mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
fix(anthropic): keep tool_search_tool_result so the search call replays with its result (#46440)
* fix(anthropic): keep tool_search_tool_result so the search call replays with its result Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): drop a client tool_result for a search call replayed with its server result * fix(anthropic): replay a server tool call the client answered as the client sent it * fix(anthropic): keep a server tool result under modify_params, count client-answered marks, and replay it from run_tool_loop modify_params no longer adds a placeholder answer for a server tool call whose result is kept, unless the call ends the request. The cache control hook counts a mark on a server tool call the client answered, since that call goes out as a plain tool_use. run_tool_loop sends provider_specific_fields back and runs only the client's tool calls. * fix(anthropic): keep a server tool result when another tool answer follows, scope run_tool_loop's replay, and keep code_interpreter_results off tool search turns * test(anthropic): audit cells for tool search result round trips, replay, cache marks, pass-through logging and chaos --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang
GitHub
parent
63c3b7a34e
commit
9ba4cc4fb3
18 changed files with 2724 additions and 71 deletions
|
|
@ -12,7 +12,7 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and
|
|||
import copy
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from collections.abc import Callable, Container, Iterable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -26,6 +26,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
cache_control_as_prompt_cache_breakpoint,
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import tool_call_ids_answered_by_client
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
is_claude_code_one_shot_subagent_request,
|
||||
supports_anthropic_cache_control,
|
||||
|
|
@ -227,12 +228,12 @@ def _as_object_list(value: object | None) -> list[object] | None:
|
|||
return _validated_object_list(value)
|
||||
|
||||
|
||||
def _tool_call_carries_cache_breakpoint(tool_call: object, message: object) -> bool:
|
||||
def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answered_tool_ids: Container[str]) -> bool:
|
||||
if _attribute_or_key(tool_call, "cache_control") is None:
|
||||
return False
|
||||
|
||||
return not tool_call_is_rebuilt_as_server_tool_use(
|
||||
_attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields")
|
||||
_attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields"), answered_tool_ids
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -424,7 +425,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
system_blocks: Final = (
|
||||
sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
|
||||
)
|
||||
return system_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages)
|
||||
message_list: Final = tuple(messages)
|
||||
answered_tool_ids: Final = tool_call_ids_answered_by_client(message_list)
|
||||
return system_blocks + sum(
|
||||
AnthropicCacheControlHook._count_cache_control_blocks(msg, answered_tool_ids) for msg in message_list
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def count_external_cache_breakpoints(
|
||||
|
|
@ -509,6 +514,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
so earlier points win when slots are scarce.
|
||||
"""
|
||||
used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
|
||||
answered_tool_ids: Final = tool_call_ids_answered_by_client(messages)
|
||||
|
||||
limit_reached = False
|
||||
for point in points:
|
||||
|
|
@ -525,14 +531,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
limit_reached = True
|
||||
break
|
||||
|
||||
if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
|
||||
if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], answered_tool_ids):
|
||||
# Client already marked this message; don't overwrite it.
|
||||
continue
|
||||
|
||||
messages[target_index] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
|
||||
messages[target_index], control, openai_dialect
|
||||
)
|
||||
if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
|
||||
if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], answered_tool_ids):
|
||||
used_blocks += 1
|
||||
|
||||
if limit_reached:
|
||||
|
|
@ -586,22 +592,26 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
return []
|
||||
|
||||
@staticmethod
|
||||
def _count_cache_control_blocks(message: object) -> int:
|
||||
def _count_cache_control_blocks(message: object, answered_tool_ids: Container[str] = frozenset()) -> int:
|
||||
message_count: Final = 1 if _carries_cache_breakpoint(message) else 0
|
||||
content: Final = _as_object_list(_attribute_or_key(message, "content"))
|
||||
content_count: Final = sum(1 for block in content if _carries_cache_breakpoint(block)) if content else 0
|
||||
tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls"))
|
||||
tool_call_count: Final = (
|
||||
sum(1 for tool_call in tool_calls if _tool_call_carries_cache_breakpoint(tool_call, message))
|
||||
sum(
|
||||
1
|
||||
for tool_call in tool_calls
|
||||
if _tool_call_carries_cache_breakpoint(tool_call, message, answered_tool_ids)
|
||||
)
|
||||
if tool_calls
|
||||
else 0
|
||||
)
|
||||
return message_count + content_count + tool_call_count
|
||||
|
||||
@staticmethod
|
||||
def _message_has_cache_control(message: AllMessageValues) -> bool:
|
||||
def _message_has_cache_control(message: AllMessageValues, answered_tool_ids: Container[str]) -> bool:
|
||||
"""Return True if the message already carries any cache_control."""
|
||||
return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0
|
||||
return AnthropicCacheControlHook._count_cache_control_blocks(message, answered_tool_ids) > 0
|
||||
|
||||
@staticmethod
|
||||
def _safe_insert_cache_control_in_message(
|
||||
|
|
|
|||
|
|
@ -20,7 +20,11 @@ import litellm.types.llms
|
|||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import anthropic_system_messages
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
|
||||
anthropic_system_messages,
|
||||
message_field,
|
||||
parts_of,
|
||||
)
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
|
||||
from litellm.types.files import get_file_extension_from_mime_type
|
||||
|
|
@ -1832,6 +1836,60 @@ def find_anthropic_server_tool_result(
|
|||
)
|
||||
|
||||
|
||||
def _tool_call_ids_in_tool_messages(messages: Sequence[object]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
tool_call_id for message in messages if isinstance(tool_call_id := message_field(message, "tool_call_id"), str)
|
||||
)
|
||||
|
||||
|
||||
def _messages_until_the_next_assistant(messages: Sequence[object], index: int) -> tuple[object, ...]:
|
||||
return tuple(
|
||||
itertools.takewhile(
|
||||
lambda message: message_field(message, "role") != "assistant",
|
||||
(messages[following_index] for following_index in range(index + 1, len(messages))),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _tool_call_ids_with_a_kept_server_result(message: object) -> frozenset[str]:
|
||||
from litellm.llms.anthropic.common_utils import tool_call_is_rebuilt_as_server_tool_use
|
||||
|
||||
provider_specific_fields: Final = message_field(message, "provider_specific_fields")
|
||||
return frozenset(
|
||||
tool_call_id
|
||||
for tool_call in parts_of(message_field(message, "tool_calls"))
|
||||
if isinstance(tool_call_id := message_field(tool_call, "id"), str)
|
||||
and tool_call_is_rebuilt_as_server_tool_use(tool_call_id, provider_specific_fields)
|
||||
)
|
||||
|
||||
|
||||
def _server_tool_answers_replaced_after(messages: Sequence[object], index: int) -> frozenset[str]:
|
||||
kept_server_result_ids: Final = _tool_call_ids_with_a_kept_server_result(messages[index])
|
||||
if not kept_server_result_ids:
|
||||
return frozenset()
|
||||
following: Final = _messages_until_the_next_assistant(messages, index)
|
||||
a_user_turn_remains: Final = any(
|
||||
message_field(message, "role") in ("user", "tool", "function")
|
||||
and message_field(message, "tool_call_id") not in kept_server_result_ids
|
||||
for message in following
|
||||
)
|
||||
return _tool_call_ids_in_tool_messages(following) & kept_server_result_ids if a_user_turn_remains else frozenset()
|
||||
|
||||
|
||||
def server_tool_answers_replaced_by_their_result(messages: Sequence[object]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
itertools.chain.from_iterable(
|
||||
_server_tool_answers_replaced_after(messages, index)
|
||||
for index, message in enumerate(messages)
|
||||
if message_field(message, "role") == "assistant"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def tool_call_ids_answered_by_client(messages: Sequence[object]) -> frozenset[str]:
|
||||
return _tool_call_ids_in_tool_messages(messages) - server_tool_answers_replaced_by_their_result(messages)
|
||||
|
||||
|
||||
_AnthropicToolInvokeItem: TypeAlias = AnthropicMessagesToolUseParam | dict[str, Any]
|
||||
|
||||
|
||||
|
|
@ -1839,6 +1897,7 @@ def convert_to_anthropic_tool_invoke(
|
|||
tool_calls: list[ChatCompletionAssistantToolCall],
|
||||
web_search_results: Sequence[object] | None = None,
|
||||
tool_results: Sequence[object] | None = None,
|
||||
answered_tool_ids: Container[str] = frozenset(),
|
||||
) -> list[_AnthropicToolInvokeItem]:
|
||||
"""
|
||||
OpenAI tool invokes:
|
||||
|
|
@ -1899,7 +1958,11 @@ def convert_to_anthropic_tool_invoke(
|
|||
context="Anthropic tool invoke",
|
||||
)
|
||||
|
||||
server_tool_result = find_anthropic_server_tool_result(tool_id, web_search_results, tool_results)
|
||||
server_tool_result = (
|
||||
None
|
||||
if tool_id in answered_tool_ids
|
||||
else find_anthropic_server_tool_result(tool_id, web_search_results, tool_results)
|
||||
)
|
||||
if server_tool_result is not None:
|
||||
anthropic_tool_invoke.append(
|
||||
{
|
||||
|
|
@ -2462,6 +2525,8 @@ def anthropic_messages_pt(
|
|||
# (Cases A/B/D in `sanitize_messages_for_tool_calling`) remains gated on
|
||||
# `modify_params` because it actually mutates conversation structure.
|
||||
messages = [_sanitize_empty_text_content(m) for m in messages]
|
||||
replaced_server_tool_answers: Final = server_tool_answers_replaced_by_their_result(messages)
|
||||
client_answered_tool_ids: Final = _tool_call_ids_in_tool_messages(messages) - replaced_server_tool_answers
|
||||
|
||||
# add role=tool support to allow function call result/error submission
|
||||
user_message_types: Final = {"user", "tool", "function"}
|
||||
|
|
@ -2583,7 +2648,9 @@ def anthropic_messages_pt(
|
|||
|
||||
user_content.append(_anthropic_content_text_element)
|
||||
|
||||
elif user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function":
|
||||
elif (
|
||||
user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function"
|
||||
) and message_field(user_message_types_block, "tool_call_id") not in replaced_server_tool_answers:
|
||||
# OpenAI's tool message content will always be a string
|
||||
user_content.append(
|
||||
convert_to_anthropic_tool_result(user_message_types_block, force_base64=force_base64)
|
||||
|
|
@ -2651,6 +2718,7 @@ def anthropic_messages_pt(
|
|||
assistant_tool_calls,
|
||||
web_search_results=_web_search_results_tc,
|
||||
tool_results=_tool_results_tc,
|
||||
answered_tool_ids=client_answered_tool_ids,
|
||||
)
|
||||
|
||||
# Group tool invoke results into (server_tool_use, result) pairs
|
||||
|
|
@ -2834,6 +2902,7 @@ def anthropic_messages_pt(
|
|||
assistant_tool_calls,
|
||||
web_search_results=_web_search_results,
|
||||
tool_results=_tool_results,
|
||||
answered_tool_ids=client_answered_tool_ids,
|
||||
)
|
||||
|
||||
# Prevent "tool_use ids must be unique" errors by filtering duplicates
|
||||
|
|
|
|||
|
|
@ -1053,13 +1053,14 @@ class ModelResponseIterator:
|
|||
# Fixes: https://github.com/BerriAI/litellm/issues/18137
|
||||
self.web_search_results.append(content_block_start["content_block"])
|
||||
provider_specific_fields["web_search_results"] = self.web_search_results
|
||||
elif content_type != "tool_search_tool_result":
|
||||
# Handle other tool results (code execution, etc.)
|
||||
# Skip tool_search_tool_result as it's internal metadata
|
||||
else:
|
||||
self.tool_results.append(content_block_start["content_block"])
|
||||
provider_specific_fields["tool_results"] = self.tool_results
|
||||
# Convert to provider-neutral code_interpreter_results
|
||||
provider_specific_fields["code_interpreter_results"] = self._build_code_interpreter_results()
|
||||
if content_type != "tool_search_tool_result":
|
||||
# Convert to provider-neutral code_interpreter_results
|
||||
provider_specific_fields["code_interpreter_results"] = (
|
||||
self._build_code_interpreter_results()
|
||||
)
|
||||
|
||||
elif type_chunk == "content_block_stop":
|
||||
ContentBlockStop(**chunk)
|
||||
|
|
|
|||
|
|
@ -2338,9 +2338,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
|
||||
## TOOL RESULTS - handle all tool result types (code execution, etc.)
|
||||
elif content["type"].endswith("_tool_result"):
|
||||
# Skip tool_search_tool_result as it's internal metadata
|
||||
if content["type"] == "tool_search_tool_result":
|
||||
continue
|
||||
# Handle web_search_tool_result separately for backwards compatibility
|
||||
if content["type"] == "web_search_tool_result" or content["type"] == "web_fetch_tool_result":
|
||||
if web_search_results is None:
|
||||
|
|
@ -2709,6 +2706,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
|
||||
if tool_results is not None:
|
||||
provider_specific_fields["tool_results"] = tool_results
|
||||
|
||||
if tool_results is not None and any(result.get("type") != "tool_search_tool_result" for result in tool_results):
|
||||
container_id: Final = (
|
||||
completion_response.get("container", {}).get("id")
|
||||
if isinstance(completion_response.get("container"), dict)
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ This file contains common utils for anthropic calls.
|
|||
import copy
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, MutableMapping, Sequence
|
||||
from collections.abc import Container, Mapping, MutableMapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
|
||||
|
|
@ -1890,9 +1890,11 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
|
|||
return None
|
||||
|
||||
|
||||
def tool_call_is_rebuilt_as_server_tool_use(tool_call_id: object, provider_specific_fields: object) -> bool:
|
||||
def tool_call_is_rebuilt_as_server_tool_use(
|
||||
tool_call_id: object, provider_specific_fields: object, answered_tool_ids: Container[str] = frozenset()
|
||||
) -> bool:
|
||||
fields: Final = _validated_claude_code_mapping(provider_specific_fields)
|
||||
if not isinstance(tool_call_id, str) or fields is None:
|
||||
if not isinstance(tool_call_id, str) or tool_call_id in answered_tool_ids or fields is None:
|
||||
return False
|
||||
return (
|
||||
find_anthropic_server_tool_result(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Client-side tool-calling loop helpers for litellm.completion."""
|
||||
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from typing import TypeAlias, cast
|
||||
from typing import Final, TypeAlias, cast
|
||||
|
||||
from typing_extensions import TypedDict, Unpack
|
||||
|
||||
|
|
@ -56,9 +56,9 @@ def _assistant_tool_call(tool_call: ChatCompletionMessageToolCall) -> ChatComple
|
|||
|
||||
|
||||
def _assistant_message(
|
||||
message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...]
|
||||
message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], replays_server_tool_results: bool
|
||||
) -> ChatCompletionAssistantMessage:
|
||||
return cast( # cast-ok: dict literal with thinking_blocks and reasoning_items spread in only when set
|
||||
return cast( # cast-ok: dict literal with thinking_blocks, reasoning_items and provider_specific_fields spread in only when set
|
||||
"ChatCompletionAssistantMessage",
|
||||
{
|
||||
"role": "assistant",
|
||||
|
|
@ -69,6 +69,10 @@ def _assistant_message(
|
|||
for key, value in (
|
||||
("thinking_blocks", getattr(message, "thinking_blocks", None)),
|
||||
("reasoning_items", getattr(message, "reasoning_items", None)),
|
||||
(
|
||||
"provider_specific_fields",
|
||||
getattr(message, "provider_specific_fields", None) if replays_server_tool_results else None,
|
||||
),
|
||||
)
|
||||
if value is not None
|
||||
},
|
||||
|
|
@ -88,6 +92,19 @@ def _function_tool_calls(message: Message) -> tuple[ChatCompletionMessageToolCal
|
|||
return tuple(_function_tool_call(tool_call) for tool_call in message.tool_calls or ())
|
||||
|
||||
|
||||
def _client_tool_calls(
|
||||
message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...]
|
||||
) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
from litellm.llms.anthropic.common_utils import tool_call_is_rebuilt_as_server_tool_use
|
||||
|
||||
provider_specific_fields: Final = getattr(message, "provider_specific_fields", None)
|
||||
return tuple(
|
||||
tool_call
|
||||
for tool_call in tool_calls
|
||||
if not tool_call_is_rebuilt_as_server_tool_use(tool_call.id, provider_specific_fields)
|
||||
)
|
||||
|
||||
|
||||
def run_tool_loop(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -123,10 +140,17 @@ def run_tool_loop(
|
|||
if not message.tool_calls:
|
||||
return message.content
|
||||
tool_calls = _function_tool_calls(message)
|
||||
client_tool_calls = _client_tool_calls(message, tool_calls)
|
||||
if not client_tool_calls:
|
||||
return message.content
|
||||
if round_number == max_rounds:
|
||||
break
|
||||
tool_results = tuple(execute_tool(tool_call) for tool_call in tool_calls)
|
||||
history = (*history, _assistant_message(message, tool_calls), *tool_results)
|
||||
tool_results = tuple(execute_tool(tool_call) for tool_call in client_tool_calls)
|
||||
history = (
|
||||
*history,
|
||||
_assistant_message(message, tool_calls, client_tool_calls != tool_calls),
|
||||
*tool_results,
|
||||
)
|
||||
raise ToolLoopMaxRoundsExceeded(max_rounds)
|
||||
|
||||
|
||||
|
|
@ -165,8 +189,15 @@ async def arun_tool_loop(
|
|||
if not message.tool_calls:
|
||||
return message.content
|
||||
tool_calls = _function_tool_calls(message)
|
||||
client_tool_calls = _client_tool_calls(message, tool_calls)
|
||||
if not client_tool_calls:
|
||||
return message.content
|
||||
if round_number == max_rounds:
|
||||
break
|
||||
tool_results = tuple([await execute_tool(tool_call) for tool_call in tool_calls])
|
||||
history = (*history, _assistant_message(message, tool_calls), *tool_results)
|
||||
tool_results = tuple([await execute_tool(tool_call) for tool_call in client_tool_calls])
|
||||
history = (
|
||||
*history,
|
||||
_assistant_message(message, tool_calls, client_tool_calls != tool_calls),
|
||||
*tool_results,
|
||||
)
|
||||
raise ToolLoopMaxRoundsExceeded(max_rounds)
|
||||
|
|
|
|||
378
tests/integration/providers/_tool_search_support.py
Normal file
378
tests/integration/providers/_tool_search_support.py
Normal file
|
|
@ -0,0 +1,378 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
import httpx
|
||||
from integration._support.anthropic_thinking import (
|
||||
JSON_LIST,
|
||||
JSON_OBJECT,
|
||||
NO_CACHE,
|
||||
chunks_of,
|
||||
deltas_of,
|
||||
stream_reply,
|
||||
streams,
|
||||
)
|
||||
from integration._support.client import Gateway, Scenario, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request
|
||||
from pydantic import JsonValue
|
||||
|
||||
Event = dict[str, JsonValue]
|
||||
Provider = Literal[
|
||||
"anthropic", "bedrock_invoke", "claude_platform", "bedrock_mantle", "vertex_ai", "azure_ai", "snowflake"
|
||||
]
|
||||
Responder = Callable[[Request], Reply]
|
||||
|
||||
MODEL: Final = "claude-opus-5-5"
|
||||
BEDROCK_MODEL: Final = "anthropic.claude-opus-5-5"
|
||||
SEARCH_ID: Final = "srvtoolu_01ScriptedSearch"
|
||||
WEATHER_ID: Final = "toolu_01ScriptedWeather"
|
||||
SEARCH_TOOL: Final = "tool_search_tool_bm25"
|
||||
SEARCH_INPUT: Final[Event] = {"query": "weather"}
|
||||
WEATHER_INPUT: Final[Event] = {"city": "Paris"}
|
||||
WEATHER_ANSWER: Final = "18C and cloudy"
|
||||
VERTEX_PROJECT: Final = "scripted-project"
|
||||
VERTEX_LOCATION: Final = "us-east5"
|
||||
VERTEX_MODEL_PATH: Final = (
|
||||
f"/v1/projects/{VERTEX_PROJECT}/locations/{VERTEX_LOCATION}/publishers/anthropic/models/{MODEL}"
|
||||
)
|
||||
TOOLS: Final[list[JsonValue]] = [
|
||||
{"type": "tool_search_tool_bm25_20251119", "name": SEARCH_TOOL},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
},
|
||||
"defer_loading": True,
|
||||
},
|
||||
]
|
||||
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
|
||||
|
||||
|
||||
def search_result(tool_use_id: str = SEARCH_ID) -> Event:
|
||||
return {
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": tool_use_id,
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
SEARCH_RESULT: Final = search_result()
|
||||
SEARCH_USE: Final[Event] = {"type": "server_tool_use", "id": SEARCH_ID, "name": SEARCH_TOOL, "input": SEARCH_INPUT}
|
||||
WEATHER_USE: Final[Event] = {"type": "tool_use", "id": WEATHER_ID, "name": "get_weather", "input": WEATHER_INPUT}
|
||||
SEARCH_TURN: Final[tuple[Event, ...]] = (SEARCH_USE, SEARCH_RESULT, WEATHER_USE)
|
||||
|
||||
|
||||
def question(marker: str) -> str:
|
||||
return f"What is the weather in Paris? marker-{marker}"
|
||||
|
||||
|
||||
def answer(marker: str) -> str:
|
||||
return f"{WEATHER_ANSWER} for marker-{marker}"
|
||||
|
||||
|
||||
def marker_of(request: Request) -> str:
|
||||
found: Final = _MARKER.findall(request.body.decode())
|
||||
assert found, request.body
|
||||
return found[0]
|
||||
|
||||
|
||||
def anthropic_message(identity: str, content: Sequence[Event], stop_reason: str | None) -> Event:
|
||||
return {
|
||||
"id": identity,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": MODEL,
|
||||
"content": list(content),
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 40, "output_tokens": 12},
|
||||
}
|
||||
|
||||
|
||||
def _block_events(index: int, block: Event) -> tuple[Event, ...]:
|
||||
stop: Final[Event] = {"type": "content_block_stop", "index": index}
|
||||
match block["type"]:
|
||||
case "server_tool_use" | "tool_use":
|
||||
return (
|
||||
{"type": "content_block_start", "index": index, "content_block": {**block, "input": {}}},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {"type": "input_json_delta", "partial_json": json.dumps(block["input"])},
|
||||
},
|
||||
stop,
|
||||
)
|
||||
case "text":
|
||||
return (
|
||||
{"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": block["text"]}},
|
||||
stop,
|
||||
)
|
||||
case _:
|
||||
return ({"type": "content_block_start", "index": index, "content_block": block}, stop)
|
||||
|
||||
|
||||
def message_events(identity: str, blocks: Sequence[Event], stop_reason: str) -> tuple[Event, ...]:
|
||||
return (
|
||||
{"type": "message_start", "message": anthropic_message(identity, (), None)},
|
||||
*chain.from_iterable(_block_events(index, block) for index, block in enumerate(blocks)),
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": {"output_tokens": 12},
|
||||
},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
|
||||
|
||||
def turn_reply(
|
||||
request: Request, identity: str, blocks: Sequence[Event], stop_reason: str, *, abort_after: int | None = None
|
||||
) -> Reply:
|
||||
if streams(request):
|
||||
return stream_reply(request, message_events(identity, blocks, stop_reason), abort_after=abort_after)
|
||||
return Reply(body=json.dumps(anthropic_message(identity, blocks, stop_reason)).encode())
|
||||
|
||||
|
||||
def request_messages(request: Request) -> tuple[Event, ...]:
|
||||
messages: Final = JSON_LIST.validate_python(JSON_OBJECT.validate_json(request.body)["messages"])
|
||||
return tuple(JSON_OBJECT.validate_python(message) for message in messages)
|
||||
|
||||
|
||||
def answer_reply(request: Request) -> Reply:
|
||||
marker: Final = marker_of(request)
|
||||
return turn_reply(request, f"msg_{marker}_answer", ({"type": "text", "text": answer(marker)},), "end_turn")
|
||||
|
||||
|
||||
def scripted_peer(first_turn: Sequence[Event] = SEARCH_TURN) -> Responder:
|
||||
def respond(request: Request) -> Reply:
|
||||
if len(request_messages(request)) > 1:
|
||||
return answer_reply(request)
|
||||
return turn_reply(request, f"msg_{marker_of(request)}_search", first_turn, "tool_use")
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def deployment(scenario: Scenario, provider: Provider, wire_url: str, upstream_url: str) -> str:
|
||||
match provider:
|
||||
case "anthropic":
|
||||
return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key="scripted-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 "bedrock_mantle":
|
||||
return scenario.model(
|
||||
model=f"bedrock_mantle/{BEDROCK_MODEL}", api_base=wire_url, api_key="scripted-mantle-key"
|
||||
)
|
||||
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=service_account_json(VERTEX_PROJECT, upstream_url.rstrip("/")),
|
||||
)
|
||||
case "azure_ai":
|
||||
return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key")
|
||||
case "snowflake":
|
||||
return scenario.model(model=f"snowflake/{MODEL}", api_base=wire_url, api_key="scripted-snowflake-key")
|
||||
|
||||
|
||||
def target(provider: Provider, stream: bool) -> str:
|
||||
match provider:
|
||||
case "anthropic" | "claude_platform":
|
||||
return "/v1/messages"
|
||||
case "bedrock_invoke":
|
||||
return f"/model/{BEDROCK_MODEL}/invoke-with-response-stream" if stream else f"/model/{BEDROCK_MODEL}/invoke"
|
||||
case "bedrock_mantle" | "azure_ai":
|
||||
return "/anthropic/v1/messages"
|
||||
case "vertex_ai":
|
||||
return f"{VERTEX_MODEL_PATH}:streamRawPredict" if stream else f"{VERTEX_MODEL_PATH}:rawPredict"
|
||||
case "snowflake":
|
||||
return "/api/v2/cortex/v1/messages"
|
||||
|
||||
|
||||
def target_path(request: Request) -> str:
|
||||
return urlsplit(unquote(request.target)).path
|
||||
|
||||
|
||||
def chat_body(model: str, messages: Sequence[JsonValue], stream: bool, **fields: JsonValue) -> Event:
|
||||
return {
|
||||
"model": model,
|
||||
"max_tokens": 64,
|
||||
"messages": list(messages),
|
||||
"tools": TOOLS,
|
||||
"stream": stream,
|
||||
**NO_CACHE,
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
def chat(gateway: Gateway, body: Mapping[str, JsonValue], *, key: str | None = None) -> httpx.Response:
|
||||
return gateway.request("POST", "/v1/chat/completions", body, key=key)
|
||||
|
||||
|
||||
def first_messages(marker: str) -> list[JsonValue]:
|
||||
return [{"role": "user", "content": question(marker)}]
|
||||
|
||||
|
||||
def tool_answer(tool_call_id: JsonValue, content: str = WEATHER_ANSWER) -> Event:
|
||||
return {"role": "tool", "tool_call_id": tool_call_id, "content": content}
|
||||
|
||||
|
||||
def replay_messages(
|
||||
marker: str, assistant: Event, answers: Sequence[JsonValue] = (tool_answer(WEATHER_ID),)
|
||||
) -> list[JsonValue]:
|
||||
return [{"role": "user", "content": question(marker)}, assistant, *answers]
|
||||
|
||||
|
||||
def _stream_tool_call_parts(deltas: Sequence[Event]) -> tuple[Event, ...]:
|
||||
calls: Final = chain.from_iterable(JSON_LIST.validate_python(delta.get("tool_calls") or []) for delta in deltas)
|
||||
return tuple(JSON_OBJECT.validate_python(call) for call in calls)
|
||||
|
||||
|
||||
def _folded_tool_call(parts: Sequence[Event], index: JsonValue) -> Event:
|
||||
mine: Final = tuple(part for part in parts if part["index"] == index)
|
||||
functions: Final = tuple(JSON_OBJECT.validate_python(part.get("function") or {}) for part in mine)
|
||||
return {
|
||||
"id": next(part["id"] for part in mine if part.get("id")),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "".join(str(function.get("name") or "") for function in functions),
|
||||
"arguments": "".join(str(function.get("arguments") or "") for function in functions),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _merged_fields(merged: Event, delta: Event) -> Event:
|
||||
return {**merged, **JSON_OBJECT.validate_python(delta.get("provider_specific_fields") or {})}
|
||||
|
||||
|
||||
def folded_assistant(deltas: Sequence[Event]) -> Event:
|
||||
parts: Final = _stream_tool_call_parts(deltas)
|
||||
indexes: Final = tuple(dict.fromkeys(part["index"] for part in parts))
|
||||
content: Final = "".join(str(delta.get("content") or "") for delta in deltas)
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": content or None,
|
||||
"tool_calls": [_folded_tool_call(parts, index) for index in indexes],
|
||||
"provider_specific_fields": reduce(_merged_fields, deltas, {}),
|
||||
}
|
||||
|
||||
|
||||
def assistant_from(response: httpx.Response, stream: bool) -> Event:
|
||||
if stream:
|
||||
return folded_assistant(deltas_of(chunks_of(response.text)))
|
||||
choices: Final = JSON_LIST.validate_python(JSON_OBJECT.validate_json(response.content)["choices"])
|
||||
return JSON_OBJECT.validate_python(JSON_OBJECT.validate_python(choices[0])["message"])
|
||||
|
||||
|
||||
def response_id(response: httpx.Response, stream: bool) -> str:
|
||||
if stream:
|
||||
return str(chunks_of(response.text)[0]["id"])
|
||||
return str(JSON_OBJECT.validate_json(response.content)["id"])
|
||||
|
||||
|
||||
def answer_text(response: httpx.Response, stream: bool) -> str:
|
||||
if stream:
|
||||
return "".join(str(delta.get("content") or "") for delta in deltas_of(chunks_of(response.text)))
|
||||
return str(assistant_from(response, stream=False)["content"])
|
||||
|
||||
|
||||
def kept_results(assistant: Event) -> JsonValue:
|
||||
return JSON_OBJECT.validate_python(assistant.get("provider_specific_fields") or {}).get("tool_results")
|
||||
|
||||
|
||||
def provider_fields(assistant: Event) -> Event:
|
||||
return JSON_OBJECT.validate_python(assistant.get("provider_specific_fields") or {})
|
||||
|
||||
|
||||
def tool_call_ids(assistant: Event) -> list[JsonValue]:
|
||||
calls: Final = JSON_LIST.validate_python(assistant.get("tool_calls") or [])
|
||||
return [JSON_OBJECT.validate_python(call)["id"] for call in calls]
|
||||
|
||||
|
||||
def wire_content(message: Event) -> list[JsonValue]:
|
||||
content: Final = message["content"]
|
||||
if isinstance(content, str):
|
||||
return [{"type": "text", "text": content}]
|
||||
return JSON_LIST.validate_python(content)
|
||||
|
||||
|
||||
def wire_assistant(request: Request, position: int = 1) -> list[JsonValue]:
|
||||
messages: Final = request_messages(request)
|
||||
assert messages[position]["role"] == "assistant", request.body
|
||||
return wire_content(messages[position])
|
||||
|
||||
|
||||
def answered_ids(request: Request) -> list[JsonValue]:
|
||||
blocks: Final = chain.from_iterable(
|
||||
wire_content(message) for message in request_messages(request) if message["role"] == "user"
|
||||
)
|
||||
objects: Final = tuple(JSON_OBJECT.validate_python(block) for block in blocks)
|
||||
return [block["tool_use_id"] for block in objects if block["type"] == "tool_result"]
|
||||
|
||||
|
||||
def user_texts(request: Request) -> list[JsonValue]:
|
||||
blocks: Final = chain.from_iterable(
|
||||
wire_content(message) for message in request_messages(request) if message["role"] == "user"
|
||||
)
|
||||
objects: Final = tuple(JSON_OBJECT.validate_python(block) for block in blocks)
|
||||
return [block["text"] for block in objects if block["type"] == "text"]
|
||||
|
||||
|
||||
def spend_row(request_id: str) -> Event:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)
|
||||
),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def hand_built_assistant(
|
||||
tool_calls: Sequence[tuple[str, str, Event]], fields: JsonValue = None, **extra: JsonValue
|
||||
) -> Event:
|
||||
calls: Final[list[JsonValue]] = [
|
||||
{"id": identity, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}}
|
||||
for identity, name, arguments in tool_calls
|
||||
]
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": calls,
|
||||
**({} if fields is None else {"provider_specific_fields": fields}),
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
SEARCH_CALL: Final[tuple[str, str, Event]] = (SEARCH_ID, SEARCH_TOOL, SEARCH_INPUT)
|
||||
WEATHER_CALL: Final[tuple[str, str, Event]] = (WEATHER_ID, "get_weather", WEATHER_INPUT)
|
||||
KEPT_FIELDS: Final[Event] = {"tool_results": [SEARCH_RESULT]}
|
||||
|
|
@ -0,0 +1,291 @@
|
|||
import asyncio
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
from integration.providers._tool_search_support import (
|
||||
MODEL,
|
||||
SEARCH_RESULT,
|
||||
SEARCH_TURN,
|
||||
WEATHER_ID,
|
||||
Event,
|
||||
answer,
|
||||
answered_ids,
|
||||
assistant_from,
|
||||
chat_body,
|
||||
deployment,
|
||||
first_messages,
|
||||
hand_built_assistant,
|
||||
KEPT_FIELDS,
|
||||
SEARCH_CALL,
|
||||
WEATHER_CALL,
|
||||
kept_results,
|
||||
marker_of,
|
||||
question,
|
||||
replay_messages,
|
||||
request_messages,
|
||||
response_id,
|
||||
scripted_peer,
|
||||
wire_assistant,
|
||||
)
|
||||
|
||||
Kind = Literal["replay", "first", "messages", "responses"]
|
||||
|
||||
_OWNED_MODEL: Final = "tool-search-chaos"
|
||||
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
|
||||
_NO_CACHE: Final[Event] = {"cache": {"no-cache": True}}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Call:
|
||||
kind: Kind
|
||||
stream: bool
|
||||
marker: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
call: _Call
|
||||
status: int
|
||||
text: str
|
||||
|
||||
|
||||
def _path(kind: Kind) -> str:
|
||||
match kind:
|
||||
case "replay" | "first":
|
||||
return "/v1/chat/completions"
|
||||
case "messages":
|
||||
return "/v1/messages"
|
||||
case "responses":
|
||||
return "/v1/responses"
|
||||
|
||||
|
||||
def _body(model: str, call: _Call) -> Event:
|
||||
match call.kind:
|
||||
case "replay":
|
||||
assistant: Final = hand_built_assistant((SEARCH_CALL, WEATHER_CALL), KEPT_FIELDS)
|
||||
return chat_body(model, replay_messages(call.marker, assistant), call.stream)
|
||||
case "first":
|
||||
return chat_body(model, first_messages(call.marker), call.stream)
|
||||
case "messages":
|
||||
return {
|
||||
"model": model,
|
||||
"max_tokens": 64,
|
||||
"stream": call.stream,
|
||||
"messages": [{"role": "user", "content": question(call.marker)}],
|
||||
}
|
||||
case "responses":
|
||||
return {"model": model, "input": question(call.marker), "max_output_tokens": 64, **_NO_CACHE}
|
||||
|
||||
|
||||
async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
|
||||
try:
|
||||
async with client.stream(
|
||||
"POST", _path(call.kind), json=_body(model, call), headers={"Authorization": f"Bearer {key}"}
|
||||
) as response:
|
||||
raw: Final = await response.aread()
|
||||
return _Served(call=call, status=response.status_code, text=raw.decode())
|
||||
except httpx.TransportError as error:
|
||||
return _Served(call=call, status=0, text=repr(error))
|
||||
|
||||
|
||||
async def _burst(base_url: str, key: str, model: str, calls: Sequence[_Call]) -> tuple[_Served, ...]:
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client:
|
||||
return tuple(await asyncio.gather(*(_send(client, key, model, call) for call in calls)))
|
||||
|
||||
|
||||
def _calls(count: int, kinds: Sequence[Kind]) -> tuple[_Call, ...]:
|
||||
return tuple(
|
||||
_Call(kind=kinds[index % len(kinds)], stream=index % 2 == 0, marker=uuid.uuid4().hex) for index in range(count)
|
||||
)
|
||||
|
||||
|
||||
def _chat_response(item: _Served) -> httpx.Response:
|
||||
return httpx.Response(item.status, content=item.text.encode())
|
||||
|
||||
|
||||
def _assert_served(item: _Served) -> None:
|
||||
assert item.status == 200, (item.call, item.text)
|
||||
match item.call.kind:
|
||||
case "first":
|
||||
assistant: Final = assistant_from(_chat_response(item), item.call.stream)
|
||||
assert kept_results(assistant) == [SEARCH_RESULT], item.text
|
||||
case "replay":
|
||||
assert answer(item.call.marker) in item.text, item.text
|
||||
case "messages":
|
||||
assert "tool_search_tool_result" in item.text, item.text
|
||||
assert "get_weather" in item.text, item.text
|
||||
case "responses":
|
||||
assert "get_weather" in item.text, item.text
|
||||
|
||||
|
||||
def _assert_upstream_saw_each_once(received: Sequence[Request], served: Sequence[_Served]) -> None:
|
||||
markers: Final = tuple(marker_of(request) for request in received)
|
||||
for item in served:
|
||||
assert markers.count(item.call.marker) == 1, (item.call, markers)
|
||||
for request in received:
|
||||
if len(request_messages(request)) > 1:
|
||||
assert wire_assistant(request) == list(SEARCH_TURN), request.body
|
||||
assert answered_ids(request) == [WEATHER_ID], request.body
|
||||
|
||||
|
||||
def _chat_ids(served: Sequence[_Served]) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
response_id(_chat_response(item), item.call.stream)
|
||||
for item in served
|
||||
if item.status == 200 and item.call.kind in ("replay", "first")
|
||||
)
|
||||
|
||||
|
||||
def _success_rows(model: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s', (model, "success")
|
||||
)
|
||||
|
||||
|
||||
async def test_upstream_outage_mid_burst_recovers_with_every_search_result_kept_and_logged_once(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
kinds: Final[tuple[Kind, ...]] = ("replay", "first", "messages", "responses", "replay", "first")
|
||||
before: Final = _calls(12, kinds)
|
||||
during: Final = _calls(6, ("replay", "first"))
|
||||
after: Final = _calls(12, kinds)
|
||||
with gateway.scenario() as scenario:
|
||||
with wire_server(scripted_peer()) as wire:
|
||||
port: Final = int(wire.url.rsplit(":", 1)[1])
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
served_before: Final = await _burst(str(gateway.client.base_url), gateway.key, model, before)
|
||||
received_before: Final = wire.drain()
|
||||
served_during: Final = await _burst(str(gateway.client.base_url), gateway.key, model, during)
|
||||
assert all(item.status == 0 or item.status >= 500 for item in served_during), served_during
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
with wire_server(scripted_peer(), port=port) as revived:
|
||||
served_after: Final = await _burst(str(gateway.client.base_url), gateway.key, model, after)
|
||||
received_after: Final = revived.drain()
|
||||
for item in (*served_before, *served_after):
|
||||
_assert_served(item)
|
||||
_assert_upstream_saw_each_once(received_before, served_before)
|
||||
_assert_upstream_saw_each_once(received_after, served_after)
|
||||
ids: Final = _chat_ids((*served_before, *served_after))
|
||||
assert len(ids) == 16, ids
|
||||
rows: Final = await asyncio.to_thread(
|
||||
eventually, lambda: _success_rows(model), lambda found: len(found) == len(before) + len(after), 70
|
||||
)
|
||||
logged: Final = tuple(str(row["request_id"]) for row in rows)
|
||||
assert len(logged) == len(before) + len(after), logged
|
||||
for request_id in ids:
|
||||
assert logged.count(request_id) == 1, (request_id, logged)
|
||||
|
||||
|
||||
async def test_slow_parallel_streams_each_keep_their_search_result(gateway: Gateway) -> None:
|
||||
def slow(request: Request) -> Reply:
|
||||
return replace(scripted_peer()(request), pause_between_chunks=0.2)
|
||||
|
||||
calls: Final = tuple(_Call(kind="first", stream=True, marker=uuid.uuid4().hex) for _ in range(10))
|
||||
with wire_server(slow) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
burst: Final = asyncio.create_task(_burst(str(gateway.client.base_url), gateway.key, model, calls))
|
||||
await asyncio.to_thread(eventually, wire.received.qsize, lambda count: count >= 1, 30)
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
served: Final = await burst
|
||||
received: Final = wire.drain()
|
||||
for item in served:
|
||||
_assert_served(item)
|
||||
_assert_upstream_saw_each_once(received, served)
|
||||
|
||||
|
||||
def _chaos_config(wire: Wire, directory: Path) -> Path:
|
||||
config: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
||||
path: Final = directory / "tool-search-chaos.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": _OWNED_MODEL,
|
||||
"litellm_params": {
|
||||
"model": f"anthropic/{MODEL}",
|
||||
"api_base": wire.url,
|
||||
"api_key": "scripted-anthropic-key",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
|
||||
async def test_worker_sigkill_mid_burst_leaves_the_survivor_replaying_the_search_result(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
calls: Final = _calls(20, ("replay",))
|
||||
release: Final = threading.Event()
|
||||
held_markers: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
held_markers.put(marker_of(request))
|
||||
assert release.wait(timeout=graceful_stop_seconds() + 60), "The burst was never released"
|
||||
return scripted_peer()(request)
|
||||
|
||||
with wire_server(held) as wire:
|
||||
path: Final = _chaos_config(wire, tmp_path)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
workers: Final = eventually(
|
||||
lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
|
||||
lambda pids: len(pids) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
burst: Final = asyncio.create_task(
|
||||
_burst(str(candidate.client.base_url), candidate.key, _OWNED_MODEL, calls)
|
||||
)
|
||||
await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
|
||||
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.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)
|
||||
release.set()
|
||||
served: Final = await burst
|
||||
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_served(item)
|
||||
follow_up: Final = _Call(kind="replay", stream=True, marker=uuid.uuid4().hex)
|
||||
(answered,) = await _burst(str(candidate.client.base_url), candidate.key, _OWNED_MODEL, (follow_up,))
|
||||
_assert_served(answered)
|
||||
received: Final = wire.drain()
|
||||
assert len(received) == 21, len(received)
|
||||
for request in received:
|
||||
assert wire_assistant(request) == list(SEARCH_TURN), request.body
|
||||
assert answered_ids(request) == [WEATHER_ID], request.body
|
||||
|
|
@ -0,0 +1,821 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.anthropic_thinking import JSON_LIST, JSON_OBJECT
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.providers._cache_control_marks_support import (
|
||||
ASK_LABEL,
|
||||
EPHEMERAL,
|
||||
POINTS,
|
||||
SYSTEM,
|
||||
SYSTEM_LABEL,
|
||||
anthropic_labels,
|
||||
anthropic_peer,
|
||||
ask,
|
||||
call_id,
|
||||
final_text,
|
||||
gateway_injected,
|
||||
post_chat,
|
||||
tool_call,
|
||||
tool_use_label,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolMessage
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
|
||||
from pydantic import JsonValue
|
||||
|
||||
from integration.providers._tool_search_support import (
|
||||
KEPT_FIELDS,
|
||||
WEATHER_ANSWER,
|
||||
SEARCH_CALL,
|
||||
SEARCH_ID,
|
||||
SEARCH_INPUT,
|
||||
SEARCH_RESULT,
|
||||
SEARCH_TOOL,
|
||||
SEARCH_TURN,
|
||||
SEARCH_USE,
|
||||
TOOLS,
|
||||
WEATHER_CALL,
|
||||
WEATHER_ID,
|
||||
WEATHER_USE,
|
||||
Event,
|
||||
Provider,
|
||||
answer,
|
||||
answer_text,
|
||||
answered_ids,
|
||||
assistant_from,
|
||||
chat,
|
||||
chat_body,
|
||||
deployment,
|
||||
first_messages,
|
||||
folded_assistant,
|
||||
hand_built_assistant,
|
||||
kept_results,
|
||||
marker_of,
|
||||
provider_fields,
|
||||
question,
|
||||
replay_messages,
|
||||
request_messages,
|
||||
response_id,
|
||||
scripted_peer,
|
||||
search_result,
|
||||
spend_row,
|
||||
target,
|
||||
target_path,
|
||||
tool_answer,
|
||||
tool_call_ids,
|
||||
turn_reply,
|
||||
user_texts,
|
||||
wire_assistant,
|
||||
)
|
||||
|
||||
_REPLAYING_PROVIDERS: Final[tuple[Provider, ...]] = (
|
||||
"anthropic",
|
||||
"bedrock_invoke",
|
||||
"claude_platform",
|
||||
"bedrock_mantle",
|
||||
"vertex_ai",
|
||||
"azure_ai",
|
||||
)
|
||||
_STREAM_MODES: Final = pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"])
|
||||
_PLAIN_SEARCH_USE: Final[Event] = {"type": "tool_use", "id": SEARCH_ID, "name": SEARCH_TOOL, "input": SEARCH_INPUT}
|
||||
_BASH_ID: Final = "srvtoolu_01ScriptedBash"
|
||||
_BASH_USE: Final[Event] = {
|
||||
"type": "server_tool_use",
|
||||
"id": _BASH_ID,
|
||||
"name": "bash_code_execution",
|
||||
"input": {"command": "echo sunny"},
|
||||
}
|
||||
_BASH_RESULT: Final[Event] = {
|
||||
"type": "bash_code_execution_tool_result",
|
||||
"tool_use_id": _BASH_ID,
|
||||
"content": {
|
||||
"type": "bash_code_execution_result",
|
||||
"stdout": "sunny\n",
|
||||
"stderr": "",
|
||||
"return_code": 0,
|
||||
"content": [],
|
||||
},
|
||||
}
|
||||
_SIGNED_THINKING: Final[Event] = {
|
||||
"type": "thinking",
|
||||
"thinking": "search for the weather tool first",
|
||||
"signature": "scripted-signature-" + "s" * 32,
|
||||
}
|
||||
|
||||
|
||||
def _round_trip(
|
||||
gateway: Gateway, provider: Provider, stream: bool, first_turn: Sequence[Event] = SEARCH_TURN
|
||||
) -> tuple[Event, httpx.Response, httpx.Response, tuple[Request, ...], str]:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer(first_turn)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, provider, wire.url, gateway.upstream_url)
|
||||
first: Final = chat(gateway, chat_body(model, first_messages(marker), stream))
|
||||
assert first.status_code == 200, first.text
|
||||
assistant: Final = assistant_from(first, stream)
|
||||
second: Final = chat(gateway, chat_body(model, replay_messages(marker, assistant), stream))
|
||||
assert second.status_code == 200, second.text
|
||||
assert answer_text(second, stream) == answer(marker), second.text
|
||||
received: Final = wire.drain()
|
||||
return assistant, first, second, received, model
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _REPLAYING_PROVIDERS)
|
||||
@_STREAM_MODES
|
||||
def test_tool_search_result_round_trips_on_every_anthropic_wire_provider(
|
||||
gateway: Gateway, provider: Provider, stream: bool
|
||||
) -> None:
|
||||
assistant, first, second, received, model = _round_trip(gateway, provider, stream)
|
||||
assert [target_path(request) for request in received] == [target(provider, stream)] * 2, received
|
||||
assert kept_results(assistant) == [SEARCH_RESULT], assistant
|
||||
assert "code_interpreter_results" not in provider_fields(assistant), assistant
|
||||
assert tool_call_ids(assistant) == [SEARCH_ID, WEATHER_ID], assistant
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
for identity in (response_id(first, stream), response_id(second, stream)):
|
||||
row = spend_row(identity)
|
||||
assert (row["model_group"], row["status"]) == (model, "success"), row
|
||||
|
||||
|
||||
@_STREAM_MODES
|
||||
def test_snowflake_claude_keeps_the_tool_search_result_on_the_response(gateway: Gateway, stream: bool) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "snowflake", wire.url, gateway.upstream_url)
|
||||
first: Final = chat(gateway, chat_body(model, first_messages(marker), stream))
|
||||
assert first.status_code == 200, first.text
|
||||
received: Final = wire.drain()
|
||||
assistant: Final = assistant_from(first, stream)
|
||||
assert [target_path(request) for request in received] == [target("snowflake", stream)], received
|
||||
assert kept_results(assistant) == [SEARCH_RESULT], assistant
|
||||
assert tool_call_ids(assistant) == [SEARCH_ID, WEATHER_ID], assistant
|
||||
assert spend_row(response_id(first, stream))["model_group"] == model
|
||||
|
||||
|
||||
def _openai_client(gateway: Gateway) -> openai.OpenAI:
|
||||
return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
|
||||
|
||||
|
||||
def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI:
|
||||
return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
|
||||
|
||||
|
||||
def _sdk_messages(messages: Sequence[JsonValue]) -> list[ChatCompletionMessageParam]:
|
||||
return cast(list[ChatCompletionMessageParam], list(messages))
|
||||
|
||||
|
||||
_SDK_TOOLS: Final = cast(list[ChatCompletionToolParam], TOOLS)
|
||||
_NO_CACHE: Final[dict[str, object]] = {"cache": {"no-cache": True}}
|
||||
|
||||
|
||||
def test_openai_sdk_sync_round_trip_replays_the_tool_search_result(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario, _openai_client(gateway) as client:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
first: Final = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=_sdk_messages(first_messages(marker)),
|
||||
tools=_SDK_TOOLS,
|
||||
max_tokens=64,
|
||||
extra_body=_NO_CACHE,
|
||||
)
|
||||
assistant: Final = JSON_OBJECT.validate_python(first.choices[0].message.model_dump(exclude_none=True))
|
||||
second: Final = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=_sdk_messages(replay_messages(marker, assistant)),
|
||||
tools=_SDK_TOOLS,
|
||||
max_tokens=64,
|
||||
extra_body=_NO_CACHE,
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert second.choices[0].message.content == answer(marker), second.model_dump()
|
||||
assert kept_results(assistant) == [SEARCH_RESULT], assistant
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
assert spend_row(second.id)["model_group"] == model
|
||||
|
||||
|
||||
async def test_openai_sdk_async_stream_round_trip_replays_the_tool_search_result(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
async with _async_openai_client(gateway) as client:
|
||||
stream: Final = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=_sdk_messages(first_messages(marker)),
|
||||
tools=_SDK_TOOLS,
|
||||
max_tokens=64,
|
||||
stream=True,
|
||||
extra_body=_NO_CACHE,
|
||||
)
|
||||
chunks: Final = [chunk async for chunk in stream]
|
||||
deltas: Final = tuple(
|
||||
JSON_OBJECT.validate_python(chunk.choices[0].delta.model_dump(exclude_none=True))
|
||||
for chunk in chunks
|
||||
if chunk.choices
|
||||
)
|
||||
assistant: Final = folded_assistant(deltas)
|
||||
second: Final = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=_sdk_messages(replay_messages(marker, assistant)),
|
||||
tools=_SDK_TOOLS,
|
||||
max_tokens=64,
|
||||
stream=True,
|
||||
extra_body=_NO_CACHE,
|
||||
)
|
||||
answered: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in second if chunk.choices])
|
||||
received: Final = wire.drain()
|
||||
assert answered == answer(marker), answered
|
||||
assert kept_results(assistant) == [SEARCH_RESULT], assistant
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
@_STREAM_MODES
|
||||
def test_code_execution_next_to_a_tool_search_keeps_both_results_and_only_the_code_run(
|
||||
gateway: Gateway, stream: bool
|
||||
) -> None:
|
||||
turn: Final = (_BASH_USE, _BASH_RESULT, SEARCH_USE, SEARCH_RESULT, WEATHER_USE)
|
||||
assistant, _, _, received, _ = _round_trip(gateway, "anthropic", stream, turn)
|
||||
fields: Final = provider_fields(assistant)
|
||||
assert fields.get("tool_results") == [_BASH_RESULT, SEARCH_RESULT], assistant
|
||||
interpreted: Final = fields.get("code_interpreter_results")
|
||||
assert isinstance(interpreted, list) and len(interpreted) == 1, assistant
|
||||
run: Final = JSON_OBJECT.validate_python(interpreted[0])
|
||||
assert (run["id"], run["code"], run["status"]) == (_BASH_ID, "echo sunny", "completed"), run
|
||||
assert wire_assistant(received[1]) == list(turn), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
def _replayed(gateway: Gateway, messages: Sequence[JsonValue]) -> tuple[httpx.Response, tuple[Request, ...]]:
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
response: Final = chat(gateway, chat_body(model, messages, stream=False))
|
||||
received: Final = wire.drain()
|
||||
return response, received
|
||||
|
||||
|
||||
def _kept_search(*calls: tuple[str, str, Event]) -> Event:
|
||||
return hand_built_assistant(calls or (SEARCH_CALL, WEATHER_CALL), KEPT_FIELDS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("answers", "assistant_turn", "answered", "texts"),
|
||||
(
|
||||
pytest.param(
|
||||
(tool_answer(SEARCH_ID, "client search"), tool_answer(WEATHER_ID)),
|
||||
(SEARCH_CALL, WEATHER_CALL),
|
||||
[WEATHER_ID],
|
||||
1,
|
||||
id="search-and-client-tool-answered",
|
||||
),
|
||||
pytest.param(
|
||||
(tool_answer(WEATHER_ID), tool_answer(SEARCH_ID, "client search")),
|
||||
(SEARCH_CALL, WEATHER_CALL),
|
||||
[WEATHER_ID],
|
||||
1,
|
||||
id="client-tool-answered-first",
|
||||
),
|
||||
pytest.param(
|
||||
(tool_answer(SEARCH_ID, "one"), tool_answer(SEARCH_ID, "two"), tool_answer(WEATHER_ID)),
|
||||
(SEARCH_CALL, WEATHER_CALL),
|
||||
[WEATHER_ID],
|
||||
1,
|
||||
id="search-answered-twice",
|
||||
),
|
||||
pytest.param(
|
||||
(tool_answer(SEARCH_ID, "client search"), {"role": "user", "content": "go on"}),
|
||||
(SEARCH_CALL,),
|
||||
[],
|
||||
2,
|
||||
id="search-answer-then-user-text",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_client_answer_to_a_kept_search_call_is_dropped_when_another_user_turn_follows(
|
||||
gateway: Gateway,
|
||||
answers: Sequence[JsonValue],
|
||||
assistant_turn: tuple[tuple[str, str, Event], ...],
|
||||
answered: list[JsonValue],
|
||||
texts: int,
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response, received = _replayed(gateway, replay_messages(marker, _kept_search(*assistant_turn), answers))
|
||||
assert response.status_code == 200, response.text
|
||||
assert answer_text(response, stream=False) == answer(marker), response.text
|
||||
assert len(received) == 1, received
|
||||
expected_turn: Final = [SEARCH_USE, SEARCH_RESULT, *([WEATHER_USE] if len(assistant_turn) == 2 else [])]
|
||||
assert wire_assistant(received[0]) == expected_turn, received[0].body
|
||||
assert answered_ids(received[0]) == answered, received[0].body
|
||||
assert len(user_texts(received[0])) == texts, received[0].body
|
||||
|
||||
|
||||
def test_client_answer_alone_replays_the_search_call_the_way_the_client_sent_it(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final = replay_messages(marker, _kept_search(SEARCH_CALL), (tool_answer(SEARCH_ID, "client search"),))
|
||||
response, received = _replayed(gateway, messages)
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
assert wire_assistant(received[0]) == [_PLAIN_SEARCH_USE], received[0].body
|
||||
assert answered_ids(received[0]) == [SEARCH_ID], received[0].body
|
||||
|
||||
|
||||
def test_interleaved_thinking_replay_drops_the_client_answer_to_the_kept_search(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
assistant: Final = hand_built_assistant(
|
||||
(SEARCH_CALL, WEATHER_CALL), KEPT_FIELDS, thinking_blocks=[_SIGNED_THINKING]
|
||||
)
|
||||
answers: Final = (tool_answer(SEARCH_ID, "client search"), tool_answer(WEATHER_ID))
|
||||
response, received = _replayed(gateway, replay_messages(marker, assistant, answers))
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
assert wire_assistant(received[0]) == [_SIGNED_THINKING, SEARCH_USE, SEARCH_RESULT, WEATHER_USE], received[0].body
|
||||
assert answered_ids(received[0]) == [WEATHER_ID], received[0].body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fields",
|
||||
(
|
||||
pytest.param({"provider_specific_fields": None}, id="null"),
|
||||
pytest.param({"provider_specific_fields": {}}, id="empty"),
|
||||
pytest.param({}, id="missing"),
|
||||
pytest.param({"provider_specific_fields": {"tool_results": []}}, id="no-results"),
|
||||
),
|
||||
)
|
||||
def test_search_call_without_a_kept_result_replays_plainly_with_both_answers(
|
||||
gateway: Gateway, fields: dict[str, JsonValue]
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
assistant: Final = {**hand_built_assistant((SEARCH_CALL, WEATHER_CALL)), **fields}
|
||||
answers: Final = (tool_answer(SEARCH_ID, "client search"), tool_answer(WEATHER_ID))
|
||||
response, received = _replayed(gateway, replay_messages(marker, assistant, answers))
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
assert wire_assistant(received[0]) == [_PLAIN_SEARCH_USE, WEATHER_USE], received[0].body
|
||||
assert answered_ids(received[0]) == [SEARCH_ID, WEATHER_ID], received[0].body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fields",
|
||||
(
|
||||
pytest.param("x" * 5120, id="5kb-string"),
|
||||
pytest.param([SEARCH_RESULT], id="list"),
|
||||
pytest.param({"tool_results": "x" * 5120}, id="results-5kb-string"),
|
||||
pytest.param({"tool_results": [1, "a", None]}, id="results-without-objects"),
|
||||
pytest.param({"tool_results": SEARCH_RESULT}, id="results-object"),
|
||||
pytest.param({"tool_results": [search_result("srvtoolu_01Other")]}, id="result-for-another-call"),
|
||||
),
|
||||
)
|
||||
def test_malformed_kept_results_replay_the_search_call_plainly(gateway: Gateway, fields: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
assistant: Final = hand_built_assistant((SEARCH_CALL, WEATHER_CALL), fields)
|
||||
response, received = _replayed(gateway, replay_messages(marker, assistant))
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
assert wire_assistant(received[0]) == [_PLAIN_SEARCH_USE, WEATHER_USE], received[0].body
|
||||
assert answered_ids(received[0]) == [WEATHER_ID], received[0].body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_call_id",
|
||||
(
|
||||
pytest.param(123, id="int"),
|
||||
pytest.param([SEARCH_ID], id="list"),
|
||||
pytest.param({"id": SEARCH_ID}, id="object"),
|
||||
),
|
||||
)
|
||||
def test_non_string_tool_answer_ids_never_crash_the_replay(gateway: Gateway, tool_call_id: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
answers: Final = (tool_answer(tool_call_id, "odd"), tool_answer(WEATHER_ID))
|
||||
response, received = _replayed(gateway, replay_messages(marker, _kept_search(), answers))
|
||||
assert response.status_code != 500, response.text
|
||||
assert '"error"' in response.text or response.status_code == 200, response.text
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
assert len(received) <= 1, received
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_call_id", (pytest.param("", id="empty-string"), pytest.param("x" * 5120, id="5kb-string"))
|
||||
)
|
||||
def test_odd_string_tool_answer_ids_keep_the_search_rebuilt(gateway: Gateway, tool_call_id: str) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
answers: Final = (tool_answer(tool_call_id, "odd"), tool_answer(WEATHER_ID))
|
||||
response, received = _replayed(gateway, replay_messages(marker, _kept_search(), answers))
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
assert wire_assistant(received[0]) == list(SEARCH_TURN), received[0].body
|
||||
forwarded: Final = answered_ids(received[0])
|
||||
assert len(forwarded) == 2, received[0].body
|
||||
assert forwarded[1] == WEATHER_ID, received[0].body
|
||||
assert SEARCH_ID not in forwarded, received[0].body
|
||||
|
||||
|
||||
_ODD_RESULTS: Final = (
|
||||
pytest.param(
|
||||
{
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": SEARCH_ID,
|
||||
"content": {"type": "tool_search_tool_result_error", "error_code": "unavailable"},
|
||||
},
|
||||
id="error-content",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": SEARCH_ID,
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": "w" * 5120}],
|
||||
},
|
||||
},
|
||||
id="5kb-reference",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("odd", _ODD_RESULTS)
|
||||
@_STREAM_MODES
|
||||
def test_unusual_search_results_from_the_upstream_are_kept_and_replayed_verbatim(
|
||||
gateway: Gateway, odd: Event, stream: bool
|
||||
) -> None:
|
||||
assistant, _, _, received, _ = _round_trip(gateway, "anthropic", stream, (SEARCH_USE, odd, WEATHER_USE))
|
||||
assert kept_results(assistant) == [odd], assistant
|
||||
assert wire_assistant(received[1]) == [SEARCH_USE, odd, WEATHER_USE], received[1].body
|
||||
|
||||
|
||||
@_STREAM_MODES
|
||||
def test_search_result_without_its_call_id_is_kept_but_not_rebuilt(gateway: Gateway, stream: bool) -> None:
|
||||
orphan: Final = {key: value for key, value in SEARCH_RESULT.items() if key != "tool_use_id"}
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer((SEARCH_USE, orphan, WEATHER_USE))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
first: Final = chat(gateway, chat_body(model, first_messages(marker), stream))
|
||||
assert first.status_code == 200, first.text
|
||||
assistant: Final = assistant_from(first, stream)
|
||||
answers: Final = (tool_answer(SEARCH_ID, "client search"), tool_answer(WEATHER_ID))
|
||||
second: Final = chat(gateway, chat_body(model, replay_messages(marker, assistant, answers), stream=False))
|
||||
assert second.status_code == 200, second.text
|
||||
received: Final = wire.drain()
|
||||
assert kept_results(assistant) == [orphan], assistant
|
||||
assert wire_assistant(received[1]) == [_PLAIN_SEARCH_USE, WEATHER_USE], received[1].body
|
||||
assert answered_ids(received[1]) == [SEARCH_ID, WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
@_STREAM_MODES
|
||||
def test_duplicated_search_result_blocks_rebuild_the_call_once(gateway: Gateway, stream: bool) -> None:
|
||||
assistant, _, _, received, _ = _round_trip(
|
||||
gateway, "anthropic", stream, (SEARCH_USE, SEARCH_RESULT, SEARCH_RESULT, WEATHER_USE)
|
||||
)
|
||||
assert kept_results(assistant) == [SEARCH_RESULT, SEARCH_RESULT], assistant
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "kind", "caller_status"),
|
||||
(
|
||||
pytest.param(400, "invalid_request_error", 400, id="400"),
|
||||
pytest.param(401, "authentication_error", 401, id="401"),
|
||||
pytest.param(529, "overloaded_error", 500, id="529"),
|
||||
),
|
||||
)
|
||||
def test_provider_error_on_the_replay_reaches_the_caller(
|
||||
gateway: Gateway, status: int, kind: str, caller_status: int
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if len(request_messages(request)) == 1:
|
||||
return scripted_peer()(request)
|
||||
body: Final = {"type": "error", "error": {"type": kind, "message": f"scripted {kind} {marker_of(request)}"}}
|
||||
return Reply(status=status, body=json.dumps(body).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
first: Final = chat(gateway, chat_body(model, first_messages(marker), stream=False))
|
||||
assert first.status_code == 200, first.text
|
||||
second: Final = chat(gateway, chat_body(model, replay_messages(marker, assistant_from(first, False)), False))
|
||||
received: Final = wire.drain()
|
||||
assert second.status_code == caller_status, second.text
|
||||
assert f"scripted {kind} {marker}" in second.text, second.text
|
||||
assert len(received) >= 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
|
||||
|
||||
def test_unauthenticated_replay_never_reaches_the_provider(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
body: Final = chat_body(model, replay_messages(marker, _kept_search()), stream=False)
|
||||
response: Final = chat(gateway, body, key=f"sk-not-a-key-{marker}")
|
||||
assert response.status_code == 401, response.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def _cut_stream_text(gateway: Gateway, body: Event) -> str:
|
||||
try:
|
||||
with gateway.client.stream(
|
||||
"POST", "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {gateway.key}"}
|
||||
) as response:
|
||||
return response.read().decode()
|
||||
except httpx.TransportError as error:
|
||||
return repr(error)
|
||||
|
||||
|
||||
def test_stream_cut_right_after_the_search_result_leaves_the_proxy_serving(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if marker_of(request) != marker:
|
||||
return scripted_peer()(request)
|
||||
return turn_reply(request, f"msg_{marker}_cut", SEARCH_TURN, "tool_use", abort_after=6)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
text: Final = _cut_stream_text(gateway, chat_body(model, first_messages(marker), stream=True))
|
||||
assert "get_weather" not in text, text
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
follow_up: Final = uuid.uuid4().hex
|
||||
recovered: Final = chat(gateway, chat_body(model, first_messages(follow_up), stream=True))
|
||||
assert recovered.status_code == 200, recovered.text
|
||||
assert kept_results(assistant_from(recovered, stream=True)) == [SEARCH_RESULT], recovered.text
|
||||
assert len(wire.drain()) == 2
|
||||
|
||||
|
||||
def test_identical_replays_each_reach_the_provider_and_land_their_own_spend_row(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
body: Final = chat_body(model, replay_messages(marker, _kept_search()), stream=False)
|
||||
responses: Final = tuple(chat(gateway, body) for _ in range(2))
|
||||
received: Final = wire.drain()
|
||||
for response in responses:
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 2, received
|
||||
assert [wire_assistant(request) for request in received] == [list(SEARCH_TURN)] * 2
|
||||
ids: Final = tuple(response_id(response, stream=False) for response in responses)
|
||||
assert len(set(ids)) == 2, ids
|
||||
for identity in ids:
|
||||
assert spend_row(identity)["model_group"] == model
|
||||
|
||||
|
||||
@_STREAM_MODES
|
||||
def test_response_cache_hit_still_carries_the_search_result_for_the_replay(gateway: Gateway, stream: bool) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
cached_body: Final = {
|
||||
key: value for key, value in chat_body(model, first_messages(marker), stream).items() if key != "cache"
|
||||
}
|
||||
first: Final = chat(gateway, cached_body)
|
||||
assert first.status_code == 200, first.text
|
||||
hit: Final = chat(gateway, cached_body)
|
||||
assert hit.status_code == 200, hit.text
|
||||
assistant: Final = assistant_from(hit, stream)
|
||||
second: Final = chat(gateway, chat_body(model, replay_messages(marker, assistant), stream=False))
|
||||
assert second.status_code == 200, second.text
|
||||
received: Final = wire.drain()
|
||||
assert kept_results(assistant_from(first, stream)) == [SEARCH_RESULT], first.text
|
||||
assert kept_results(assistant) == [SEARCH_RESULT], hit.text
|
||||
assert len(received) == 2, [request.body for request in received]
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
def _messages_turn(model: str, marker: str, stream: bool, replay: bool) -> Event:
|
||||
turns: Final[list[JsonValue]] = [{"role": "user", "content": question(marker)}]
|
||||
replayed: Final[list[JsonValue]] = [
|
||||
{"role": "assistant", "content": list(SEARCH_TURN)},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": WEATHER_ID, "content": "18C"}]},
|
||||
]
|
||||
return {
|
||||
"model": model,
|
||||
"max_tokens": 64,
|
||||
"stream": stream,
|
||||
"tools": [
|
||||
{"type": "tool_search_tool_bm25_20251119", "name": SEARCH_TOOL},
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a city",
|
||||
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}},
|
||||
"defer_loading": True,
|
||||
},
|
||||
],
|
||||
"messages": [*turns, *(replayed if replay else [])],
|
||||
}
|
||||
|
||||
|
||||
def _streamed_message_blocks(text: str) -> list[JsonValue]:
|
||||
events: Final = tuple(
|
||||
JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: ")
|
||||
)
|
||||
return [event["content_block"] for event in events if event["type"] == "content_block_start"]
|
||||
|
||||
|
||||
@_STREAM_MODES
|
||||
def test_messages_endpoint_passes_the_search_result_and_its_replay_through(gateway: Gateway, stream: bool) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
first: Final = gateway.request("POST", "/v1/messages", _messages_turn(model, marker, stream, replay=False))
|
||||
second: Final = gateway.request("POST", "/v1/messages", _messages_turn(model, marker, stream, replay=True))
|
||||
received: Final = wire.drain()
|
||||
assert first.status_code == 200, first.text
|
||||
assert second.status_code == 200, second.text
|
||||
blocks: Final = (
|
||||
_streamed_message_blocks(first.text)
|
||||
if stream
|
||||
else JSON_LIST.validate_python(JSON_OBJECT.validate_json(first.content)["content"])
|
||||
)
|
||||
assert SEARCH_RESULT in blocks, first.text
|
||||
assert answer(marker) in second.text, second.text
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
|
||||
|
||||
_RESPONSES_TOOLS: Final[list[JsonValue]] = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_responses_bridge_reports_the_search_turn_the_same_way(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{
|
||||
"model": model,
|
||||
"input": question(marker),
|
||||
"tools": _RESPONSES_TOOLS,
|
||||
"max_output_tokens": 64,
|
||||
**_NO_CACHE,
|
||||
},
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
output: Final = JSON_OBJECT.validate_json(response.content)["output"]
|
||||
assert isinstance(output, list), response.text
|
||||
calls: Final = [
|
||||
JSON_OBJECT.validate_python(item)["name"]
|
||||
for item in output
|
||||
if JSON_OBJECT.validate_python(item)["type"] == "function_call"
|
||||
]
|
||||
assert "get_weather" in calls, response.text
|
||||
|
||||
|
||||
def test_run_tool_loop_through_the_proxy_executes_only_the_client_tool(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
executed: Final[list[str]] = [] # mutable-ok: the executor records each tool it ran for the assertion
|
||||
|
||||
def execute(call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
executed.append(str(call.function.name)) # mutable-ok: see above
|
||||
return {"role": "tool", "tool_call_id": call.id, "content": WEATHER_ANSWER}
|
||||
|
||||
with wire_server(scripted_peer()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
content: Final = litellm.run_tool_loop(
|
||||
model=f"litellm_proxy/{model}",
|
||||
messages=cast(list[AllMessageValues], first_messages(marker)),
|
||||
tools=cast(list[ChatCompletionToolParam], TOOLS),
|
||||
execute_tool=execute,
|
||||
api_base=str(gateway.client.base_url),
|
||||
api_key=gateway.key,
|
||||
max_tokens=64,
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert content == answer(marker), content
|
||||
assert executed == ["get_weather"], executed
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
def _marked_search_call() -> Event:
|
||||
return {
|
||||
"id": SEARCH_ID,
|
||||
"type": "function",
|
||||
"function": {"name": SEARCH_TOOL, "arguments": json.dumps(SEARCH_INPUT)},
|
||||
"cache_control": EPHEMERAL,
|
||||
}
|
||||
|
||||
|
||||
def _kept_marked_search() -> Event:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [_marked_search_call()],
|
||||
"provider_specific_fields": KEPT_FIELDS,
|
||||
}
|
||||
|
||||
|
||||
_SEARCH_LABEL: Final = f"assistant:tool_use:{SEARCH_ID}"
|
||||
_CITY_CALLS: Final = ("London", "Paris")
|
||||
|
||||
|
||||
def _answered_cities() -> list[JsonValue]:
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [tool_call(city, cache_control=EPHEMERAL) for city in _CITY_CALLS],
|
||||
},
|
||||
*({"role": "tool", "tool_call_id": call_id(city), "content": "sunny."} for city in _CITY_CALLS),
|
||||
]
|
||||
|
||||
|
||||
def test_client_answered_search_call_mark_counts_toward_the_cap(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "system", "content": SYSTEM},
|
||||
ask(marked=True),
|
||||
*_answered_cities(),
|
||||
_kept_marked_search(),
|
||||
tool_answer(SEARCH_ID, f"client search marker-{marker}"),
|
||||
]
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
status, _, text = post_chat(
|
||||
gateway, chat_body(model, messages, stream=False, cache_control_injection_points=POINTS)
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert status == 200, text
|
||||
assert len(received) == 1, received
|
||||
assert anthropic_labels(received[0]) == [
|
||||
ASK_LABEL,
|
||||
*(tool_use_label(city) for city in _CITY_CALLS),
|
||||
_SEARCH_LABEL,
|
||||
]
|
||||
assert answered_ids(received[0]) == [*(call_id(city) for city in _CITY_CALLS), SEARCH_ID]
|
||||
|
||||
|
||||
def test_search_call_mark_is_not_counted_when_its_client_answer_is_dropped(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "system", "content": SYSTEM},
|
||||
ask(marked=True),
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [*(tool_call(city, cache_control=EPHEMERAL) for city in _CITY_CALLS), _marked_search_call()],
|
||||
"provider_specific_fields": KEPT_FIELDS,
|
||||
},
|
||||
*({"role": "tool", "tool_call_id": call_id(city), "content": "sunny."} for city in _CITY_CALLS),
|
||||
tool_answer(SEARCH_ID, "client search"),
|
||||
{"role": "user", "content": final_text(marker)},
|
||||
]
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
status, _, text = post_chat(
|
||||
gateway, chat_body(model, messages, stream=False, cache_control_injection_points=POINTS)
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert status == 200, text
|
||||
assert len(received) == 1, received
|
||||
assert anthropic_labels(received[0]) == [SYSTEM_LABEL, ASK_LABEL, *(tool_use_label(city) for city in _CITY_CALLS)]
|
||||
assert answered_ids(received[0]) == [call_id(city) for city in _CITY_CALLS]
|
||||
assert wire_assistant(received[0])[-2:] == [SEARCH_USE, SEARCH_RESULT], received[0].body
|
||||
|
||||
|
||||
def test_response_cache_hit_serves_the_client_answered_search_request_once(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "system", "content": SYSTEM},
|
||||
ask(marked=True),
|
||||
_kept_marked_search(),
|
||||
tool_answer(SEARCH_ID, f"client search marker-{marker}"),
|
||||
]
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
|
||||
body: Final = {
|
||||
key: value
|
||||
for key, value in chat_body(model, messages, stream=False, cache_control_injection_points=POINTS).items()
|
||||
if key != "cache"
|
||||
}
|
||||
first: Final = post_chat(gateway, body)
|
||||
second: Final = post_chat(gateway, body)
|
||||
received: Final = wire.drain()
|
||||
assert (first[0], second[0]) == (200, 200), (first[2], second[2])
|
||||
assert first[1] == second[1], (first[2], second[2])
|
||||
assert [anthropic_labels(request) for request in received] == [
|
||||
[SYSTEM_LABEL, ASK_LABEL, _SEARCH_LABEL, f"user:tool_result:{SEARCH_ID}"]
|
||||
]
|
||||
assert gateway_injected(first[1]) is True
|
||||
|
|
@ -0,0 +1,150 @@
|
|||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, gateway_from_environment, object_value
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.wire import wire_server
|
||||
from integration.providers._cache_control_marks_support import (
|
||||
ASK_LABEL,
|
||||
EPHEMERAL,
|
||||
POINTS,
|
||||
SYSTEM,
|
||||
SYSTEM_LABEL,
|
||||
anthropic_labels,
|
||||
anthropic_peer,
|
||||
ask,
|
||||
call_id,
|
||||
final_text,
|
||||
post_chat,
|
||||
tool_call,
|
||||
tool_use_label,
|
||||
)
|
||||
from integration.providers._tool_search_support import (
|
||||
KEPT_FIELDS,
|
||||
SEARCH_CALL,
|
||||
SEARCH_ID,
|
||||
SEARCH_TURN,
|
||||
SEARCH_USE,
|
||||
SEARCH_RESULT,
|
||||
WEATHER_CALL,
|
||||
WEATHER_ID,
|
||||
answer,
|
||||
answer_text,
|
||||
answered_ids,
|
||||
assistant_from,
|
||||
chat,
|
||||
chat_body,
|
||||
deployment,
|
||||
first_messages,
|
||||
hand_built_assistant,
|
||||
replay_messages,
|
||||
scripted_peer,
|
||||
wire_assistant,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_LANE_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_CITIES: Final = ("London", "Paris")
|
||||
|
||||
|
||||
def _modify_params_config(directory: Path) -> Path:
|
||||
lane: Final = object_value(yaml.safe_load(_LANE_CONFIG.read_text()))
|
||||
settings: Final = object_value(lane["litellm_settings"])
|
||||
path: Final = directory / "proxy_config_modify_params.yaml"
|
||||
path.write_text(yaml.safe_dump({**lane, "litellm_settings": {**settings, "modify_params": True}}, sort_keys=False))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def modify_params_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as lane:
|
||||
directory: Final = tmp_path_factory.mktemp("modify_params")
|
||||
with owned_proxy(lane, directory, {}, config=_modify_params_config(directory)) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"])
|
||||
def test_modify_params_round_trip_keeps_the_search_result_and_no_placeholder_answer(
|
||||
modify_params_gateway: Gateway, stream: bool
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(scripted_peer()) as wire, modify_params_gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, modify_params_gateway.upstream_url)
|
||||
first: Final = chat(modify_params_gateway, chat_body(model, first_messages(marker), stream))
|
||||
assert first.status_code == 200, first.text
|
||||
second: Final = chat(
|
||||
modify_params_gateway, chat_body(model, replay_messages(marker, assistant_from(first, stream)), stream)
|
||||
)
|
||||
assert second.status_code == 200, second.text
|
||||
assert answer_text(second, stream) == answer(marker), second.text
|
||||
received: Final = wire.drain()
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
def test_modify_params_placeholder_for_the_kept_search_is_dropped_when_no_answer_was_sent(
|
||||
modify_params_gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
assistant: Final = hand_built_assistant((SEARCH_CALL, WEATHER_CALL), KEPT_FIELDS)
|
||||
with wire_server(scripted_peer()) as wire, modify_params_gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, modify_params_gateway.upstream_url)
|
||||
response: Final = chat(
|
||||
modify_params_gateway,
|
||||
chat_body(model, replay_messages(marker, assistant, ({"role": "user", "content": "go on"},)), False),
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(received) == 1, received
|
||||
assert wire_assistant(received[0])[:2] == [SEARCH_USE, SEARCH_RESULT], received[0].body
|
||||
assert SEARCH_ID not in answered_ids(received[0]), received[0].body
|
||||
|
||||
|
||||
def _marked_search_turn() -> dict[str, JsonValue]:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": SEARCH_ID,
|
||||
"type": "function",
|
||||
"function": {"name": "tool_search_tool_bm25", "arguments": '{"query": "weather"}'},
|
||||
"cache_control": EPHEMERAL,
|
||||
}
|
||||
],
|
||||
"provider_specific_fields": KEPT_FIELDS,
|
||||
}
|
||||
|
||||
|
||||
def test_modify_params_trailing_search_turn_stays_within_the_cache_mark_cap(modify_params_gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "system", "content": SYSTEM},
|
||||
ask(marked=True),
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [tool_call(city, cache_control=EPHEMERAL) for city in _CITIES],
|
||||
},
|
||||
*({"role": "tool", "tool_call_id": call_id(city), "content": "sunny."} for city in _CITIES),
|
||||
{"role": "user", "content": final_text(marker)},
|
||||
_marked_search_turn(),
|
||||
]
|
||||
with wire_server(anthropic_peer) as wire, modify_params_gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, "anthropic", wire.url, modify_params_gateway.upstream_url)
|
||||
status, _, text = post_chat(
|
||||
modify_params_gateway, chat_body(model, messages, stream=False, cache_control_injection_points=POINTS)
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert status == 200, text
|
||||
assert len(received) == 1, received
|
||||
labels: Final = anthropic_labels(received[0])
|
||||
assert len(labels) <= 4, labels
|
||||
assert labels[:4] == [SYSTEM_LABEL, ASK_LABEL, *(tool_use_label(city) for city in _CITIES)], labels
|
||||
|
|
@ -0,0 +1,169 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_thinking import JSON_LIST, JSON_OBJECT, aws_chunks, sse_chunks
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers._tool_search_support import (
|
||||
BEDROCK_MODEL,
|
||||
MODEL,
|
||||
SEARCH_RESULT,
|
||||
SEARCH_TOOL,
|
||||
SEARCH_TURN,
|
||||
Event,
|
||||
anthropic_message,
|
||||
marker_of,
|
||||
message_events,
|
||||
question,
|
||||
turn_reply,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_ANTHROPIC_KEY: Final = "scripted-anthropic-key"
|
||||
_ANTHROPIC_BASE: Final = "http://api.anthropic.com"
|
||||
_BY_REQUEST_ID: Final = 'SELECT response FROM "LiteLLM_SpendLogs" WHERE request_id=%s AND status=%s'
|
||||
_BY_DEPLOYMENT: Final = 'SELECT response FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s'
|
||||
_NATIVE_TOOLS: Final[list[JsonValue]] = [{"type": "tool_search_tool_bm25_20251119", "name": SEARCH_TOOL}]
|
||||
|
||||
|
||||
def _identity(marker: str) -> str:
|
||||
return f"msg_{marker}_search"
|
||||
|
||||
|
||||
def _search_turn(request: Request) -> Reply:
|
||||
return turn_reply(request, _identity(marker_of(request)), SEARCH_TURN, "tool_use")
|
||||
|
||||
|
||||
def _events(marker: str) -> tuple[Event, ...]:
|
||||
return message_events(_identity(marker), SEARCH_TURN, "tool_use")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig() -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as gateway:
|
||||
yield gateway
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wire() -> Iterator[Wire]:
|
||||
with wire_server(_search_turn) as served:
|
||||
yield served
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drained_wire(wire: Wire) -> None:
|
||||
wire.drain()
|
||||
|
||||
|
||||
def _config_storing_prompts(directory: Path) -> Path:
|
||||
config: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
||||
path: Final = directory / "store-prompts.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{**config, "general_settings": {**config["general_settings"], "store_prompts_in_spend_logs": True}}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def logged(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("tool-search-passthrough")
|
||||
overrides: Final = {
|
||||
"ANTHROPIC_API_BASE": _ANTHROPIC_BASE,
|
||||
"ANTHROPIC_API_KEY": _ANTHROPIC_KEY,
|
||||
"AIOHTTP_TRUST_ENV": "True",
|
||||
"HTTP_PROXY": wire.url,
|
||||
"NO_PROXY": "127.0.0.1,localhost",
|
||||
}
|
||||
with owned_proxy(rig, directory, overrides, config=_config_storing_prompts(directory), workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _logged_response(query: str, value: str) -> Event:
|
||||
rows: Final = eventually(lambda: read_rows(query, (value, "success")), lambda found: len(found) == 1, seconds=70)
|
||||
return JSON_OBJECT.validate_python(rows[0]["response"])
|
||||
|
||||
|
||||
def _logged_search_results(response: Mapping[str, JsonValue]) -> JsonValue:
|
||||
if "choices" not in response:
|
||||
return [block for block in JSON_LIST.validate_python(response.get("content") or []) if _is_search_result(block)]
|
||||
choice: Final = JSON_OBJECT.validate_python(JSON_LIST.validate_python(response["choices"])[0])
|
||||
message: Final = JSON_OBJECT.validate_python(choice.get("message") or {})
|
||||
return JSON_OBJECT.validate_python(message.get("provider_specific_fields") or {}).get("tool_results")
|
||||
|
||||
|
||||
def _is_search_result(block: JsonValue) -> bool:
|
||||
return JSON_OBJECT.validate_python(block).get("type") == "tool_search_tool_result"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"])
|
||||
def test_anthropic_passthrough_relays_the_search_turn_and_logs_its_result(
|
||||
logged: Gateway, wire: Wire, stream: bool
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
body: Final[Event] = {
|
||||
"model": MODEL,
|
||||
"max_tokens": 64,
|
||||
"stream": stream,
|
||||
"tools": _NATIVE_TOOLS,
|
||||
"messages": [{"role": "user", "content": question(marker)}],
|
||||
}
|
||||
response: Final = logged.request("POST", "/anthropic/v1/messages", body)
|
||||
assert response.status_code == 200, response.text
|
||||
if stream:
|
||||
assert response.content == b"".join(sse_chunks(_events(marker))), response.text
|
||||
else:
|
||||
assert JSON_OBJECT.validate_json(response.content) == anthropic_message(
|
||||
_identity(marker), SEARCH_TURN, "tool_use"
|
||||
), response.text
|
||||
received: Final = wire.drain()
|
||||
assert [request.target for request in received] == [f"{_ANTHROPIC_BASE}/v1/messages"], response.text
|
||||
assert json.loads(received[0].body) == body, received[0].body
|
||||
stored: Final = _logged_response(_BY_REQUEST_ID, _identity(marker))
|
||||
assert _logged_search_results(stored) == [SEARCH_RESULT], stored
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["invoke", "invoke-with-response-stream"])
|
||||
def test_bedrock_invoke_passthrough_relays_the_search_turn_and_logs_its_result(
|
||||
logged: Gateway, wire: Wire, stream: bool
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
action: Final = "invoke-with-response-stream" if stream else "invoke"
|
||||
with logged.scenario() as scenario:
|
||||
deployment: Final = scenario.model(
|
||||
model=f"bedrock/{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,
|
||||
)
|
||||
body: Final[Event] = {
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 64,
|
||||
"tools": _NATIVE_TOOLS,
|
||||
"messages": [{"role": "user", "content": question(marker)}],
|
||||
}
|
||||
response: Final = logged.request("POST", f"/bedrock/model/{deployment}/{action}", body)
|
||||
assert response.status_code == 200, response.text
|
||||
if stream:
|
||||
assert response.content == b"".join(aws_chunks(_events(marker))), response.text
|
||||
else:
|
||||
assert JSON_OBJECT.validate_json(response.content) == anthropic_message(
|
||||
_identity(marker), SEARCH_TURN, "tool_use"
|
||||
), response.text
|
||||
targets: Final = [unquote(request.target) for request in wire.drain()]
|
||||
assert targets == [f"/model/{BEDROCK_MODEL}/{action}"], targets
|
||||
stored: Final = _logged_response(_BY_DEPLOYMENT, deployment)
|
||||
assert _logged_search_results(stored) == [SEARCH_RESULT], stored
|
||||
|
|
@ -0,0 +1,212 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
|
||||
from integration._support.anthropic_thinking import sse_chunks
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
_MODEL: Final = "claude-opus-5-5"
|
||||
_API_KEY: Final = "synthetic-anthropic-key"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_SEARCH_ID: Final = "srvtoolu_01ScriptedSearch"
|
||||
_WEATHER_ID: Final = "toolu_01ScriptedWeather"
|
||||
_SEARCH_INPUT: Final = {"query": "weather"}
|
||||
_SEARCH_RESULT: Final = {
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": _SEARCH_ID,
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
|
||||
},
|
||||
}
|
||||
_TOOLS: Final = [
|
||||
{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
},
|
||||
"defer_loading": True,
|
||||
},
|
||||
]
|
||||
_USAGE: Final = {"input_tokens": 40, "output_tokens": 12}
|
||||
|
||||
|
||||
def _search_turn_blocks() -> list[JsonValue]:
|
||||
return [
|
||||
{"type": "server_tool_use", "id": _SEARCH_ID, "name": "tool_search_tool_bm25", "input": _SEARCH_INPUT},
|
||||
_SEARCH_RESULT,
|
||||
{"type": "tool_use", "id": _WEATHER_ID, "name": "get_weather", "input": {"city": "Paris"}},
|
||||
]
|
||||
|
||||
|
||||
def _message(identity: str, content: Sequence[JsonValue], stop_reason: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": identity,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": _MODEL,
|
||||
"content": list(content),
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": None,
|
||||
"usage": _USAGE,
|
||||
}
|
||||
|
||||
|
||||
def _search_turn_events(identity: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return (
|
||||
{"type": "message_start", "message": {**_message(identity, [], "tool_use"), "stop_reason": None}},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": "server_tool_use",
|
||||
"id": _SEARCH_ID,
|
||||
"name": "tool_search_tool_bm25",
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": json.dumps(_SEARCH_INPUT)},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "content_block_start", "index": 1, "content_block": _SEARCH_RESULT},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 2,
|
||||
"content_block": {"type": "tool_use", "id": _WEATHER_ID, "name": "get_weather", "input": {}},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 2,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 2},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": _USAGE},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
|
||||
|
||||
def _assistant_turn_on_the_wire(body: dict[str, JsonValue]) -> list[JsonValue]:
|
||||
messages: Final = body["messages"]
|
||||
assert isinstance(messages, list), body
|
||||
assistant: Final = messages[1]
|
||||
assert isinstance(assistant, dict) and assistant["role"] == "assistant", body
|
||||
content: Final = assistant["content"]
|
||||
assert isinstance(content, list), body
|
||||
return content
|
||||
|
||||
|
||||
def _second_turn_messages(assistant: dict[str, JsonValue]) -> list[JsonValue]:
|
||||
return [
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
assistant,
|
||||
{"role": "tool", "tool_call_id": _WEATHER_ID, "content": "18C and cloudy"},
|
||||
]
|
||||
|
||||
|
||||
def _respond(identity: str, stream: bool, captured: list[dict[str, JsonValue]]):
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/v1/messages", request.target
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
captured.append(body) # mutable-ok: the scripted upstream records each request it served
|
||||
if len(captured) == 2:
|
||||
return Reply(body=json.dumps(_message(identity, [{"type": "text", "text": "18C"}], "end_turn")).encode())
|
||||
if stream:
|
||||
return Reply(content_type="text/event-stream", chunks=sse_chunks(_search_turn_events(identity)))
|
||||
return Reply(body=json.dumps(_message(identity, _search_turn_blocks(), "tool_use")).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _streamed_assistant(lines: Sequence[str]) -> dict[str, JsonValue]:
|
||||
deltas: Final = tuple(
|
||||
chunk["choices"][0]["delta"]
|
||||
for line in lines
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
for chunk in (json.loads(line.removeprefix("data: ")),)
|
||||
if chunk["choices"]
|
||||
)
|
||||
tool_calls: Final = {} # mutable-ok: stream deltas are folded into one tool call per index
|
||||
for delta in deltas:
|
||||
for call in delta.get("tool_calls") or ():
|
||||
entry = tool_calls.setdefault(call["index"], {"id": None, "type": "function", "name": "", "arguments": ""})
|
||||
entry["id"] = call.get("id") or entry["id"]
|
||||
entry["name"] += call["function"].get("name") or ""
|
||||
entry["arguments"] += call["function"].get("arguments") or ""
|
||||
fields: Final = next(
|
||||
(
|
||||
delta["provider_specific_fields"]
|
||||
for delta in reversed(deltas)
|
||||
if (delta.get("provider_specific_fields") or {}).get("tool_results")
|
||||
),
|
||||
{},
|
||||
)
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": entry["id"],
|
||||
"type": "function",
|
||||
"function": {"name": entry["name"], "arguments": entry["arguments"]},
|
||||
}
|
||||
for entry in tool_calls.values()
|
||||
],
|
||||
"provider_specific_fields": fields,
|
||||
}
|
||||
|
||||
|
||||
def _round_trip(gateway: Gateway, stream: bool) -> tuple[dict[str, JsonValue], list[JsonValue]]:
|
||||
identity: Final = f"anthropic-tool-search-{uuid.uuid4().hex}"
|
||||
captured: Final[list[dict[str, JsonValue]]] = []
|
||||
with wire_server(_respond(identity, stream, captured)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
first_request: Final = {
|
||||
"model": model,
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "What is the weather in Paris?"}],
|
||||
"tools": _TOOLS,
|
||||
"stream": stream,
|
||||
}
|
||||
first: Final = gateway.request("POST", "/v1/chat/completions", first_request)
|
||||
assert first.status_code == 200, first.text
|
||||
assistant: Final = (
|
||||
_streamed_assistant(first.text.splitlines()) if stream else first.json()["choices"][0]["message"]
|
||||
)
|
||||
second: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{**first_request, "stream": False, "messages": _second_turn_messages(assistant)},
|
||||
)
|
||||
assert second.status_code == 200, second.text
|
||||
assert len(wire.drain()) == 2
|
||||
return assistant, _assistant_turn_on_the_wire(captured[1])
|
||||
|
||||
|
||||
def _assert_search_replayed_in_place(assistant: dict[str, JsonValue], wire_content: list[JsonValue]) -> None:
|
||||
fields: Final = assistant["provider_specific_fields"]
|
||||
assert isinstance(fields, dict) and fields.get("tool_results") == [_SEARCH_RESULT], assistant
|
||||
assert wire_content == [
|
||||
{"type": "server_tool_use", "id": _SEARCH_ID, "name": "tool_search_tool_bm25", "input": _SEARCH_INPUT},
|
||||
_SEARCH_RESULT,
|
||||
{"type": "tool_use", "id": _WEATHER_ID, "name": "get_weather", "input": {"city": "Paris"}},
|
||||
], wire_content
|
||||
|
||||
|
||||
def test_chat_completions_replays_a_tool_search_result_right_after_its_search_call(gateway: Gateway) -> None:
|
||||
assistant, wire_content = _round_trip(gateway, stream=False)
|
||||
_assert_search_replayed_in_place(assistant, wire_content)
|
||||
|
||||
|
||||
def test_streamed_chat_completions_replays_a_tool_search_result_right_after_its_search_call(gateway: Gateway) -> None:
|
||||
assistant, wire_content = _round_trip(gateway, stream=True)
|
||||
_assert_search_replayed_in_place(assistant, wire_content)
|
||||
|
|
@ -5,6 +5,7 @@ import socket
|
|||
import socketserver
|
||||
import ssl
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Generator, Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -24,13 +25,23 @@ from cryptography.x509.oid import NameOID
|
|||
from google import genai
|
||||
from google.genai import types
|
||||
from google.oauth2.credentials import Credentials
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.anthropic_thinking import sse_chunks
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers._tool_search_support import (
|
||||
SEARCH_TOOL,
|
||||
SEARCH_TURN,
|
||||
marker_of,
|
||||
message_events,
|
||||
question,
|
||||
turn_reply,
|
||||
)
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_PROJECT: Final = "scripted-project"
|
||||
_TOKEN: Final = "scripted-vertex-token"
|
||||
|
|
@ -167,6 +178,8 @@ def _vertex_tunnel(destination: Wire) -> Generator[_Tunnel, None, None]:
|
|||
|
||||
|
||||
def _vertex_peer(request: Request) -> Reply:
|
||||
if SEARCH_TOOL.encode() in request.body:
|
||||
return turn_reply(request, _search_identity(marker_of(request)), SEARCH_TURN, "tool_use")
|
||||
if "enerateContent" in request.target:
|
||||
if ":streamGenerateContent" in request.target:
|
||||
return Reply(chunks=_GEMINI_SSE, content_type="text/event-stream")
|
||||
|
|
@ -176,6 +189,10 @@ def _vertex_peer(request: Request) -> Reply:
|
|||
return Reply(body=json.dumps(_CLAUDE_REPLY).encode())
|
||||
|
||||
|
||||
def _search_identity(marker: str) -> str:
|
||||
return f"msg_{marker}_search"
|
||||
|
||||
|
||||
def _token_peer(request: Request) -> Reply:
|
||||
assert request.target == "/_oauth/token", request.target
|
||||
return Reply(body=json.dumps({"access_token": _TOKEN, "expires_in": 3600, "token_type": "Bearer"}).encode())
|
||||
|
|
@ -514,3 +531,36 @@ def test_credential_less_vertex_project_forwards_the_caller_bearer_token_when_th
|
|||
assert upstream.headers["x-caller-trace"] == "trace-1", upstream.headers
|
||||
assert key not in upstream.target, upstream.target
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("action", ["rawPredict", "streamRawPredict"])
|
||||
def test_vertex_partner_stream_relays_a_tool_search_turn_verbatim_and_logs_it_once(rig: _Rig, action: str) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model_path: Final = f"/v1/projects/{_PROJECT}/locations/global/publishers/anthropic/models/claude-sonnet-4-5"
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
"max_tokens": 64,
|
||||
"stream": True,
|
||||
"tools": [{"type": "tool_search_tool_bm25_20251119", "name": SEARCH_TOOL}],
|
||||
"messages": [{"role": "user", "content": question(marker)}],
|
||||
}
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
response: Final = rig.proxy.client.post(
|
||||
f"/vertex_ai{model_path}:{action}", json=body, headers={"Authorization": f"Bearer {key}"}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
expected: Final = b"".join(sse_chunks(message_events(_search_identity(marker), SEARCH_TURN, "tool_use")))
|
||||
assert response.content == expected, response.text
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"{model_path}:{action}?alt=sse")
|
||||
], received
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.headers["x-litellm-call-id"],)
|
||||
),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert rows[0]["status"] == "success", rows
|
||||
|
|
|
|||
130
tests/integration/sdk/test_tool_loop_server_tool_replay_sdk.py
Normal file
130
tests/integration/sdk/test_tool_loop_server_tool_replay_sdk.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from queue import SimpleQueue
|
||||
from typing import Final, cast
|
||||
|
||||
import litellm
|
||||
from integration._support.wire import wire_server
|
||||
from integration.providers._tool_search_support import (
|
||||
MODEL,
|
||||
SEARCH_RESULT,
|
||||
SEARCH_TURN,
|
||||
SEARCH_USE,
|
||||
TOOLS,
|
||||
WEATHER_ANSWER,
|
||||
WEATHER_ID,
|
||||
WEATHER_USE,
|
||||
Event,
|
||||
answer,
|
||||
answered_ids,
|
||||
first_messages,
|
||||
scripted_peer,
|
||||
wire_assistant,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolMessage, ChatCompletionToolParam
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
|
||||
_API_KEY: Final = "scripted-anthropic-key"
|
||||
_TOOLS: Final = cast(list[ChatCompletionToolParam], TOOLS)
|
||||
|
||||
|
||||
def _messages(marker: str) -> list[AllMessageValues]:
|
||||
return cast(list[AllMessageValues], first_messages(marker))
|
||||
|
||||
|
||||
def _answer(call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
return {"role": "tool", "tool_call_id": call.id, "content": WEATHER_ANSWER}
|
||||
|
||||
|
||||
def _recording(executed: SimpleQueue[str]) -> Callable[[ChatCompletionMessageToolCall], ChatCompletionToolMessage]:
|
||||
def execute(call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
executed.put(str(call.function.name))
|
||||
return _answer(call)
|
||||
|
||||
return execute
|
||||
|
||||
|
||||
def _async_recording(
|
||||
executed: SimpleQueue[str],
|
||||
) -> Callable[[ChatCompletionMessageToolCall], Awaitable[ChatCompletionToolMessage]]:
|
||||
async def execute(call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
executed.put(str(call.function.name))
|
||||
return _answer(call)
|
||||
|
||||
return execute
|
||||
|
||||
|
||||
def _drained(executed: SimpleQueue[str]) -> tuple[str, ...]:
|
||||
return tuple(executed.get() for _ in range(executed.qsize()))
|
||||
|
||||
|
||||
def _run(marker: str, wire_url: str, executed: SimpleQueue[str]) -> str | None:
|
||||
return litellm.run_tool_loop(
|
||||
model=f"anthropic/{MODEL}",
|
||||
messages=_messages(marker),
|
||||
tools=_TOOLS,
|
||||
execute_tool=_recording(executed),
|
||||
api_base=wire_url,
|
||||
api_key=_API_KEY,
|
||||
max_tokens=64,
|
||||
)
|
||||
|
||||
|
||||
def test_run_tool_loop_executes_only_the_client_tool_and_replays_the_kept_search() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
executed: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
with wire_server(scripted_peer()) as wire:
|
||||
content: Final = _run(marker, wire.url, executed)
|
||||
received: Final = wire.drain()
|
||||
assert content == answer(marker), content
|
||||
assert _drained(executed) == ("get_weather",)
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
async def test_arun_tool_loop_executes_only_the_client_tool_and_replays_the_kept_search() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
executed: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
with wire_server(scripted_peer()) as wire:
|
||||
content: Final = await litellm.arun_tool_loop(
|
||||
model=f"anthropic/{MODEL}",
|
||||
messages=_messages(marker),
|
||||
tools=_TOOLS,
|
||||
execute_tool=_async_recording(executed),
|
||||
api_base=wire.url,
|
||||
api_key=_API_KEY,
|
||||
max_tokens=64,
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert content == answer(marker), content
|
||||
assert _drained(executed) == ("get_weather",)
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
||||
|
||||
def test_run_tool_loop_returns_a_turn_that_only_used_the_server_search() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
closing: Final = f"No weather tool was needed for marker-{marker}"
|
||||
turn: Final[Sequence[Event]] = (SEARCH_USE, SEARCH_RESULT, {"type": "text", "text": closing})
|
||||
executed: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
with wire_server(scripted_peer(turn)) as wire:
|
||||
content: Final = _run(marker, wire.url, executed)
|
||||
received: Final = wire.drain()
|
||||
assert content == closing, content
|
||||
assert _drained(executed) == ()
|
||||
assert len(received) == 1, received
|
||||
|
||||
|
||||
def test_run_tool_loop_replays_a_client_only_turn_as_a_plain_tool_call() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
executed: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
with wire_server(scripted_peer((WEATHER_USE,))) as wire:
|
||||
content: Final = _run(marker, wire.url, executed)
|
||||
received: Final = wire.drain()
|
||||
assert content == answer(marker), content
|
||||
assert _drained(executed) == ("get_weather",)
|
||||
assert len(received) == 2, received
|
||||
assert wire_assistant(received[1]) == [WEATHER_USE], received[1].body
|
||||
assert answered_ids(received[1]) == [WEATHER_ID], received[1].body
|
||||
|
|
@ -1166,6 +1166,48 @@ def test_cache_control_hook_counts_tool_call_marks_except_answered_server_tool_c
|
|||
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("follow_ups", "expected_breakpoints"),
|
||||
[
|
||||
pytest.param((), 0, id="rebuilt_from_its_kept_result"),
|
||||
pytest.param(
|
||||
({"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "found get_weather"},),
|
||||
1,
|
||||
id="replayed_with_the_client_answer",
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
{"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "found get_weather"},
|
||||
{"role": "user", "content": "Go on"},
|
||||
),
|
||||
0,
|
||||
id="client_answer_replaced_by_the_kept_result",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_cache_control_hook_counts_a_server_tool_call_mark_only_when_the_call_replays_as_tool_use(
|
||||
follow_ups: tuple[AllMessageValues, ...], expected_breakpoints: int
|
||||
):
|
||||
search_call: Final = cast(
|
||||
AllMessageValues,
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "srvtoolu_01Search",
|
||||
"type": "function",
|
||||
"function": {"name": "tool_search_tool_bm25", "arguments": "{}"},
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
"provider_specific_fields": {"tool_results": [{"tool_use_id": "srvtoolu_01Search"}]},
|
||||
},
|
||||
)
|
||||
|
||||
assert AnthropicCacheControlHook.count_request_cache_breakpoints([search_call, *follow_ups]) == expected_breakpoints
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool_call_type", ["function", "custom", None])
|
||||
@pytest.mark.parametrize(
|
||||
"mark",
|
||||
|
|
|
|||
|
|
@ -1363,6 +1363,24 @@ def test_multiple_tools_streaming_has_index_zero():
|
|||
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
|
||||
|
||||
|
||||
def test_streamed_tool_search_tool_result_is_kept_in_tool_results():
|
||||
search_result: Final = {
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01Search",
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
|
||||
},
|
||||
}
|
||||
iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
parsed: Final = iterator.chunk_parser(
|
||||
chunk={"type": "content_block_start", "index": 1, "content_block": search_result}
|
||||
)
|
||||
|
||||
assert parsed.choices[0].delta.provider_specific_fields == {"tool_results": [search_result]}
|
||||
|
||||
|
||||
def test_streaming_chunks_have_stable_ids():
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
|
||||
first_chunk = {
|
||||
|
|
|
|||
|
|
@ -3503,49 +3503,203 @@ def test_code_execution_tool_results_in_hidden_params():
|
|||
assert hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"] == "hello\n"
|
||||
|
||||
|
||||
def test_tool_search_tool_result_not_in_tool_results():
|
||||
"""
|
||||
Test that tool_search_tool_result is NOT included in tool_results
|
||||
since it's internal metadata, not actual tool execution results.
|
||||
"""
|
||||
import httpx
|
||||
def _tool_search_turn(tool_references: list[str], last_block: dict[str, object]) -> list[dict[str, object]]:
|
||||
return [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01Search",
|
||||
"name": "tool_search_tool_bm25",
|
||||
"input": {"query": "weather"},
|
||||
},
|
||||
{
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01Search",
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": name} for name in tool_references],
|
||||
},
|
||||
},
|
||||
last_block,
|
||||
]
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
config = AnthropicConfig()
|
||||
|
||||
mock_anthropic_response = {
|
||||
def _parsed_assistant_message(content: list[dict[str, object]], stop_reason: str):
|
||||
anthropic_response: Final = {
|
||||
"id": "msg_01XYZ",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
"content": [
|
||||
{"type": "text", "text": "Found tools."},
|
||||
{"type": "tool_search_tool_result", "tool_references": ["tool1", "tool2"]},
|
||||
],
|
||||
"stop_reason": "stop",
|
||||
"model": "claude-opus-5-5",
|
||||
"content": content,
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 100, "output_tokens": 50},
|
||||
}
|
||||
|
||||
mock_raw_response = MagicMock(spec=httpx.Response)
|
||||
mock_raw_response.json.return_value = mock_anthropic_response
|
||||
mock_raw_response.status_code = 200
|
||||
mock_raw_response.headers = {}
|
||||
|
||||
model_response = ModelResponse()
|
||||
|
||||
transformed_response = config.transform_parsed_response(
|
||||
completion_response=mock_anthropic_response,
|
||||
raw_response=mock_raw_response,
|
||||
model_response=model_response,
|
||||
json_mode=False,
|
||||
prefix_prompt=None,
|
||||
raw_response: Final = MagicMock(spec=httpx.Response)
|
||||
raw_response.json.return_value = anthropic_response
|
||||
raw_response.status_code = 200
|
||||
raw_response.headers = {}
|
||||
return (
|
||||
AnthropicConfig()
|
||||
.transform_parsed_response(
|
||||
completion_response=anthropic_response,
|
||||
raw_response=raw_response,
|
||||
model_response=ModelResponse(),
|
||||
json_mode=False,
|
||||
prefix_prompt=None,
|
||||
)
|
||||
.choices[0]
|
||||
.message
|
||||
)
|
||||
|
||||
# Verify tool_search_tool_result is NOT in tool_results
|
||||
provider_fields = transformed_response.choices[0].message.provider_specific_fields
|
||||
assert provider_fields.get("tool_results") is None
|
||||
|
||||
_WEATHER_CALL: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}}
|
||||
|
||||
|
||||
def _web_search_turn(last_block: dict[str, object]) -> list[dict[str, object]]:
|
||||
return [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01Web", "name": "web_search", "input": {"query": "Paris facts"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01Web",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://example.com/paris",
|
||||
"title": "Paris",
|
||||
"encrypted_content": "EqgfCioIARgB",
|
||||
"page_age": None,
|
||||
}
|
||||
],
|
||||
},
|
||||
last_block,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("server_call_answered", "client_call_answered", "modify_params"),
|
||||
[
|
||||
pytest.param(False, True, False, id="client_answers_its_own_call"),
|
||||
pytest.param(False, True, True, id="client_answers_its_own_call_with_modify_params"),
|
||||
pytest.param(True, True, False, id="client_also_answers_the_server_call"),
|
||||
pytest.param(False, False, True, id="modify_params_answers_the_unanswered_client_call"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
[_tool_search_turn(["get_weather"], _WEATHER_CALL), _web_search_turn(_WEATHER_CALL)],
|
||||
ids=["tool_search", "web_search"],
|
||||
)
|
||||
def test_server_tool_result_round_trips_right_after_its_call(
|
||||
content: list[dict[str, object]],
|
||||
server_call_answered: bool,
|
||||
client_call_answered: bool,
|
||||
modify_params: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
|
||||
|
||||
message: Final = _parsed_assistant_message(content, "tool_use")
|
||||
monkeypatch.setattr(litellm, "modify_params", modify_params)
|
||||
server_answer: Final = {"role": "tool", "tool_call_id": content[0]["id"], "content": "Error: unknown tool"}
|
||||
client_answer: Final = {"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}
|
||||
|
||||
replayed: Final = anthropic_messages_pt(
|
||||
messages=[
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
message.model_dump(),
|
||||
*((server_answer,) if server_call_answered else ()),
|
||||
*((client_answer,) if client_call_answered else ()),
|
||||
],
|
||||
model="claude-opus-5-5",
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
assert replayed[1]["content"] == content
|
||||
assert [turn["role"] for turn in replayed] == ["user", "assistant", "user"]
|
||||
assert [block["tool_use_id"] for block in replayed[2]["content"]] == ["toolu_01Weather"]
|
||||
|
||||
|
||||
def _block_ids(message: dict[str, object], block_type: str, id_key: str) -> list[object]:
|
||||
return sorted(block[id_key] for block in message["content"] if block["type"] == block_type)
|
||||
|
||||
|
||||
_SEARCH_THAT_FOUND_NOTHING: Final = _tool_search_turn([], {"type": "text", "text": "I have no tool that books flights."})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("follow_ups", "modify_params", "server_tool_use_ids"),
|
||||
[
|
||||
pytest.param(
|
||||
({"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "no tool found"},),
|
||||
False,
|
||||
[],
|
||||
id="client_answers_a_search_that_found_nothing",
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
{"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "no tool found"},
|
||||
{"role": "user", "content": "Book it anyway"},
|
||||
),
|
||||
False,
|
||||
["srvtoolu_01Search"],
|
||||
id="client_answers_a_search_that_found_nothing_and_writes_again",
|
||||
),
|
||||
pytest.param((), True, [], id="modify_params_with_nothing_after_a_search_that_found_nothing"),
|
||||
],
|
||||
)
|
||||
def test_replayed_search_that_found_nothing_answers_every_tool_use_and_ends_on_a_user_turn(
|
||||
follow_ups: tuple[dict[str, str], ...],
|
||||
modify_params: bool,
|
||||
server_tool_use_ids: list[str],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
|
||||
|
||||
message: Final = _parsed_assistant_message(_SEARCH_THAT_FOUND_NOTHING, "end_turn")
|
||||
monkeypatch.setattr(litellm, "modify_params", modify_params)
|
||||
|
||||
replayed: Final = anthropic_messages_pt(
|
||||
messages=[{"role": "user", "content": "Find a tool and use it"}, message.model_dump(), *follow_ups],
|
||||
model="claude-opus-5-5",
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
assert [turn["role"] for turn in replayed] == ["user", "assistant", "user"]
|
||||
assert _block_ids(replayed[2], "tool_result", "tool_use_id") == _block_ids(replayed[1], "tool_use", "id")
|
||||
assert _block_ids(replayed[1], "server_tool_use", "id") == server_tool_use_ids
|
||||
|
||||
|
||||
_CODE_EXECUTION_BLOCKS: Final = [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01Bash",
|
||||
"name": "bash_code_execution",
|
||||
"input": {"command": "python3 -c 'print(2 + 2)'"},
|
||||
},
|
||||
{
|
||||
"type": "bash_code_execution_tool_result",
|
||||
"tool_use_id": "srvtoolu_01Bash",
|
||||
"content": {"type": "bash_code_execution_result", "stdout": "4\n", "stderr": "", "return_code": 0},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("code_execution_blocks", "code_interpreter_call_ids"),
|
||||
[
|
||||
pytest.param([], None, id="tool_search_only"),
|
||||
pytest.param(_CODE_EXECUTION_BLOCKS, ["srvtoolu_01Bash"], id="tool_search_and_code_execution"),
|
||||
],
|
||||
)
|
||||
def test_parsed_tool_search_turn_adds_code_interpreter_results_only_for_code_execution(
|
||||
code_execution_blocks: list[dict[str, object]], code_interpreter_call_ids: list[str] | None
|
||||
):
|
||||
search_turn: Final = _tool_search_turn(["get_weather"], {"type": "text", "text": "2 + 2 is 4"})
|
||||
message: Final = _parsed_assistant_message([*code_execution_blocks, *search_turn], "end_turn")
|
||||
code_interpreter_results: Final = message.provider_specific_fields.get("code_interpreter_results")
|
||||
|
||||
assert search_turn[1] in message.provider_specific_fields["tool_results"]
|
||||
assert (
|
||||
None if code_interpreter_results is None else [call.id for call in code_interpreter_results]
|
||||
) == code_interpreter_call_ids
|
||||
|
||||
|
||||
def test_web_search_tool_result_backwards_compatibility():
|
||||
|
|
@ -4345,9 +4499,7 @@ def test_service_tier_without_anthropic_equivalent_raises_with_drop_params_hint(
|
|||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
|
||||
with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params") as exc_info:
|
||||
litellm.utils.get_optional_params(
|
||||
model="claude-sonnet-4-6", custom_llm_provider="anthropic", service_tier=tier
|
||||
)
|
||||
litellm.utils.get_optional_params(model="claude-sonnet-4-6", custom_llm_provider="anthropic", service_tier=tier)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ WEATHER_TOOLS: Final = (
|
|||
)
|
||||
|
||||
|
||||
def _openai_response(content: str | None, tool_calls: list | None = None) -> dict:
|
||||
def _openai_response(content: str | None, tool_calls: list | None = None, **message_fields: object) -> dict:
|
||||
return {
|
||||
"id": "chatcmpl-tool-loop",
|
||||
"object": "chat.completion",
|
||||
|
|
@ -41,6 +41,7 @@ def _openai_response(content: str | None, tool_calls: list | None = None) -> dic
|
|||
"role": "assistant",
|
||||
"content": content,
|
||||
"tool_calls": tool_calls,
|
||||
**message_fields,
|
||||
},
|
||||
}
|
||||
],
|
||||
|
|
@ -418,3 +419,130 @@ async def test_arun_tool_loop_two_rounds(respx_mock: respx.MockRouter, monkeypat
|
|||
{"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_2"},
|
||||
]
|
||||
assert messages == messages_snapshot
|
||||
|
||||
|
||||
ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages"
|
||||
TOOL_SEARCH_CALL: Final = {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01Search",
|
||||
"name": "tool_search_tool_bm25",
|
||||
"input": {"query": "weather"},
|
||||
}
|
||||
WEATHER_CALL: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}}
|
||||
|
||||
|
||||
def _tool_search_result(tool_names: list[str]) -> dict:
|
||||
return {
|
||||
"type": "tool_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01Search",
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": name} for name in tool_names],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _anthropic_response(content: list[dict], stop_reason: str) -> dict:
|
||||
return {
|
||||
"id": "msg_tool_loop",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-opus-5-5",
|
||||
"content": content,
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
|
||||
|
||||
def test_tool_search_call_is_replayed_with_its_result_instead_of_executed(respx_mock: respx.MockRouter) -> None:
|
||||
search_turn: Final = [TOOL_SEARCH_CALL, _tool_search_result(["get_weather"]), WEATHER_CALL]
|
||||
respx_mock.post(ANTHROPIC_MESSAGES_URL).mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_anthropic_response(search_turn, "tool_use")),
|
||||
httpx.Response(200, json=_anthropic_response([{"type": "text", "text": "Paris is 72F"}], "end_turn")),
|
||||
]
|
||||
)
|
||||
executed: Final = []
|
||||
|
||||
def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
executed.append(tc.function.name)
|
||||
return _tool_result(tc)
|
||||
|
||||
answer: Final = litellm.run_tool_loop(
|
||||
model="anthropic/claude-opus-5-5",
|
||||
messages=[{"role": "user", "content": "weather in Paris?"}],
|
||||
tools=WEATHER_TOOLS,
|
||||
execute_tool=executor,
|
||||
api_key="sk-ant-test",
|
||||
)
|
||||
|
||||
assert answer == "Paris is 72F"
|
||||
assert executed == ["get_weather"]
|
||||
second_messages: Final = _request_bodies(respx_mock)[1]["messages"]
|
||||
assert second_messages[1]["content"] == search_turn
|
||||
assert [block["tool_use_id"] for block in second_messages[2]["content"]] == ["toolu_01Weather"]
|
||||
|
||||
|
||||
async def test_arun_tool_loop_answers_after_a_search_that_found_nothing(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler())
|
||||
route: Final = respx_mock.post(ANTHROPIC_MESSAGES_URL).mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json=_anthropic_response(
|
||||
[TOOL_SEARCH_CALL, _tool_search_result([]), {"type": "text", "text": "I have no flight tool."}],
|
||||
"end_turn",
|
||||
),
|
||||
)
|
||||
)
|
||||
executed: Final = []
|
||||
|
||||
async def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
executed.append(tc.function.name)
|
||||
return _tool_result(tc)
|
||||
|
||||
answer: Final = await litellm.arun_tool_loop(
|
||||
model="anthropic/claude-opus-5-5",
|
||||
messages=[{"role": "user", "content": "book me a flight"}],
|
||||
tools=WEATHER_TOOLS,
|
||||
execute_tool=executor,
|
||||
api_key="sk-ant-test",
|
||||
)
|
||||
|
||||
assert answer == "I have no flight tool."
|
||||
assert executed == []
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
def test_client_tool_turn_is_replayed_without_its_provider_specific_fields(respx_mock: respx.MockRouter) -> None:
|
||||
respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
|
||||
side_effect=[
|
||||
httpx.Response(
|
||||
200,
|
||||
json=_openai_response(
|
||||
None,
|
||||
[_tool_call("call_1", "get_weather", {"city": "Paris"})],
|
||||
reasoning_details=[{"type": "reasoning.text", "text": "Check the weather tool"}],
|
||||
),
|
||||
),
|
||||
httpx.Response(200, json=_openai_response("Paris is 72F")),
|
||||
]
|
||||
)
|
||||
|
||||
answer: Final = litellm.run_tool_loop(
|
||||
model="openai/gpt-5-mini",
|
||||
messages=[{"role": "user", "content": "weather in Paris?"}],
|
||||
tools=WEATHER_TOOLS,
|
||||
execute_tool=_tool_result,
|
||||
api_key="sk-test",
|
||||
)
|
||||
|
||||
assert answer == "Paris is 72F"
|
||||
replayed_assistant: Final = _request_bodies(respx_mock)[1]["messages"][1]
|
||||
assert replayed_assistant["tool_calls"][0]["id"] == "call_1"
|
||||
assert "provider_specific_fields" not in replayed_assistant
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue