diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py index c7fb6a542ff..c93306cdced 100644 --- a/tests/base_sdk_tests/check_sdk_http.py +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -6,12 +6,17 @@ import json import os import struct import zlib +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence +from functools import partial from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import Queue from threading import Event, Thread -from typing import Final +from typing import Final, cast +from unittest import SkipTest, skipIf from unittest.mock import patch +from pydantic import BaseModel, JsonValue, TypeAdapter + ENVIRONMENT: Final = {"LITELLM_LOCAL_MODEL_COST_MAP": "True", "PYTHON_DOTENV_DISABLED": "1"} with patch.dict(os.environ, ENVIRONMENT): @@ -21,6 +26,7 @@ RESPONSES: Final[Queue[tuple[int, bytes]]] = Queue() REQUESTS: Final[Queue[tuple[str, dict[str, str], bytes]]] = Queue() ARRIVED: Final = Event() RELEASE: Final = Event() +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) MESSAGES: Final = [{"role": "user", "content": "ping"}] CHAT: Final = { "id": "chat-http-check", @@ -45,7 +51,14 @@ class RecordingUpstream(BaseHTTPRequestHandler): RELEASE.wait(timeout=10) return self.send_response(status) - self.send_header("Content-Type", "text/event-stream" if body.startswith(b"data:") else "application/vnd.amazon.eventstream" if body[:1] == b"\x00" else "application/json") + self.send_header( + "Content-Type", + "text/event-stream" + if body.startswith((b"data:", b"event:")) + else "application/vnd.amazon.eventstream" + if body[:1] == b"\x00" + else "application/json", + ) self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) @@ -71,6 +84,189 @@ def bedrock_event(event: str, payload: object) -> bytes: return frame + struct.pack(">I", zlib.crc32(frame)) +def payload(response: object) -> Mapping[str, JsonValue]: + assert isinstance(response, BaseModel), response + return JSON_OBJECT.validate_json(response.model_dump_json(exclude_none=True)) + + +@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 completion(tools=...) imports optional mcp") +def check_tools(base: str, asynchronous: bool = False) -> None: + tools: Final = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + call: Final = { + "id": "call-http-check", + "type": "function", + "function": {"name": "lookup", "arguments": '{"city":"Paris"}'}, + } + arguments: Final = dict(model="openai/gpt-6", messages=MESSAGES, api_key="test-key", api_base=base + "/v1") + choices: Final = [ + {"index": 0, "message": {"role": "assistant", "tool_calls": [call]}, "finish_reason": "tool_calls"} + ] + enqueue({**CHAT, "choices": choices}) + result: Final = ( + asyncio.run(litellm.acompletion(**arguments, tools=tools)) + if asynchronous + else litellm.completion(**arguments, tools=tools) + ) + assert payload(result)["choices"] == [{**choices[0], "provider_specific_fields": {}}] + path, headers, body = REQUESTS.get(timeout=10) + assert path == "/v1/chat/completions" + assert headers["authorization"] == "Bearer test-key" + assert json.loads(body) == {"model": "gpt-6", "messages": MESSAGES, "tools": tools} + + +@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 responses() imports optional mcp") +def check_responses(base: str, asynchronous: bool = False) -> None: + output: Final[Sequence[JsonValue]] = [ + { + "id": "msg-http-check", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + } + ] + enqueue( + { + "id": "resp-http-check", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-6", + "output": output, + "usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, + } + ) + arguments: Final = dict(model="openai/gpt-6", input="ping", api_key="test-key", api_base=base + "/v1") + response: Final = asyncio.run(litellm.aresponses(**arguments)) if asynchronous else litellm.responses(**arguments) + assert payload(response)["output"] == output + assert payload(response)["usage"] == {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5} + path, headers, body = REQUESTS.get(timeout=10) + assert path == "/v1/responses" + assert headers["authorization"] == "Bearer test-key" + assert json.loads(body) == {"model": "gpt-6", "input": "ping"} + + +@skipIf(importlib.util.find_spec("fastapi") is None, "BUG: LIT-9333 cache hits import proxy-only fastapi") +def check_cache_hit(base: str) -> None: + from litellm.caching.caching import Cache + from litellm.types.caching import LiteLLMCacheType + + arguments: Final = dict( + model="openai/gpt-6", + messages=MESSAGES, + api_key="test-key", + api_base=base + "/v1", + caching=True, + timeout=3, + num_retries=0, + max_retries=0, + ) + with patch.object(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)): + enqueue(CHAT) + first: Final = litellm.completion(**arguments) + assert REQUESTS.get(timeout=10)[0] == "/v1/chat/completions" + cached: Final = litellm.completion(**arguments) + assert payload(cached) == payload(first) + assert payload(cached)["choices"] == [{ + "index": 0, "message": {"role": "assistant", "content": "pong"}, + "finish_reason": "stop", "provider_specific_fields": {}, + }] + assert payload(cached)["usage"] == CHAT["usage"] + assert REQUESTS.empty(), "the second SDK call missed the response cache" + + +@skipIf(importlib.util.find_spec("fastapi") is None, "BUG: LIT-9333 Messages streaming imports proxy-only fastapi") +def check_messages_stream_and_callback(base: str) -> None: + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + successes: Final[Queue[Mapping[str, object]]] = Queue() + failures: Final[Queue[object]] = Queue() + + class Probe(CustomLogger): + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + if kwargs.get("litellm_call_id") == "base-sdk-messages-check": + successes.put(kwargs) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + if kwargs.get("litellm_call_id") == "base-sdk-messages-check": + failures.put(response_obj) + + frames: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "message_start", + "message": { + "id": "msg-stream-http-check", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ) + RESPONSES.put( + (200, b"".join(f"event: {frame['type']}\ndata: {json.dumps(frame)}\n\n".encode() for frame in frames)) + ) + + async def check_async() -> None: + with patch.object(litellm, "callbacks", [Probe()]): + create: Final = cast( # cast-ok: SDK facade exposes untyped kwargs and iterator items + Callable[..., Awaitable[object]], litellm.anthropic.messages.acreate + ) + stream: Final = await create( + model="anthropic/claude-sonnet-5-5", + messages=MESSAGES, + api_key="test-key", + api_base=base, + max_tokens=16, + stream=True, + litellm_call_id="base-sdk-messages-check", + ) + assert isinstance(stream, AsyncIterable), stream + iterator: Final = cast(AsyncIterable[object], stream) # cast-ok: SDK returns an unparameterized iterator + chunks: Final = tuple([frame async for frame in iterator]) + assert all(isinstance(frame, bytes) for frame in chunks), chunks + received: Final = b"".join(frame for frame in chunks if isinstance(frame, bytes)) + assert tuple(json.loads(line[6:]) for line in received.splitlines() if line.startswith(b"data: ")) == frames + success: Final = await asyncio.to_thread(successes.get, True, 10) + await GLOBAL_LOGGING_WORKER.flush() + logging: Final = success["standard_logging_object"] + assert isinstance(logging, dict) + assert (logging["prompt_tokens"], logging["completion_tokens"], logging["total_tokens"]) == (3, 2, 5) + assert successes.empty(), "a successful stream fired more than one success callback" + assert failures.empty(), "a successful stream fired a failure callback" + + asyncio.run(check_async()) + path, headers, body = REQUESTS.get(timeout=10) + assert path == "/v1/messages" + assert headers["x-api-key"] == "test-key" + assert json.loads(body) == {"model": "claude-sonnet-5-5", "messages": MESSAGES, "max_tokens": 16, "stream": True} + + def check_http(base: str) -> None: arguments: Final = dict(model="openai/gpt-4o", messages=MESSAGES, api_key="test-key", api_base=base + "/v1") enqueue(CHAT) @@ -191,6 +387,18 @@ def main() -> None: worker.start() try: check_http(f"http://127.0.0.1:{server.server_port}") + for check in ( + check_tools, + partial(check_tools, asynchronous=True), + check_responses, + partial(check_responses, asynchronous=True), + check_cache_hit, + check_messages_stream_and_callback, + ): + try: + check(f"http://127.0.0.1:{server.server_port}") + except SkipTest as skipped: + print(f"SKIP: {skipped}") assert RESPONSES.empty() and REQUESTS.empty(), "unconsumed HTTP exchanges" finally: RELEASE.set() diff --git a/tests/integration/caching/test_cache_max_messages.py b/tests/integration/caching/test_cache_max_messages.py index 89401b7c4ec..2079761d489 100644 --- a/tests/integration/caching/test_cache_max_messages.py +++ b/tests/integration/caching/test_cache_max_messages.py @@ -1,19 +1,42 @@ +from __future__ import annotations + import json import os import uuid -from collections.abc import Callable -from typing import Final +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Final, cast import pytest +import yaml from integration._support.anthropic_sse import message_json -from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.process import owned_proxy from integration._support.provider import SharedProvider -from integration._support.wire import Reply -from pydantic import JsonValue +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter from redis import Redis +from redis.exceptions import ResponseError _Turns = Callable[[str], tuple[list[JsonValue], list[JsonValue]]] +_INDEXES: Final = TypeAdapter(tuple[bytes, ...]) +_COUNT: Final = TypeAdapter(tuple[int]) + + +def _redis_command(redis: Redis[bytes], *arguments: str | int) -> object: + execute: Final = cast( + Callable[..., object], redis.execute_command + ) # cast-ok: Redis stubs omit command return types + return execute(*arguments) + + +def _search_indexes(redis: Redis[bytes]) -> tuple[bytes, ...]: + return _INDEXES.validate_python(_redis_command(redis, "FT._LIST")) + + +def _index_size(redis: Redis[bytes], index: str) -> int: + return _COUNT.validate_python(_redis_command(redis, "FT.SEARCH", index, "*", "LIMIT", 0, 0))[0] def _claude_code_turns(task: str) -> tuple[list[JsonValue], list[JsonValue]]: @@ -108,26 +131,27 @@ def _first(value: JsonValue) -> dict[str, JsonValue]: return object_value(value[0]) -def _cached_responses(redis: Redis) -> frozenset[bytes]: +def _cached_responses(redis: Redis[bytes]) -> frozenset[bytes]: digests: Final = tuple(key for key in redis.scan_iter() if len(key) == 64) return frozenset(key for key in digests if b'"response"' in (redis.get(key) or b"")) -@pytest.mark.parametrize( - ("path", "model", "turns"), - [ - pytest.param("/v1/messages", "anthropic/claude-sonnet-5-5", _claude_code_turns, id="messages"), - pytest.param("/v1/chat/completions", "openai/gpt-5.4", _agent_turns, id="chat-completions"), - pytest.param("/v1/responses", "openai/responses/gpt-5.6-sol", _responses_turns, id="responses"), - ], +_SCENARIOS: Final = ( + pytest.param("/v1/messages", "anthropic/claude-sonnet-5-5", _claude_code_turns, id="messages"), + pytest.param("/v1/chat/completions", "openai/gpt-5.4", _agent_turns, id="chat-completions"), + pytest.param("/v1/responses", "openai/responses/gpt-5.6-sol", _responses_turns, id="responses"), ) + + +@pytest.mark.parametrize(("path", "model", "turns"), _SCENARIOS) def test_cache_serves_a_turn_under_max_messages_and_skips_one_past_it( gateway: Gateway, provider: SharedProvider, path: str, model: str, turns: _Turns ) -> None: under_cap, past_cap = turns(f"update the config {uuid.uuid4().hex}") provider.expect(_reply(path, "first answer"), _reply(path, "second answer"), _reply(path, "third answer")) - with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as redis: + redis_client: Final[Redis[bytes]] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + with redis_client as redis: cached_before: Final = _cached_responses(redis) gateway.post(path, _body(path, model, under_cap)) eventually(lambda: _cached_responses(redis) - cached_before, lambda written: len(written) == 1) @@ -142,3 +166,149 @@ def test_cache_serves_a_turn_under_max_messages_and_skips_one_past_it( "a turn past max_messages was served from the cache" ) assert len(provider.received()) == 3 + + +def _tool_reply(path: str, assistant: Mapping[str, JsonValue]) -> Reply: + base: Final = JSON_OBJECT.validate_json(_reply(path, "unused").body) + if path == "/v1/messages": + return Reply(body=json.dumps({**base, "content": assistant["content"], "stop_reason": "tool_use"}).encode()) + if path == "/v1/responses": + return Reply(body=json.dumps({**base, "output": [{**assistant, "id": "fc_cached"}]}).encode()) + return Reply( + body=json.dumps( + { + **base, + "choices": [ + { + "index": 0, + "message": assistant, + "finish_reason": "tool_calls", + } + ], + } + ).encode() + ) + + +def _stable_payload(path: str, payload: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + if path == "/v1/responses": + assert string_value(payload["id"]).startswith("resp_") + return {key: value for key, value in payload.items() if key != "id"} + return payload + + +@pytest.mark.parametrize(("path", "model", "turns"), _SCENARIOS) +def test_semantic_cache_embeds_tool_results_and_does_not_replay_the_previous_turn( + gateway: Gateway, + provider: SharedProvider, + tmp_path: Path, + path: str, + model: str, + turns: _Turns, +) -> None: + redis_client: Final[Redis[bytes]] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + with redis_client as redis: + try: + _search_indexes(redis) + except ResponseError: + pytest.skip("Redis Search is required for semantic caching") + + task: Final = f"update the config {uuid.uuid4().hex}" + first_turn, past_cap = turns(task) + equivalent_first, equivalent_past_cap = turns(task + " please") + tool_turn: Final = past_cap[:-2] + tool_result: Final = "calc.py test_calc.py" if path == "/v1/messages" else "wrote a.yaml" + index: Final = f"semantic-tool-turn-{uuid.uuid4().hex}" + + def embed(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/v1/embeddings") + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["model"] == "text-embedding-3-small" + prompt: Final = string_value(body["input"]) + vector: Final = ([0.0, 1.0] if tool_result in prompt else [1.0, 0.0]) + [0.0] * 1534 + return Reply( + body=json.dumps( + { + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": vector}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ).encode() + ) + + with wire_server(embed) as embeddings: + config: Final = JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + models: Final = config["model_list"] + assert isinstance(models, list) + settings: Final = object_value(config["litellm_settings"]) + config_path: Final = tmp_path / "semantic-cache.yaml" + config_path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + *models, + { + "model_name": "semantic-embedding", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": embeddings.url + "/v1", + "api_key": "scripted-embedding-key", + }, + }, + ], + "litellm_settings": { + **settings, + "cache_params": { + "type": "redis-semantic", + "redis_url": f"redis://{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}", + "similarity_threshold": 0.8, + "redis_semantic_cache_embedding_model": "semantic-embedding", + "redis_semantic_cache_index_name": index, + }, + }, + } + ) + ) + first_reply: Final = _tool_reply(path, object_value(tool_turn[len(first_turn)])) + provider.expect( + first_reply, + _reply(path, "tool result used"), + _reply(path, "second answer"), + _reply(path, "third answer"), + ) + try: + with owned_proxy(gateway, tmp_path, {}, config=config_path) as proxy: + initial: Final = proxy.post(path, _body(path, model, first_turn)) + eventually(lambda: _index_size(redis, index), lambda count: count == 1) + cached: Final = proxy.post(path, _body(path, model, equivalent_first)) + assert _stable_payload(path, cached) == _stable_payload(path, initial), ( + "semantic cache lost the provider-returned tool call" + ) + assert len(provider.received()) == 1 + + follow_up: Final = proxy.post(path, _body(path, model, tool_turn)) + assert _answer(path, follow_up) == "tool result used", "semantic cache replayed the prior tool call" + eventually(lambda: _index_size(redis, index), lambda count: count == 2) + assert ( + _answer(path, proxy.post(path, _body(path, model, equivalent_past_cap[:-2]))) + == "tool result used" + ) + assert len(provider.received()) == 1 + + past_cap_twice: Final = tuple(proxy.post(path, _body(path, model, past_cap)) for _ in range(2)) + assert tuple(_answer(path, answer) for answer in past_cap_twice) == ( + "second answer", + "third answer", + ) + assert len(provider.received()) == 2 + assert any( + tool_result in string_value(JSON_OBJECT.validate_json(request.body)["input"]) + for request in embeddings.drain() + ) + finally: + if index.encode() in _search_indexes(redis): + _redis_command(redis, "FT.DROPINDEX", index, "DD") diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_live_lifecycle_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_live_lifecycle_wire.py index cb7043c0362..ff4c24f92eb 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_live_lifecycle_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_live_lifecycle_wire.py @@ -1,51 +1,61 @@ -import json import threading import uuid +from collections.abc import Mapping from typing import Final -from integration._support.client import Gateway +import pytest +from integration._support.anthropic_sse import CONTENT_TAIL, message_start, sse +from integration._support.client import JSON_OBJECT, Gateway from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue -_MODEL: Final = "claude-sonnet-4-5-20250929" +_MODEL: Final = "claude-sonnet-5-5" _API_KEY: Final = "synthetic-anthropic-key" -def _sse(event: str, payload: dict[str, object]) -> bytes: - return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode() - - -def test_messages_stream_message_start_reaches_client_before_content_without_fallback( +@pytest.mark.parametrize("block_type", ("text", "thinking")) +def test_messages_stream_lifecycle_reaches_client_before_content_without_fallback( gateway: Gateway, + block_type: str, ) -> None: - """With no fallback able to take over, the proxy must not hold lifecycle - frames back for a retry that cannot happen: message_start reaches the - client while the upstream is still thinking.""" gate: Final = threading.Event() - head: Final = _sse("message_start", {"type": "message_start", "message": {"id": "msg_live_1"}}) + _sse( - "content_block_start", - {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + block_start: Final[Mapping[str, JsonValue]] = { + "type": "content_block_start", + "index": 0, + "content_block": {"type": block_type, block_type: ""}, + } + head: Final = ( + message_start("msg_live_1", _MODEL) + + sse( + "content_block_start", + block_start, + ) + + sse("ping", {"type": "ping"}) ) + delta: Final[Mapping[str, JsonValue]] = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": f"{block_type}_delta", block_type: "Hello"}, + } tail: Final = ( - _sse( + sse( "content_block_delta", - {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + delta, ) - + _sse("content_block_stop", {"type": "content_block_stop", "index": 0}) - + _sse( - "message_delta", - {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 3}}, - ) - + _sse("message_stop", {"type": "message_stop"}) + + CONTENT_TAIL ) prompt: Final = "live-lifecycle-" + uuid.uuid4().hex def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == "/v1/messages" assert request.headers["x-api-key"] == _API_KEY - body: Final = json.loads(request.body) - assert body["model"] == _MODEL - assert body["stream"] is True - assert body["messages"] == [{"role": "user", "content": prompt}] + body: Final = JSON_OBJECT.validate_json(request.body) + assert body == { + "model": _MODEL, + "stream": True, + "max_tokens": 16, + "messages": [{"role": "user", "content": prompt}], + } return Reply(content_type="text/event-stream", chunks=(head, tail), gate_after_first=gate) with wire_server(respond) as wire, gateway.scenario() as scenario: @@ -60,23 +70,27 @@ def test_messages_stream_message_start_reaches_client_before_content_without_fal "messages": [{"role": "user", "content": prompt}], }, headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=10, ) as response: assert response.status_code == 200, response.read().decode() - lines = response.iter_lines() - first_event: Final = next( - json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ") + events: Final = ( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: ") ) - assert first_event["type"] == "message_start" - gate.set() - events: Final = (first_event,) + tuple( - json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ") - ) - assert tuple(event["type"] for event in events) == ( - "message_start", - "content_block_start", + try: + start: Final = next(events) + assert start["type"] == "message_start" + assert next(events) == block_start + assert next(events) == {"type": "ping"} + finally: + gate.set() + remaining: Final = tuple(events) + assert tuple(event["type"] for event in remaining) == ( "content_block_delta", "content_block_stop", "message_delta", "message_stop", - ), f"observed events: {events!r}" + ), f"observed events: {remaining!r}" + assert remaining[0] == delta assert [request.target for request in wire.drain()] == ["/v1/messages"] diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py index fb0a9d24c0f..96929ff15c4 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py @@ -1,10 +1,14 @@ import json import time import uuid +from contextlib import nullcontext +from pathlib import Path from typing import Final import pytest +import yaml from integration._support.client import JSON_OBJECT, Gateway +from integration._support.process import owned_proxy from integration._support.timing import assert_cut_short, timed from integration._support.wire import Reply, Request, wire_server @@ -13,17 +17,23 @@ _CONFIGURED_TIMEOUT_SECONDS: Final = 1.0 @pytest.mark.covers("providers.anthropic_messages.configured_timeout_aborts_stalled_upstream") -def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream(gateway: Gateway) -> None: +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("timeout_source", ("deployment", "request_timeout")) +def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream( + gateway: Gateway, tmp_path: Path, stream: bool, timeout_source: str +) -> None: prompt: Final = "stall-" + uuid.uuid4().hex def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == "/v1/messages" assert request.headers["x-api-key"] == "synthetic-anthropic-key" body: Final = JSON_OBJECT.validate_json(request.body) - assert body["model"] == "claude-sonnet-4-5-20250929" - assert body["messages"] == [{"role": "user", "content": prompt}] - assert body["max_tokens"] == 16 - assert "timeout" not in body + assert body == { + "model": "claude-sonnet-5-5", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 16, + "stream": stream, + } time.sleep(_UPSTREAM_STALL_SECONDS) return Reply( body=json.dumps( @@ -31,7 +41,7 @@ def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream(ga "id": "msg_stalled", "type": "message", "role": "assistant", - "model": "claude-sonnet-4-5-20250929", + "model": "claude-sonnet-5-5", "content": [{"type": "text", "text": "too late"}], "stop_reason": "end_turn", "stop_sequence": None, @@ -40,18 +50,35 @@ def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream(ga ).encode() ) - with wire_server(respond) as wire, gateway.scenario() as scenario: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = JSON_OBJECT.validate_python(config["litellm_settings"]) + config_path: Final = tmp_path / "request-timeout.yaml" + config_path.write_text( + yaml.safe_dump( + { + **config, + "litellm_settings": {**settings, "request_timeout": _CONFIGURED_TIMEOUT_SECONDS}, + } + ) + ) + context: Final = ( + owned_proxy(gateway, tmp_path, {}, config=config_path) + if timeout_source == "request_timeout" + else nullcontext(gateway) + ) + with context as proxy, wire_server(respond) as wire, proxy.scenario() as scenario: model: Final = scenario.model( - model="anthropic/claude-sonnet-4-5-20250929", + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key="synthetic-anthropic-key", - timeout=_CONFIGURED_TIMEOUT_SECONDS, + model_info={}, + **({"timeout": _CONFIGURED_TIMEOUT_SECONDS} if timeout_source == "deployment" else {}), ) response: Final = timed( - lambda: gateway.request( + lambda: proxy.request( "POST", "/v1/messages", - {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + {"model": model, "max_tokens": 16, "stream": stream, "messages": [{"role": "user", "content": prompt}]}, ) ) assert response.value.status_code == 408, response.value.text diff --git a/tests/integration/sdk/test_anthropic_messages_bridge_mid_stream_failure_sdk.py b/tests/integration/sdk/test_anthropic_messages_bridge_mid_stream_failure_sdk.py index f719b8fbbb2..98c42220aab 100644 --- a/tests/integration/sdk/test_anthropic_messages_bridge_mid_stream_failure_sdk.py +++ b/tests/integration/sdk/test_anthropic_messages_bridge_mid_stream_failure_sdk.py @@ -10,6 +10,7 @@ from integration._support.anthropic_sse import ( delta_text, dropping_reply, event_types, + message_stream, parse_sse, stream_reply, user_prompt, @@ -17,10 +18,12 @@ from integration._support.anthropic_sse import ( from integration._support.client import eventually, object_value from integration._support.openai_wire import answering_model_discovery, chat_stream, posted_targets from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import TypeAdapter import litellm from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER _BACKEND: Final = "gpt-4o-mini" _PROVIDER_KEY: Final = "integration-provider-key" @@ -47,7 +50,7 @@ class _CallbackProbe(CustomLogger): self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object ) -> None: with self._lock: - self._successes.append(kwargs.get("litellm_call_id")) + self._successes.append((kwargs.get("litellm_call_id"), kwargs.get("standard_logging_object"))) def failures(self) -> tuple[object, ...]: with self._lock: @@ -74,14 +77,15 @@ def _upstream(marker: str) -> Callable[[Request], Reply]: return answering_model_discovery(respond) -async def _stream(wire: Wire, prompt: str) -> tuple[SseEvent, ...]: +async def _stream(wire: Wire, prompt: str, model: str = f"hosted_vllm/{_BACKEND}") -> tuple[SseEvent, ...]: response: Final = await litellm.anthropic.messages.acreate( - model=f"hosted_vllm/{_BACKEND}", - api_base=wire.url + "/v1", + model=model, + api_base=wire.url if model.startswith("anthropic/") else wire.url + "/v1", api_key=_PROVIDER_KEY, max_tokens=16, stream=True, messages=[{"role": "user", "content": prompt}], + litellm_call_id=prompt, ) assert isinstance(response, AsyncIterable), response frames: Final = [frame async for frame in response] @@ -111,3 +115,44 @@ async def test_sdk_bridged_stream_failing_after_content_reports_the_provider_err assert len(failures) == 1, failures assert isinstance(failures[0], Exception), failures assert not isinstance(failures[0], MidStreamFallbackError), failures + assert len(probe.successes()) == 1, probe.successes() + + +async def test_sdk_native_messages_stream_delivers_one_success_callback_with_usage() -> None: + probe: Final = _CallbackProbe() + litellm.callbacks.append(probe) + prompt: Final = uuid.uuid4().hex + model: Final = "claude-sonnet-5-5" + + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/v1/messages") + assert request.headers["x-api-key"] == _PROVIDER_KEY + assert json.loads(request.body) == { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "max_tokens": 16, + } + return stream_reply(message_stream("msg_success", model, _TEXT)) + + with wire_server(respond) as wire: + events: Final = await _stream(wire, prompt, f"anthropic/{model}") + assert delta_text(events) == _TEXT + assert event_types(events) == ( + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ) + await _settled(probe.successes) + await GLOBAL_LOGGING_WORKER.flush() + successes: Final = probe.successes() + assert len(successes) == 1, successes + success: Final = TypeAdapter(tuple[object, object]).validate_python(successes[0]) + assert success[0] == prompt + logging: Final = TypeAdapter(dict[str, object]).validate_python(success[1]) + assert (logging["prompt_tokens"], logging["completion_tokens"], logging["total_tokens"]) == (5, 3, 8) + assert probe.failures() == () + assert posted_targets(wire) == ("/v1/messages",)