diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 6f16e1e3fb4..ad33c47fc51 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -6,7 +6,7 @@ import traceback from collections.abc import Mapping, Sequence from typing import Final, Literal, cast -from pydantic import ConfigDict, TypeAdapter, with_config +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, TypedDict import litellm @@ -222,6 +222,9 @@ async def convert_to_streaming_response_async( role=choice["message"]["role"], function_call=choice["message"].get("function_call", None), tool_calls=choice["message"].get("tool_calls", None), + provider_specific_fields=_cached_provider_specific_fields( + _CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(choice["message"]) + ), ) finish_reason = choice.get("finish_reason", None) @@ -243,6 +246,9 @@ async def convert_to_streaming_response_async( completion_tokens=response_object["usage"].get("completion_tokens", 0), prompt_tokens=response_object["usage"].get("prompt_tokens", 0), total_tokens=response_object["usage"].get("total_tokens", 0), + completion_tokens_details=_cached_completion_tokens_details( + _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + ), ), ) @@ -334,6 +340,9 @@ def convert_to_streaming_response( model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) + model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details( + _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + ) if "id" in response_object: model_response_object.id = response_object["id"] @@ -390,6 +399,35 @@ class _ParallelToolUseArgs(TypedDict): _PARALLEL_TOOL_USE_ARGS: Final = TypeAdapter(_ParallelToolUseArgs) +class _CachedMessageProviderFields(TypedDict, total=False): + provider_specific_fields: ReadOnly[object] + + +class _CachedUsageDetails(TypedDict, total=False): + completion_tokens_details: ReadOnly[object] + + +_CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields) +_CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails) +_FIELD_MAP: Final = TypeAdapter(dict[str, object]) +_COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper) + + +def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None: + fields: Final = message.get("provider_specific_fields") + return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else None + + +def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionTokensDetailsWrapper | None: + details: Final = usage.get("completion_tokens_details") + if not isinstance(details, dict): + return None + try: + return _COMPLETION_TOKENS_DETAILS.validate_python(details) + except ValidationError: + return None + + def handle_invalid_parallel_tool_calls( tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall], ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 236286a493b..484978bae04 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -280,6 +280,7 @@ class CustomStreamWrapper: self._response_headers = _response_headers self.response_id: str | None = None + self.received_model_id: str | None = None self.logging_loop = None self.rules = Rules() self.stream_options = stream_options or getattr(logging_obj, "stream_options", None) @@ -838,6 +839,7 @@ class CustomStreamWrapper: self.response_id = id if id and isinstance(id, str) and id.strip(): + self.received_model_id = id model_response.hidden_params["received_model_id"] = id if self.response_id is not None and isinstance(self.response_id, str): @@ -1756,7 +1758,9 @@ class CustomStreamWrapper: self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) def finish_reason_handler(self): - model_response: Final = self.model_response_creator() + model_response: Final = self.model_response_creator( + hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} + ) _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason if _finish_reason is not None: model_response.choices[0].finish_reason = _finish_reason diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index dff10473ae7..4e5b761b9a3 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -66,7 +66,7 @@ from litellm.types.utils import ( ) from ...base import BaseLLM -from ..common_utils import AnthropicError, process_anthropic_headers +from ..common_utils import AnthropicError, process_anthropic_headers, refusal_provider_specific_fields from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig, JsonModeRequest if TYPE_CHECKING: @@ -1102,7 +1102,8 @@ class ModelResponseIterator: # These are automatically handled by Anthropic API, we just pass them through pass elif type_chunk == "message_delta": - finish_reason, usage, container = self._handle_message_delta(chunk) + finish_reason, usage, container, refusal_fields = self._handle_message_delta(chunk) + provider_specific_fields.update(refusal_fields) text = self._release_answer_close() + self._drain_split_answer_blocks() if container: provider_specific_fields["container"] = container @@ -1297,7 +1298,9 @@ class ModelResponseIterator: return "" return _split_answer_text(blocks, self.response_format_schema) - def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]: + def _handle_message_delta( + self, chunk: dict + ) -> tuple[str, Usage | None, dict[str, Any] | None, Mapping[str, object]]: """ Handle message_delta event for finish_reason, usage, and container. @@ -1319,7 +1322,7 @@ class ModelResponseIterator: self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None ) container: Final = message_delta["delta"].get("container") - return finish_reason, usage, container + return finish_reason, usage, container, refusal_provider_specific_fields(message_delta["delta"]) def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None: """ diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index f4e50bf3acd..487f231e1db 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -104,10 +104,12 @@ from litellm.utils import ( ) from ..common_utils import ( + REFUSAL_STOP_FIELDS, AnthropicError, AnthropicModelInfo, eager_input_streaming_flag, process_anthropic_headers, + refusal_provider_specific_fields, requires_native_compaction_beta, strip_advisor_blocks_from_messages, ) @@ -2723,7 +2725,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if compaction_blocks is not None: provider_specific_fields["compaction_blocks"] = compaction_blocks - return provider_specific_fields + refusal_fields: Final = refusal_provider_specific_fields( + REFUSAL_STOP_FIELDS.validate_python(completion_response) + ) + return {**provider_specific_fields, **refusal_fields} def transform_parsed_response( self, @@ -2808,7 +2813,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _message.provider_specific_fields = provider_specific_fields if json_mode_message is not None: + refusal_fields: Final = refusal_provider_specific_fields( + REFUSAL_STOP_FIELDS.validate_python(completion_response) + ) completion_response["stop_reason"] = "stop" + if refusal_fields: + json_mode_message.provider_specific_fields = dict(refusal_fields) _message = json_mode_message model_response.choices[0].message = _message diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 47a82aa097c..dbcb7ea4962 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -8,11 +8,12 @@ import re from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import Any, ClassVar, Final, Literal, TypeVar +from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar from urllib.parse import quote, urlparse import httpx -from pydantic import ConfigDict, Field, StrictBool, TypeAdapter, ValidationError +from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -2150,3 +2151,35 @@ def create_anthropic_model_list_response( "first_id": data[0]["id"] if data else None, "last_id": data[-1]["id"] if data else None, } + + +class RefusalStopFields(TypedDict, total=False): + stop_reason: ReadOnly[object] + stop_details: ReadOnly[object] + + +def _is_dict(value: object) -> bool: + return isinstance(value, dict) + + +def _no_stop_fields() -> Mapping[str, object]: + return {} + + +def _dict_or_empty(value: object) -> object: + return value if _is_dict(value) else _no_stop_fields() + + +REFUSAL_STOP_FIELDS: Final[TypeAdapter[RefusalStopFields]] = TypeAdapter( + Annotated[RefusalStopFields, BeforeValidator(_dict_or_empty)] +) + + +def refusal_provider_specific_fields(message: Mapping[str, object]) -> Mapping[str, object]: + stop_reason: Final = message.get("stop_reason") + stop_details: Final = message.get("stop_details") + if stop_details is not None: + return {"stop_reason": stop_reason, "stop_details": stop_details} + if stop_reason == "refusal": + return {"stop_reason": stop_reason} + return {} diff --git a/tests/integration/_support/anthropic_refusal.py b/tests/integration/_support/anthropic_refusal.py new file mode 100644 index 00000000000..6d5fd236074 --- /dev/null +++ b/tests/integration/_support/anthropic_refusal.py @@ -0,0 +1,533 @@ +import hashlib +import json +import os +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final, TypeAlias + +import yaml +from integration._support.anthropic_thinking import ( + JSON_LIST, + JSON_OBJECT, + Event, + marker_of, + stream_reply, + streams, +) +from integration._support.client import eventually +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response" +DOCS_STREAMING_REFUSAL_URL: Final = ( + "https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals" +) + +# Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12. +DOCS_REFUSAL: Final[Event] = { + "id": "msg_01XFUDYJgAACzvnptvVoYEL", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", + }, + "usage": {"input_tokens": 412, "output_tokens": 0}, +} +DOCS_PARTIAL_REFUSAL: Final[Event] = { + "role": "assistant", + "content": [{"type": "text", "text": "Hello.."}], + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", + }, +} + +# Captured from api.anthropic.com (claude-sonnet-5-5, prompt "Reply with the single word: ready") on 2026-10-13. +CAPTURED_TURN: Final[Event] = { + "model": "claude-sonnet-5-5", + "id": "msg_011CfyXFGze1cDkk8o7gDRRS", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ready"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, +} +CAPTURED_EVENTS: Final[tuple[Event, ...]] = ( + { + "type": "message_start", + "message": { + "model": "claude-sonnet-5-5", + "id": "msg_011CfyXFLwXUxsp7nNAm3gcu", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ready"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 4, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, +) + +DOCS_STOP_DETAILS: Final = DOCS_REFUSAL["stop_details"] +REFUSAL_FIELDS: Final[Event] = {"stop_reason": "refusal", "stop_details": DOCS_STOP_DETAILS} +STOP_FIELDS: Final = ("stop_reason", "stop_details") +LONG_STOP_DETAILS: Final = "d" * 5120 +MALFORMED_STOP_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"int": 7, "list": ["refusal"], "empty": "", "long": LONG_STOP_DETAILS} +) +_FIXTURE: Final = re.compile(r"refusal-fixture-([a-z]+)") + +Location: TypeAlias = tuple[str | int, ...] + + +@dataclass(frozen=True, slots=True) +class Fixture: + status: int + body: bytes + events: tuple[Event, ...] = () + + +def _text_block(index: int, block: JsonValue) -> tuple[Event, ...]: + text: Final = JSON_OBJECT.validate_python(block)["text"] + return ( + {"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": index}, + ) + + +def refusal_events(message: Mapping[str, JsonValue]) -> tuple[Event, ...]: + usage: Final = JSON_OBJECT.validate_python(message["usage"]) + start: Final[Event] = { + "type": "message_start", + "message": { + "id": message["id"], + "type": "message", + "role": "assistant", + "model": message["model"], + "content": [], + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": usage, + }, + } + blocks: Final = chain.from_iterable( + _text_block(index, block) for index, block in enumerate(JSON_LIST.validate_python(message["content"])) + ) + delta: Final[Event] = { + "type": "message_delta", + "delta": { + "stop_reason": message["stop_reason"], + "stop_sequence": None, + "stop_details": message["stop_details"], + }, + "usage": {name: usage[name] for name in ("output_tokens", "output_tokens_details") if name in usage}, + } + return (start, *blocks, delta, {"type": "message_stop"}) + + +def _without_stop_details(value: Mapping[str, JsonValue]) -> Event: + return {name: item for name, item in value.items() if name != "stop_details"} + + +def _captured_event(event: Event, stop_details: JsonValue, *, keep: bool) -> Event: + if event["type"] == "message_start": + message: Final = JSON_OBJECT.validate_python(event["message"]) + return {**event, "message": message if keep else _without_stop_details(message)} + if event["type"] != "message_delta": + return event + delta: Final = JSON_OBJECT.validate_python(event["delta"]) + replaced: Final = {**delta, "stop_details": stop_details} if keep else _without_stop_details(delta) + return {**event, "delta": replaced} + + +def _captured(stop_details: JsonValue, *, keep: bool = True) -> Fixture: + body: Final = {**CAPTURED_TURN, "stop_details": stop_details} if keep else _without_stop_details(CAPTURED_TURN) + return Fixture( + 200, + json.dumps(body).encode(), + tuple(_captured_event(event, stop_details, keep=keep) for event in CAPTURED_EVENTS), + ) + + +def _duplicated_stop_details() -> bytes: + pair: Final = f'"stop_details": {json.dumps(DOCS_STOP_DETAILS)}' + return json.dumps(DOCS_REFUSAL).replace(pair, f"{pair}, {pair}").encode() + + +def _anthropic_error(kind: str, message: str) -> bytes: + return json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode() + + +_PARTIAL: Final[Event] = {**DOCS_REFUSAL, **DOCS_PARTIAL_REFUSAL} +_REASONED: Final[Event] = { + **_PARTIAL, + "usage": {"input_tokens": 412, "output_tokens": 10, "output_tokens_details": {"thinking_tokens": 4}}, +} + +JSON_TOOL_NAME: Final = "json_tool_call" +JSON_ANSWER_SCHEMA: Final[Event] = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + "additionalProperties": False, +} +JSON_ANSWER_FORMAT: Final[Event] = { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": JSON_ANSWER_SCHEMA, "strict": True}, +} +DOCS_PARTIAL_TEXT: Final = "Hello.." +_DECLARED_INPUT: Final[Event] = {"answer": DOCS_PARTIAL_TEXT} +_ENVELOPED_INPUT: Final[Event] = {"parameters": _DECLARED_INPUT} +DECLARED_PARTIAL_JSON: Final = '{"answer": "Hello..' +ENVELOPED_PARTIAL_JSON: Final = '{"parameters": {"answer": "Hello..' +_JSON_TOOL_USE_ID: Final = "toolu_01RefusalFixtureJsonTool" + + +def _json_tool_refusal(tool_input: Event, partial_json: str) -> Fixture: + message: Final[Event] = { + **DOCS_REFUSAL, + "content": [{"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": tool_input}], + } + start, delta, stop = refusal_events(DOCS_REFUSAL) + tool_block: Final[tuple[Event, ...]] = ( + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": {}}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": partial_json}, + }, + {"type": "content_block_stop", "index": 0}, + ) + return Fixture(200, json.dumps(message).encode(), (start, *tool_block, delta, stop)) + + +FIXTURES: Final[Mapping[str, Fixture]] = MappingProxyType( + { + "refusal": Fixture(200, json.dumps(DOCS_REFUSAL).encode(), refusal_events(DOCS_REFUSAL)), + "partial": Fixture(200, json.dumps(_PARTIAL).encode(), refusal_events(_PARTIAL)), + "reasoned": Fixture(200, json.dumps(_REASONED).encode(), refusal_events(_REASONED)), + "jsontool": _json_tool_refusal(_DECLARED_INPUT, DECLARED_PARTIAL_JSON), + "jsonheld": _json_tool_refusal(_ENVELOPED_INPUT, ENVELOPED_PARTIAL_JSON), + "duplicate": Fixture(200, _duplicated_stop_details()), + "null": _captured(None), + "missing": _captured(None, keep=False), + **{name: _captured(value) for name, value in MALFORMED_STOP_DETAILS.items()}, + "nondict": Fixture(200, b"[]"), + "unauthorized": Fixture(401, _anthropic_error("authentication_error", "invalid x-api-key")), + "invalid": Fixture(400, _anthropic_error("invalid_request_error", "max_tokens: Field required")), + } +) + + +def refusal_prompt(fixture: str, marker: str) -> str: + return f"refusal-fixture-{fixture} marker-{marker}" + + +CITED_FIELDS: Final[Event] = {"citations": ["https://example.com/source"]} +_PLAIN_DELTA: Final[Event] = {"role": "assistant", "content": "ready"} +_CITED_DELTA: Final[Event] = {**_PLAIN_DELTA, "provider_specific_fields": CITED_FIELDS} +_OPENAI_USAGE: Final[Event] = { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + "completion_tokens_details": {"reasoning_tokens": 4, "text_tokens": 3}, +} +UNVALIDATED_USAGE_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "text": {"reasoning_tokens": "abc"}, + "fraction": {"reasoning_tokens": 1.5}, + "object": {"reasoning_tokens": {"a": 1}}, + } +) + + +@dataclass(frozen=True, slots=True) +class OpenAIFixture: + first_delta: Event + usage: tuple[Event, ...] = () + + +_OPENAI_FIXTURES: Final[Mapping[str, OpenAIFixture]] = MappingProxyType( + { + "plain": OpenAIFixture(_PLAIN_DELTA), + "cited": OpenAIFixture(_CITED_DELTA), + "plainusage": OpenAIFixture(_PLAIN_DELTA, (_OPENAI_USAGE,)), + "citedusage": OpenAIFixture(_CITED_DELTA, (_OPENAI_USAGE,)), + } +) + + +def openai_id(marker: str) -> str: + return f"chatcmpl-{marker}" + + +def _openai_frame(marker: str, body: Event) -> bytes: + chunk: Final[Event] = { + "id": openai_id(marker), + "object": "chat.completion.chunk", + "created": 1760000000, + "model": "gpt-4o-mini", + **body, + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +def _openai_choice(delta: Event, finish_reason: str | None) -> Event: + return {"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]} + + +def _openai_stream(name: str, request: Request) -> Reply: + assert streams(request), request.body + marker: Final = marker_of(request) + fixture: Final = _OPENAI_FIXTURES[name] + return Reply( + content_type="text/event-stream", + chunks=( + _openai_frame(marker, _openai_choice(fixture.first_delta, None)), + _openai_frame(marker, _openai_choice({"content": " now"}, "stop")), + *(_openai_frame(marker, {"choices": [], "usage": usage}) for usage in fixture.usage), + b"data: [DONE]\n\n", + ), + ) + + +def respond(request: Request) -> Reply: + found: Final[Sequence[str]] = _FIXTURE.findall(request.body.decode()) + assert found, request.body + if request.target.endswith("/chat/completions"): + return _openai_stream(found[-1], request) + fixture: Final = FIXTURES[found[-1]] + if fixture.status != 200 or not fixture.events or not streams(request): + return Reply(status=fixture.status, body=fixture.body) + return stream_reply(request, fixture.events) + + +def _stop_fields(fields: JsonValue) -> Event: + if not isinstance(fields, dict): + return {} + return {name: fields[name] for name in STOP_FIELDS if name in fields} + + +def relayed(value: JsonValue, location: Location = ()) -> tuple[tuple[Location, Event], ...]: + if isinstance(value, list): + return tuple(chain.from_iterable(relayed(item, (*location, index)) for index, item in enumerate(value))) + if not isinstance(value, dict): + return () + found: Final = _stop_fields(value.get("provider_specific_fields")) + here: Final = (((*location, "provider_specific_fields"), found),) if found else () + return (*here, *chain.from_iterable(relayed(item, (*location, name)) for name, item in value.items())) + + +_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") +_STORED: Final = TypeAdapter(bytes | None) +STALE_ID_SUFFIX: Final = "-stale" + + +def _stored(store: Redis, key: str) -> bytes: + if not _CACHE_KEY.fullmatch(key) or store.type(key) != b"string": + return b"" + return _STORED.validate_python(store.get(key)) or b"" + + +def _entry_key(store: Redis, needle: str) -> str | None: + keys: Final = (raw.decode() for raw in store.scan_iter(count=1000)) + return next((key for key in keys if needle.encode() in _stored(store, key)), None) + + +def _with_message_fields(response: Event, fields: JsonValue) -> Event: + choice: Final = JSON_OBJECT.validate_python(JSON_LIST.validate_python(response["choices"])[0]) + message: Final = JSON_OBJECT.validate_python(choice["message"]) + return {**response, "choices": [{**choice, "message": {**message, "provider_specific_fields": fields}}]} + + +def _with_usage_details(response: Event, details: JsonValue) -> Event: + usage: Final = JSON_OBJECT.validate_python(response["usage"]) + return {**response, "usage": {**usage, "completion_tokens_details": details}} + + +def _rewrite_cached_response(needle: str, rewrite: Callable[[Event], Event]) -> str: + store: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + key: Final = eventually(lambda: _entry_key(store, needle), lambda found: found is not None) + assert key is not None, needle + entry: Final = JSON_OBJECT.validate_json(_stored(store, key)) + stored: Final = entry["response"] + encoded: Final = isinstance(stored, str) + response: Final = ( + JSON_OBJECT.validate_json(stored) if isinstance(stored, str) else JSON_OBJECT.validate_python(stored) + ) + rewritten: Final = {**rewrite(response), "id": f"{response['id']}{STALE_ID_SUFFIX}"} + store.set(key, json.dumps({**entry, "response": json.dumps(rewritten) if encoded else rewritten}), ex=60) + return str(rewritten["id"]) + + +def rewrite_cached_message_fields(needle: str, fields: JsonValue) -> str: + return _rewrite_cached_response(needle, lambda response: _with_message_fields(response, fields)) + + +def rewrite_cached_usage_details(needle: str, details: JsonValue) -> str: + return _rewrite_cached_response(needle, lambda response: _with_usage_details(response, details)) + + +def key_hash(key: str) -> str: + return hashlib.sha256(key.encode()).hexdigest() + + +def _key_hash_of(event: Mapping[str, JsonValue]) -> JsonValue: + return JSON_OBJECT.validate_python(event.get("metadata") or {}).get("user_api_key_hash") + + +class Ledger: + def __init__(self, wire: Wire) -> None: + self.wire: Final = wire + self.batches: tuple[Request, ...] = () + + def events(self) -> tuple[Event, ...]: + self.batches = (*self.batches, *self.wire.drain()) + decoded: Final = chain.from_iterable(JSON_LIST.validate_json(batch.body) for batch in self.batches) + return tuple(JSON_OBJECT.validate_python(event) for event in decoded) + + def for_key(self, key: str) -> tuple[Event, ...]: + hashed: Final = key_hash(key) + return tuple(event for event in self.events() if _key_hash_of(event) == hashed) + + def await_key(self, key: str, count: int = 1, seconds: float = 30) -> tuple[Event, ...]: + return eventually(lambda: self.for_key(key), lambda found: len(found) >= count, seconds=seconds) + + +_HIDDEN_PARAMS_PROBE: Final = """ +import httpx +from litellm.integrations.custom_logger import CustomLogger + + +class HiddenParamsProbe(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + payload = kwargs.get("standard_logging_object") or {{}} + hidden = getattr(response_obj, "_hidden_params", None) or {{}} + event = {{ + "metadata": {{"user_api_key_hash": (payload.get("metadata") or {{}}).get("user_api_key_hash")}}, + "response_id": getattr(response_obj, "id", None), + "received_model_id": hidden.get("received_model_id"), + }} + async with httpx.AsyncClient(trust_env=False) as client: + await client.post({url!r}, json=[event]) + + +instance = HiddenParamsProbe() +""" + + +def hidden_params_probe(directory: Path, url: str) -> str: + (directory / "hidden_params_probe.py").write_text(_HIDDEN_PARAMS_PROBE.format(url=url)) + return "hidden_params_probe.instance" + + +def sink_config( + directory: Path, + sinks: Mapping[str, str], + *, + retries: int = 0, + extra_callbacks: Sequence[str] = (), + store_prompts: bool = False, +) -> Path: + shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = { + **JSON_OBJECT.validate_python(shared["litellm_settings"]), + "callbacks": [*sinks, *extra_callbacks], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + } + callback_settings: Final = { + name: { + "callback_type": "generic_api", + "endpoint": url, + "headers": {"x-refusal-sink": name}, + "max_retries": retries, + "retry_delay": 0.5, + } + for name, url in sinks.items() + } + general_settings: Final = { + **JSON_OBJECT.validate_python(shared["general_settings"]), + "store_prompts_in_spend_logs": store_prompts, + } + path: Final = directory / "anthropic-refusal-sinks.yaml" + path.write_text( + yaml.safe_dump( + { + **shared, + "general_settings": general_settings, + "litellm_settings": settings, + "callback_settings": callback_settings, + } + ) + ) + return path + + +def frames(text: str) -> tuple[Event, ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: {") + ) + + +def frame_types(events: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + return tuple(event.get("type") for event in events) diff --git a/tests/integration/caching/response_cache_case.py b/tests/integration/caching/response_cache_case.py index 9df34721202..8dfb6a0a479 100644 --- a/tests/integration/caching/response_cache_case.py +++ b/tests/integration/caching/response_cache_case.py @@ -336,6 +336,29 @@ def flip_entry_to_empty(store: Redis, key: str, field: str) -> str: return stale_id +def rewrite_usage_details(store: Redis, key: str, details: JsonValue) -> str: + """Rewrite the stored usage's ``completion_tokens_details`` to ``details`` and stamp the id ``-stale``, keeping the + encoding, and return the stale id.""" + raw: Final = store.get(key) + assert raw is not None, key + entry, is_json = _decode_entry(raw) + response: Final = entry["response"] + stored: Final = object_value(json.loads(response) if isinstance(response, str) else response) + stale_id: Final = f"{string_value(stored['id'])}{STALE_ID_SUFFIX}" + usage: Final = object_value(stored["usage"]) + rewritten_response: Final[dict[str, JsonValue]] = { + **stored, + "usage": {**usage, "completion_tokens_details": details}, + "id": stale_id, + } + rewritten: Final[dict[str, JsonValue]] = { + **entry, + "response": json.dumps(rewritten_response) if isinstance(response, str) else rewritten_response, + } + store.set(key, json.dumps(rewritten) if is_json else str(rewritten), ex=STALE_ENTRY_TTL_SECONDS) + return stale_id + + def chat_body(model: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]: return {"model": model, "messages": [{"role": "user", "content": text}], **extra} diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py new file mode 100644 index 00000000000..dd4b7602b96 --- /dev/null +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py @@ -0,0 +1,337 @@ +import asyncio +import json +import re +import threading +import time +import uuid +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.anthropic_refusal import ( + REFUSAL_FIELDS, + Ledger, + Location, + refusal_prompt, + relayed, + respond, + sink_config, +) +from integration._support.anthropic_thinking import JSON_OBJECT, MODEL, Event, chunks_of +from integration._support.client import Gateway, eventually, held +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.timing import assert_not_delayed_by +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark: Final = [ + pytest.mark.timeout(2 * graceful_stop_seconds() + 120), + pytest.mark.xdist_group("anthropic_refusal_stop_details_chaos"), +] + +_MODEL_NAME: Final = "anthropic-refusal-chaos" +_SLOW_SINK_SECONDS: Final = 4 +_SINK: Final = "refusal_chaos_sink" +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields") +_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields") +_PATHS: Final[Mapping[str, str]] = MappingProxyType( + {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"} +) + +Surface: TypeAlias = Literal["chat", "messages", "responses"] +_RecordProperty: TypeAlias = Callable[[str, object], None] + + +@dataclass(frozen=True, slots=True) +class _Call: + surface: Surface + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +@dataclass(frozen=True, slots=True) +class _Held: + release: threading.Event + markers: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + self.markers.put(_markers_in(request.body.decode())[0]) + assert self.release.wait(timeout=90), "The burst was never released" + return respond(request) + + +def _markers_in(text: str) -> tuple[str, ...]: + return tuple(dict.fromkeys(_MARKER.findall(text))) + + +def _calls(count: int) -> tuple[_Call, ...]: + surfaces: Final[tuple[Surface, ...]] = ("chat", "messages", "responses") + return tuple( + _Call(surfaces[index % 3], stream=(index // 3) % 2 == 0, marker=uuid.uuid4().hex) for index in range(count) + ) + + +def _body(call: _Call) -> Event: + text: Final = refusal_prompt("refusal", call.marker) + shared: Final[Event] = {"model": _MODEL_NAME, "stream": call.stream, "cache": {"no-cache": True}} + match call.surface: + case "chat": + return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64} + case "messages": + return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64} + case "responses": + return {**shared, "input": text, "max_output_tokens": 64} + + +async def _send(client: httpx.AsyncClient, key: str, call: _Call) -> _Served: + try: + response: Final = await client.post( + _PATHS[call.surface], json=_body(call), headers={"Authorization": f"Bearer {key}"} + ) + except httpx.TransportError as error: + return _Served(call, 0, repr(error)) + return _Served(call, response.status_code, response.text) + + +async def _burst(gateway: Gateway, calls: Sequence[_Call]) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client: + return tuple(await asyncio.gather(*(_send(client, gateway.key, call) for call in calls))) + + +def _chaos_config(directory: Path, upstream: Wire, sink: str, *, retries: int) -> Path: + base: Final = sink_config(directory, {_SINK: sink}, retries=retries) + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(base.read_text())) + model_list: Final = [ + { + "model_name": _MODEL_NAME, + "litellm_params": {"model": f"anthropic/{MODEL}", "api_base": upstream.url, "api_key": "scripted-key"}, + } + ] + path: Final = directory / "anthropic-refusal-chaos.yaml" + path.write_text(yaml.safe_dump({**config, "model_list": model_list})) + return path + + +def _ledger_events(ledgers: Sequence[Ledger]) -> tuple[Event, ...]: + return tuple(chain.from_iterable(ledger.events() for ledger in ledgers)) + + +def _marker_counts(ledgers: Sequence[Ledger], markers: frozenset[str]) -> Mapping[str, int]: + found: Final = chain.from_iterable(_markers_in(json.dumps(event)) for event in _ledger_events(ledgers)) + return MappingProxyType(Counter(marker for marker in found if marker in markers)) + + +def _events_for(ledgers: Sequence[Ledger], marker: str) -> tuple[Event, ...]: + return tuple(event for event in _ledger_events(ledgers) if marker in json.dumps(event)) + + +def _assert_answered(item: _Served) -> None: + assert item.status == 200, (item.call, item.text) + if item.call.surface != "chat": + return + if item.call.stream: + found: Final = tuple(chain.from_iterable(relayed(chunk) for chunk in chunks_of(item.text))) + assert found == ((_DELTA_FIELDS, REFUSAL_FIELDS),), item.text + return + assert relayed(JSON_OBJECT.validate_json(item.text)) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), item.text + + +def _assert_logged_chat_relays(ledgers: Sequence[Ledger], calls: Sequence[_Call]) -> None: + for call in calls: + if call.surface != "chat": + continue + (event,) = _events_for(ledgers, call.marker) + assert relayed(event["response"]) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), event + + +def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: + return eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _port(wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + assert port is not None, wire.url + return port + + +def _accept(_request: Request) -> Reply: + return Reply() + + +async def _timed_chat(gateway: Gateway) -> float: + started: Final = time.monotonic() + (served,) = await _burst(gateway, (_Call("chat", stream=False, marker=uuid.uuid4().hex),)) + _assert_answered(served) + return time.monotonic() - started + + +async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(30) + markers: Final = frozenset(call.marker for call in calls) + with wire_server(respond) as upstream, ExitStack() as first_sink: + sink: Final = first_sink.enter_context(wire_server(_accept)) + first: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=6) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + _worker_pids(owned) + served_first: Final = await _burst(owned.gateway, calls[:10]) + await asyncio.to_thread( + eventually, lambda: sum(_marker_counts((first,), markers).values()), lambda total: total >= 1, 30 + ) + first.events() + first_sink.close() + served_during: Final = await _burst(owned.gateway, calls[10:20]) + with wire_server(_accept, port=_port(sink)) as restarted: + second: Final = Ledger(restarted) + served_after: Final = await _burst(owned.gateway, calls[20:]) + for item in (*served_first, *served_during, *served_after): + _assert_answered(item) + counts: Final = await asyncio.to_thread( + held, + lambda: _marker_counts((first, second), markers), + lambda found: set(found) == markers and set(found.values()) == {1}, + holding=3, + seconds=60, + ) + assert sum(counts.values()) == 30, counts + _assert_logged_chat_relays((first, second), calls) + assert len(upstream.drain()) == 30 + + +async def test_slow_sink_neither_deadlocks_nor_uploads_a_refusal_event_twice(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(30) + markers: Final = frozenset(call.marker for call in calls) + + def slow(_request: Request) -> Reply: + time.sleep(_SLOW_SINK_SECONDS) + return Reply() + + with wire_server(respond) as upstream, wire_server(slow) as sink: + ledger: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + _worker_pids(owned) + control: Final = await _timed_chat(owned.gateway) + served: Final = await _burst(owned.gateway, calls) + for item in served: + _assert_answered(item) + follow_up: Final = await _timed_chat(owned.gateway) + assert_not_delayed_by((follow_up,), control=(control,), injected_delay=_SLOW_SINK_SECONDS) + counts: Final = await asyncio.to_thread( + held, + lambda: _marker_counts((ledger,), markers), + lambda found: set(found) == markers and set(found.values()) == {1}, + holding=5, + seconds=90, + ) + assert sum(counts.values()) == 30, counts + _assert_logged_chat_relays((ledger,), calls) + + +async def test_worker_sigkill_mid_burst_leaves_the_sibling_relaying_and_logging_once( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(20) + hold: Final = _Held(threading.Event(), SimpleQueue()) + with wire_server(hold.respond) as upstream, wire_server(_accept) as sink: + ledger: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + workers: Final = _worker_pids(owned) + burst: Final = asyncio.create_task(_burst(owned.gateway, calls)) + await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, upstream.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + hold.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + completed: Final = tuple(item for item in served if item.status == 200) + assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served]) + for item in completed: + _assert_answered(item) + follow_up: Final = _Call("chat", stream=True, marker=uuid.uuid4().hex) + (answered,) = await _burst(owned.gateway, (follow_up,)) + _assert_answered(answered) + survivors: Final = frozenset((*(item.call.marker for item in completed), follow_up.marker)) + counts: Final = await asyncio.to_thread( + held, + lambda: _marker_counts((ledger,), frozenset(call.marker for call in (*calls, follow_up))), + lambda found: set(found) == survivors and set(found.values()) == {1}, + holding=3, + seconds=60, + ) + assert sum(counts.values()) == len(survivors), counts + _assert_logged_chat_relays((ledger,), tuple(item.call for item in completed)) + + +async def test_proxy_restart_mid_burst_never_logs_a_refusal_twice( + gateway: Gateway, tmp_path: Path, record_property: _RecordProperty +) -> None: + calls: Final = _calls(20) + markers: Final = frozenset(call.marker for call in calls) + hold: Final = _Held(threading.Event(), SimpleQueue()) + with wire_server(hold.respond) as upstream, wire_server(_accept) as sink: + ledger: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first: + _worker_pids(first) + burst: Final = asyncio.create_task(_burst(first.gateway, calls)) + await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60) + first.begin_shutdown() + hold.release.set() + served: Final = await burst + for item in served: + _assert_answered(item) + second_directory: Final = tmp_path / "second" + second_directory.mkdir() + with owned_proxy_process(gateway, second_directory, {}, config=config, workers=2) as second: + follow_up: Final = _Call("chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(second.gateway, (follow_up,)) + _assert_answered(answered) + await asyncio.to_thread( + eventually, + lambda: _marker_counts((ledger,), frozenset({follow_up.marker})), + lambda found: found[follow_up.marker] == 1, + 30, + ) + counts: Final = _marker_counts((ledger,), markers) + assert set(counts.values()) <= {1}, counts + _assert_logged_chat_relays((ledger,), tuple(call for call in calls if call.marker in counts)) + record_property("served", len(served)) + record_property("landed", sum(counts.values())) diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py new file mode 100644 index 00000000000..461839ccbcc --- /dev/null +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py @@ -0,0 +1,1356 @@ +import asyncio +import json +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from typing import Final, Literal, TypeAlias +from urllib.parse import unquote, urlsplit + +import anthropic +import httpx +import openai +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.anthropic_refusal import ( + CITED_FIELDS, + DECLARED_PARTIAL_JSON, + DOCS_PARTIAL_TEXT, + DOCS_REFUSAL, + DOCS_STOP_DETAILS, + ENVELOPED_PARTIAL_JSON, + FIXTURES, + JSON_ANSWER_FORMAT, + JSON_TOOL_NAME, + MALFORMED_STOP_DETAILS, + REFUSAL_FIELDS, + UNVALIDATED_USAGE_DETAILS, + Ledger, + Location, + frame_types, + frames, + hidden_params_probe, + openai_id, + refusal_events, + refusal_prompt, + relayed, + respond, + rewrite_cached_message_fields, + rewrite_cached_usage_details, + sink_config, +) +from integration._support.anthropic_thinking import ( + BEDROCK_MODEL, + JSON_LIST, + JSON_OBJECT, + MODEL, + NO_CACHE, + Event, + aws_chunks, + chunks_of, + sse_chunks, +) +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, held +from integration._support.forward_proxy import tunnelling_forward_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Wire, wire_server +from pydantic import BaseModel, JsonValue + +pytestmark: Final = [ + pytest.mark.timeout(2 * graceful_stop_seconds() + 120), + pytest.mark.xdist_group("anthropic_refusal_stop_details_wire"), +] + +_ANTHROPIC_BASE: Final = "http://api.anthropic.com" +_ANTHROPIC_KEY: Final = "scripted-anthropic-key" +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_HOST: Final = f"{_VERTEX_LOCATION}-aiplatform.googleapis.com" +_VERTEX_MODEL_PATH: Final = ( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{MODEL}" +) +_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields") +_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields") +_RESPONSE_FIELDS: Final[Location] = ("provider_specific_fields",) +_STREAMED_RESPONSE_FIELDS: Final[Location] = ("response", "provider_specific_fields") + +Client: TypeAlias = Literal["sdk", "async_sdk", "httpx"] +Provider: TypeAlias = Literal[ + "anthropic", + "bedrock_invoke", + "claude_platform", + "vertex_ai", + "azure_ai", + "bedrock_mantle", + "snowflake", + "openai_compatible", +] +Surface: TypeAlias = Literal["chat", "chat_stream", "messages", "messages_stream", "responses", "responses_stream"] + +CLIENTS: Final[tuple[Client, ...]] = ("sdk", "async_sdk", "httpx") +FAMILY: Final[tuple[Provider, ...]] = ( + "bedrock_invoke", + "claude_platform", + "vertex_ai", + "azure_ai", + "bedrock_mantle", + "snowflake", +) +_STREAM_TARGETS: Final[Mapping[Provider, tuple[str, str]]] = { + "anthropic": ("/v1/messages", "/v1/messages"), + "bedrock_invoke": (f"/model/{BEDROCK_MODEL}/invoke", f"/model/{BEDROCK_MODEL}/invoke-with-response-stream"), + "claude_platform": ("/v1/messages", "/v1/messages"), + "vertex_ai": (f"{_VERTEX_MODEL_PATH}:rawPredict", f"{_VERTEX_MODEL_PATH}:streamRawPredict"), + "azure_ai": ("/anthropic/v1/messages", "/anthropic/v1/messages"), + "bedrock_mantle": ("/anthropic/v1/messages", "/anthropic/v1/messages"), + "snowflake": ("/api/v2/cortex/v1/messages", "/api/v2/cortex/v1/messages"), +} + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + wire: Wire + vertex: Wire + sink: Ledger + forbidden: Ledger + absent: Ledger + hidden: Ledger + credentials: str + + +def _service_account_json(token_url: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": _VERTEX_PROJECT, + "private_key_id": "scripted", + "private_key": private_key, + "client_email": f"scripted@{_VERTEX_PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{token_url}/_oauth/authorize", + "token_uri": f"{token_url}/_oauth/token", + } + ) + + +def _port(wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + assert port is not None, wire.url + return port + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("anthropic-refusal-stop-details") + certificate: Final = write_self_signed_cert(directory, (_VERTEX_HOST,)) + with ( + gateway_from_environment() as parent, + wire_server(respond) as wire, + wire_server(respond, tls=server_context(*certificate)) as vertex, + tunnelling_forward_proxy(_port(vertex)) as tunnel, + wire_server(lambda _request: Reply()) as sink, + wire_server(lambda _request: Reply(status=403)) as forbidden, + wire_server(lambda _request: Reply(status=404)) as absent, + wire_server(lambda _request: Reply()) as hidden, + ): + credentials: Final = _service_account_json(parent.upstream_url.rstrip("/")) + credentials_path: Final = directory / "vertex-service-account.json" + credentials_path.write_text(credentials) + config: Final = sink_config( + directory, + {"refusal_sink": sink.url, "refusal_sink_403": forbidden.url, "refusal_sink_404": absent.url}, + extra_callbacks=(hidden_params_probe(directory, hidden.url),), + store_prompts=True, + ) + overrides: Final = { + "ANTHROPIC_API_BASE": _ANTHROPIC_BASE, + "ANTHROPIC_API_KEY": _ANTHROPIC_KEY, + "AIOHTTP_TRUST_ENV": "True", + "HTTP_PROXY": wire.url, + "http_proxy": wire.url, + "HTTPS_PROXY": tunnel.url, + "https_proxy": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(certificate[0]), + "DEFAULT_VERTEXAI_PROJECT": _VERTEX_PROJECT, + "DEFAULT_VERTEXAI_LOCATION": _VERTEX_LOCATION, + "DEFAULT_GOOGLE_APPLICATION_CREDENTIALS": str(credentials_path), + } + with owned_proxy(parent, directory, overrides, config=config, workers=2) as owned: + yield Rig(owned, wire, vertex, Ledger(sink), Ledger(forbidden), Ledger(absent), Ledger(hidden), credentials) + + +@pytest.fixture(autouse=True) +def _drained_upstreams(rig: Rig) -> None: + rig.wire.drain() + rig.vertex.drain() + + +def _marker() -> str: + return uuid.uuid4().hex + + +def _deployment(scenario: Scenario, provider: Provider, rig: Rig) -> str: + wire_url: Final = rig.wire.url + match provider: + case "anthropic": + return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key=_ANTHROPIC_KEY) + case "bedrock_invoke": + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "claude_platform": + return scenario.model( + model=f"bedrock/claude_platform/{MODEL}", + api_base=wire_url, + api_key="scripted-platform-key", + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + case "vertex_ai": + return scenario.model( + model=f"vertex_ai/{MODEL}", + api_base=f"{wire_url}{_VERTEX_MODEL_PATH}", + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=rig.credentials, + ) + case "azure_ai": + return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key") + case "bedrock_mantle": + return scenario.model( + model=f"bedrock_mantle/{BEDROCK_MODEL}", + api_base=wire_url, + api_key="scripted-mantle-key", + aws_region_name="us-east-1", + ) + case "snowflake": + return scenario.model(model=f"snowflake/{MODEL}", api_base=wire_url, api_key="scripted-snowflake-key") + case "openai_compatible": + return scenario.model(model="openai/gpt-4o-mini", api_base=f"{wire_url}/v1", api_key="scripted-openai-key") + + +def _base_url(rig: Rig) -> str: + return str(rig.gateway.client.base_url).rstrip("/") + + +def _dump(model: BaseModel) -> Event: + return JSON_OBJECT.validate_python(model.model_dump()) + + +def _post(rig: Rig, path: str, body: Mapping[str, JsonValue], key: str) -> bytes: + response: Final = rig.gateway.request("POST", path, body, key=key) + assert response.status_code == 200, response.text + return response.content + + +def _chat_request(model: str, text: str, *, stream: bool, cached: bool) -> Event: + request: Final[Event] = { + "model": model, + "messages": [{"role": "user", "content": text}], + "stream": stream, + "max_tokens": 64, + } + return request if cached else {**request, **NO_CACHE} + + +def _chat(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event: + match client: + case "httpx": + return JSON_OBJECT.validate_json( + _post(rig, "/v1/chat/completions", _chat_request(model, text, stream=False, cached=cached), key) + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump( + sdk.chat.completions.create( + model=model, messages=[{"role": "user", "content": text}], max_tokens=64, extra_body=NO_CACHE + ) + ) + case "async_sdk": + + async def call() -> Event: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump( + await sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": text}], + max_tokens=64, + extra_body=NO_CACHE, + ) + ) + + return asyncio.run(call()) + + +def _chat_stream( + rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False +) -> tuple[Event, ...]: + match client: + case "httpx": + return chunks_of( + _post(rig, "/v1/chat/completions", _chat_request(model, text, stream=True, cached=cached), key).decode() + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return tuple( + _dump(chunk) + for chunk in sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": text}], + max_tokens=64, + stream=True, + extra_body=NO_CACHE, + ) + ) + case "async_sdk": + + async def call() -> tuple[Event, ...]: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + stream: Final = await sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": text}], + max_tokens=64, + stream=True, + extra_body=NO_CACHE, + ) + return tuple([_dump(chunk) async for chunk in stream]) + + return asyncio.run(call()) + + +def _messages_request(model: str, text: str, *, stream: bool, cached: bool) -> Event: + request: Final[Event] = { + "model": model, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": text}], + } + return request if cached else {**request, **NO_CACHE} + + +def _messages(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event: + match client: + case "httpx": + return JSON_OBJECT.validate_json( + _post(rig, "/v1/messages", _messages_request(model, text, stream=False, cached=cached), key) + ) + case "sdk": + with anthropic.Anthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + return _dump( + sdk.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": text}], extra_body=NO_CACHE + ) + ) + case "async_sdk": + + async def call() -> Event: + async with anthropic.AsyncAnthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + return _dump( + await sdk.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": text}], + extra_body=NO_CACHE, + ) + ) + + return asyncio.run(call()) + + +def _messages_stream( + rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False +) -> tuple[Event, ...]: + match client: + case "httpx": + return frames( + _post(rig, "/v1/messages", _messages_request(model, text, stream=True, cached=cached), key).decode() + ) + case "sdk": + with anthropic.Anthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + return tuple( + _dump(event) + for event in sdk.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": text}], + stream=True, + extra_body=NO_CACHE, + ) + ) + case "async_sdk": + + async def call() -> tuple[Event, ...]: + async with anthropic.AsyncAnthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + stream: Final = await sdk.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": text}], + stream=True, + extra_body=NO_CACHE, + ) + return tuple([_dump(event) async for event in stream]) + + return asyncio.run(call()) + + +def _responses_request(model: str, text: str, *, stream: bool, cached: bool) -> Event: + request: Final[Event] = {"model": model, "input": text, "max_output_tokens": 64, "stream": stream} + return request if cached else {**request, **NO_CACHE} + + +def _responses(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event: + match client: + case "httpx": + return JSON_OBJECT.validate_json( + _post(rig, "/v1/responses", _responses_request(model, text, stream=False, cached=cached), key) + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump(sdk.responses.create(model=model, input=text, max_output_tokens=64, extra_body=NO_CACHE)) + case "async_sdk": + + async def call() -> Event: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump( + await sdk.responses.create(model=model, input=text, max_output_tokens=64, extra_body=NO_CACHE) + ) + + return asyncio.run(call()) + + +def _responses_stream(rig: Rig, key: str, model: str, text: str, client: Client = "httpx") -> tuple[Event, ...]: + match client: + case "httpx": + return frames( + _post(rig, "/v1/responses", _responses_request(model, text, stream=True, cached=False), key).decode() + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return tuple( + _dump(event) + for event in sdk.responses.create( + model=model, input=text, max_output_tokens=64, stream=True, extra_body=NO_CACHE + ) + ) + case "async_sdk": + + async def call() -> tuple[Event, ...]: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + stream: Final = await sdk.responses.create( + model=model, input=text, max_output_tokens=64, stream=True, extra_body=NO_CACHE + ) + return tuple([_dump(event) async for event in stream]) + + return asyncio.run(call()) + + +def _relayed_in(events: Sequence[JsonValue]) -> tuple[tuple[Location, Event], ...]: + return tuple(chain.from_iterable(relayed(event) for event in events)) + + +def _choice(chunk: Mapping[str, JsonValue]) -> Event: + choices: Final = JSON_LIST.validate_python(chunk.get("choices") or []) + return JSON_OBJECT.validate_python(choices[0]) if choices else {} + + +def _finish_reasons(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + return tuple(_choice(chunk)["finish_reason"] for chunk in chunks if _choice(chunk).get("finish_reason")) + + +def _streamed_text(chunks: Sequence[Mapping[str, JsonValue]]) -> str: + return "".join( + str(JSON_OBJECT.validate_python(_choice(chunk).get("delta") or {}).get("content") or "") for chunk in chunks + ) + + +def _message_text(body: Mapping[str, JsonValue]) -> JsonValue: + return JSON_OBJECT.validate_python(_choice(body)["message"]).get("content") + + +def _native_stops(events: Sequence[Mapping[str, JsonValue]]) -> tuple[Event, ...]: + deltas: Final = tuple( + JSON_OBJECT.validate_python(event["delta"]) for event in events if event.get("type") == "message_delta" + ) + return tuple({name: delta.get(name) for name in ("stop_reason", "stop_details")} for delta in deltas) + + +def _native_deltas(events: Sequence[Mapping[str, JsonValue]]) -> tuple[Event, ...]: + deltas: Final = tuple( + JSON_OBJECT.validate_python(event["delta"]) for event in events if event.get("type") == "message_delta" + ) + return tuple({name: delta[name] for name in ("stop_reason", "stop_details") if name in delta} for delta in deltas) + + +def _targets(wire: Wire) -> tuple[str, ...]: + return tuple(unquote(request.target) for request in wire.drain()) + + +def _paths(wire: Wire) -> tuple[str, ...]: + return tuple(urlsplit(target).path for target in _targets(wire)) + + +def _logged(rig: Rig, key: str, count: int = 1) -> tuple[Event, ...]: + events: Final = rig.sink.await_key(key, count) + assert len(events) == count, events + return events + + +def _logged_response(rig: Rig, key: str) -> JsonValue: + return _logged(rig, key)[0]["response"] + + +@pytest.mark.parametrize("client", CLIENTS) +def test_chat_refusal_relays_the_docs_stop_details_to_the_caller_and_the_log(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body + assert _choice(body)["finish_reason"] == "content_filter", body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_chat_stream_refusal_relays_the_docs_stop_details_in_one_chunk_and_the_log(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks + assert _finish_reasons(chunks) == ("content_filter",), chunks + assert _streamed_text(chunks) == "", chunks + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_messages_refusal_returns_the_docs_body_and_logs_the_stop_details(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _messages(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert (body["stop_reason"], body["stop_details"], body["content"]) == ("refusal", DOCS_STOP_DETAILS, []), body + assert relayed(body) == (), body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_messages_stream_refusal_keeps_the_client_frames_and_logs_the_stop_details(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert frame_types(events) == frame_types(FIXTURES["refusal"].events), events + assert _native_stops(events) == (REFUSAL_FIELDS,), events + assert _relayed_in(events) == (), events + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_responses_refusal_relays_the_docs_stop_details_at_the_top_level(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _responses(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert relayed(body) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), body + assert (body["status"], body["incomplete_details"]) == ("incomplete", {"reason": "content_filter"}), body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_responses_stream_refusal_relays_the_docs_stop_details_on_the_incomplete_event( + rig: Rig, client: Client +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + events: Final = _responses_stream(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert _relayed_in(events) == ((_STREAMED_RESPONSE_FIELDS, REFUSAL_FIELDS),), events + assert frame_types(events)[-1] == "response.incomplete", events + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +def test_anthropic_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + body: Final = { + "model": MODEL, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + content: Final = _post(rig, "/anthropic/v1/messages", body, key) + if stream: + assert content == b"".join(sse_chunks(refusal_events(DOCS_REFUSAL))), content + else: + assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content + assert _targets(rig.wire) == (f"{_ANTHROPIC_BASE}/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("provider", FAMILY) +def test_family_chat_refusal_relays_the_docs_stop_details(rig: Rig, provider: Provider) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker())) + expected: Final = () if provider == "snowflake" else ((_CHAT_FIELDS, REFUSAL_FIELDS),) + assert relayed(body) == expected, body + assert _choice(body)["finish_reason"] == ("stop" if provider == "snowflake" else "content_filter"), body + assert _paths(rig.wire) == (_STREAM_TARGETS[provider][0],) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == expected, logged + + +@pytest.mark.parametrize("provider", FAMILY) +def test_family_chat_stream_refusal_relays_the_docs_stop_details(rig: Rig, provider: Provider) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, rig) + key: Final = scenario.key() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt("refusal", _marker())) + assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks + assert _finish_reasons(chunks) == ("content_filter",), chunks + assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("provider", FAMILY) +def test_family_messages_stream_refusal_keeps_frames_and_logs_the_stop_details(rig: Rig, provider: Provider) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, rig) + key: Final = scenario.key() + events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker())) + native: Final = {"stop_reason": "refusal"} if provider == "snowflake" else REFUSAL_FIELDS + assert _native_deltas(events) == (native,), events + assert _relayed_in(events) == (), events + assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _bedrock_passthrough_deployment(scenario: Scenario, rig: Rig) -> str: + return scenario.model( + model=f"bedrock/{BEDROCK_MODEL}", + api_base=rig.wire.url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=rig.wire.url, + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["invoke", "invoke-with-response-stream"]) +def test_bedrock_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None: + with rig.gateway.scenario() as scenario: + deployment: Final = _bedrock_passthrough_deployment(scenario, rig) + key: Final = scenario.key() + action: Final = "invoke-with-response-stream" if stream else "invoke" + body: Final = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + content: Final = _post(rig, f"/bedrock/model/{deployment}/{action}", body, key) + if stream: + assert content == b"".join(aws_chunks(refusal_events(DOCS_REFUSAL))), content + else: + assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content + assert _paths(rig.wire) == (f"/model/{BEDROCK_MODEL}/{action}",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("stream", [False, True], ids=["rawPredict", "streamRawPredict"]) +def test_vertex_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + action: Final = "streamRawPredict" if stream else "rawPredict" + body: Final = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + content: Final = _post(rig, f"/vertex_ai{_VERTEX_MODEL_PATH}:{action}", body, key) + if stream: + assert content == b"".join(sse_chunks(refusal_events(DOCS_REFUSAL))), content + else: + assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content + assert _paths(rig.vertex) == (f"{_VERTEX_MODEL_PATH}:{action}",) + assert _targets(rig.wire) == () + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _surface_call(rig: Rig, surface: Surface, key: str, model: str, text: str) -> JsonValue: + match surface: + case "chat": + return _chat(rig, key, model, text, cached=True) + case "chat_stream": + return list(_chat_stream(rig, key, model, text, cached=True)) + case "messages": + return _messages(rig, key, model, text, cached=True) + case "messages_stream": + return list(_messages_stream(rig, key, model, text, cached=True)) + case "responses": + return _responses(rig, key, model, text, cached=True) + case "responses_stream": + return list(_responses_stream(rig, key, model, text)) + + +@dataclass(frozen=True, slots=True) +class CacheTwin: + fields: tuple[tuple[Location, Event], ...] + served_from_cache: tuple[bool, bool] + + +_CACHE_TWINS: Final[Mapping[Surface, CacheTwin]] = { + "chat": CacheTwin(((_CHAT_FIELDS, REFUSAL_FIELDS),), (False, True)), + "chat_stream": CacheTwin(((_DELTA_FIELDS, REFUSAL_FIELDS),), (False, True)), + "messages": CacheTwin((), (False, False)), + "messages_stream": CacheTwin((), (False, False)), +} + + +def _relayed_anywhere(value: JsonValue) -> tuple[tuple[Location, Event], ...]: + return _relayed_in(JSON_LIST.validate_python(value)) if isinstance(value, list) else relayed(value) + + +@pytest.mark.parametrize("surface", tuple(_CACHE_TWINS)) +def test_cache_hit_twin_carries_the_same_stop_details_as_the_miss(rig: Rig, surface: Surface) -> None: + twin: Final = _CACHE_TWINS[surface] + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("refusal", _marker()) + first: Final = _surface_call(rig, surface, key, model, text) + _logged(rig, key) + second: Final = _surface_call(rig, surface, key, model, text) + logged: Final = _logged(rig, key, 2) + upstream: Final = _paths(rig.wire) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "upstream": upstream, "first": first, "second": second} + assert (served_from_cache, len(upstream)) == (twin.served_from_cache, twin.served_from_cache.count(False)), ( + evidence + ) + assert (_relayed_anywhere(first), _relayed_anywhere(second)) == (twin.fields, twin.fields), evidence + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, ( + served_from_cache, + logged, + ) + + +def test_responses_cache_hit_logs_the_same_stop_details_as_the_miss(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("refusal", _marker()) + first: Final = _surface_call(rig, "responses", key, model, text) + _logged(rig, key) + _surface_call(rig, "responses", key, model, text) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + assert (len(_paths(rig.wire)), served_from_cache.count(True)) == (1, 1), (served_from_cache, logged) + assert relayed(first) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), first + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, ( + served_from_cache, + logged, + ) + + +def test_passthrough_is_not_served_from_the_response_cache(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + body: Final = { + "model": MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + first: Final = _post(rig, "/anthropic/v1/messages", body, key) + second: Final = _post(rig, "/anthropic/v1/messages", body, key) + assert JSON_OBJECT.validate_json(first) == JSON_OBJECT.validate_json(second) == DOCS_REFUSAL + assert _targets(rig.wire) == (f"{_ANTHROPIC_BASE}/v1/messages",) * 2 + logged: Final = _logged(rig, key, 2) + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2 + + +@pytest.mark.parametrize("variant", tuple(MALFORMED_STOP_DETAILS)) +def test_malformed_stop_details_are_relayed_as_sent_without_breaking_the_turn(rig: Rig, variant: str) -> None: + value: Final = MALFORMED_STOP_DETAILS[variant] + expected: Final[Event] = {"stop_reason": "end_turn", "stop_details": value} + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt(variant, _marker())) + assert relayed(body) == ((_CHAT_FIELDS, expected),), body + assert (_choice(body)["finish_reason"], _message_text(body)) == ("stop", "ready"), body + chunks: Final = _chat_stream(rig, key, model, refusal_prompt(variant, _marker())) + assert _relayed_in(chunks) == ((_DELTA_FIELDS, expected),), chunks + assert (_finish_reasons(chunks), _streamed_text(chunks)) == (("stop",), "ready"), chunks + assert len(_paths(rig.wire)) == 2 + logged: Final = _logged(rig, key, 2) + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, expected),),) * 2, logged + assert rig.gateway.client.get("/health/liveliness").status_code == 200 + + +def test_duplicated_stop_details_key_relays_the_docs_object_once(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("duplicate", _marker())) + assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def test_non_object_upstream_body_fails_the_call_and_leaves_other_keys_serving(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + broken_key: Final = scenario.key() + healthy_key: Final = scenario.key() + broken: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + _chat_request(model, refusal_prompt("nondict", _marker()), stream=False, cached=False), + key=broken_key, + ) + assert broken.status_code == 500, broken.text + assert "error" in JSON_OBJECT.validate_json(broken.content), broken.text + healthy: Final = _chat(rig, healthy_key, model, refusal_prompt("refusal", _marker())) + assert relayed(healthy) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), healthy + assert len(_paths(rig.wire)) == 2 + failure: Final = _logged(rig, broken_key)[0] + assert (failure["status"], relayed(failure["response"])) == ("failure", ()), failure + + +@pytest.mark.parametrize( + ("fixture", "status", "message"), + [("unauthorized", 401, "invalid x-api-key"), ("invalid", 400, "max_tokens: Field required")], +) +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +def test_provider_errors_reach_the_caller_and_log_a_failure_without_stop_fields( + rig: Rig, fixture: str, status: int, message: str, stream: bool +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + _chat_request(model, refusal_prompt(fixture, _marker()), stream=stream, cached=False), + key=key, + ) + assert response.status_code == status, response.text + assert message in response.text, response.text + assert _paths(rig.wire) == ("/v1/messages",) + failure: Final = _logged(rig, key)[0] + assert (failure["status"], relayed(failure["response"])) == ("failure", ()), failure + + +def test_failing_sinks_get_one_attempt_each_and_the_proxy_stays_healthy(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker())) + assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body + for ledger in (rig.sink, rig.forbidden, rig.absent): + events = ledger.await_key(key) + assert tuple(relayed(event["response"]) for event in events) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) + held( + lambda: tuple(len(ledger.for_key(key)) for ledger in (rig.sink, rig.forbidden, rig.absent)), + lambda counts: counts == (1, 1, 1), + holding=3, + ) + readiness: Final = rig.gateway.client.get("/health/readiness") + assert (readiness.status_code, readiness.json()["status"]) == (200, "healthy"), readiness.text + + +def test_unauthenticated_refusal_request_never_reaches_the_upstream_or_the_sink(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + stranger: Final = "sk-not-a-virtual-key-" + _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + _chat_request(model, refusal_prompt("refusal", _marker()), stream=False, cached=False), + key=stranger, + ) + assert response.status_code == 401, response.text + assert _targets(rig.wire) == () + assert rig.sink.for_key(stranger) == () + + +@pytest.mark.parametrize("fixture", ["null", "missing"]) +def test_normal_turn_adds_no_stop_fields_on_any_surface(rig: Rig, fixture: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt(fixture, _marker())) + assert relayed(body) == (), body + assert (_choice(body)["finish_reason"], _message_text(body)) == ("stop", "ready"), body + chunks: Final = _chat_stream(rig, key, model, refusal_prompt(fixture, _marker())) + assert _relayed_in(chunks) == (), chunks + assert (_finish_reasons(chunks), _streamed_text(chunks)) == (("stop",), "ready"), chunks + events: Final = _messages_stream(rig, key, model, refusal_prompt(fixture, _marker())) + expected_native: Final = {"stop_reason": "end_turn", "stop_details": None} + assert _native_stops(events) == (expected_native,), events + assert len(_paths(rig.wire)) == 3 + logged: Final = _logged(rig, key, 3) + assert tuple(relayed(event["response"]) for event in logged) == ((), (), ()), logged + + +def test_refusal_after_partial_text_keeps_the_text_and_relays_the_stop_details(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + chat_key: Final = scenario.key() + messages_key: Final = scenario.key() + chunks: Final = _chat_stream(rig, chat_key, model, refusal_prompt("partial", _marker())) + assert _streamed_text(chunks) == "Hello..", chunks + assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks + assert _finish_reasons(chunks) == ("content_filter",), chunks + logged_chat: Final = _logged_response(rig, chat_key) + assert relayed(logged_chat) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged_chat + assert _message_text(JSON_OBJECT.validate_python(logged_chat)) == "Hello..", logged_chat + events: Final = _messages_stream(rig, messages_key, model, refusal_prompt("partial", _marker())) + assert _native_stops(events) == (REFUSAL_FIELDS,), events + assert frame_types(events) == frame_types(FIXTURES["partial"].events), events + logged_messages: Final = _logged_response(rig, messages_key) + assert relayed(logged_messages) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged_messages + assert len(_paths(rig.wire)) == 2 + + +def test_repeated_identical_refusals_log_once_per_request(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("refusal", _marker()) + bodies: Final = (_chat(rig, key, model, text), _chat(rig, key, model, text)) + assert tuple(relayed(body) for body in bodies) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, bodies + assert len(_paths(rig.wire)) == 2 + logged: Final = _logged(rig, key, 2) + assert len({event["litellm_call_id"] for event in logged}) == 2, logged + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2 + + +@pytest.mark.parametrize("fixture", ["refusal", "int", "null"]) +def test_chat_stream_logged_response_keeps_the_received_model_id(rig: Rig, fixture: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt(fixture, _marker())) + assert len({chunk["id"] for chunk in chunks}) == 1, chunks + probed: Final = rig.hidden.await_key(key) + assert len(probed) == 1, probed + assert (probed[0]["response_id"], probed[0]["received_model_id"]) == (chunks[0]["id"], chunks[0]["id"]), probed + + +def _delta_fields(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + deltas: Final = tuple(JSON_OBJECT.validate_python(_choice(chunk).get("delta") or {}) for chunk in chunks) + return tuple( + delta["provider_specific_fields"] for delta in deltas if delta.get("provider_specific_fields") is not None + ) + + +def _message_fields(response: JsonValue) -> tuple[JsonValue, ...]: + message: Final = JSON_OBJECT.validate_python(_choice(JSON_OBJECT.validate_python(response))["message"]) + fields: Final = message.get("provider_specific_fields") + return () if fields is None else (fields,) + + +def test_openai_compatible_stream_logged_response_keeps_the_received_model_id(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + marker: Final = _marker() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt("plain", marker)) + assert (_streamed_text(chunks), _finish_reasons(chunks), _delta_fields(chunks)) == ("ready now", ("stop",), ()) + assert {chunk["id"] for chunk in chunks} == {openai_id(marker)}, chunks + assert _paths(rig.wire) == ("/v1/chat/completions",) + probed: Final = rig.hidden.await_key(key) + assert len(probed) == 1, probed + assert (probed[0]["response_id"], probed[0]["received_model_id"]) == (openai_id(marker),) * 2, probed + + +@pytest.mark.parametrize(("fixture", "fields"), [("plain", ()), ("cited", (CITED_FIELDS,))], ids=["plain", "cited"]) +def test_openai_compatible_stream_cache_hit_replays_the_provider_fields_of_the_miss( + rig: Rig, fixture: str, fields: tuple[JsonValue, ...] +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + text: Final = refusal_prompt(fixture, _marker()) + first: Final = _chat_stream(rig, key, model, text, cached=True) + _logged(rig, key) + second: Final = _chat_stream(rig, key, model, text, cached=True) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence + streamed: Final = tuple( + (_streamed_text(chunks), _finish_reasons(chunks), _delta_fields(chunks)) for chunks in (first, second) + ) + assert streamed == (("ready now", ("stop",), fields),) * 2, evidence + assert tuple(_message_fields(event["response"]) for event in logged) == (fields, fields), evidence + + +def _spend_rows(rig: Rig, key: str) -> tuple[Event, ...]: + response: Final = rig.gateway.request("GET", "/spend/logs", params={"api_key": key}) + assert response.status_code == 200, response.text + return tuple(JSON_OBJECT.validate_python(row) for row in JSON_LIST.validate_json(response.content)) + + +def test_spend_logs_endpoint_reports_the_stop_details_of_a_refusal(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker())) + rows: Final = eventually(lambda: _spend_rows(rig, key), lambda found: len(found) >= 1, seconds=70) + assert len(rows) == 1, rows + reported: Final = (rows[0]["request_id"], relayed(rows[0]["response"])) + expected: Final = (body["id"], ((_CHAT_FIELDS, REFUSAL_FIELDS),)) + assert ((body["id"], relayed(body)), reported) == (expected, expected), rows + + +def _usage_counts(usage: JsonValue, total: str, details: str) -> tuple[JsonValue, JsonValue]: + counted: Final = JSON_OBJECT.validate_python(usage) + return (counted[total], JSON_OBJECT.validate_python(counted[details])["reasoning_tokens"]) + + +def _chat_usage_counts(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]: + return tuple( + _usage_counts(chunk["usage"], "completion_tokens", "completion_tokens_details") + for chunk in chunks + if chunk.get("usage") + ) + + +def _cached_stream_usage_counts( + rig: Rig, surface: Literal["chat_stream", "responses_stream"], key: str, model: str, text: str +) -> tuple[tuple[JsonValue, JsonValue], ...]: + match surface: + case "chat_stream": + request: Final = { + **_chat_request(model, text, stream=True, cached=True), + "stream_options": {"include_usage": True}, + } + return _chat_usage_counts(chunks_of(_post(rig, "/v1/chat/completions", request, key).decode())) + case "responses_stream": + events: Final = frames( + _post(rig, "/v1/responses", _responses_request(model, text, stream=True, cached=True), key).decode() + ) + responses: Final = tuple( + JSON_OBJECT.validate_python(event["response"]) + for event in events + if isinstance(event.get("response"), dict) + ) + return tuple( + _usage_counts(response["usage"], "output_tokens", "output_tokens_details") + for response in responses + if response.get("usage") + ) + + +def _logged_usage_counts(logged: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]: + return tuple( + _usage_counts( + JSON_OBJECT.validate_python(event["response"])["usage"], "completion_tokens", "completion_tokens_details" + ) + for event in logged + ) + + +@pytest.mark.parametrize("fixture", ["plainusage", "citedusage"], ids=["plain", "cited"]) +def test_openai_compatible_stream_cache_hit_replays_the_usage_details_of_the_miss(rig: Rig, fixture: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + text: Final = refusal_prompt(fixture, _marker()) + first: Final = _cached_stream_usage_counts(rig, "chat_stream", key, model, text) + _logged(rig, key) + second: Final = _cached_stream_usage_counts(rig, "chat_stream", key, model, text) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence + assert (first, second, _logged_usage_counts(logged)) == (((7, 4),), ((7, 4),), ((7, 4), (7, 4))), evidence + + +def test_text_completion_stream_cache_hit_replays_the_usage_details_of_the_miss(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + request: Final[Event] = { + "model": model, + "prompt": refusal_prompt("plainusage", _marker()), + "max_tokens": 64, + "stream": True, + "stream_options": {"include_usage": True}, + } + first: Final = _chat_usage_counts(chunks_of(_post(rig, "/v1/completions", request, key).decode())) + _logged(rig, key) + second: Final = _chat_usage_counts(chunks_of(_post(rig, "/v1/completions", request, key).decode())) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence + assert (first, second, _logged_usage_counts(logged)) == (((7, 4),), ((7, 4),), ((7, 4), (7, 4))), evidence + + +@pytest.mark.parametrize("surface", ["chat_stream", "responses_stream"]) +def test_anthropic_stream_cache_hit_keeps_the_usage_details_of_the_miss( + rig: Rig, surface: Literal["chat_stream", "responses_stream"] +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("reasoned", _marker()) + first: Final = _cached_stream_usage_counts(rig, surface, key, model, text) + _logged(rig, key) + second: Final = _cached_stream_usage_counts(rig, surface, key, model, text) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/messages",)), evidence + assert (first, second, _logged_usage_counts(logged)) == (((10, 4),), ((10, 4),), ((10, 4), (10, 4))), evidence + + +def _response_id(event: Mapping[str, JsonValue]) -> JsonValue: + return JSON_OBJECT.validate_python(event["response"]).get("id") + + +@pytest.mark.parametrize("variant", tuple(MALFORMED_STOP_DETAILS)) +def test_stream_cache_hit_skips_cached_provider_fields_that_are_not_an_object(rig: Rig, variant: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + marker: Final = _marker() + text: Final = refusal_prompt("plain", marker) + filled: Final = chunks_of( + _post( + rig, + "/v1/chat/completions", + {**_chat_request(model, text, stream=True, cached=True), "cache": {"ttl": 2}}, + key, + ).decode() + ) + assert _streamed_text(filled) == "ready now", filled + _logged(rig, key) + stale_id: Final = rewrite_cached_message_fields(openai_id(marker), MALFORMED_STOP_DETAILS[variant]) + replayed: Final = eventually( + lambda: _chat_stream(rig, key, model, text, cached=True), + lambda chunks: {chunk["id"] for chunk in chunks} == {stale_id}, + seconds=30, + ) + streamed: Final = (_streamed_text(replayed), _finish_reasons(replayed), _delta_fields(replayed)) + assert streamed == ("ready now", ("stop",), ()), replayed + assert _paths(rig.wire) == ("/v1/chat/completions",) + stale_logs: Final = eventually( + lambda: tuple(event for event in rig.sink.for_key(key) if _response_id(event) == stale_id), + lambda found: len(found) >= 1, + ) + assert tuple(_message_fields(event["response"]) for event in stale_logs) == ((),) * len(stale_logs), stale_logs + + +JsonModeRoute: TypeAlias = Literal["anthropic_native", "anthropic_tool", "bedrock_invoke"] +_UNMAPPED_MODEL: Final = "claude-scripted-9" +_JSON_ANSWER: Final = json.dumps({"answer": DOCS_PARTIAL_TEXT}) + + +@dataclass(frozen=True, slots=True) +class JsonModeCase: + route: JsonModeRoute + fixture: str + streamed_text: str + message_text: str + streamed_finish_reason: str + message_finish_reason: str + requested: Literal["json_tool", "native"] + + +_JSON_MODE_CASES: Final[Mapping[str, JsonModeCase]] = { + "anthropic_native": JsonModeCase( + "anthropic_native", + "partial", + DOCS_PARTIAL_TEXT, + DOCS_PARTIAL_TEXT, + "content_filter", + "content_filter", + "native", + ), + "anthropic_tool_declared": JsonModeCase( + "anthropic_tool", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), + "anthropic_tool_enveloped": JsonModeCase( + "anthropic_tool", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), + "bedrock_invoke_declared": JsonModeCase( + "bedrock_invoke", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), + "bedrock_invoke_enveloped": JsonModeCase( + "bedrock_invoke", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), +} + + +def _json_mode_deployment(scenario: Scenario, route: JsonModeRoute, rig: Rig) -> tuple[str, Provider]: + match route: + case "anthropic_native": + return _deployment(scenario, "anthropic", rig), "anthropic" + case "anthropic_tool": + return ( + scenario.model(model=f"anthropic/{_UNMAPPED_MODEL}", api_base=rig.wire.url, api_key=_ANTHROPIC_KEY), + "anthropic", + ) + case "bedrock_invoke": + return _deployment(scenario, "bedrock_invoke", rig), "bedrock_invoke" + + +def _requested_format(body: Mapping[str, JsonValue]) -> tuple[bool, bool]: + tools: Final = JSON_LIST.validate_python(body.get("tools") or []) + output_config: Final = JSON_OBJECT.validate_python(body.get("output_config") or {}) + json_tool: Final = any(JSON_OBJECT.validate_python(tool).get("name") == JSON_TOOL_NAME for tool in tools) + return json_tool, "output_format" in body or "format" in output_config + + +def _received(wire: Wire) -> tuple[tuple[str, tuple[bool, bool]], ...]: + return tuple( + (urlsplit(unquote(request.target)).path, _requested_format(JSON_OBJECT.validate_json(request.body))) + for request in wire.drain() + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +@pytest.mark.parametrize("case", tuple(_JSON_MODE_CASES)) +def test_json_mode_refusal_relays_the_docs_stop_details_through_the_response_format_path( + rig: Rig, case: str, stream: bool +) -> None: + expected: Final = _JSON_MODE_CASES[case] + with rig.gateway.scenario() as scenario: + model, provider = _json_mode_deployment(scenario, expected.route, rig) + key: Final = scenario.key() + request: Final = { + **_chat_request(model, refusal_prompt(expected.fixture, _marker()), stream=stream, cached=False), + "response_format": JSON_ANSWER_FORMAT, + } + content: Final = _post(rig, "/v1/chat/completions", request, key) + requested: Final = (expected.requested == "json_tool", expected.requested == "native") + assert _received(rig.wire) == ((_STREAM_TARGETS[provider][int(stream)], requested),) + if stream: + chunks: Final = chunks_of(content.decode()) + streamed: Final = (_streamed_text(chunks), _finish_reasons(chunks), _relayed_in(chunks)) + assert streamed == ( + expected.streamed_text, + (expected.streamed_finish_reason,), + ((_DELTA_FIELDS, REFUSAL_FIELDS),), + ), chunks + else: + body: Final = JSON_OBJECT.validate_json(content) + answered: Final = (_message_text(body), _choice(body)["finish_reason"], relayed(body)) + assert answered == ( + expected.message_text, + expected.message_finish_reason, + ((_CHAT_FIELDS, REFUSAL_FIELDS),), + ), body + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _serialized_choice(body: Mapping[str, JsonValue]) -> Event: + choice: Final = _choice(body) + return {**choice, "message": JSON_OBJECT.validate_python(choice["message"])} + + +_ANTHROPIC_MESSAGE_FIELDS: Final[Event] = {"citations": None, "thinking_blocks": None} + + +@pytest.mark.parametrize(("fixture", "content"), [("refusal", None), ("partial", DOCS_PARTIAL_TEXT)]) +def test_chat_refusal_body_serializes_the_stop_details_beside_the_existing_message_fields( + rig: Rig, fixture: str, content: str | None +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt(fixture, _marker())) + expected: Final[Event] = { + "finish_reason": "content_filter", + "index": 0, + "message": { + "content": content, + "role": "assistant", + "provider_specific_fields": {**_ANTHROPIC_MESSAGE_FIELDS, **REFUSAL_FIELDS}, + }, + } + assert _serialized_choice(body) == expected, body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _usage_of(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]: + usages: Final = tuple(JSON_OBJECT.validate_python(chunk["usage"]) for chunk in chunks if chunk.get("usage")) + return tuple((usage["completion_tokens"], usage.get("completion_tokens_details")) for usage in usages) + + +def _served_ids(response: httpx.Response) -> frozenset[JsonValue]: + return frozenset(chunk["id"] for chunk in chunks_of(response.text)) if response.status_code == 200 else frozenset() + + +@pytest.mark.parametrize("variant", tuple(UNVALIDATED_USAGE_DETAILS)) +def test_stream_cache_hit_skips_cached_usage_details_that_fail_validation(rig: Rig, variant: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + marker: Final = _marker() + request: Final = { + **_chat_request(model, refusal_prompt("plainusage", marker), stream=True, cached=True), + "stream_options": {"include_usage": True}, + } + filled: Final = chunks_of(_post(rig, "/v1/chat/completions", {**request, "cache": {"ttl": 2}}, key).decode()) + assert (_streamed_text(filled), _chat_usage_counts(filled)) == ("ready now", ((7, 4),)), filled + _logged(rig, key) + stale_id: Final = rewrite_cached_usage_details(openai_id(marker), UNVALIDATED_USAGE_DETAILS[variant]) + replayed: Final = eventually( + lambda: rig.gateway.request("POST", "/v1/chat/completions", request, key=key), + lambda response: response.status_code != 200 or _served_ids(response) == {stale_id}, + seconds=30, + ) + assert replayed.status_code == 200, replayed.text + chunks: Final = chunks_of(replayed.text) + streamed: Final = (_streamed_text(chunks), _finish_reasons(chunks), _usage_of(chunks)) + assert streamed == ("ready now", ("stop",), ((7, {"reasoning_tokens": 0, "text_tokens": 7}),)), chunks + assert _paths(rig.wire) == ("/v1/chat/completions",) + stale_logs: Final = eventually( + lambda: tuple(event for event in rig.sink.for_key(key) if _response_id(event) == stale_id), + lambda found: len(found) >= 1, + ) + assert tuple(event["status"] for event in stale_logs) == ("success",) * len(stale_logs), stale_logs diff --git a/tests/integration/sdk/test_response_cache_sync_sdk.py b/tests/integration/sdk/test_response_cache_sync_sdk.py index 54d78adbaba..1fe73ec9012 100644 --- a/tests/integration/sdk/test_response_cache_sync_sdk.py +++ b/tests/integration/sdk/test_response_cache_sync_sdk.py @@ -9,17 +9,23 @@ from __future__ import annotations import os from collections.abc import Iterator +from itertools import chain from typing import Final import pytest +from pydantic import JsonValue, TypeAdapter import litellm from litellm.caching.caching import Cache +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.caching import LiteLLMCacheType from litellm.types.utils import ModelResponse, TextCompletionResponse +from tests.integration._support.client import JSON_OBJECT from tests.integration.caching.response_cache_case import ( CHAT, + CHAT_CHUNK, CHAT_MODEL, + CHAT_STREAM, PROVIDER_KEY, TEXT, TEXT_MODEL, @@ -31,8 +37,10 @@ from tests.integration.caching.response_cache_case import ( prompt, redis_store, rescript, + rewrite_usage_details, scripted_in_process, ) +from tests.integration.cost_calculation.cost_tracking_case import SseResponse @pytest.fixture @@ -123,3 +131,77 @@ def test_sync_text_completion_with_empty_choices_reaches_upstream_again() -> Non hit: Final = _text_completion(handle.api_base(), text) assert hit.id == refilled.id, hit.model_dump_json() assert upstream.calls(handle.scenario_id) == 2 + + +_REASONING_USAGE_FRAME: Final = ( + CHAT_CHUNK + '[],"usage":{"prompt_tokens":5,"completion_tokens":7,"total_tokens":12,' + '"completion_tokens_details":{"reasoning_tokens":4}}}' +) +_REASONING_STREAM: Final = SseResponse( + content_type="text/event-stream", + frames=(*CHAT_STREAM.frames[:-1], _REASONING_USAGE_FRAME, CHAT_STREAM.frames[-1]), +) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_UNVALIDATED_USAGE_DETAILS: Final[dict[str, JsonValue]] = { + "text": {"reasoning_tokens": "abc"}, + "fraction": {"reasoning_tokens": 1.5}, + "object": {"reasoning_tokens": {"a": 1}}, +} + + +def _stream(api_base: str, text: str) -> tuple[dict[str, JsonValue], ...]: + stream: Final = litellm.completion( + model=CHAT_MODEL, + messages=[{"role": "user", "content": text}], + api_base=api_base, + api_key=PROVIDER_KEY, + stream=True, + stream_options={"include_usage": True}, + ) + assert isinstance(stream, CustomStreamWrapper), type(stream) + return tuple(JSON_OBJECT.validate_python(chunk.model_dump(mode="json")) for chunk in stream) + + +def _replayed(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[frozenset[JsonValue], str, tuple[JsonValue, ...]]: + choices: Final = tuple( + JSON_OBJECT.validate_python(choice) + for choice in chain.from_iterable(_JSON_LIST.validate_python(chunk.get("choices") or []) for chunk in chunks) + ) + text: Final = "".join(str(JSON_OBJECT.validate_python(choice["delta"]).get("content") or "") for choice in choices) + usages: Final = tuple(JSON_OBJECT.validate_python(chunk["usage"]) for chunk in chunks if chunk.get("usage")) + reasoning: Final = tuple( + JSON_OBJECT.validate_python(usage.get("completion_tokens_details") or {}).get("reasoning_tokens") + for usage in usages + ) + return frozenset(chunk["id"] for chunk in chunks), text, reasoning + + +@pytest.mark.usefixtures("sdk_cache") +def test_sync_stream_cache_hit_replays_the_usage_details_of_the_miss() -> None: + upstream: Final = Upstream() + text: Final = prompt() + with scripted_in_process(_REASONING_STREAM) as handle, redis_store() as store: + first: Final = _replayed(_stream(handle.api_base(), text)) + (first_id,) = first[0] + assert first[1:] == ("streamed response", (4,)), first + await_entry(store, str(first_id)) + hit: Final = _replayed(_stream(handle.api_base(), text)) + assert hit == first, hit + assert upstream.calls(handle.scenario_id) == 1 + + +@pytest.mark.usefixtures("sdk_cache") +@pytest.mark.parametrize("variant", tuple(_UNVALIDATED_USAGE_DETAILS)) +def test_sync_stream_cache_hit_skips_cached_usage_details_that_fail_validation(variant: str) -> None: + upstream: Final = Upstream() + text: Final = prompt() + with scripted_in_process(_REASONING_STREAM) as handle, redis_store() as store: + first: Final = _replayed(_stream(handle.api_base(), text)) + (first_id,) = first[0] + assert first[1:] == ("streamed response", (4,)), first + stale_id: Final = rewrite_usage_details( + store, await_entry(store, str(first_id)), _UNVALIDATED_USAGE_DETAILS[variant] + ) + hit: Final = _replayed(_stream(handle.api_base(), text)) + assert hit == (frozenset({stale_id}), "streamed response", (0,)), hit + assert upstream.calls(handle.scenario_id) == 1 diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 2c3dae21058..e9d96f057ae 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -11,6 +11,7 @@ content frame, restoring per-token cadence on cache hits. import pytest +import litellm from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _split_assembled_content_for_replay, convert_to_streaming_response, @@ -87,6 +88,71 @@ async def test_async_yields_multiple_content_chunks_with_lossless_join(): assert reassembled == text +@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"]) +@pytest.mark.asyncio +async def test_async_and_sync_replays_carry_the_cached_provider_specific_fields(content): + provider_specific_fields = {"stop_reason": "refusal", "stop_details": {"category": "cyber"}} + + def payload(): + cached = _async_payload(content=content) + cached["choices"][0]["message"]["provider_specific_fields"] = provider_specific_fields + return cached + + async_chunks = await _collect_async(payload()) + sync_chunks = list(convert_to_streaming_response(payload())) + + assert [c.choices[0].delta.provider_specific_fields for c in async_chunks] == [ + c.choices[0].delta.provider_specific_fields for c in sync_chunks + ] + assert async_chunks[0].choices[0].delta.provider_specific_fields == provider_specific_fields + + +@pytest.mark.parametrize("provider_fields", [{"thought_signatures": ["sig"]}, None], ids=["provider_fields", "text"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_the_cached_completion_tokens_details(replay, provider_fields): + cached = _async_payload() + cached["choices"][0]["message"]["provider_specific_fields"] = provider_fields + cached["usage"]["completion_tokens_details"] = {"reasoning_tokens": 4, "text_tokens": 3} + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.get("provider_specific_fields") == provider_fields + assert rebuilt.usage.completion_tokens_details.reasoning_tokens == 4 + assert rebuilt.usage.completion_tokens_details.text_tokens == 3 + + +@pytest.mark.parametrize("cached_fields", [7, [1, 2], "", "x" * 5000], ids=["int", "list", "empty", "long"]) +@pytest.mark.asyncio +async def test_the_async_replay_ignores_cached_provider_fields_that_are_not_an_object(cached_fields): + cached = _async_payload() + cached["choices"][0]["message"]["provider_specific_fields"] = cached_fields + + chunks = await _collect_async(cached) + + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "Hello world! How are you?" + assert all(c.choices[0].delta.provider_specific_fields is None for c in chunks) + + +@pytest.mark.parametrize( + "details", + [{"reasoning_tokens": "abc"}, {"reasoning_tokens": 1.5}, {"reasoning_tokens": {"a": 1}}], + ids=["text", "fraction", "object"], +) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_not_validate(replay, details): + cached = _async_payload() + cached["usage"]["completion_tokens_details"] = details + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.content == "Hello world! How are you?" + assert rebuilt.usage.completion_tokens == 7 + + @pytest.mark.asyncio async def test_async_finish_reason_only_on_last_chunk(): chunks = await _collect_async(_async_payload()) @@ -204,9 +270,7 @@ def test_sync_multi_chunk_and_lossless_join(): def test_sync_delta_and_choice_metadata_only_on_first_chunk(): - thinking_blocks = [ - {"type": "thinking", "thinking": "cached thinking", "signature": "sig"} - ] + thinking_blocks = [{"type": "thinking", "thinking": "cached thinking", "signature": "sig"}] payload = { "id": "chatcmpl-test-sync", "object": "chat.completion", diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 0172f0e9aeb..a604e61b9c2 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -732,6 +732,42 @@ def test_set_response_id_propagation_valid_to_invalid( assert initialized_custom_stream_wrapper.response_id == "first-valid-id" +@pytest.mark.parametrize( + "last_delta", + [ + Delta(content="Hello"), + Delta(provider_specific_fields={"stop_reason": "refusal"}), + ], + ids=["content", "provider_specific_fields"], +) +def test_logged_response_keeps_received_model_id_when_the_finish_chunk_is_synthesized( + logging_obj: Logging, last_delta: Delta +): + upstream_id: Final = "chatcmpl-upstream-id" + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[ + ModelResponseStream( + id=upstream_id, + model="gpt-4o-mini", + choices=[StreamingChoices(index=0, delta=last_delta, finish_reason="content_filter")], + ) + ] + ), + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + emitted: Final = list(wrapper) + logged: Final = litellm.stream_chunk_builder(chunks=wrapper.chunks) + + assert len(emitted) == 2 + assert [chunk._hidden_params.get("received_model_id") for chunk in emitted] == [upstream_id, upstream_id] + assert logged is not None + assert logged._hidden_params.get("received_model_id") == upstream_id + + @pytest.mark.asyncio async def test_streaming_completion_start_time(logging_obj: Logging): """Test that the start time is set correctly""" diff --git a/tests/unit/llms/anthropic/anthropic_docs_refusal.py b/tests/unit/llms/anthropic/anthropic_docs_refusal.py new file mode 100644 index 00000000000..2bd75b4aab5 --- /dev/null +++ b/tests/unit/llms/anthropic/anthropic_docs_refusal.py @@ -0,0 +1,19 @@ +from typing import Final + +ANTHROPIC_DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response" + +# Copied verbatim from ANTHROPIC_DOCS_REFUSAL_URL on 2026-10-12. Streams carry the same stop_details on message_delta. +ANTHROPIC_DOCS_REFUSAL_MESSAGE: Final = { + "id": "msg_01XFUDYJgAACzvnptvVoYEL", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", + }, + "usage": {"input_tokens": 412, "output_tokens": 0}, +} diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 7d1d2a3c434..a072010bfa9 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -22,6 +22,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import OutputCodeInterpreterCall from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @pytest.mark.asyncio @@ -3366,3 +3367,50 @@ def test_anthropic_streaming_invalid_key_maps_to_authentication_error() -> None: assert exc_info.value.status_code == 401 assert exc_info.value.llm_provider == "anthropic" assert route.calls.last.request.headers["x-api-key"] == "bad-anthropic-key" + + +def _message_stream_events(stop_reason: str, stop_details: dict | None) -> tuple[dict, ...]: + return ( + { + "type": "message_start", + "message": { + "id": ANTHROPIC_DOCS_REFUSAL_MESSAGE["id"], + "type": "message", + "role": "assistant", + "model": ANTHROPIC_DOCS_REFUSAL_MESSAGE["model"], + "content": [], + "usage": {"input_tokens": 412, "output_tokens": 0}, + }, + }, + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None, "stop_details": stop_details}, + "usage": {"output_tokens": 0}, + }, + {"type": "message_stop"}, + ) + + +def _assemble_stream(events: tuple[dict, ...]) -> litellm.ModelResponse: + iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + chunks: Final = [iterator.chunk_parser(event) for event in events] + return litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + +def test_streamed_refusal_stop_details_reach_the_assembled_message(): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + assembled: Final = _assemble_stream(_message_stream_events("refusal", stop_details)) + + assert assembled.choices[0].finish_reason == "content_filter" + provider_fields: Final = assembled.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == stop_details + + +def test_streamed_end_turn_carries_no_stop_details(): + assembled: Final = _assemble_stream(_message_stream_events("end_turn", None)) + + assert assembled.choices[0].finish_reason == "stop" + provider_fields: Final = assembled.choices[0].message.provider_specific_fields or {} + assert "stop_details" not in provider_fields + assert "stop_reason" not in provider_fields diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index bbcc8629a10..5f4ecc4440b 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -39,6 +39,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE def test_response_format_transformation_unit_test(): @@ -9082,3 +9083,93 @@ def test_completion_surfaces_response_citations_in_provider_specific_fields(resp assert response.choices[0].message.provider_specific_fields["citations"] == [ [{**citation, "supported_text": "the grass is green"}] ] + + +def _refusal_message_body(stop_reason: str, stop_details: dict | None) -> dict: + return {**ANTHROPIC_DOCS_REFUSAL_MESSAGE, "stop_reason": stop_reason, "stop_details": stop_details} + + +def _transform_message_body(body: dict) -> litellm.ModelResponse: + return AnthropicConfig().transform_parsed_response( + completion_response=body, + raw_response=httpx.Response(status_code=200, headers={}), + model_response=litellm.ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + + +def test_refusal_stop_details_kept_in_provider_specific_fields(): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + result: Final = _transform_message_body(ANTHROPIC_DOCS_REFUSAL_MESSAGE) + + assert result.choices[0].finish_reason == "content_filter" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == stop_details + assert result._hidden_params["provider_specific_fields"]["stop_details"] == stop_details + + +def test_refusal_without_stop_details_keeps_the_native_stop_reason(): + result: Final = _transform_message_body(_refusal_message_body("refusal", None)) + + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert "stop_details" not in provider_fields + + +def test_end_turn_response_carries_no_stop_details(): + result: Final = _transform_message_body(_refusal_message_body("end_turn", None)) + + assert result.choices[0].finish_reason == "stop" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert "stop_details" not in provider_fields + assert "stop_reason" not in provider_fields + + +def _json_mode_message_body(stop_reason: str, stop_details: dict | None) -> dict: + return { + **_refusal_message_body(stop_reason, stop_details), + "content": [ + {"type": "tool_use", "id": "toolu_json", "name": RESPONSE_FORMAT_TOOL_NAME, "input": {"answer": "Hello"}} + ], + } + + +def _transform_json_mode_body(body: dict) -> litellm.ModelResponse: + return AnthropicConfig().transform_parsed_response( + completion_response=body, + raw_response=httpx.Response(status_code=200, headers={}), + model_response=litellm.ModelResponse(), + json_mode=True, + prefix_prompt=None, + ) + + +def test_json_mode_refusal_keeps_stop_details_on_the_message(): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + result: Final = _transform_json_mode_body(_json_mode_message_body("refusal", stop_details)) + + message: Final = result.choices[0].message + assert json.loads(message.content) == {"answer": "Hello"} + assert message.provider_specific_fields == {"stop_reason": "refusal", "stop_details": stop_details} + + +def test_json_mode_tool_use_turn_adds_no_stop_fields_to_the_message(): + result: Final = _transform_json_mode_body(_json_mode_message_body("tool_use", None)) + + message: Final = result.choices[0].message + assert json.loads(message.content) == {"answer": "Hello"} + assert not message.provider_specific_fields + + +def test_non_dict_upstream_body_still_transforms_without_refusal_fields() -> None: + result = AnthropicConfig().transform_parsed_response( + completion_response=MagicMock(), + raw_response=httpx.Response(200, headers={}), + model_response=litellm.ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + provider_specific_fields = result.choices[0].message.provider_specific_fields or {} + assert "stop_details" not in provider_specific_fields diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 859b57c7cf3..8cb9ad9b73c 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -4,7 +4,7 @@ import json import uuid from types import SimpleNamespace from typing import Final -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx import pytest @@ -17,6 +17,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE ONE_PIXEL_PNG = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" @@ -1689,3 +1690,23 @@ class TestBedrockAnthropicCombinedRegressions: user_msg = result["messages"][0] has_cache_control = any(isinstance(c, dict) and "cache_control" in c for c in user_msg["content"]) assert has_cache_control + + +def test_refusal_stop_details_kept_through_bedrock_invoke_transform(): + body: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE + result: Final = AmazonAnthropicClaudeConfig().transform_response( + model="anthropic.claude-sonnet-5-5-v1:0", + raw_response=httpx.Response(status_code=200, headers={}, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].finish_reason == "content_filter" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 56e76df25a4..ace24c8a3f4 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -3,6 +3,7 @@ import copy import json from collections.abc import Mapping from typing import Final +from unittest.mock import MagicMock import httpx import pytest @@ -19,6 +20,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im ) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VertexPartnerProvider +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @pytest.mark.parametrize( @@ -1125,3 +1127,23 @@ def test_chat_flagged_model_replays_a_byte_identical_prefix_around_a_mid_convers _assert_prefix_stable(requests) assert [m["role"] for m in requests[1]["messages"]] == ["user", "assistant", "user", "system"] assert [m["role"] for m in requests[2]["messages"]] == ["user", "assistant", "user", "system", "assistant", "user"] + + +def test_refusal_stop_details_kept_through_vertex_anthropic_transform(): + body: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE + result: Final = VertexAIAnthropicConfig().transform_response( + model="claude-sonnet-5-5", + raw_response=httpx.Response(status_code=200, headers={}, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].finish_reason == "content_filter" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 973d6bc9f7b..88c5e2d2c1b 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -16,6 +16,7 @@ from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE async def _drain_tasks(): @@ -2951,3 +2952,58 @@ def test_get_user_from_metadata(end_user_id): ) assert response == "test" + + +def _anthropic_sse_stream(stop_reason: str, stop_details: dict | None) -> list[str]: + events: Final = ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": ANTHROPIC_DOCS_REFUSAL_MESSAGE["id"], + "type": "message", + "role": "assistant", + "model": ANTHROPIC_DOCS_REFUSAL_MESSAGE["model"], + "content": [], + "usage": {"input_tokens": 412, "output_tokens": 0}, + }, + }, + ), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None, "stop_details": stop_details}, + "usage": {"output_tokens": 0}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + return [f"event: {name}\ndata: {json.dumps(data)}\n\n" for name, data in events] + + +def test_streamed_messages_refusal_stop_details_on_the_logged_response(mock_logging_obj): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + logged: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response( + all_chunks=_anthropic_sse_stream("refusal", stop_details), + litellm_logging_obj=mock_logging_obj, + model="claude-sonnet-5-5", + ) + + assert logged.choices[0].finish_reason == "content_filter" + provider_fields: Final = logged.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == stop_details + + +def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj): + logged: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response( + all_chunks=_anthropic_sse_stream("end_turn", None), + litellm_logging_obj=mock_logging_obj, + model="claude-sonnet-5-5", + ) + + assert logged.choices[0].finish_reason == "stop" + provider_fields: Final = logged.choices[0].message.provider_specific_fields or {} + assert "stop_details" not in provider_fields