mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix(responses bridge): round-trip tool_search_call and tool_search_output items between turns (#46444)
* fix(responses bridge): keep defer_loading on function tools sent to the Responses API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): give the defer_loading casts a cast-ok reason Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(responses bridge): build defer_loading outside the tool loop so Final stays valid Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): round-trip tool_search_call and tool_search_output items between turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(responses bridge): read replayed tool-search and namespace fields through typed helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): keep tool_search items on a turn with no message or tool call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): keep every reasoning item on a search-only turn Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): leave content unset on a search-only turn so replay adds no empty message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): replay tool search items only with tool_search, count replayed tools in the allowlist, keep namespace off guardrails * fix(responses bridge): honor drop_params for defer_loading and route tool_search to Responses Under drop_params, defer_loading is dropped when no tool_search tool is sent, so the 400 it would cause upstream goes away as before. A non-bool value answers 400 without drop_params and is dropped with it. A null tool-level value falls back to the nested function value. A request carrying a tool_search tool now bridges to Responses on OpenAI and Azure endpoints even with reasoning_effort none, since chat completions rejects the tool_search type. * fix(responses bridge): replay tool search items in output order and count replayed tools in tool policy Store the ordered output prefix through the last tool search item, reasoning included, so a second reasoning item no longer makes OpenAI reject the replay. The shared chat-to-Responses converter now replays tool search items and namespace only when the bridge opts in, so Noma, the complexity router and guardrail payloads keep their shape. The tool policy guardrail and the allowlist count replayed tools only when the request can send them. * fix(responses bridge): forward defer_loading only beside a tool_search tool defer_loading only means something with tool_search, and backends without it reject the key (xAI answers 403). A true value with no tool_search tool now answers 400 naming drop_params and is dropped under drop_params; a false value there is a no-op and is dropped. * fix(tool_loop): keep tool search items and namespace between rounds * fix(responses bridge): check replayed tool search loads when a deployment adds tool_search * fix(chat): drop litellm response fields before strict OpenAI-compatible providers A chat reply from the Responses bridge carries provider_specific_fields on the message and on each tool call, plus reasoning_items. Groq, Fireworks, Mistral and Together reject unknown fields, so sending that reply back to them as history failed with a 400 or 422. Strip them in the shared helper and use it in those providers. * fix(caching): keep reasoning_items and tool search items on response cache hits A cached bridge response lost reasoning_items on a non-streaming hit and lost both reasoning_items and provider_specific_fields (tool_search_items) on a streamed hit, so the next turn could not replay the search or the reasoning around it * test(integration): cover tool search item replay on the chat to Responses bridge Proxy, SDK, authorization and chaos cells for returning and replaying tool_search_call and tool_search_output items and the tool call namespace * test(integration): add response cache twins for the tool search guardrail cells Each guardrail cell now has a twin where the answer is already in the Redis response cache: a key opted out of global guardrails caches a replay that loads a policy-blocked tool, the same replay without guardrails is cached before block mode sees it, rewrite mode is not served the unrewritten conversation's cached answer, and an allowed replay is served from the cache. * fix(responses bridge): forward defer_loading beside a client-executed tool_search The bridge now replays a client tool_search_call and the client's tool_search_output through provider_specific_fields.tool_search_items, so a deferred function beside a client tool_search keeps defer_loading like it does beside a hosted one * test(responses bridge): pin the client tool_search turn and its replay, streamed and not * fix(responses bridge): strip nested namespace for strict providers and keep streamed reasoning_items on text replies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses bridge): keep dropping defer_loading beside a client-executed tool_search Forwarding defer_loading beside a client tool_search made GPT stop with a search call the chat client has to answer by hand, so a plain question that used to get tool calls got none. Keep main's behavior of dropping it there. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang
GitHub
parent
1311894845
commit
52fc37888b
34 changed files with 3273 additions and 89 deletions
|
|
@ -304,6 +304,136 @@ def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[
|
|||
)
|
||||
|
||||
|
||||
_TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"})
|
||||
_TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_ITEM_TYPES})
|
||||
_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items"
|
||||
|
||||
|
||||
def _json_field(value: object, key: str) -> object:
|
||||
return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON
|
||||
|
||||
|
||||
def _tool_search_replay_item(item: object) -> Mapping[str, object] | None:
|
||||
raw: Final[object] = item.model_dump(exclude_unset=True) if isinstance(item, BaseModel) else item
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
fields: Final = cast(dict[str, object], raw) # cast-ok: provider JSON item, keys forwarded verbatim
|
||||
if fields.get("type") not in _TOOL_SEARCH_REPLAY_ITEM_TYPES:
|
||||
return None
|
||||
return {key: value for key, value in fields.items() if value is not None}
|
||||
|
||||
|
||||
def _tool_search_items_from_output_items(output_items: object) -> tuple[Mapping[str, object], ...]:
|
||||
if not isinstance(output_items, list):
|
||||
return ()
|
||||
entries: Final = cast(Sequence[object], output_items) # cast-ok: each entry is re-validated
|
||||
items: Final = tuple(item for item in (_tool_search_replay_item(entry) for entry in entries) if item is not None)
|
||||
last_search: Final = max(
|
||||
(index for index, item in enumerate(items) if item.get("type") in _TOOL_SEARCH_ITEM_TYPES), default=-1
|
||||
)
|
||||
return items[: last_search + 1]
|
||||
|
||||
|
||||
def _tool_search_provider_fields(
|
||||
items: Sequence[Mapping[str, object]],
|
||||
) -> dict[str, object] | None: # mutable-ok: Message.provider_specific_fields is typed as a dict
|
||||
return {_TOOL_SEARCH_ITEMS_KEY: [dict(item) for item in items]} if items else None
|
||||
|
||||
|
||||
def _tool_search_input_items(msg: object) -> tuple[Mapping[str, object], ...]:
|
||||
return _tool_search_items_from_output_items(
|
||||
_json_field(_json_field(msg, "provider_specific_fields"), _TOOL_SEARCH_ITEMS_KEY)
|
||||
)
|
||||
|
||||
|
||||
def carries_tool_search_items(provider_specific_fields: object) -> bool:
|
||||
return bool(_json_field(provider_specific_fields, _TOOL_SEARCH_ITEMS_KEY))
|
||||
|
||||
|
||||
def _tool_definition_name(tool: object) -> tuple[str, ...]:
|
||||
name: Final = _json_field(tool, "server_label" if _json_field(tool, "type") == "mcp" else "name")
|
||||
return (name,) if isinstance(name, str) else ()
|
||||
|
||||
|
||||
def _namespace_members(tool: object) -> Sequence[object]:
|
||||
members: Final = _json_field(tool, "tools")
|
||||
if _json_field(tool, "type") != "namespace" or not isinstance(members, list):
|
||||
return (tool,)
|
||||
return cast(Sequence[object], members) # cast-ok: untyped JSON, each member re-checked
|
||||
|
||||
|
||||
def _loaded_tools(item: Mapping[str, object]) -> Sequence[object]:
|
||||
tools: Final = item.get("tools")
|
||||
if item.get("type") != "tool_search_output" or not isinstance(tools, list):
|
||||
return ()
|
||||
return cast(Sequence[object], tools) # cast-ok: untyped JSON, each tool re-checked
|
||||
|
||||
|
||||
def _json_list(value: object) -> Sequence[object] | None:
|
||||
return cast(Sequence[object], value) if isinstance(value, list) else None # cast-ok: untyped JSON
|
||||
|
||||
|
||||
def _message_loaded_tools(message: object) -> Iterator[object]:
|
||||
loaded: Final = chain.from_iterable(_loaded_tools(item) for item in _tool_search_input_items(message))
|
||||
return chain.from_iterable(_namespace_members(tool) for tool in loaded)
|
||||
|
||||
|
||||
def replayed_tool_search_tools(request: object) -> tuple[object, ...]:
|
||||
messages: Final = _json_list(_json_field(request, "messages")) or ()
|
||||
return tuple(chain.from_iterable(_message_loaded_tools(message) for message in messages))
|
||||
|
||||
|
||||
def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]:
|
||||
return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request)))
|
||||
|
||||
|
||||
def _loads_any_tool(message: object, names: frozenset[str]) -> bool:
|
||||
loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message))
|
||||
return any(name in names for name in loaded_names)
|
||||
|
||||
|
||||
def _without_tool_search_items(message: object) -> object:
|
||||
fields: Final = _json_field(message, "provider_specific_fields")
|
||||
if not isinstance(fields, dict):
|
||||
return message
|
||||
kept: Final = {
|
||||
key: value
|
||||
for key, value in cast(Mapping[str, object], fields).items() # cast-ok: untyped JSON
|
||||
if key != _TOOL_SEARCH_ITEMS_KEY
|
||||
}
|
||||
return {**cast(Mapping[str, object], message), "provider_specific_fields": kept} # cast-ok: dict checked above
|
||||
|
||||
|
||||
def without_tool_search_items_loading(
|
||||
request: object, names: frozenset[str]
|
||||
) -> list[object] | None: # mutable-ok: request messages stay a list for every downstream handler
|
||||
messages: Final = _json_list(_json_field(request, "messages"))
|
||||
if messages is None:
|
||||
return None
|
||||
return [_without_tool_search_items(message) if _loads_any_tool(message, names) else message for message in messages]
|
||||
|
||||
|
||||
def _effective_tools(request: object) -> Sequence[object]:
|
||||
override: Final = _json_list(_json_field(_json_field(request, "extra_body"), "tools"))
|
||||
return override if override is not None else (_json_list(_json_field(request, "tools")) or ())
|
||||
|
||||
|
||||
def _requests_tool_search(request: object) -> bool:
|
||||
return any(_json_field(tool, "type") == "tool_search" for tool in _effective_tools(request))
|
||||
|
||||
|
||||
def _reasoning_ids(items: Sequence[Mapping[str, object]]) -> frozenset[str]:
|
||||
ids: Final = (item.get("id") for item in items if item.get("type") == "reasoning")
|
||||
return frozenset(item_id for item_id in ids if isinstance(item_id, str))
|
||||
|
||||
|
||||
def _assistant_replay_items(msg: "AllMessageValues", replay_tool_search: bool) -> tuple[Mapping[str, object], ...]:
|
||||
searched: Final = _tool_search_input_items(msg) if replay_tool_search else ()
|
||||
searched_reasoning_ids: Final = _reasoning_ids(searched)
|
||||
reasoning: Final = (item for item in _reasoning_input_items(msg) if item.get("id") not in searched_reasoning_ids)
|
||||
return (*searched, *reasoning)
|
||||
|
||||
|
||||
def _as_chat_reasoning_items(
|
||||
reasoning_items: Sequence[_BuiltReasoningItem | ChatCompletionReasoningItem],
|
||||
) -> list[ChatCompletionReasoningItem] | None:
|
||||
|
|
@ -346,6 +476,7 @@ def _fold_choices_into_one(
|
|||
tool_calls: Sequence[Mapping[str, object]],
|
||||
trailing_reasoning_content: str | None,
|
||||
trailing_reasoning_item: _BuiltReasoningItem | None,
|
||||
tool_search_items: Sequence[Mapping[str, object]],
|
||||
structured_output: bool = False,
|
||||
) -> "Choices":
|
||||
"""Clients that only read choices[0] must still see every text and tool call.
|
||||
|
|
@ -373,6 +504,7 @@ def _fold_choices_into_one(
|
|||
tool_calls=list(tool_calls) or None,
|
||||
reasoning_content=" ".join(value for value in reasoning_contents if value) or None,
|
||||
reasoning_items=_as_chat_reasoning_items(reasoning_items),
|
||||
provider_specific_fields=_tool_search_provider_fields(tool_search_items),
|
||||
),
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
index=0,
|
||||
|
|
@ -457,7 +589,9 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
|
|||
provider_specific_fields: Mapping[str, object]
|
||||
|
||||
|
||||
def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
|
||||
def tool_call_dict_from_output_item(
|
||||
item: Mapping[str, Any], index: int, *, include_namespace: bool = False
|
||||
) -> _ChatToolCallDict:
|
||||
"""Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
|
||||
completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
|
||||
string payload in ``input`` rather than ``arguments``; both map to
|
||||
|
|
@ -486,12 +620,24 @@ def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Cha
|
|||
provider_specific_fields = vars(raw_provider_fields)
|
||||
else:
|
||||
provider_specific_fields = None
|
||||
if provider_specific_fields:
|
||||
tool_call_dict["provider_specific_fields"] = provider_specific_fields
|
||||
function_chunk["provider_specific_fields"] = provider_specific_fields
|
||||
namespace: Final = item.get("namespace") if include_namespace else None
|
||||
merged_provider_fields: Final = {
|
||||
**(provider_specific_fields or {}),
|
||||
**({"namespace": namespace} if isinstance(namespace, str) else {}),
|
||||
}
|
||||
if merged_provider_fields:
|
||||
tool_call_dict["provider_specific_fields"] = merged_provider_fields
|
||||
function_chunk["provider_specific_fields"] = merged_provider_fields
|
||||
return tool_call_dict
|
||||
|
||||
|
||||
def _tool_call_namespace(tool_call: object) -> str | None:
|
||||
top_level: Final = _json_field(_json_field(tool_call, "provider_specific_fields"), "namespace")
|
||||
function_fields: Final = _json_field(_json_field(tool_call, "function"), "provider_specific_fields")
|
||||
nested: Final = _json_field(function_fields, "namespace")
|
||||
return top_level if isinstance(top_level, str) else nested if isinstance(nested, str) else None
|
||||
|
||||
|
||||
def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFunctionParam | ToolChoiceCustomParam:
|
||||
if choice_type == "custom":
|
||||
return ToolChoiceCustomParam(type="custom", name=name)
|
||||
|
|
@ -641,9 +787,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
*,
|
||||
drop_params: bool = False,
|
||||
keep_prompt_cache_breakpoints: bool = False,
|
||||
replay_tool_search: bool = False,
|
||||
include_namespace: bool = False,
|
||||
) -> tuple[list[object], str | None]:
|
||||
converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
|
||||
messages, drop_params=drop_params
|
||||
messages,
|
||||
drop_params=drop_params,
|
||||
replay_tool_search=replay_tool_search,
|
||||
include_namespace=include_namespace,
|
||||
)
|
||||
return (
|
||||
converted_input_items
|
||||
|
|
@ -653,7 +804,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
)
|
||||
|
||||
def _convert_chat_completion_messages_to_responses_input(
|
||||
self, messages: list["AllMessageValues"], *, drop_params: bool = False
|
||||
self,
|
||||
messages: list["AllMessageValues"],
|
||||
*,
|
||||
drop_params: bool = False,
|
||||
replay_tool_search: bool = False,
|
||||
include_namespace: bool = False,
|
||||
) -> tuple[list[object], str | None]:
|
||||
input_items: Final[list[object]] = []
|
||||
instructions: str | None = None
|
||||
|
|
@ -735,7 +891,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
)
|
||||
)
|
||||
elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
input_items.extend(_assistant_replay_items(msg, replay_tool_search))
|
||||
if content:
|
||||
input_items.append(
|
||||
{
|
||||
|
|
@ -758,6 +914,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
input_tool_call["name"] = function["name"]
|
||||
if "arguments" in function:
|
||||
input_tool_call["arguments"] = function["arguments"]
|
||||
if include_namespace and (
|
||||
namespace := _tool_call_namespace(cast(object, tool_call)) # cast-ok: untyped
|
||||
):
|
||||
input_tool_call["namespace"] = namespace
|
||||
input_items.append(input_tool_call)
|
||||
elif isinstance(custom, dict):
|
||||
input_items.append(
|
||||
|
|
@ -772,7 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
raise ValueError(f"tool call not supported: {tool_call}")
|
||||
elif content is not None:
|
||||
if role == "assistant":
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
input_items.extend(_assistant_replay_items(msg, replay_tool_search))
|
||||
input_items.append(
|
||||
{
|
||||
"type": "message",
|
||||
|
|
@ -783,7 +943,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
}
|
||||
)
|
||||
elif role == "assistant":
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
input_items.extend(_assistant_replay_items(msg, replay_tool_search))
|
||||
|
||||
return input_items, instructions
|
||||
|
||||
|
|
@ -890,6 +1050,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
client: object | None = None,
|
||||
) -> dict:
|
||||
optional_params = self._extract_extra_body_params(optional_params)
|
||||
provider: Final = _served_provider(litellm_params)
|
||||
base_model: Final = litellm_params.get("base_model")
|
||||
supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model, provider) or (
|
||||
|
|
@ -901,6 +1062,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
messages,
|
||||
drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
|
||||
keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
|
||||
replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast
|
||||
include_namespace=True,
|
||||
)
|
||||
# OpenAI's Responses API rejects an empty input. For a system-only
|
||||
# request, carry the system message as a system-role input item instead
|
||||
|
|
@ -925,8 +1088,6 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
)
|
||||
instructions: Final = None if is_system_only_request else converted_instructions
|
||||
|
||||
optional_params = self._extract_extra_body_params(optional_params)
|
||||
|
||||
# Build responses API request using the reverse transformation logic
|
||||
responses_api_request: Final = ResponsesAPIOptionalRequestParams()
|
||||
|
||||
|
|
@ -1100,7 +1261,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
# Tool calls accumulate into the single trailing tool_calls choice
|
||||
# like the typed branches above; a choice per call would hide every
|
||||
# call after choices[0] from chat clients
|
||||
accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
|
||||
accumulated_tool_calls.append(
|
||||
tool_call_dict_from_output_item(raw_item, tool_call_index, include_namespace=True)
|
||||
)
|
||||
tool_call_index += 1
|
||||
elif handle_raw_dict_callback is not None:
|
||||
choice, index = handle_raw_dict_callback(item=raw_item, index=index)
|
||||
|
|
@ -1109,22 +1272,31 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
else:
|
||||
pass # don't fail request if item in list is not supported
|
||||
|
||||
tool_search_items: Final = _tool_search_items_from_output_items(output_items)
|
||||
if not choices and not accumulated_tool_calls:
|
||||
return choices
|
||||
if not tool_search_items:
|
||||
return choices
|
||||
return [
|
||||
LiteLLMResponsesTransformationHandler._build_empty_choice(
|
||||
output_items, "stop", tool_search_items=tool_search_items
|
||||
)
|
||||
]
|
||||
return [
|
||||
_fold_choices_into_one(
|
||||
choices=choices,
|
||||
tool_calls=accumulated_tool_calls,
|
||||
trailing_reasoning_content=reasoning_content,
|
||||
trailing_reasoning_item=pending_reasoning_item,
|
||||
tool_search_items=tool_search_items,
|
||||
structured_output=structured_output,
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _build_empty_incomplete_choice(
|
||||
def _build_empty_choice(
|
||||
output_items: Sequence[object],
|
||||
finish_reason: Literal["length", "content_filter"],
|
||||
finish_reason: Literal["stop", "length", "content_filter"],
|
||||
tool_search_items: Sequence[Mapping[str, object]] = (),
|
||||
) -> "Choices":
|
||||
from litellm.types.utils import Choices, Message
|
||||
|
||||
|
|
@ -1136,9 +1308,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
if summary_block.get("text")
|
||||
)
|
||||
message: Final = Message(
|
||||
content="",
|
||||
content=None if tool_search_items else "",
|
||||
reasoning_content=reasoning_content if reasoning_content else None,
|
||||
reasoning_items=_as_chat_reasoning_items(reasoning_items),
|
||||
provider_specific_fields=_tool_search_provider_fields(tool_search_items),
|
||||
)
|
||||
return Choices(message=message, finish_reason=finish_reason, index=0)
|
||||
|
||||
|
|
@ -1265,7 +1438,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None
|
||||
)
|
||||
if len(choices) == 0:
|
||||
choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason))
|
||||
choices.append(self._build_empty_choice(output_items, incomplete_finish_reason))
|
||||
else:
|
||||
for choice in choices:
|
||||
choice.finish_reason = incomplete_finish_reason
|
||||
|
|
@ -1802,7 +1975,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
# New output item added
|
||||
output_item = parsed_chunk.get("item", {})
|
||||
if output_item.get("type") in ("function_call", "custom_tool_call"):
|
||||
converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
|
||||
converted: Final = tool_call_dict_from_output_item(
|
||||
output_item, parsed_chunk.get("output_index", 0), include_namespace=True
|
||||
)
|
||||
provider_specific_fields: Final = converted.get("provider_specific_fields")
|
||||
|
||||
function_chunk: Final = ChatCompletionToolCallFunctionChunk(
|
||||
|
|
@ -1815,14 +1990,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index(
|
||||
tool_call_index_map, parsed_chunk.get("output_index", 0)
|
||||
)
|
||||
tool_call_chunk: Final = ChatCompletionToolCallChunk(
|
||||
tool_call_chunk: Final = _ChatToolCallDict(
|
||||
id=converted["id"],
|
||||
index=tool_call_index,
|
||||
type="function",
|
||||
function=function_chunk,
|
||||
**({"provider_specific_fields": provider_specific_fields} if provider_specific_fields else {}),
|
||||
)
|
||||
if provider_specific_fields:
|
||||
tool_call_chunk.provider_specific_fields = provider_specific_fields
|
||||
|
||||
return ModelResponseStream(
|
||||
choices=[
|
||||
|
|
@ -1878,7 +2052,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
delta=Delta(
|
||||
tool_calls=(
|
||||
tool_call_dict_from_output_item(
|
||||
output_item, parsed_chunk.get("output_index", 0)
|
||||
output_item, parsed_chunk.get("output_index", 0), include_namespace=True
|
||||
),
|
||||
)
|
||||
),
|
||||
|
|
@ -1973,6 +2147,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
delta=Delta(
|
||||
content="",
|
||||
reasoning_items=terminal_reasoning_items_typed,
|
||||
provider_specific_fields=_tool_search_provider_fields(
|
||||
_tool_search_items_from_output_items(cast(object, output_items)) # cast-ok: untyped
|
||||
),
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
)
|
||||
from litellm.types.llms.databricks import DatabricksTool
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionReasoningItem,
|
||||
ChatCompletionThinkingBlock,
|
||||
ImageURLListItem,
|
||||
OpenAIModerationResponse,
|
||||
|
|
@ -156,6 +157,21 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
|
|||
del choice.enhancements
|
||||
|
||||
|
||||
def _first_replay_delta(
|
||||
message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None
|
||||
) -> Delta:
|
||||
return Delta(
|
||||
content=message.get("content"),
|
||||
role=message["role"],
|
||||
function_call=message.get("function_call"),
|
||||
tool_calls=message.get("tool_calls"),
|
||||
reasoning_items=None if reasoning_items is None else list(reasoning_items),
|
||||
provider_specific_fields=_cached_provider_specific_fields(
|
||||
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _invalid_choices_message(response_object: Mapping[str, object]) -> str:
|
||||
raw_keys: Final = list(response_object.keys())
|
||||
if "choices" not in response_object:
|
||||
|
|
@ -217,15 +233,7 @@ async def convert_to_streaming_response_async(
|
|||
t["index"] = index
|
||||
pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t))
|
||||
choice["message"]["tool_calls"] = pydantic_tool_calls
|
||||
delta = Delta(
|
||||
content=choice["message"].get("content", None),
|
||||
role=choice["message"]["role"],
|
||||
function_call=choice["message"].get("function_call", None),
|
||||
tool_calls=choice["message"].get("tool_calls", None),
|
||||
provider_specific_fields=_cached_provider_specific_fields(
|
||||
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(choice["message"])
|
||||
),
|
||||
)
|
||||
delta = _first_replay_delta(choice["message"], choice["message"].get("reasoning_items"))
|
||||
finish_reason = choice.get("finish_reason", None)
|
||||
|
||||
if finish_reason is None:
|
||||
|
|
@ -754,6 +762,7 @@ def convert_to_model_response_object(
|
|||
provider_specific_fields=provider_specific_fields,
|
||||
reasoning_content=reasoning_content,
|
||||
thinking_blocks=thinking_blocks,
|
||||
reasoning_items=choice["message"].get("reasoning_items", None),
|
||||
annotations=choice["message"].get("annotations", None),
|
||||
images=_normalize_images_for_message(choice["message"].get("images", None)),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1807,18 +1807,62 @@ def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT:
|
|||
)
|
||||
|
||||
|
||||
LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"})
|
||||
LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS: Final = frozenset(
|
||||
{"thinking_blocks", "reasoning_items", "provider_specific_fields"}
|
||||
)
|
||||
LITELLM_INTERNAL_MESSAGE_FIELDS: Final = LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS | {"reasoning_content"}
|
||||
|
||||
|
||||
def _tool_call_without_provider_specific_fields(tool_call: object) -> object:
|
||||
if not isinstance(tool_call, dict):
|
||||
return tool_call
|
||||
fields: Final = cast(dict[str, object], tool_call) # cast-ok: untyped JSON tool call, keys forwarded verbatim
|
||||
function_fields: Final = (
|
||||
cast("dict[str, object]", fields.get("function")) # cast-ok: untyped JSON function chunk
|
||||
if isinstance(fields.get("function"), dict)
|
||||
else None
|
||||
)
|
||||
if "provider_specific_fields" not in fields and (
|
||||
function_fields is None or "provider_specific_fields" not in function_fields
|
||||
):
|
||||
return cast("dict[str, object]", tool_call) # cast-ok: untyped JSON tool call, no strip needed
|
||||
return {
|
||||
key: (
|
||||
{fkey: fvalue for fkey, fvalue in function_fields.items() if fkey != "provider_specific_fields"}
|
||||
if key == "function" and function_fields is not None
|
||||
else value
|
||||
)
|
||||
for key, value in fields.items()
|
||||
if key != "provider_specific_fields"
|
||||
}
|
||||
|
||||
|
||||
def _message_without_fields(message: AllMessageValues, dropped: frozenset[str]) -> AllMessageValues:
|
||||
kept: Final = {key: value for key, value in message.items() if key not in dropped}
|
||||
tool_calls: Final = message.get("tool_calls")
|
||||
if not isinstance(tool_calls, list):
|
||||
return cast(AllMessageValues, kept) # cast-ok: same TypedDict minus dropped keys
|
||||
return cast( # cast-ok: same TypedDict minus dropped keys, tool calls minus their provider_specific_fields
|
||||
AllMessageValues,
|
||||
{
|
||||
**kept,
|
||||
"tool_calls": [
|
||||
_tool_call_without_provider_specific_fields(tool_call)
|
||||
for tool_call in cast(list[object], tool_calls) # cast-ok: untyped JSON tool calls
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def strip_litellm_response_only_fields(message: AllMessageValues) -> AllMessageValues:
|
||||
"""Drop the fields litellm attaches to a returned message (e.g. Anthropic thinking blocks or Responses
|
||||
API items) that no OpenAI-compatible endpoint accepts back, so a reply can be sent to a strict one."""
|
||||
return _message_without_fields(message, LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS)
|
||||
|
||||
|
||||
def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessageValues:
|
||||
"""Drop the fields litellm attaches to assistant messages (e.g. when translating Anthropic thinking
|
||||
blocks) that OpenAI-compatible endpoints with strict schemas reject as extra inputs."""
|
||||
if LITELLM_INTERNAL_MESSAGE_FIELDS.isdisjoint(message):
|
||||
return message
|
||||
return cast( # cast-ok: same TypedDict minus internal keys
|
||||
AllMessageValues,
|
||||
{key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS},
|
||||
)
|
||||
"""``strip_litellm_response_only_fields`` plus ``reasoning_content``, for endpoints that reject it too."""
|
||||
return _message_without_fields(message, LITELLM_INTERNAL_MESSAGE_FIELDS)
|
||||
|
||||
|
||||
def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import base64
|
|||
import time
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from itertools import groupby
|
||||
from itertools import chain, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Optional, TypeAlias, TypedDict, Union, cast
|
||||
|
||||
|
|
@ -1330,10 +1330,29 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o
|
|||
_set_hidden_response_cost(response, float(computed_cost))
|
||||
|
||||
|
||||
def _field_of(value: object, name: str) -> object:
|
||||
if isinstance(value, dict):
|
||||
return cast(Mapping[str, object], value).get(name) # cast-ok: chunk dicts are untyped JSON
|
||||
return getattr(value, name, None)
|
||||
|
||||
|
||||
def _delta_reasoning_items(delta: object) -> Sequence[object]:
|
||||
items: Final = _field_of(delta, "reasoning_items")
|
||||
return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim
|
||||
|
||||
|
||||
def _chunk_reasoning_items(chunk: object) -> Sequence[object]:
|
||||
choices: Final = _field_of(chunk, "choices")
|
||||
if not isinstance(choices, list) or not choices:
|
||||
return ()
|
||||
return _delta_reasoning_items(_field_of(cast(Sequence[object], choices)[0], "delta")) # cast-ok: untyped chunk
|
||||
|
||||
|
||||
_NON_TEXT_DELTA_FIELDS: Final = (
|
||||
"tool_calls",
|
||||
"function_call",
|
||||
"reasoning_content",
|
||||
"reasoning_items",
|
||||
"thinking_blocks",
|
||||
"annotations",
|
||||
"audio",
|
||||
|
|
@ -1510,6 +1529,11 @@ def stream_chunk_builder(
|
|||
reasoning_chunks
|
||||
)
|
||||
|
||||
typed_chunks: Final = cast(Sequence[object], chunks) # cast-ok: chunks are untyped upstream
|
||||
reasoning_items: Final = tuple(chain.from_iterable(_chunk_reasoning_items(chunk) for chunk in typed_chunks))
|
||||
if reasoning_items:
|
||||
response["choices"][0]["message"]["reasoning_items"] = list(reasoning_items)
|
||||
|
||||
annotation_chunks: Final = [
|
||||
chunk
|
||||
for chunk in chunks
|
||||
|
|
|
|||
|
|
@ -13,7 +13,10 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
||||
get_response_headers,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_legacy_defs
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
strip_litellm_response_only_fields,
|
||||
unpack_legacy_defs,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -520,15 +523,8 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
|
|||
llm_provider="fireworks_ai",
|
||||
)
|
||||
filter_value_from_dict(cast(dict, message), "cache_control")
|
||||
# Remove fields not permitted by FireworksAI (additionalProperties: false
|
||||
# on their ChatMessage schema) that may cause:
|
||||
# "Extra inputs are not permitted, field: 'messages[n].<field>'"
|
||||
if isinstance(message, dict):
|
||||
m = cast(dict, message)
|
||||
m.pop("provider_specific_fields", None)
|
||||
m.pop("thinking_blocks", None)
|
||||
|
||||
return messages
|
||||
return [strip_litellm_response_only_fields(message) for message in messages]
|
||||
|
||||
# Cached index of fireworks_ai/* entries from litellm.model_cost. Building
|
||||
# this index requires a full scan of model_cost (tens of thousands of
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields
|
||||
from litellm.llms.openai.chat.gpt_transformation import (
|
||||
OpenAIChatCompletionStreamingHandler,
|
||||
)
|
||||
|
|
@ -152,7 +153,7 @@ class GroqChatConfig(OpenAILikeChatConfig):
|
|||
for k, v in _message.items():
|
||||
if v is not None:
|
||||
new_message[k] = v
|
||||
messages[idx] = new_message
|
||||
messages[idx] = strip_litellm_response_only_fields(new_message)
|
||||
|
||||
if is_async:
|
||||
return super()._transform_messages(messages=messages, model=model, is_async=True)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
handle_messages_with_content_list_to_str_conversion,
|
||||
strip_litellm_internal_message_fields,
|
||||
strip_none_values_from_message,
|
||||
)
|
||||
from litellm.llms.openai.chat.gpt_transformation import (
|
||||
|
|
@ -424,17 +425,14 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
@classmethod
|
||||
def _strip_output_only_fields(cls, message: AllMessageValues) -> AllMessageValues:
|
||||
"""
|
||||
``reasoning_content`` and ``thinking_blocks`` are output-only fields that
|
||||
LiteLLM attaches to assistant responses. Mistral's input schema forbids
|
||||
LiteLLM attaches output-only fields (``reasoning_content``, ``thinking_blocks``,
|
||||
``provider_specific_fields``) to assistant responses. Mistral's input schema forbids
|
||||
unknown fields, so replaying them verbatim in a follow-up turn triggers a
|
||||
422 ``extra_forbidden``. Drop them before the request is sent.
|
||||
"""
|
||||
if message["role"] != "assistant":
|
||||
return message
|
||||
return cast(
|
||||
AllMessageValues,
|
||||
{k: v for k, v in message.items() if k not in ("reasoning_content", "thinking_blocks")},
|
||||
)
|
||||
return strip_litellm_internal_message_fields(message)
|
||||
|
||||
@classmethod
|
||||
def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues:
|
||||
|
|
|
|||
|
|
@ -25,6 +25,9 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
replayed_tool_search_tool_names,
|
||||
)
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
|
|
@ -256,7 +259,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
return None if images else "no scannable content"
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
|
||||
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name,
|
||||
and the tools a replayed tool_search_output loads)."""
|
||||
names: Final[list[str]] = []
|
||||
for tool in data.get("tools") or []:
|
||||
if isinstance(tool, dict) and tool.get("type") == "function":
|
||||
|
|
@ -266,7 +270,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
for fn in data.get("functions") or []:
|
||||
if isinstance(fn, dict) and fn.get("name"):
|
||||
names.append(str(fn["name"]))
|
||||
return names
|
||||
return [*names, *replayed_tool_search_tool_names(cast(object, data))] # cast-ok: upcast, read field by field
|
||||
|
||||
def _extract_inputs(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ from types import MappingProxyType
|
|||
from typing import (
|
||||
Final,
|
||||
Literal,
|
||||
cast, # noqa: TID251 # rebuilding a TypedDict minus keys has no checked spelling
|
||||
overload,
|
||||
)
|
||||
|
||||
|
|
@ -18,6 +17,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields
|
||||
from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.utils import supports_function_calling, supports_reasoning, supports_response_schema
|
||||
|
|
@ -25,7 +25,6 @@ from litellm.utils import supports_function_calling, supports_reasoning, support
|
|||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call")
|
||||
LITELLM_INTERNAL_ASSISTANT_FIELDS: Final = frozenset({"thinking_blocks", "provider_specific_fields"})
|
||||
FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling"
|
||||
STRUCTURED_OUTPUTS_DOCS_URL: Final = "https://docs.together.ai/docs/inference/chat/structured-outputs"
|
||||
|
||||
|
|
@ -174,12 +173,9 @@ def _drop_response_format(passed_params: Container[str], model: str, drop_params
|
|||
|
||||
|
||||
def _without_litellm_internal_fields(message: AllMessageValues) -> AllMessageValues:
|
||||
if message["role"] != "assistant" or LITELLM_INTERNAL_ASSISTANT_FIELDS.isdisjoint(message):
|
||||
if message["role"] != "assistant":
|
||||
return message
|
||||
return cast( # cast-ok: rebuilding the same TypedDict minus internal keys loses the narrowed type
|
||||
"AllMessageValues",
|
||||
{key: value for key, value in message.items() if key not in LITELLM_INTERNAL_ASSISTANT_FIELDS},
|
||||
)
|
||||
return strip_litellm_response_only_fields(message)
|
||||
|
||||
|
||||
class TogetherAIChatConfig(OpenAIGPTConfig):
|
||||
|
|
|
|||
|
|
@ -4,11 +4,16 @@ from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence
|
|||
from typing import Any, Final, Literal, TypedDict
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import ReadOnly, Required
|
||||
|
||||
from litellm import ChatCompletionToolParam
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
replayed_tool_search_tools,
|
||||
without_tool_search_items_loading,
|
||||
)
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
|
|
@ -40,6 +45,7 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
GUARDRAIL_NAME: Final = "tool_permission"
|
||||
_REQUEST_DATA: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
def _object_mapping(value: object) -> Mapping[str, object] | None:
|
||||
|
|
@ -52,6 +58,12 @@ def _object_list(value: object) -> Sequence[object] | None:
|
|||
return value if isinstance(value, list) else None
|
||||
|
||||
|
||||
def _function_tool_name(tool: object) -> str | None:
|
||||
fields: Final = _object_mapping(tool)
|
||||
name: Final = fields.get("name") if fields is not None and fields.get("type") == "function" else None
|
||||
return name if isinstance(name, str) else None
|
||||
|
||||
|
||||
class _ToolPermissionRuleFields(TypedDict, total=False):
|
||||
"""The config-file shape a :class:`ToolPermissionRule` is built from."""
|
||||
|
||||
|
|
@ -653,6 +665,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
if forced_tool_name is not None:
|
||||
request_tools.append((forced_tool_name, "function"))
|
||||
|
||||
replayed: Final = map(_function_tool_name, replayed_tool_search_tools(_REQUEST_DATA.validate_python(data)))
|
||||
request_tools.extend((name, "function") for name in replayed if name is not None)
|
||||
|
||||
return request_tools
|
||||
|
||||
def _modify_request_with_permission_errors(
|
||||
|
|
@ -701,6 +716,12 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
if named_function_call in error_tool_names:
|
||||
data["function_call"] = "none"
|
||||
|
||||
messages: Final = without_tool_search_items_loading(
|
||||
_REQUEST_DATA.validate_python(data), frozenset(denied_tool_names)
|
||||
)
|
||||
if messages is not None:
|
||||
data["messages"] = messages # rebind-ok: rewrite mode edits the request in place
|
||||
|
||||
return data
|
||||
|
||||
def _create_permission_error_result(
|
||||
|
|
|
|||
|
|
@ -23,9 +23,11 @@ Configuration in proxy config YAML:
|
|||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import TypedDict, Unpack
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import replayed_tool_search_tool_names
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
|
|
@ -38,6 +40,7 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
GUARDRAIL_NAME: Final = "tool_policy"
|
||||
_REQUEST_DATA: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||
|
|
@ -133,15 +136,18 @@ class ToolPolicyGuardrail(CustomGuardrail):
|
|||
"""
|
||||
if input_type == "request":
|
||||
tools: Final = inputs.get("tools") or []
|
||||
tool_names = [
|
||||
function_tool_names: Final = [
|
||||
t["function"]["name"]
|
||||
for t in tools
|
||||
if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name")
|
||||
]
|
||||
if not tool_names:
|
||||
route: Final = _get_request_route_from_data(request_data)
|
||||
if route:
|
||||
tool_names = extract_request_tool_names(route, request_data)
|
||||
route: Final = _get_request_route_from_data(request_data)
|
||||
replayed_tool_names: Final = replayed_tool_search_tool_names(_REQUEST_DATA.validate_python(request_data))
|
||||
tool_names = (
|
||||
extract_request_tool_names(route, request_data)
|
||||
if route and not function_tool_names
|
||||
else [*function_tool_names, *replayed_tool_names]
|
||||
)
|
||||
else:
|
||||
tool_calls: Final = inputs.get("tool_calls") or []
|
||||
tool_names = []
|
||||
|
|
|
|||
|
|
@ -2868,13 +2868,18 @@ class LiteLLMCompletionResponsesConfig:
|
|||
else _attribute_fields(provider_fields)
|
||||
)
|
||||
|
||||
namespace: Final = getattr(tool_call_item, "namespace", None)
|
||||
merged_provider_fields: Final = {
|
||||
**(provider_specific_fields if isinstance(provider_specific_fields, dict) else {}),
|
||||
**({"namespace": namespace} if isinstance(namespace, str) else {}),
|
||||
}
|
||||
function_dict: Final[dict[str, object]] = {
|
||||
"name": item.name,
|
||||
"arguments": item.arguments,
|
||||
}
|
||||
|
||||
if provider_specific_fields:
|
||||
function_dict["provider_specific_fields"] = provider_specific_fields
|
||||
if merged_provider_fields:
|
||||
function_dict["provider_specific_fields"] = merged_provider_fields
|
||||
|
||||
tool_call_dict: Final[dict[str, object]] = {
|
||||
"id": LiteLLMCompletionResponsesConfig.tool_call_id_from_responses_item(
|
||||
|
|
@ -2886,8 +2891,8 @@ class LiteLLMCompletionResponsesConfig:
|
|||
"index": index,
|
||||
}
|
||||
|
||||
if provider_specific_fields:
|
||||
tool_call_dict["provider_specific_fields"] = provider_specific_fields
|
||||
if merged_provider_fields:
|
||||
tool_call_dict["provider_specific_fields"] = merged_provider_fields
|
||||
|
||||
return tool_call_dict
|
||||
|
||||
|
|
|
|||
|
|
@ -46,19 +46,24 @@ def _expect_model_response(response: object) -> ModelResponse:
|
|||
|
||||
|
||||
def _assistant_tool_call(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionAssistantToolCall:
|
||||
return ChatCompletionAssistantToolCall(
|
||||
id=tool_call.id,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=tool_call.function.name, arguments=tool_call.function.arguments
|
||||
),
|
||||
provider_fields: Final[object] = getattr(tool_call, "provider_specific_fields", None)
|
||||
return cast( # cast-ok: dict literal with provider_specific_fields spread in only when set
|
||||
"ChatCompletionAssistantToolCall",
|
||||
{
|
||||
"id": tool_call.id,
|
||||
"type": "function",
|
||||
"function": ChatCompletionToolCallFunctionChunk(
|
||||
name=tool_call.function.name, arguments=tool_call.function.arguments
|
||||
),
|
||||
**({"provider_specific_fields": provider_fields} if provider_fields else {}),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _assistant_message(
|
||||
message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], replays_server_tool_results: bool
|
||||
message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], replays_provider_specific_fields: bool
|
||||
) -> ChatCompletionAssistantMessage:
|
||||
return cast( # cast-ok: dict literal with thinking_blocks, reasoning_items and provider_specific_fields spread in only when set
|
||||
return cast( # cast-ok: dict literal with the optional message fields spread in only when set
|
||||
"ChatCompletionAssistantMessage",
|
||||
{
|
||||
"role": "assistant",
|
||||
|
|
@ -71,7 +76,9 @@ def _assistant_message(
|
|||
("reasoning_items", getattr(message, "reasoning_items", None)),
|
||||
(
|
||||
"provider_specific_fields",
|
||||
getattr(message, "provider_specific_fields", None) if replays_server_tool_results else None,
|
||||
getattr(message, "provider_specific_fields", None)
|
||||
if replays_provider_specific_fields
|
||||
else None,
|
||||
),
|
||||
)
|
||||
if value is not None
|
||||
|
|
@ -105,6 +112,18 @@ def _client_tool_calls(
|
|||
)
|
||||
|
||||
|
||||
def _replays_provider_specific_fields(
|
||||
message: Message,
|
||||
tool_calls: tuple[ChatCompletionMessageToolCall, ...],
|
||||
client_tool_calls: tuple[ChatCompletionMessageToolCall, ...],
|
||||
) -> bool:
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import carries_tool_search_items
|
||||
|
||||
return client_tool_calls != tool_calls or carries_tool_search_items(
|
||||
getattr(message, "provider_specific_fields", None)
|
||||
)
|
||||
|
||||
|
||||
def run_tool_loop(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -148,7 +167,9 @@ def run_tool_loop(
|
|||
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),
|
||||
_assistant_message(
|
||||
message, tool_calls, _replays_provider_specific_fields(message, tool_calls, client_tool_calls)
|
||||
),
|
||||
*tool_results,
|
||||
)
|
||||
raise ToolLoopMaxRoundsExceeded(max_rounds)
|
||||
|
|
@ -197,7 +218,9 @@ async def arun_tool_loop(
|
|||
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),
|
||||
_assistant_message(
|
||||
message, tool_calls, _replays_provider_specific_fields(message, tool_calls, client_tool_calls)
|
||||
),
|
||||
*tool_results,
|
||||
)
|
||||
raise ToolLoopMaxRoundsExceeded(max_rounds)
|
||||
|
|
|
|||
340
tests/integration/_support/tool_search.py
Normal file
340
tests/integration/_support/tool_search.py
Normal file
|
|
@ -0,0 +1,340 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.client import string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.openai_wire import answering_model_discovery
|
||||
from integration._support.wire import Reply, Request
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
WEATHER: Final = "get_weather"
|
||||
TIME: Final = "get_time"
|
||||
ONLY_SEARCH: Final = "only-search"
|
||||
TWO_SEARCHES: Final = "two-searches"
|
||||
REJECT_REPLAY: Final = "reject-replay"
|
||||
REJECTED: Final = "scripted: the replayed tool_search_output was rejected"
|
||||
TOOL_SEARCH: Final[dict[str, JsonValue]] = {"type": "tool_search"}
|
||||
PARAMETERS: Final[dict[str, JsonValue]] = {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
}
|
||||
REPLAY_TYPES: Final = frozenset({"reasoning", "tool_search_call", "tool_search_output", "function_call"})
|
||||
|
||||
|
||||
def function(name: str) -> dict[str, JsonValue]:
|
||||
return {"name": name, "description": f"Look up {name}", "parameters": PARAMETERS}
|
||||
|
||||
|
||||
def deferred(name: str) -> dict[str, JsonValue]:
|
||||
return {"type": "function", "function": function(name), "defer_loading": True}
|
||||
|
||||
|
||||
def chat_tools(*names: str) -> list[JsonValue]:
|
||||
return [TOOL_SEARCH, *(deferred(name) for name in names)]
|
||||
|
||||
|
||||
def loaded(name: str) -> dict[str, JsonValue]:
|
||||
return {"type": "function", **function(name), "strict": False, "defer_loading": True}
|
||||
|
||||
|
||||
def namespace(name: str, *members: str) -> dict[str, JsonValue]:
|
||||
return {"type": "namespace", "name": name, "description": f"{name} tools", "tools": [loaded(m) for m in members]}
|
||||
|
||||
|
||||
def reasoning(marker: str, label: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"rs_{label}_{marker}",
|
||||
"type": "reasoning",
|
||||
"summary": [],
|
||||
"encrypted_content": f"enc-{label}-{marker}",
|
||||
}
|
||||
|
||||
|
||||
def search_call(marker: str, label: str = "a") -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"tsc_{label}_{marker}",
|
||||
"type": "tool_search_call",
|
||||
"arguments": {"query": f"tools {label}"},
|
||||
"execution": "server",
|
||||
"status": "completed",
|
||||
}
|
||||
|
||||
|
||||
def search_output(marker: str, tools: Sequence[JsonValue], label: str = "a") -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"tso_{label}_{marker}",
|
||||
"type": "tool_search_output",
|
||||
"execution": "server",
|
||||
"status": "completed",
|
||||
"tools": list(tools),
|
||||
}
|
||||
|
||||
|
||||
def call_id(marker: str, name: str = WEATHER) -> str:
|
||||
return f"call_{name}_{marker}"
|
||||
|
||||
|
||||
def replayed_call(marker: str, name: str = WEATHER) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"type": "function_call",
|
||||
"call_id": call_id(marker, name),
|
||||
"name": name,
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"namespace": name,
|
||||
}
|
||||
|
||||
|
||||
def call_item(marker: str, name: str = WEATHER) -> dict[str, JsonValue]:
|
||||
return {**replayed_call(marker, name), "id": f"fc_{name}_{marker}", "status": "completed"}
|
||||
|
||||
|
||||
def searched(marker: str, label: str, name: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return (reasoning(marker, label), search_call(marker, label), search_output(marker, [loaded(name)], label))
|
||||
|
||||
|
||||
def turn_one_output(marker: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return (*searched(marker, "a", WEATHER), reasoning(marker, "b"), call_item(marker))
|
||||
|
||||
|
||||
def turn_two_output(marker: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return (*searched(marker, "c", TIME), call_item(marker, TIME))
|
||||
|
||||
|
||||
def only_search_output(marker: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return searched(marker, "a", WEATHER)
|
||||
|
||||
|
||||
def on_the_wire(item: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {**item, "created_by": None} if item.get("type") == "tool_search_call" else dict(item)
|
||||
|
||||
|
||||
def _message_item(identity: str, text: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"msg_{identity}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": text, "annotations": []}],
|
||||
}
|
||||
|
||||
|
||||
def _response(identity: str, model: str, output: Sequence[JsonValue], status: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": status,
|
||||
"model": model,
|
||||
"output": list(output),
|
||||
"usage": rv.USAGE,
|
||||
}
|
||||
|
||||
|
||||
def _item_events(index: int, item: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
|
||||
added: Final[dict[str, JsonValue]] = {"type": "response.output_item.added", "output_index": index, "item": {**item}}
|
||||
done: Final[dict[str, JsonValue]] = {"type": "response.output_item.done", "output_index": index, "item": {**item}}
|
||||
if item.get("type") == "function_call":
|
||||
return (
|
||||
{**added, "item": {**item, "arguments": "", "status": "in_progress"}},
|
||||
{
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"output_index": index,
|
||||
"item_id": item["id"],
|
||||
"delta": item["arguments"],
|
||||
},
|
||||
done,
|
||||
)
|
||||
if item.get("type") == "message":
|
||||
contents: Final = rv.ITEMS.validate_python(item["content"])
|
||||
return (
|
||||
{**added, "item": {**item, "content": [], "status": "in_progress"}},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"output_index": index,
|
||||
"item_id": item["id"],
|
||||
"content_index": 0,
|
||||
"delta": contents[0]["text"],
|
||||
},
|
||||
done,
|
||||
)
|
||||
return (added, done)
|
||||
|
||||
|
||||
def stream_chunks(identity: str, model: str, output: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]:
|
||||
created: Final[dict[str, JsonValue]] = {
|
||||
"type": "response.created",
|
||||
"response": _response(identity, model, [], "in_progress"),
|
||||
}
|
||||
items: Final = tuple(chain.from_iterable(_item_events(index, item) for index, item in enumerate(output)))
|
||||
completed: Final[dict[str, JsonValue]] = {
|
||||
"type": "response.completed",
|
||||
"response": _response(identity, model, [dict(item) for item in output], "completed"),
|
||||
}
|
||||
events: Final = (created, *items, completed)
|
||||
return tuple(rv.sse({**event, "sequence_number": number}) for number, event in enumerate(events))
|
||||
|
||||
|
||||
def reply(
|
||||
identity: str, model: str, output: Sequence[Mapping[str, JsonValue]], *, stream: bool, pause: float = 0
|
||||
) -> Reply:
|
||||
if stream:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=stream_chunks(identity, model, output),
|
||||
pause_between_chunks=pause,
|
||||
)
|
||||
return Reply(body=json.dumps(_response(identity, model, [dict(item) for item in output], "completed")).encode())
|
||||
|
||||
|
||||
def _count(items: Sequence[Mapping[str, JsonValue]], kind: str) -> int:
|
||||
return sum(1 for item in items if item.get("type") == kind)
|
||||
|
||||
|
||||
def _scripted_output(text: str, marker: str, items: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]:
|
||||
results: Final = _count(items, "function_call_output")
|
||||
if TWO_SEARCHES in text and results == 1:
|
||||
return tuple(on_the_wire(item) for item in turn_two_output(marker))
|
||||
if results or (ONLY_SEARCH in text and _count(items, "tool_search_output")):
|
||||
return (_message_item(f"{marker}-{results}", rv.answer(marker)),)
|
||||
chosen: Final = only_search_output(marker) if ONLY_SEARCH in text else turn_one_output(marker)
|
||||
return tuple(on_the_wire(item) for item in chosen)
|
||||
|
||||
|
||||
def scripted_with_pause(pause: float) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
body: Final = pcb.body_of(request)
|
||||
text: Final = request.body.decode()
|
||||
marker: Final = rv.newest_marker(text) or uuid.uuid4().hex
|
||||
items: Final = pcb.input_items(request)
|
||||
if REJECT_REPLAY in text and _count(items, "tool_search_output"):
|
||||
return rv.error(400, REJECTED, "invalid_request_error")
|
||||
return reply(
|
||||
f"resp_{marker}-{uuid.uuid4().hex}",
|
||||
string_value(body["model"]),
|
||||
[rv.JSON_OBJECT.validate_python(item) for item in _scripted_output(text, marker, items)],
|
||||
stream=body.get("stream") is True,
|
||||
pause=pause,
|
||||
)
|
||||
|
||||
return answering_model_discovery(respond)
|
||||
|
||||
|
||||
scripted: Final = scripted_with_pause(0)
|
||||
|
||||
|
||||
def replayed(request: Request) -> list[dict[str, JsonValue]]:
|
||||
return [item for item in pcb.input_items(request) if item.get("type") in REPLAY_TYPES]
|
||||
|
||||
|
||||
def first_message(text: str) -> dict[str, JsonValue]:
|
||||
choices: Final = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(text)["choices"])
|
||||
return rv.JSON_OBJECT.validate_python(choices[0]["message"])
|
||||
|
||||
|
||||
def provider_fields(value: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
fields: Final = value.get("provider_specific_fields")
|
||||
return rv.JSON_OBJECT.validate_python(fields) if isinstance(fields, dict) else {}
|
||||
|
||||
|
||||
def tool_calls(message: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]:
|
||||
return rv.ITEMS.validate_python(message.get("tool_calls") or [])
|
||||
|
||||
|
||||
def namespaces(message: Mapping[str, JsonValue]) -> list[JsonValue]:
|
||||
return [provider_fields(call).get("namespace") for call in tool_calls(message)]
|
||||
|
||||
|
||||
def chunks_of(text: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
rv.JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: {")
|
||||
)
|
||||
|
||||
|
||||
def _deltas(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]:
|
||||
choices: Final = tuple(rv.ITEMS.validate_python(chunk.get("choices") or []) for chunk in chunks)
|
||||
return tuple(rv.JSON_OBJECT.validate_python(found[0].get("delta") or {}) for found in choices if found)
|
||||
|
||||
|
||||
def _merged(left: Mapping[str, JsonValue], right: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {**left, **{key: value for key, value in right.items() if value is not None}}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CallParts:
|
||||
index: int
|
||||
id: str | None
|
||||
name: str | None
|
||||
arguments: str
|
||||
fields: Mapping[str, JsonValue]
|
||||
|
||||
|
||||
def _call_parts(delta: Mapping[str, JsonValue]) -> tuple[_CallParts, ...]:
|
||||
calls: Final = rv.ITEMS.validate_python(delta.get("tool_calls") or [])
|
||||
functions: Final = tuple(rv.JSON_OBJECT.validate_python(call.get("function") or {}) for call in calls)
|
||||
return tuple(
|
||||
_CallParts(
|
||||
index=int(str(call.get("index") or 0)),
|
||||
id=call["id"] if isinstance(call.get("id"), str) else None,
|
||||
name=fn["name"] if isinstance(fn.get("name"), str) else None,
|
||||
arguments=str(fn.get("arguments") or ""),
|
||||
fields=_merged(provider_fields(fn), provider_fields(call)),
|
||||
)
|
||||
for call, fn in zip(calls, functions, strict=True)
|
||||
)
|
||||
|
||||
|
||||
def _rebuilt_call(parts: Sequence[_CallParts]) -> dict[str, JsonValue]:
|
||||
fields: Final = reduce(_merged, (part.fields for part in parts), {})
|
||||
ids: Final = tuple(part.id for part in parts if part.id)
|
||||
names: Final = tuple(part.name for part in parts if part.name)
|
||||
return {
|
||||
"id": ids[0],
|
||||
"type": "function",
|
||||
"function": {"name": names[0], "arguments": "".join(part.arguments for part in parts)},
|
||||
**({"provider_specific_fields": fields} if fields else {}),
|
||||
}
|
||||
|
||||
|
||||
def assistant_from_stream(chunks: Sequence[Mapping[str, JsonValue]]) -> dict[str, JsonValue]:
|
||||
deltas: Final = _deltas(chunks)
|
||||
parts: Final = tuple(chain.from_iterable(_call_parts(delta) for delta in deltas))
|
||||
indexes: Final = sorted({part.index for part in parts})
|
||||
calls: Final[list[JsonValue]] = [
|
||||
_rebuilt_call([part for part in parts if part.index == index]) for index in indexes
|
||||
]
|
||||
fields: Final = reduce(_merged, (provider_fields(delta) for delta in deltas), {})
|
||||
reasoning_items: Final[list[JsonValue]] = list(
|
||||
chain.from_iterable(rv.ITEMS.validate_python(delta.get("reasoning_items") or []) for delta in deltas)
|
||||
)
|
||||
content: Final = "".join(str(delta.get("content") or "") for delta in deltas)
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": content or None,
|
||||
**({"tool_calls": calls} if calls else {}),
|
||||
**({"provider_specific_fields": fields} if fields else {}),
|
||||
**({"reasoning_items": reasoning_items} if reasoning_items else {}),
|
||||
}
|
||||
|
||||
|
||||
def tool_result(marker: str, name: str = WEATHER) -> dict[str, JsonValue]:
|
||||
return {"role": "tool", "tool_call_id": call_id(marker, name), "content": "sunny"}
|
||||
|
||||
|
||||
def spend_rows_of(model: str, response_id: str) -> list[dict[str, JsonValue]]:
|
||||
rows: Final = read_rows(
|
||||
'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)
|
||||
)
|
||||
return [row for row in rows if rv.same_response(string_value(row["request_id"]), response_id)]
|
||||
|
|
@ -0,0 +1,166 @@
|
|||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support import tool_search as ts
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.openai_wire import answering_model_discovery, chat_reply
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
from tests.integration.caching.response_cache_case import await_entry, redis_store
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(120)
|
||||
|
||||
|
||||
def _shell() -> str:
|
||||
return f"run_shell_{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
def _replay_body(model: str, marker: str, loads: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]:
|
||||
call: Final[dict[str, JsonValue]] = {
|
||||
"id": ts.call_id(marker),
|
||||
"type": "function",
|
||||
"function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": ts.WEATHER},
|
||||
}
|
||||
items: Final[list[JsonValue]] = [ts.search_call(marker), ts.search_output(marker, loads)]
|
||||
return {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "user", "content": pcb.prompt(marker)},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [call],
|
||||
"provider_specific_fields": {"tool_search_items": items},
|
||||
},
|
||||
ts.tool_result(marker),
|
||||
],
|
||||
"tools": ts.chat_tools(ts.WEATHER),
|
||||
**pcb.NO_CACHE,
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
def _bridge(scenario: Scenario, wire: Wire, **parameters: JsonValue) -> str:
|
||||
return scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1", **parameters)
|
||||
|
||||
|
||||
def _posts(wire: Wire, marker: str) -> int:
|
||||
return len(pcb.with_marker(pcb.drained_posts(wire), marker))
|
||||
|
||||
|
||||
def _assert_denied(response: httpx.Response, wire: Wire, marker: str, tool: str) -> None:
|
||||
assert response.status_code == 403, response.text
|
||||
error: Final = rv.JSON_OBJECT.validate_python(rv.JSON_OBJECT.validate_json(response.text)["error"])
|
||||
assert error["type"] == "tool_access_denied", response.text
|
||||
assert tool in str(error["message"]) and ts.WEATHER not in str(error["message"]), response.text
|
||||
assert _posts(wire, marker) == 0, "a denied replay reached the upstream"
|
||||
|
||||
|
||||
def _assert_served(response: httpx.Response, wire: Wire, marker: str) -> None:
|
||||
assert response.status_code == 200, response.text
|
||||
assert ts.first_message(response.text)["content"] == rv.answer(marker), response.text
|
||||
assert _posts(wire, marker) == 1
|
||||
|
||||
|
||||
def test_key_without_an_allowlist_still_replays_any_loaded_tool(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bridge(scenario, wire)
|
||||
key: Final = scenario.key(models=[model])
|
||||
body: Final = _replay_body(model, marker, [ts.loaded(ts.WEATHER), ts.loaded(_shell())])
|
||||
_assert_served(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("placement", ["tools", "extra_body", "deployment", "absent"])
|
||||
def test_allowlisted_key_cannot_replay_a_search_that_loaded_another_tool(gateway: Gateway, placement: str) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
shell: Final = _shell()
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bridge(scenario, wire, **({"tools": [ts.TOOL_SEARCH]} if placement == "deployment" else {}))
|
||||
key: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]})
|
||||
loads: Final[list[JsonValue]] = [ts.loaded(ts.WEATHER), ts.loaded(shell)]
|
||||
overrides: Final[dict[str, dict[str, JsonValue]]] = {
|
||||
"tools": {},
|
||||
"extra_body": {
|
||||
"tools": [ts.deferred(ts.WEATHER)],
|
||||
"extra_body": {"tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)]},
|
||||
},
|
||||
"deployment": {"tools": [ts.deferred(ts.WEATHER)]},
|
||||
"absent": {"tools": [{"type": "function", "function": ts.function(ts.WEATHER)}]},
|
||||
}
|
||||
body: Final = {**_replay_body(model, marker, loads), **overrides[placement]}
|
||||
_assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker, shell)
|
||||
|
||||
|
||||
def test_allowlisted_key_cannot_replay_a_denied_namespace_member(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
shell: Final = _shell()
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bridge(scenario, wire)
|
||||
key: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]})
|
||||
body: Final = _replay_body(model, marker, [ts.namespace("ops", ts.WEATHER, shell)])
|
||||
_assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker, shell)
|
||||
|
||||
|
||||
def _chat_upstream(request: Request) -> Reply:
|
||||
marker: Final = rv.newest_marker(request.body.decode())
|
||||
return chat_reply(f"chatcmpl-{marker}", "gpt-4o-mini", rv.answer(marker), stream=False)
|
||||
|
||||
|
||||
def test_plain_chat_deployment_also_counts_replayed_search_loads(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
shell: Final = _shell()
|
||||
with wire_server(answering_model_discovery(_chat_upstream)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=f"{wire.url}/v1")
|
||||
key: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]})
|
||||
body: Final = _replay_body(model, marker, [ts.loaded(shell)])
|
||||
_assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker, shell)
|
||||
|
||||
|
||||
def test_team_allowlist_applies_and_a_key_allowlist_overrides_it(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
shell: Final = _shell()
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bridge(scenario, wire)
|
||||
team: Final = scenario.team(models=[model], metadata={"allowed_tools": [ts.WEATHER]})
|
||||
team_key: Final = scenario.key(team_id=team, models=[model])
|
||||
wider_key: Final = scenario.key(team_id=team, models=[model], metadata={"allowed_tools": [ts.WEATHER, shell]})
|
||||
body: Final = _replay_body(model, marker, [ts.loaded(ts.WEATHER), ts.loaded(shell)])
|
||||
_assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=team_key), wire, marker, shell)
|
||||
_assert_served(gateway.request("POST", "/v1/chat/completions", body, key=wider_key), wire, marker)
|
||||
|
||||
|
||||
def test_allowed_replay_is_served_and_its_cache_hit_still_checks_the_allowlist(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
shell: Final = _shell()
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bridge(scenario, wire)
|
||||
allowed: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]})
|
||||
cached: Final = {
|
||||
key: value
|
||||
for key, value in _replay_body(model, marker, [ts.loaded(ts.WEATHER)]).items()
|
||||
if key != "cache"
|
||||
}
|
||||
first: Final = gateway.request("POST", "/v1/chat/completions", cached, key=allowed)
|
||||
_assert_served(first, wire, marker)
|
||||
await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(first.text)["id"]))
|
||||
hit: Final = gateway.request("POST", "/v1/chat/completions", cached, key=allowed)
|
||||
assert hit.status_code == 200 and hit.headers.get("x-litellm-cache-key"), (hit.text, dict(hit.headers))
|
||||
assert ts.first_message(hit.text)["content"] == rv.answer(marker), hit.text
|
||||
shell_body: Final = {
|
||||
key: value
|
||||
for key, value in _replay_body(model, marker, [ts.loaded(shell)]).items()
|
||||
if key != "cache"
|
||||
}
|
||||
populated: Final = gateway.request("POST", "/v1/chat/completions", shell_body)
|
||||
assert populated.status_code == 200, populated.text
|
||||
await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(populated.text)["id"]))
|
||||
assert _posts(wire, marker) == 1
|
||||
_assert_denied(gateway.request("POST", "/v1/chat/completions", shell_body, key=allowed), wire, marker, shell)
|
||||
|
|
@ -0,0 +1,276 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support import tool_search as ts
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.database import write_rows
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.wire import Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
from tests.integration.caching.response_cache_case import await_entry, redis_store
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_POLICY: Final = "integration-tool-policy"
|
||||
_BLOCK: Final = "tool-permission-block"
|
||||
_REWRITE: Final = "tool-permission-rewrite"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Guarded:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
model: str
|
||||
shell: str
|
||||
blocked: str
|
||||
|
||||
|
||||
def _permission(name: str, shell: str, action: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "tool_permission",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"default_action": "allow",
|
||||
"on_disallowed_action": action,
|
||||
"rules": [{"id": "deny_shell", "tool_name": f"^{shell}$", "decision": "deny"}],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _config(directory: Path, model: str, upstream: str, shell: str) -> Path:
|
||||
path: Final = directory / "tool_search_guardrails.yaml"
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": model,
|
||||
"litellm_params": {"model": pcb.MODEL, "api_base": f"{upstream}/v1", "api_key": "sk-upstream"},
|
||||
}
|
||||
],
|
||||
"general_settings": {
|
||||
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
||||
"database_url": "os.environ/DATABASE_URL",
|
||||
"store_model_in_db": True,
|
||||
},
|
||||
"router_settings": {"disable_cooldowns": True},
|
||||
"litellm_settings": {
|
||||
"cache": True,
|
||||
"cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"},
|
||||
},
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": _POLICY,
|
||||
"litellm_params": {"guardrail": "tool_policy", "mode": "pre_call", "default_on": True},
|
||||
},
|
||||
_permission(_BLOCK, shell, "block"),
|
||||
_permission(_REWRITE, shell, "rewrite"),
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def guarded(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Guarded]:
|
||||
suffix: Final = uuid.uuid4().hex[:12]
|
||||
shell: Final = f"run_shell_{suffix}"
|
||||
blocked: Final = f"wipe_disk_{suffix}"
|
||||
model: Final = f"tool-search-guarded-{suffix}"
|
||||
policy: Final = session_gateway.request(
|
||||
"POST", "/v1/tool/policy", {"tool_name": blocked, "input_policy": "blocked"}
|
||||
)
|
||||
assert policy.status_code == 200, policy.text
|
||||
directory: Final = tmp_path_factory.mktemp("tool_search_guardrails")
|
||||
try:
|
||||
with (
|
||||
wire_server(ts.scripted) as wire,
|
||||
owned_proxy(session_gateway, directory, {}, config=_config(directory, model, wire.url, shell), workers=2)
|
||||
as proxy,
|
||||
):
|
||||
yield _Guarded(proxy, wire, model, shell, blocked)
|
||||
finally:
|
||||
write_rows('DELETE FROM "LiteLLM_ToolTable" WHERE tool_name = %s', (blocked,))
|
||||
|
||||
|
||||
def _call(marker: str, name: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": ts.call_id(marker, name),
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": name},
|
||||
}
|
||||
|
||||
|
||||
def _searched_turn(marker: str, name: str, label: str, loads: list[JsonValue]) -> list[JsonValue]:
|
||||
items: Final[list[JsonValue]] = [ts.search_call(marker, label), ts.search_output(marker, loads, label)]
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [_call(marker, name)],
|
||||
"provider_specific_fields": {"tool_search_items": items},
|
||||
},
|
||||
ts.tool_result(marker, name),
|
||||
]
|
||||
|
||||
|
||||
def _two_turn_body(guarded: _Guarded, marker: str, **fields: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": guarded.model,
|
||||
"messages": [
|
||||
{"role": "user", "content": pcb.prompt(marker)},
|
||||
*_searched_turn(marker, ts.WEATHER, "a", [ts.loaded(ts.WEATHER)]),
|
||||
*_searched_turn(marker, ts.TIME, "b", [ts.loaded(ts.TIME), ts.loaded(guarded.shell)]),
|
||||
],
|
||||
"tools": ts.chat_tools(ts.WEATHER, ts.TIME),
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
def _posts(wire: Wire, marker: str) -> tuple[bytes, ...]:
|
||||
return tuple(request.body for request in pcb.with_marker(pcb.drained_posts(wire), marker))
|
||||
|
||||
|
||||
def _assert_refused(response: httpx.Response, wire: Wire, marker: str, message: str, tool: str) -> None:
|
||||
assert response.status_code == 400, response.text
|
||||
assert message in response.text and tool in response.text, response.text
|
||||
assert _posts(wire, marker) == (), "a refused replay reached the upstream"
|
||||
|
||||
|
||||
def _one_search_body(
|
||||
guarded: _Guarded, marker: str, loads: list[JsonValue], **fields: JsonValue
|
||||
) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": guarded.model,
|
||||
"messages": [{"role": "user", "content": pcb.prompt(marker)}, *_searched_turn(marker, ts.WEATHER, "a", loads)],
|
||||
"tools": ts.chat_tools(ts.WEATHER),
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
def _blocked_load_body(guarded: _Guarded, marker: str) -> dict[str, JsonValue]:
|
||||
return _one_search_body(guarded, marker, [ts.loaded(ts.WEATHER), ts.loaded(guarded.blocked)])
|
||||
|
||||
|
||||
def _cache(guarded: _Guarded, marker: str, body: dict[str, JsonValue], key: str | None = None) -> None:
|
||||
populated: Final = guarded.proxy.request("POST", "/v1/chat/completions", body, key=key)
|
||||
assert populated.status_code == 200, populated.text
|
||||
await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(populated.text)["id"]))
|
||||
hit: Final = guarded.proxy.request("POST", "/v1/chat/completions", body, key=key)
|
||||
assert hit.status_code == 200 and hit.headers.get("x-litellm-cache-key"), (hit.text, dict(hit.headers))
|
||||
assert ts.first_message(hit.text)["content"] == rv.answer(marker), hit.text
|
||||
assert len(_posts(guarded.wire, marker)) == 1
|
||||
|
||||
|
||||
def _assert_policy_refuses_the_replay(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
refused: Final = guarded.proxy.request("POST", "/v1/chat/completions", _blocked_load_body(guarded, marker))
|
||||
_assert_refused(refused, guarded.wire, marker, "Violated tool policy", guarded.blocked)
|
||||
detail: Final = rv.JSON_OBJECT.validate_python(rv.JSON_OBJECT.validate_json(refused.text)["error"])
|
||||
assert rv.JSON_OBJECT.validate_python(detail["provider_specific_fields"])["blocked_tools"] == [guarded.blocked]
|
||||
|
||||
|
||||
def test_tool_policy_blocks_a_replayed_search_that_loaded_a_blocked_tool(guarded: _Guarded) -> None:
|
||||
control_marker: Final = uuid.uuid4().hex
|
||||
control: Final = guarded.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": guarded.model,
|
||||
"messages": [{"role": "user", "content": pcb.prompt(control_marker)}],
|
||||
"tools": [{"type": "function", "function": ts.function(guarded.blocked)}],
|
||||
},
|
||||
)
|
||||
_assert_refused(control, guarded.wire, control_marker, "Violated tool policy", guarded.blocked)
|
||||
_assert_policy_refuses_the_replay(guarded)
|
||||
_assert_policy_refuses_the_replay(guarded)
|
||||
|
||||
|
||||
def test_tool_permission_block_mode_refuses_a_replayed_search_that_loaded_a_denied_tool(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = guarded.proxy.request(
|
||||
"POST", "/v1/chat/completions", _two_turn_body(guarded, marker, guardrails=[_BLOCK])
|
||||
)
|
||||
_assert_refused(response, guarded.wire, marker, "Violated guardrail policy", guarded.shell)
|
||||
|
||||
|
||||
def test_tool_permission_rewrite_mode_drops_only_the_search_that_loaded_the_denied_tool(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = guarded.proxy.request(
|
||||
"POST", "/v1/chat/completions", _two_turn_body(guarded, marker, guardrails=[_REWRITE])
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert ts.first_message(response.text)["content"] == rv.answer(marker), response.text
|
||||
(posted,) = pcb.with_marker(pcb.drained_posts(guarded.wire), marker)
|
||||
assert ts.replayed(posted) == [
|
||||
ts.search_call(marker, "a"),
|
||||
ts.search_output(marker, [ts.loaded(ts.WEATHER)], "a"),
|
||||
ts.replayed_call(marker, ts.WEATHER),
|
||||
ts.replayed_call(marker, ts.TIME),
|
||||
]
|
||||
assert guarded.shell not in posted.body.decode()
|
||||
|
||||
|
||||
def test_guardrails_leave_an_allowed_replay_untouched(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
body: Final = _one_search_body(guarded, marker, [ts.loaded(ts.WEATHER)], guardrails=[_BLOCK])
|
||||
response: Final = guarded.proxy.request("POST", "/v1/chat/completions", body)
|
||||
assert response.status_code == 200, response.text
|
||||
(posted,) = pcb.with_marker(pcb.drained_posts(guarded.wire), marker)
|
||||
assert ts.replayed(posted) == [
|
||||
ts.search_call(marker, "a"),
|
||||
ts.search_output(marker, [ts.loaded(ts.WEATHER)], "a"),
|
||||
ts.replayed_call(marker, ts.WEATHER),
|
||||
]
|
||||
|
||||
|
||||
def test_guardrails_serve_an_allowed_replay_from_the_cache(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_cache(guarded, marker, _one_search_body(guarded, marker, [ts.loaded(ts.WEATHER)], guardrails=[_BLOCK]))
|
||||
|
||||
|
||||
def test_tool_policy_refuses_a_replay_whose_answer_an_exempt_key_cached(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
body: Final = _blocked_load_body(guarded, marker)
|
||||
with guarded.proxy.scenario() as scenario:
|
||||
exempt: Final = scenario.key(models=[guarded.model], metadata={"disable_global_guardrails": True})
|
||||
_cache(guarded, marker, body, key=exempt)
|
||||
refused: Final = guarded.proxy.request("POST", "/v1/chat/completions", body)
|
||||
_assert_refused(refused, guarded.wire, marker, "Violated tool policy", guarded.blocked)
|
||||
|
||||
|
||||
def test_tool_permission_block_mode_refuses_a_replay_whose_answer_is_cached(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_cache(guarded, marker, _two_turn_body(guarded, marker))
|
||||
response: Final = guarded.proxy.request(
|
||||
"POST", "/v1/chat/completions", _two_turn_body(guarded, marker, guardrails=[_BLOCK])
|
||||
)
|
||||
_assert_refused(response, guarded.wire, marker, "Violated guardrail policy", guarded.shell)
|
||||
|
||||
|
||||
def test_tool_permission_rewrite_mode_is_not_served_the_cached_unrewritten_answer(guarded: _Guarded) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_cache(guarded, marker, _two_turn_body(guarded, marker))
|
||||
rewritten: Final = _two_turn_body(guarded, marker, guardrails=[_REWRITE])
|
||||
first: Final = guarded.proxy.request("POST", "/v1/chat/completions", rewritten)
|
||||
assert first.status_code == 200, first.text
|
||||
(posted,) = pcb.with_marker(pcb.drained_posts(guarded.wire), marker)
|
||||
assert guarded.shell not in posted.body.decode()
|
||||
await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(first.text)["id"]))
|
||||
hit: Final = guarded.proxy.request("POST", "/v1/chat/completions", rewritten)
|
||||
assert hit.status_code == 200 and hit.headers.get("x-litellm-cache-key"), (hit.text, dict(hit.headers))
|
||||
assert ts.first_message(hit.text)["content"] == rv.answer(marker), hit.text
|
||||
assert _posts(guarded.wire, marker) == ()
|
||||
|
|
@ -19,9 +19,12 @@ import pytest
|
|||
import yaml
|
||||
from _pytest.mark.structures import ParameterSet
|
||||
from integration._support import claude_code as cc
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support import tool_search as ts
|
||||
from integration._support.anthropic_thinking import BEDROCK_MODEL, EVENT_STREAM, Event, aws_chunks, sse_chunks
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, string_value
|
||||
from integration._support.openai_wire import answering_model_discovery
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
|
@ -1008,3 +1011,311 @@ def test_a_global_drop_params_drops_a_malformed_defer_loading_and_serves_the_sea
|
|||
("get_weather", True),
|
||||
("lookup_order", None),
|
||||
], tools
|
||||
|
||||
|
||||
_BRIDGE_OUTAGE: Final = "scripted: upstream outage"
|
||||
|
||||
BridgeEndpoint: TypeAlias = Literal["chat", "messages", "responses"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeCall:
|
||||
endpoint: BridgeEndpoint
|
||||
marker: str
|
||||
stream: bool
|
||||
replay: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeServed:
|
||||
call: _BridgeCall
|
||||
status: int
|
||||
text: str
|
||||
|
||||
|
||||
_BRIDGE_ENDPOINTS: Final[tuple[BridgeEndpoint, ...]] = ("chat", "chat", "messages", "responses")
|
||||
|
||||
|
||||
def _bridge_calls(
|
||||
count: int, endpoint: Callable[[int], BridgeEndpoint], stream: Callable[[int], bool]
|
||||
) -> tuple[_BridgeCall, ...]:
|
||||
return tuple(
|
||||
_BridgeCall(endpoint(index), uuid.uuid4().hex, stream(index), endpoint(index) != "messages" and index % 3 != 0)
|
||||
for index in range(count)
|
||||
)
|
||||
|
||||
|
||||
def _bridge_chat_messages(call: _BridgeCall) -> list[JsonValue]:
|
||||
user: Final[dict[str, JsonValue]] = {"role": "user", "content": pcb.prompt(call.marker)}
|
||||
if not call.replay:
|
||||
return [user]
|
||||
items: Final[list[JsonValue]] = [
|
||||
ts.search_call(call.marker),
|
||||
ts.search_output(call.marker, [ts.loaded(ts.WEATHER)]),
|
||||
]
|
||||
tool_call: Final[dict[str, JsonValue]] = {
|
||||
"id": ts.call_id(call.marker),
|
||||
"type": "function",
|
||||
"function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": ts.WEATHER},
|
||||
}
|
||||
return [
|
||||
user,
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [tool_call],
|
||||
"provider_specific_fields": {"tool_search_items": items},
|
||||
},
|
||||
ts.tool_result(call.marker),
|
||||
]
|
||||
|
||||
|
||||
def _bridge_responses_input(call: _BridgeCall) -> list[JsonValue]:
|
||||
user: Final[dict[str, JsonValue]] = {"type": "message", "role": "user", "content": pcb.prompt(call.marker)}
|
||||
if not call.replay:
|
||||
return [user]
|
||||
return [
|
||||
user,
|
||||
*ts.searched(call.marker, "a", ts.WEATHER),
|
||||
ts.replayed_call(call.marker),
|
||||
{"type": "function_call_output", "call_id": ts.call_id(call.marker), "output": "sunny"},
|
||||
]
|
||||
|
||||
|
||||
def _bridge_expected_replay(call: _BridgeCall) -> list[dict[str, JsonValue]]:
|
||||
if not call.replay:
|
||||
return []
|
||||
match call.endpoint:
|
||||
case "chat":
|
||||
return [
|
||||
ts.search_call(call.marker),
|
||||
ts.search_output(call.marker, [ts.loaded(ts.WEATHER)]),
|
||||
ts.replayed_call(call.marker),
|
||||
]
|
||||
case "responses":
|
||||
return [*ts.searched(call.marker, "a", ts.WEATHER), ts.replayed_call(call.marker)]
|
||||
case "messages":
|
||||
return []
|
||||
|
||||
|
||||
def _bridge_request(call: _BridgeCall, model: str) -> tuple[str, dict[str, JsonValue]]:
|
||||
common: Final[dict[str, JsonValue]] = {"model": model, "stream": call.stream, "num_retries": 0, **pcb.NO_CACHE}
|
||||
match call.endpoint:
|
||||
case "chat":
|
||||
return "/v1/chat/completions", {
|
||||
**common,
|
||||
"messages": _bridge_chat_messages(call),
|
||||
"tools": ts.chat_tools(ts.WEATHER),
|
||||
}
|
||||
case "messages":
|
||||
return "/v1/messages", {
|
||||
**common,
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": pcb.prompt(call.marker)}],
|
||||
"tools": [{"name": ts.WEATHER, "description": f"Look up {ts.WEATHER}", "input_schema": ts.PARAMETERS}],
|
||||
}
|
||||
case "responses":
|
||||
return "/v1/responses", {
|
||||
**common,
|
||||
"input": _bridge_responses_input(call),
|
||||
"tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)],
|
||||
}
|
||||
|
||||
|
||||
async def _bridge_send(client: httpx.AsyncClient, key: str, model: str, call: _BridgeCall) -> _BridgeServed:
|
||||
route, body = _bridge_request(call, model)
|
||||
headers: Final = {"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}
|
||||
async with client.stream("POST", route, json=body, headers=headers) as response:
|
||||
raw: Final = await response.aread()
|
||||
return _BridgeServed(call, response.status_code, raw.decode())
|
||||
|
||||
|
||||
async def _bridge_burst(
|
||||
gateway: Gateway, model: str, calls: tuple[_BridgeCall, ...], *, tolerate_transport_errors: bool = False
|
||||
) -> tuple[_BridgeServed, ...]:
|
||||
async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(_bridge_send(client, gateway.key, model, call) for call in calls),
|
||||
return_exceptions=tolerate_transport_errors,
|
||||
)
|
||||
unexpected: Final = [
|
||||
result
|
||||
for result in results
|
||||
if isinstance(result, BaseException) and not isinstance(result, httpx.TransportError)
|
||||
]
|
||||
assert unexpected == [], unexpected
|
||||
return tuple(result for result in results if isinstance(result, _BridgeServed))
|
||||
|
||||
|
||||
def _bridge_assert_chat_answered(served: _BridgeServed) -> None:
|
||||
assistant: Final = (
|
||||
ts.assistant_from_stream(ts.chunks_of(served.text)) if served.call.stream else ts.first_message(served.text)
|
||||
)
|
||||
if served.call.replay:
|
||||
assert assistant["content"] == rv.answer(served.call.marker), served.text
|
||||
return
|
||||
assert ts.provider_fields(assistant).get("tool_search_items") == [
|
||||
*ts.searched(served.call.marker, "a", ts.WEATHER)
|
||||
], served.text
|
||||
assert ts.namespaces(assistant) == [ts.WEATHER], served.text
|
||||
|
||||
|
||||
def _bridge_assert_answered(served: _BridgeServed) -> None:
|
||||
assert served.status == 200, served.text
|
||||
match served.call.endpoint:
|
||||
case "chat":
|
||||
_bridge_assert_chat_answered(served)
|
||||
case "messages":
|
||||
assert '"tool_use"' in served.text and ts.call_id(served.call.marker) in served.text, served.text
|
||||
assert "tool_search" not in served.text, served.text
|
||||
case "responses":
|
||||
expected: Final = rv.answer(served.call.marker) if served.call.replay else f"tso_a_{served.call.marker}"
|
||||
assert expected in served.text, served.text
|
||||
|
||||
|
||||
def _bridge_assert_upstream_saw_each_call_once(received: tuple[Request, ...], calls: tuple[_BridgeCall, ...]) -> None:
|
||||
posts: Final = tuple(request for request in received if request.method == "POST")
|
||||
assert sorted(rv.newest_marker(request.body.decode()) or "" for request in posts) == sorted(
|
||||
call.marker for call in calls
|
||||
), [request.body for request in posts]
|
||||
expected: Final = {call.marker: _bridge_expected_replay(call) for call in calls}
|
||||
assert [ts.replayed(request) for request in posts] == [
|
||||
expected[rv.newest_marker(request.body.decode()) or ""] for request in posts
|
||||
], [request.body for request in posts]
|
||||
|
||||
|
||||
def _bridge_caller_id(served: _BridgeServed) -> str:
|
||||
if not served.call.stream:
|
||||
return string_value(rv.JSON_OBJECT.validate_json(served.text)["id"])
|
||||
first: Final = ts.chunks_of(served.text)[0]
|
||||
match served.call.endpoint:
|
||||
case "chat":
|
||||
return string_value(first["id"])
|
||||
case "messages":
|
||||
return string_value(rv.JSON_OBJECT.validate_python(first["message"])["id"])
|
||||
case "responses":
|
||||
return string_value(rv.JSON_OBJECT.validate_python(first["response"])["id"])
|
||||
|
||||
|
||||
def _bridge_assert_one_success_row(model: str, served: _BridgeServed) -> None:
|
||||
caller_id: Final = _bridge_caller_id(served)
|
||||
rows: Final = eventually(lambda: ts.spend_rows_of(model, caller_id), lambda found: len(found) >= 1, seconds=60)
|
||||
assert [row["status"] for row in rows] == ["success"], rows
|
||||
|
||||
|
||||
def _bridge_outage_for(down: frozenset[str]) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if rv.newest_marker(request.body.decode()) in down:
|
||||
return rv.error(503, _BRIDGE_OUTAGE, "server_error")
|
||||
return ts.scripted(request)
|
||||
|
||||
return answering_model_discovery(respond)
|
||||
|
||||
|
||||
async def test_upstream_outage_mid_burst_fails_only_its_calls_and_every_replay_lands_once(gateway: Gateway) -> None:
|
||||
calls: Final = _bridge_calls(24, lambda index: _BRIDGE_ENDPOINTS[index % 4], lambda index: (index // 4) % 2 == 0)
|
||||
down: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 5 == 1)
|
||||
with wire_server(_bridge_outage_for(down)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
served: Final = await _bridge_burst(gateway, model, calls)
|
||||
assert len(served) == len(calls)
|
||||
for item in served:
|
||||
if item.call.marker in down:
|
||||
assert item.status == 503 and _BRIDGE_OUTAGE in item.text, item.text
|
||||
assert rv.answer(item.call.marker) not in item.text, item.text
|
||||
else:
|
||||
_bridge_assert_answered(item)
|
||||
recovery: Final = _BridgeCall("chat", uuid.uuid4().hex, stream=True, replay=True)
|
||||
(recovered,) = await _bridge_burst(gateway, model, (recovery,))
|
||||
_bridge_assert_answered(recovered)
|
||||
_bridge_assert_upstream_saw_each_call_once(wire.drain(), (*calls, recovery))
|
||||
for item in (*served, recovered):
|
||||
if item.call.marker not in down:
|
||||
_bridge_assert_one_success_row(model, item)
|
||||
|
||||
|
||||
async def test_slow_upstream_streams_keep_items_intact_without_a_duplicate_upstream_call(gateway: Gateway) -> None:
|
||||
calls: Final = _bridge_calls(12, lambda index: _BRIDGE_ENDPOINTS[index % 4], lambda index: (index // 4) % 2 == 0)
|
||||
with wire_server(ts.scripted_with_pause(0.3)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
served: Final = await _bridge_burst(gateway, model, calls)
|
||||
assert len(served) == len(calls)
|
||||
for item in served:
|
||||
_bridge_assert_answered(item)
|
||||
_bridge_assert_upstream_saw_each_call_once(wire.drain(), calls)
|
||||
for item in served:
|
||||
_bridge_assert_one_success_row(model, item)
|
||||
|
||||
|
||||
def _bridge_chaos_config(wire: Wire, directory: Path) -> Path:
|
||||
template: Final = Path("tests/integration/proxy_config.yaml").read_text()
|
||||
base: Final = rv.JSON_OBJECT.validate_python(yaml.safe_load(template))
|
||||
config: Final = {
|
||||
**base,
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": _CONFIG_MODEL,
|
||||
"litellm_params": {"model": pcb.MODEL, "api_base": f"{wire.url}/v1", "api_key": "sk-upstream"},
|
||||
}
|
||||
],
|
||||
}
|
||||
path: Final = directory / "responses-bridge-tool-search-chaos.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def _bridge_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(2 * graceful_stop_seconds() + 120)
|
||||
async def test_worker_kill_mid_replay_burst_leaves_the_sibling_replaying_search_items(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
calls: Final = _bridge_calls(20, lambda index: _BRIDGE_ENDPOINTS[index % 4], lambda _: False)
|
||||
release: Final = threading.Event()
|
||||
held_markers: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
held_markers.put(rv.newest_marker(request.body.decode()) or "")
|
||||
assert release.wait(timeout=90), "The burst was never released"
|
||||
return ts.scripted(request)
|
||||
|
||||
with wire_server(answering_model_discovery(held)) as wire:
|
||||
with owned_proxy_process(
|
||||
gateway, tmp_path, {}, config=_bridge_chaos_config(wire, tmp_path), workers=2
|
||||
) as owned:
|
||||
workers: Final = eventually(
|
||||
lambda: pcb.started_worker_pids(owned.log),
|
||||
lambda pids: len(pids) == 2,
|
||||
seconds=graceful_stop_seconds() + 60,
|
||||
)
|
||||
burst: Final = asyncio.create_task(
|
||||
_bridge_burst(owned.gateway, _CONFIG_MODEL, calls, tolerate_transport_errors=True)
|
||||
)
|
||||
await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == len(calls), 60)
|
||||
held_by: Final = MappingProxyType(
|
||||
{pid: _bridge_open_upstream_connections(pid, wire.url) for pid in workers}
|
||||
)
|
||||
assert sum(held_by.values()) == len(calls), 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
|
||||
assert len(served) == held_by[survivor_pid], (held_by, len(served))
|
||||
for item in served:
|
||||
_bridge_assert_answered(item)
|
||||
follow_up: Final = _BridgeCall("chat", uuid.uuid4().hex, stream=True, replay=True)
|
||||
(answered,) = await _bridge_burst(owned.gateway, _CONFIG_MODEL, (follow_up,))
|
||||
_bridge_assert_answered(answered)
|
||||
_bridge_assert_upstream_saw_each_call_once(wire.drain(), (*calls, follow_up))
|
||||
for item in (*served, answered):
|
||||
_bridge_assert_one_success_row(_CONFIG_MODEL, item)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,496 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from typing import Final, Literal, TypeAlias, cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from anthropic import Anthropic
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support import tool_search as ts
|
||||
from integration._support.client import Gateway, Scenario, eventually, string_value
|
||||
from integration._support.wire import Wire, wire_server
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolUnionParam
|
||||
from pydantic import JsonValue
|
||||
|
||||
from tests.integration.caching.response_cache_case import await_entry, redis_store
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(120)
|
||||
|
||||
Client: TypeAlias = Literal["httpx", "openai_sync", "openai_async"]
|
||||
_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {
|
||||
"name": ts.WEATHER,
|
||||
"description": "Look up get_weather",
|
||||
"input_schema": ts.PARAMETERS,
|
||||
}
|
||||
|
||||
|
||||
def _base_url(gateway: Gateway) -> str:
|
||||
return str(gateway.client.base_url).rstrip("/")
|
||||
|
||||
|
||||
def _sdk_body(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {key: value for key, value in body.items() if key not in ("model", "messages", "tools", "stream")}
|
||||
|
||||
|
||||
def _sync_sdk(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[object, ...]:
|
||||
messages: Final = cast(Iterable[ChatCompletionMessageParam], body["messages"])
|
||||
tools: Final = cast(Iterable[ChatCompletionToolUnionParam], body["tools"])
|
||||
model: Final = string_value(body["model"])
|
||||
with OpenAI(api_key=gateway.key, base_url=f"{_base_url(gateway)}/v1", max_retries=0) as client:
|
||||
if body.get("stream") is True:
|
||||
chunks: Final = client.chat.completions.create(
|
||||
model=model, messages=messages, tools=tools, stream=True, extra_body=_sdk_body(body)
|
||||
)
|
||||
return tuple(chunk.model_dump() for chunk in chunks)
|
||||
completion: Final = client.chat.completions.create(
|
||||
model=model, messages=messages, tools=tools, stream=False, extra_body=_sdk_body(body)
|
||||
)
|
||||
return (completion.model_dump(),)
|
||||
|
||||
|
||||
async def _async_sdk(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[object, ...]:
|
||||
messages: Final = cast(Iterable[ChatCompletionMessageParam], body["messages"])
|
||||
tools: Final = cast(Iterable[ChatCompletionToolUnionParam], body["tools"])
|
||||
model: Final = string_value(body["model"])
|
||||
async with AsyncOpenAI(api_key=gateway.key, base_url=f"{_base_url(gateway)}/v1", max_retries=0) as client:
|
||||
if body.get("stream") is True:
|
||||
chunks: Final = await client.chat.completions.create(
|
||||
model=model, messages=messages, tools=tools, stream=True, extra_body=_sdk_body(body)
|
||||
)
|
||||
return tuple([chunk.model_dump() async for chunk in chunks])
|
||||
completion: Final = await client.chat.completions.create(
|
||||
model=model, messages=messages, tools=tools, stream=False, extra_body=_sdk_body(body)
|
||||
)
|
||||
return (completion.model_dump(),)
|
||||
|
||||
|
||||
def _send(gateway: Gateway, client: Client, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
|
||||
match client:
|
||||
case "httpx":
|
||||
response: Final = gateway.request("POST", "/v1/chat/completions", dict(body))
|
||||
assert response.status_code == 200, response.text
|
||||
if body.get("stream") is True:
|
||||
return ts.chunks_of(response.text)
|
||||
return (rv.JSON_OBJECT.validate_json(response.text),)
|
||||
case "openai_sync":
|
||||
return tuple(rv.JSON_OBJECT.validate_python(item) for item in _sync_sdk(gateway, body))
|
||||
case "openai_async":
|
||||
return tuple(rv.JSON_OBJECT.validate_python(item) for item in asyncio.run(_async_sdk(gateway, body)))
|
||||
|
||||
|
||||
def _assistant(served: Sequence[Mapping[str, JsonValue]], stream: bool) -> dict[str, JsonValue]:
|
||||
if stream:
|
||||
return ts.assistant_from_stream(served)
|
||||
choices: Final = rv.ITEMS.validate_python(served[0]["choices"])
|
||||
return rv.JSON_OBJECT.validate_python(choices[0]["message"])
|
||||
|
||||
|
||||
def _sent_message(message: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
kept: Final = ("role", "content", "tool_calls", "provider_specific_fields", "reasoning_items")
|
||||
return {key: value for key, value in message.items() if key in kept and value is not None}
|
||||
|
||||
|
||||
def _search_items(marker: str) -> list[JsonValue]:
|
||||
return [*ts.searched(marker, "a", ts.WEATHER)]
|
||||
|
||||
|
||||
def _replay(marker: str) -> list[dict[str, JsonValue]]:
|
||||
return [*ts.searched(marker, "a", ts.WEATHER), ts.reasoning(marker, "b"), ts.replayed_call(marker)]
|
||||
|
||||
|
||||
def _body(
|
||||
model: str, messages: Sequence[JsonValue], *, stream: bool = False, **extra: JsonValue
|
||||
) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": model,
|
||||
"messages": list(messages),
|
||||
"tools": ts.chat_tools(ts.WEATHER),
|
||||
"stream": stream,
|
||||
**pcb.NO_CACHE,
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def _user(marker: str, *words: str) -> dict[str, JsonValue]:
|
||||
return {"role": "user", "content": " ".join((pcb.prompt(marker), *words))}
|
||||
|
||||
|
||||
def _two_turns(
|
||||
gateway: Gateway, wire: Wire, model: str, client: Client, stream: bool, marker: str
|
||||
) -> tuple[dict[str, JsonValue], list[dict[str, JsonValue]], tuple[str, str]]:
|
||||
user: Final = _user(marker)
|
||||
first: Final = _send(gateway, client, _body(model, [user], stream=stream))
|
||||
assistant: Final = _assistant(first, stream)
|
||||
pcb.posted(wire, marker)
|
||||
history: Final = [user, _sent_message(assistant), ts.tool_result(marker)]
|
||||
second: Final = _send(gateway, client, _body(model, history, stream=stream))
|
||||
answered: Final = _assistant(second, stream)
|
||||
assert answered["content"] == rv.answer(marker), answered
|
||||
identities: Final = (string_value(first[0]["id"]), string_value(second[0]["id"]))
|
||||
return assistant, ts.replayed(pcb.posted(wire, marker)), identities
|
||||
|
||||
|
||||
def _assert_turn_one(assistant: Mapping[str, JsonValue], marker: str) -> None:
|
||||
assert ts.provider_fields(assistant).get("tool_search_items") == _search_items(marker), assistant
|
||||
assert [call["id"] for call in ts.tool_calls(assistant)] == [ts.call_id(marker)], assistant
|
||||
assert ts.namespaces(assistant) == [ts.WEATHER], assistant
|
||||
|
||||
|
||||
def _assert_spend(model: str, response_ids: Sequence[str]) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: tuple(ts.spend_rows_of(model, response_id) for response_id in response_ids),
|
||||
lambda found: all(found),
|
||||
seconds=70,
|
||||
)
|
||||
assert [[row["status"] for row in found] for found in rows] == [["success"]] * len(response_ids), rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ["httpx", "openai_sync", "openai_async"])
|
||||
def test_chat_bridge_returns_and_replays_tool_search_items_non_streaming(gateway: Gateway, client: Client) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
assistant, replayed, response_ids = _two_turns(gateway, wire, model, client, False, marker)
|
||||
_assert_turn_one(assistant, marker)
|
||||
assert replayed == _replay(marker)
|
||||
_assert_spend(model, response_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ["httpx", "openai_async"])
|
||||
def test_chat_bridge_streams_tool_search_items_and_replays_the_rebuilt_message(
|
||||
gateway: Gateway, client: Client
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
assistant, replayed, response_ids = _two_turns(gateway, wire, model, client, True, marker)
|
||||
_assert_turn_one(assistant, marker)
|
||||
assert replayed == _replay(marker)
|
||||
_assert_spend(model, response_ids)
|
||||
|
||||
|
||||
def _history(marker: str, **assistant_fields: JsonValue) -> list[JsonValue]:
|
||||
call: Final[dict[str, JsonValue]] = {
|
||||
"id": ts.call_id(marker),
|
||||
"type": "function",
|
||||
"function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'},
|
||||
}
|
||||
return [
|
||||
_user(marker),
|
||||
{"role": "assistant", "content": None, "tool_calls": [call], **assistant_fields},
|
||||
ts.tool_result(marker),
|
||||
]
|
||||
|
||||
|
||||
def _replayed_history(marker: str) -> list[JsonValue]:
|
||||
return _history(
|
||||
marker,
|
||||
tool_calls=[
|
||||
{
|
||||
"id": ts.call_id(marker),
|
||||
"type": "function",
|
||||
"function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": ts.WEATHER},
|
||||
}
|
||||
],
|
||||
provider_specific_fields={"tool_search_items": _search_items(marker)},
|
||||
reasoning_items=[ts.reasoning(marker, "a"), ts.reasoning(marker, "b")],
|
||||
)
|
||||
|
||||
|
||||
def _replayed_upstream(
|
||||
gateway: Gateway,
|
||||
scenario: Scenario,
|
||||
wire: Wire,
|
||||
marker: str,
|
||||
body: Mapping[str, JsonValue],
|
||||
**model_params: JsonValue,
|
||||
) -> list[dict[str, JsonValue]]:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1", **model_params)
|
||||
response: Final = gateway.request("POST", "/v1/chat/completions", {**body, "model": model})
|
||||
assert response.status_code == 200, response.text
|
||||
assert ts.first_message(response.text)["content"] == rv.answer(marker), response.text
|
||||
return ts.replayed(pcb.posted(wire, marker))
|
||||
|
||||
|
||||
def test_turn_without_tool_search_strips_search_items_but_keeps_namespace_and_reasoning(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
body: Final = {
|
||||
**_body("", _replayed_history(marker)),
|
||||
"tools": [{"type": "function", "function": ts.function(ts.WEATHER)}],
|
||||
}
|
||||
assert _replayed_upstream(gateway, scenario, wire, marker, body) == [
|
||||
ts.reasoning(marker, "a"),
|
||||
ts.reasoning(marker, "b"),
|
||||
ts.replayed_call(marker),
|
||||
]
|
||||
|
||||
|
||||
def test_tool_search_only_in_extra_body_tools_still_replays(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
body: Final = {
|
||||
**_body("", _replayed_history(marker)),
|
||||
"tools": [ts.deferred(ts.WEATHER)],
|
||||
"extra_body": {"tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)]},
|
||||
}
|
||||
assert _replayed_upstream(gateway, scenario, wire, marker, body) == _replay(marker)
|
||||
|
||||
|
||||
def test_extra_body_tools_without_tool_search_win_over_tools_and_skip_the_replay(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
body: Final = {
|
||||
**_body("", _replayed_history(marker)),
|
||||
"extra_body": {"tools": [{"type": "function", **ts.function(ts.WEATHER)}]},
|
||||
}
|
||||
assert _replayed_upstream(gateway, scenario, wire, marker, body) == [
|
||||
ts.reasoning(marker, "a"),
|
||||
ts.reasoning(marker, "b"),
|
||||
ts.replayed_call(marker),
|
||||
]
|
||||
|
||||
|
||||
def test_deployment_tool_search_replays_items_the_request_did_not_ask_for(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
body: Final = {**_body("", _replayed_history(marker)), "tools": [ts.deferred(ts.WEATHER)]}
|
||||
replayed: Final = _replayed_upstream(gateway, scenario, wire, marker, body, tools=[ts.TOOL_SEARCH])
|
||||
assert replayed == _replay(marker)
|
||||
|
||||
|
||||
_MALFORMED_ITEMS: Final[dict[str, JsonValue]] = {
|
||||
"int": 7,
|
||||
"string": "tool_search_call",
|
||||
"empty-object": {},
|
||||
"empty-list": [],
|
||||
"null": None,
|
||||
"junk-list": [1, "x", None, {"type": "message", "content": "hi"}, {"no": "type"}],
|
||||
"5kb-string": "x" * 5120,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("items", list(_MALFORMED_ITEMS.values()), ids=list(_MALFORMED_ITEMS))
|
||||
def test_malformed_tool_search_items_are_ignored_without_a_server_error(gateway: Gateway, items: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
body: Final = _body("", _history(marker, provider_specific_fields={"tool_search_items": items}))
|
||||
assert _replayed_upstream(gateway, scenario, wire, marker, body) == [
|
||||
{k: v for k, v in ts.replayed_call(marker).items() if k != "namespace"}
|
||||
]
|
||||
|
||||
|
||||
def test_search_output_with_non_list_tools_and_items_repeated_on_two_turns_replay_verbatim(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
odd_output: Final = {**ts.search_output(marker, []), "tools": "not-a-list"}
|
||||
items: Final[list[JsonValue]] = [ts.search_call(marker), odd_output]
|
||||
first_call: Final = {
|
||||
"id": ts.call_id(marker),
|
||||
"type": "function",
|
||||
"function": {"name": ts.WEATHER, "arguments": "{}"},
|
||||
}
|
||||
messages: Final[list[JsonValue]] = [
|
||||
_user(marker),
|
||||
{"role": "assistant", "content": "one", "provider_specific_fields": {"tool_search_items": items}},
|
||||
{"role": "user", "content": "again"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [first_call],
|
||||
"provider_specific_fields": {"tool_search_items": items},
|
||||
},
|
||||
ts.tool_result(marker),
|
||||
]
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
replayed: Final = _replayed_upstream(gateway, scenario, wire, marker, _body("", messages))
|
||||
assert replayed == [
|
||||
ts.search_call(marker),
|
||||
odd_output,
|
||||
ts.search_call(marker),
|
||||
odd_output,
|
||||
{"type": "function_call", "call_id": ts.call_id(marker), "name": ts.WEATHER, "arguments": "{}"},
|
||||
]
|
||||
|
||||
|
||||
_NAMESPACES: Final[dict[str, JsonValue]] = {"int": 3, "list": ["get_weather"], "empty": "", "5kb": "n" * 5120}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", list(_NAMESPACES.values()), ids=list(_NAMESPACES))
|
||||
def test_malformed_namespace_is_dropped_and_a_long_one_is_forwarded(gateway: Gateway, value: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
call: Final[dict[str, JsonValue]] = {
|
||||
"id": ts.call_id(marker),
|
||||
"type": "function",
|
||||
"function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": value},
|
||||
}
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
body: Final = _body("", _history(marker, tool_calls=[call]))
|
||||
bare: Final = {k: v for k, v in ts.replayed_call(marker).items() if k != "namespace"}
|
||||
expected: Final = {**bare, "namespace": value} if isinstance(value, str) and value else bare
|
||||
assert _replayed_upstream(gateway, scenario, wire, marker, body) == [expected]
|
||||
|
||||
|
||||
def test_upstream_rejecting_the_replay_reaches_the_caller(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
messages: Final = [_user(marker, ts.REJECT_REPLAY), *_replayed_history(marker)[1:]]
|
||||
response: Final = gateway.request("POST", "/v1/chat/completions", {**_body(model, messages)})
|
||||
assert response.status_code == 400, response.text
|
||||
assert ts.REJECTED in response.text, response.text
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker)
|
||||
healthy: Final = gateway.request("POST", "/v1/chat/completions", _body(model, [_user(uuid.uuid4().hex)]))
|
||||
assert healthy.status_code == 200, healthy.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
|
||||
def test_turn_that_only_searched_returns_the_items_and_continues_the_search(gateway: Gateway, stream: bool) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
user: Final = _user(marker, ts.ONLY_SEARCH)
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
first: Final = _send(gateway, "httpx", _body(model, [user], stream=stream))
|
||||
assistant: Final = _assistant(first, stream)
|
||||
choices: Final = chain.from_iterable(rv.ITEMS.validate_python(chunk["choices"]) for chunk in first)
|
||||
finish: Final = [choice["finish_reason"] for choice in choices if choice.get("finish_reason")]
|
||||
assert finish == ["stop"], first
|
||||
assert assistant.get("content") in (None, ""), assistant
|
||||
assert ts.provider_fields(assistant).get("tool_search_items") == _search_items(marker), assistant
|
||||
assert "tool_calls" not in assistant or not assistant["tool_calls"], assistant
|
||||
pcb.posted(wire, marker)
|
||||
second: Final = _send(gateway, "httpx", _body(model, [user, _sent_message(assistant)], stream=stream))
|
||||
assert _assistant(second, stream)["content"] == rv.answer(marker), second
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == _search_items(marker)
|
||||
|
||||
|
||||
def test_two_searches_in_two_turns_replay_in_order(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
user: Final = _user(marker, ts.TWO_SEARCHES)
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
tools: Final = ts.chat_tools(ts.WEATHER, ts.TIME)
|
||||
first: Final = _assistant(_send(gateway, "httpx", {**_body(model, [user]), "tools": tools}), False)
|
||||
pcb.posted(wire, marker)
|
||||
history: Final[list[JsonValue]] = [user, _sent_message(first), ts.tool_result(marker)]
|
||||
second: Final = _assistant(_send(gateway, "httpx", {**_body(model, history), "tools": tools}), False)
|
||||
assert ts.provider_fields(second).get("tool_search_items") == [*ts.searched(marker, "c", ts.TIME)], second
|
||||
assert ts.namespaces(second) == [ts.TIME], second
|
||||
pcb.posted(wire, marker)
|
||||
final: Final = [*history, _sent_message(second), ts.tool_result(marker, ts.TIME)]
|
||||
third: Final = _assistant(_send(gateway, "httpx", {**_body(model, final), "tools": tools}), False)
|
||||
assert third["content"] == rv.answer(marker), third
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == [
|
||||
*_replay(marker),
|
||||
*ts.searched(marker, "c", ts.TIME),
|
||||
ts.replayed_call(marker, ts.TIME),
|
||||
]
|
||||
|
||||
|
||||
def _cached_twin(
|
||||
gateway: Gateway, body: Mapping[str, JsonValue], stream: bool
|
||||
) -> tuple[httpx.Response, httpx.Response]:
|
||||
first: Final = gateway.request("POST", "/v1/chat/completions", dict(body))
|
||||
assert first.status_code == 200, first.text
|
||||
first_frame: Final = ts.chunks_of(first.text)[0] if stream else rv.JSON_OBJECT.validate_json(first.text)
|
||||
first_id: Final = string_value(first_frame["id"])
|
||||
await_entry(redis_store(), first_id)
|
||||
second: Final = gateway.request("POST", "/v1/chat/completions", dict(body))
|
||||
assert second.status_code == 200, second.text
|
||||
assert second.headers.get("x-litellm-cache-key"), dict(second.headers)
|
||||
return first, second
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
|
||||
def test_cache_hit_turn_carries_the_items_and_replays_them(gateway: Gateway, stream: bool) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
user: Final = _user(marker)
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
cached_body: Final = {
|
||||
key: value for key, value in _body(model, [user], stream=stream).items() if key != "cache"
|
||||
}
|
||||
first, hit = _cached_twin(gateway, cached_body, stream)
|
||||
pcb.posted(wire, marker)
|
||||
served: Final = ts.chunks_of(hit.text) if stream else (rv.JSON_OBJECT.validate_json(hit.text),)
|
||||
assistant: Final = _assistant(served, stream)
|
||||
_assert_turn_one(assistant, marker)
|
||||
second: Final = _send(gateway, "httpx", _body(model, [user, _sent_message(assistant), ts.tool_result(marker)]))
|
||||
assert _assistant(second, False)["content"] == rv.answer(marker), second
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker)
|
||||
first_frames: Final = ts.chunks_of(first.text) if stream else (rv.JSON_OBJECT.validate_json(first.text),)
|
||||
assert served[0]["id"] == first_frames[0]["id"]
|
||||
|
||||
|
||||
def test_messages_endpoint_sees_the_same_tool_use_on_both_sides(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
with Anthropic(api_key=gateway.key, base_url=_base_url(gateway), max_retries=0) as client:
|
||||
message: Final = client.messages.create(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
messages=[{"role": "user", "content": pcb.prompt(marker)}],
|
||||
tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # JSON tool definition
|
||||
)
|
||||
pcb.posted(wire, marker)
|
||||
uses: Final = [block for block in message.content if block.type == "tool_use"]
|
||||
assert [(block.name, block.input) for block in uses] == [(ts.WEATHER, {"city": "Paris"})], message
|
||||
assert "tool_search" not in message.model_dump_json(), message.model_dump_json()
|
||||
with client.messages.stream(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
messages=[{"role": "user", "content": pcb.prompt(marker)}],
|
||||
tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # JSON tool definition
|
||||
) as stream:
|
||||
streamed: Final = stream.get_final_message()
|
||||
pcb.posted(wire, marker)
|
||||
streamed_uses: Final = [block for block in streamed.content if block.type == "tool_use"]
|
||||
assert [(block.name, block.input) for block in streamed_uses] == [(ts.WEATHER, {"city": "Paris"})]
|
||||
assert "tool_search" not in streamed.model_dump_json(), streamed.model_dump_json()
|
||||
|
||||
|
||||
def test_messages_endpoint_turn_that_only_searched_is_answered(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": f"{pcb.prompt(marker)} {ts.ONLY_SEARCH}"}],
|
||||
"tools": [_ANTHROPIC_TOOL],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = rv.JSON_OBJECT.validate_json(response.text)
|
||||
assert body["stop_reason"] == "end_turn", body
|
||||
assert [block for block in rv.ITEMS.validate_python(body["content"]) if block.get("type") == "tool_use"] == []
|
||||
pcb.posted(wire, marker)
|
||||
|
||||
|
||||
def test_native_responses_forwards_tool_search_items_and_namespace_verbatim(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
history: Final[list[JsonValue]] = [
|
||||
{"type": "message", "role": "user", "content": pcb.prompt(marker)},
|
||||
*ts.searched(marker, "a", ts.WEATHER),
|
||||
ts.replayed_call(marker),
|
||||
{"type": "function_call_output", "call_id": ts.call_id(marker), "output": "sunny"},
|
||||
]
|
||||
with wire_server(ts.scripted) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": model, "input": history, "tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)], **pcb.NO_CACHE},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert rv.answer(marker) in response.text, response.text
|
||||
assert pcb.input_items(pcb.posted(wire, marker)) == history
|
||||
|
|
@ -19,8 +19,9 @@ from integration._support import claude_code as cc
|
|||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.anthropic_thinking import BEDROCK_MODEL, EVENT_STREAM, MODEL, Event, aws_chunks, sse_chunks
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.openai_wire import answering_model_discovery, responses_reply
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers.test_bedrock_auth_wire import RESPONSE as CONVERSE_RESPONSE
|
||||
|
|
@ -1324,3 +1325,113 @@ def test_chat_bridge_forwards_defer_loading_so_tool_search_has_a_deferred_tool(
|
|||
_defer_wire_function("get_time"),
|
||||
_defer_wire_function("get_news", defer_loading=False),
|
||||
]
|
||||
|
||||
|
||||
_OURS_PARAMETERS: Final[dict[str, JsonValue]] = {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
}
|
||||
_OURS_TOOL_SEARCH: Final[dict[str, JsonValue]] = {"type": "tool_search"}
|
||||
|
||||
|
||||
def _ours_function(name: str, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"name": name, "description": f"Look up {name}", "parameters": _OURS_PARAMETERS, **extra}
|
||||
|
||||
|
||||
def _ours_wire_function(name: str, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"type": "function", "strict": None, **_ours_function(name), **extra}
|
||||
|
||||
|
||||
_OURS_CALL_ID: Final = "call_weather"
|
||||
_OURS_TOOL_SEARCH_CALL: Final[dict[str, JsonValue]] = {
|
||||
"id": "tsc_weather",
|
||||
"type": "tool_search_call",
|
||||
"arguments": {"query": "weather"},
|
||||
"execution": "server",
|
||||
"status": "completed",
|
||||
}
|
||||
_OURS_TOOL_SEARCH_OUTPUT: Final[dict[str, JsonValue]] = {
|
||||
"id": "tso_weather",
|
||||
"type": "tool_search_output",
|
||||
"execution": "server",
|
||||
"status": "completed",
|
||||
"tools": [_ours_wire_function("get_weather", defer_loading=True)],
|
||||
}
|
||||
_OURS_WEATHER_CALL: Final[dict[str, JsonValue]] = {
|
||||
"type": "function_call",
|
||||
"call_id": _OURS_CALL_ID,
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"namespace": "get_weather",
|
||||
}
|
||||
|
||||
|
||||
def _ours_searches_then_answers(request: Request) -> Reply:
|
||||
body: Final = pcb.body_of(request)
|
||||
marker: Final = rv.newest_marker(request.body.decode())
|
||||
identity: Final = f"resp_{marker}-{uuid.uuid4().hex}"
|
||||
if any(item.get("type") == "function_call_output" for item in pcb.input_items(request)):
|
||||
return responses_reply(identity, string_value(body["model"]), rv.answer(marker), stream=False)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": body["model"],
|
||||
"output": [
|
||||
_OURS_TOOL_SEARCH_CALL,
|
||||
_OURS_TOOL_SEARCH_OUTPUT,
|
||||
{**_OURS_WEATHER_CALL, "id": "fc_weather", "status": "completed"},
|
||||
],
|
||||
"usage": rv.USAGE,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def test_chat_bridge_replays_tool_search_items_on_the_next_turn(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
tools: Final = [
|
||||
_OURS_TOOL_SEARCH,
|
||||
{"type": "function", "function": _ours_function("get_weather"), "defer_loading": True},
|
||||
]
|
||||
user_turn: Final[dict[str, JsonValue]] = {"role": "user", "content": pcb.prompt(marker)}
|
||||
with wire_server(answering_model_discovery(_ours_searches_then_answers)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1")
|
||||
first: Final = gateway.request(
|
||||
"POST", "/v1/chat/completions", {"model": model, "messages": [user_turn], "tools": tools, **pcb.NO_CACHE}
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
assistant_turn: Final = rv.JSON_OBJECT.validate_python(
|
||||
rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(first.text)["choices"])[0]["message"]
|
||||
)
|
||||
assert rv.JSON_OBJECT.validate_python(assistant_turn["provider_specific_fields"])["tool_search_items"] == [
|
||||
_OURS_TOOL_SEARCH_CALL,
|
||||
_OURS_TOOL_SEARCH_OUTPUT,
|
||||
], first.text
|
||||
pcb.posted(wire, marker)
|
||||
|
||||
second: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [
|
||||
user_turn,
|
||||
assistant_turn,
|
||||
{"role": "tool", "tool_call_id": _OURS_CALL_ID, "content": "sunny"},
|
||||
],
|
||||
"tools": tools,
|
||||
**pcb.NO_CACHE,
|
||||
},
|
||||
)
|
||||
assert second.status_code == 200, second.text
|
||||
replayed: Final = [
|
||||
item
|
||||
for item in pcb.input_items(pcb.posted(wire, marker))
|
||||
if item["type"] not in ("message", "function_call_output")
|
||||
]
|
||||
assert replayed == [_OURS_TOOL_SEARCH_CALL, _OURS_TOOL_SEARCH_OUTPUT, _OURS_WEATHER_CALL]
|
||||
|
|
|
|||
201
tests/integration/sdk/test_responses_bridge_tool_search_sdk.py
Normal file
201
tests/integration/sdk/test_responses_bridge_tool_search_sdk.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support import tool_search as ts
|
||||
from integration._support.wire import Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolMessage, ChatCompletionToolParam
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Message, ModelResponse
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(120)
|
||||
|
||||
_API_KEY: Final = "sk-scripted-provider"
|
||||
|
||||
|
||||
def _tools(*tools: JsonValue) -> list[ChatCompletionToolParam]:
|
||||
return cast(list[ChatCompletionToolParam], list(tools))
|
||||
|
||||
|
||||
def _messages(*messages: object) -> list[AllMessageValues]:
|
||||
return cast(list[AllMessageValues], list(messages))
|
||||
|
||||
|
||||
def _user(marker: str, *words: str) -> dict[str, JsonValue]:
|
||||
return {"role": "user", "content": " ".join((pcb.prompt(marker), *words))}
|
||||
|
||||
|
||||
def _replay(marker: str) -> list[dict[str, JsonValue]]:
|
||||
return [*ts.searched(marker, "a", ts.WEATHER), ts.reasoning(marker, "b"), ts.replayed_call(marker)]
|
||||
|
||||
|
||||
def _message(response: object) -> Message:
|
||||
assert isinstance(response, ModelResponse), type(response)
|
||||
choice: Final = response.choices[0]
|
||||
assert isinstance(choice, Choices), type(choice)
|
||||
return choice.message
|
||||
|
||||
|
||||
def _assert_turn_one(message: Message, marker: str) -> None:
|
||||
fields: Final = rv.JSON_OBJECT.validate_python(message.provider_specific_fields or {})
|
||||
assert fields.get("tool_search_items") == [*ts.searched(marker, "a", ts.WEATHER)], message
|
||||
calls: Final = message.tool_calls or []
|
||||
call_fields: Final = [
|
||||
rv.JSON_OBJECT.validate_python(getattr(call, "provider_specific_fields", None) or {}) for call in calls
|
||||
]
|
||||
assert call_fields == [{"namespace": ts.WEATHER}], message
|
||||
|
||||
|
||||
def _kwargs(wire: Wire) -> dict[str, str]:
|
||||
return {"api_base": f"{wire.url}/v1", "api_key": _API_KEY}
|
||||
|
||||
|
||||
def test_completion_returns_and_replays_tool_search_items() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
tools: Final = _tools(*ts.chat_tools(ts.WEATHER))
|
||||
with wire_server(ts.scripted) as wire:
|
||||
user: Final = _user(marker)
|
||||
first: Final = _message(
|
||||
litellm.completion(model=pcb.MODEL, messages=_messages(user), tools=tools, **_kwargs(wire))
|
||||
)
|
||||
_assert_turn_one(first, marker)
|
||||
pcb.posted(wire, marker)
|
||||
second: Final = _message(
|
||||
litellm.completion(
|
||||
model=pcb.MODEL,
|
||||
messages=_messages(user, first, ts.tool_result(marker)),
|
||||
tools=tools,
|
||||
**_kwargs(wire),
|
||||
)
|
||||
)
|
||||
assert second.content == rv.answer(marker), second
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker)
|
||||
|
||||
|
||||
def test_streamed_completion_rebuilt_by_stream_chunk_builder_replays_tool_search_items() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
tools: Final = _tools(*ts.chat_tools(ts.WEATHER))
|
||||
|
||||
async def run(wire: Wire) -> tuple[Message, Message]:
|
||||
user: Final = _user(marker)
|
||||
stream: Final = await litellm.acompletion(
|
||||
model=pcb.MODEL, messages=_messages(user), tools=tools, stream=True, **_kwargs(wire)
|
||||
)
|
||||
chunks: Final = [chunk async for chunk in stream] # pyright: ignore[reportGeneralTypeIssues] # stream wrapper
|
||||
built: Final = _message(litellm.stream_chunk_builder(chunks, messages=_messages(user)))
|
||||
pcb.posted(wire, marker)
|
||||
answered: Final = _message(
|
||||
await litellm.acompletion(
|
||||
model=pcb.MODEL,
|
||||
messages=_messages(user, built.model_dump(), ts.tool_result(marker)),
|
||||
tools=tools,
|
||||
**_kwargs(wire),
|
||||
)
|
||||
)
|
||||
return built, answered
|
||||
|
||||
with wire_server(ts.scripted) as wire:
|
||||
built, answered = asyncio.run(run(wire))
|
||||
_assert_turn_one(built, marker)
|
||||
assert [item["id"] for item in built.reasoning_items or []] == [f"rs_a_{marker}", f"rs_b_{marker}"], built
|
||||
assert answered.content == rv.answer(marker), answered
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker)
|
||||
|
||||
|
||||
def test_streamed_turn_without_tool_search_replays_its_reasoning() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
tools: Final = _tools({"type": "function", "function": ts.function(ts.WEATHER)})
|
||||
with wire_server(ts.scripted) as wire:
|
||||
user: Final = _user(marker)
|
||||
stream: Final = litellm.completion(
|
||||
model=pcb.MODEL, messages=_messages(user), tools=tools, stream=True, **_kwargs(wire)
|
||||
)
|
||||
chunks: Final = list(stream) # pyright: ignore[reportArgumentType] # stream wrapper
|
||||
built: Final = _message(litellm.stream_chunk_builder(chunks, messages=_messages(user)))
|
||||
pcb.posted(wire, marker)
|
||||
answered: Final = _message(
|
||||
litellm.completion(
|
||||
model=pcb.MODEL,
|
||||
messages=_messages(user, built.model_dump(), ts.tool_result(marker)),
|
||||
tools=tools,
|
||||
**_kwargs(wire),
|
||||
)
|
||||
)
|
||||
assert answered.content == rv.answer(marker), answered
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == [
|
||||
ts.reasoning(marker, "a"),
|
||||
ts.reasoning(marker, "b"),
|
||||
ts.replayed_call(marker),
|
||||
]
|
||||
|
||||
|
||||
def test_router_completion_replays_tool_search_items() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
tools: Final = _tools(*ts.chat_tools(ts.WEATHER))
|
||||
with wire_server(ts.scripted) as wire:
|
||||
router: Final = Router(
|
||||
model_list=[{"model_name": "bridged", "litellm_params": {"model": pcb.MODEL, **_kwargs(wire)}}],
|
||||
num_retries=0,
|
||||
)
|
||||
user: Final = _user(marker)
|
||||
first: Final = _message(router.completion(model="bridged", messages=_messages(user), tools=tools))
|
||||
_assert_turn_one(first, marker)
|
||||
pcb.posted(wire, marker)
|
||||
second: Final = _message(
|
||||
router.completion(model="bridged", messages=_messages(user, first, ts.tool_result(marker)), tools=tools)
|
||||
)
|
||||
assert second.content == rv.answer(marker), second
|
||||
assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker)
|
||||
|
||||
|
||||
def _sunny(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
return {"role": "tool", "tool_call_id": tool_call.id, "content": "sunny"}
|
||||
|
||||
|
||||
async def _async_sunny(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
|
||||
return _sunny(tool_call)
|
||||
|
||||
|
||||
def _posted_in_order(wire: Wire, marker: str) -> Sequence[list[dict[str, JsonValue]]]:
|
||||
return [ts.replayed(request) for request in pcb.with_marker(pcb.drained_posts(wire), marker)]
|
||||
|
||||
|
||||
def test_run_tool_loop_keeps_tool_search_items_and_namespace_between_rounds() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire:
|
||||
answer: Final = litellm.run_tool_loop(
|
||||
model=pcb.MODEL,
|
||||
messages=_messages(_user(marker)),
|
||||
tools=_tools(*ts.chat_tools(ts.WEATHER)),
|
||||
execute_tool=_sunny,
|
||||
**_kwargs(wire),
|
||||
)
|
||||
assert answer == rv.answer(marker)
|
||||
assert _posted_in_order(wire, marker) == [[], _replay(marker)]
|
||||
|
||||
|
||||
def test_async_run_tool_loop_keeps_two_searches_in_order() -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(ts.scripted) as wire:
|
||||
answer: Final = asyncio.run(
|
||||
litellm.arun_tool_loop(
|
||||
model=pcb.MODEL,
|
||||
messages=_messages(_user(marker, ts.TWO_SEARCHES)),
|
||||
tools=_tools(*ts.chat_tools(ts.WEATHER, ts.TIME)),
|
||||
execute_tool=_async_sunny,
|
||||
**_kwargs(wire),
|
||||
)
|
||||
)
|
||||
assert answer == rv.answer(marker)
|
||||
assert _posted_in_order(wire, marker) == [
|
||||
[],
|
||||
_replay(marker),
|
||||
[*_replay(marker), *ts.searched(marker, "c", ts.TIME), ts.replayed_call(marker, ts.TIME)],
|
||||
]
|
||||
|
|
@ -331,6 +331,187 @@ def test_chunk_parser_function_call_added_produces_tool_use():
|
|||
assert choice.finish_reason is None
|
||||
|
||||
|
||||
def test_streamed_function_call_namespace_replays_on_next_turn():
|
||||
from unittest.mock import Mock
|
||||
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
OpenAiResponsesToChatCompletionStreamIterator,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
||||
|
||||
iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
|
||||
streamed = iterator.chunk_parser(
|
||||
{
|
||||
"type": ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
"item": {
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"call_id": "call-ns",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"namespace": "get_weather",
|
||||
},
|
||||
}
|
||||
)
|
||||
tool_call = streamed.choices[0].delta.tool_calls[0].model_dump(exclude_none=True)
|
||||
|
||||
request = LiteLLMResponsesTransformationHandler().transform_request(
|
||||
model="gpt-5.5",
|
||||
messages=[
|
||||
{"role": "user", "content": "weather in Paris"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
|
||||
{"role": "tool", "tool_call_id": "call-ns", "content": "sunny"},
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
|
||||
function_calls = [item for item in request["input"] if item.get("type") == "function_call"]
|
||||
assert [call["namespace"] for call in function_calls] == ["get_weather"]
|
||||
|
||||
|
||||
def test_function_level_namespace_replays_when_tool_call_level_is_missing():
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
)
|
||||
|
||||
tool_call = {
|
||||
"id": "call-ns",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"provider_specific_fields": {"namespace": "get_weather"},
|
||||
},
|
||||
}
|
||||
|
||||
input_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
|
||||
[
|
||||
{"role": "user", "content": "weather in Paris"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
|
||||
{"role": "tool", "tool_call_id": "call-ns", "content": "sunny"},
|
||||
],
|
||||
include_namespace=True,
|
||||
)
|
||||
|
||||
function_calls = [item for item in input_items if item.get("type") == "function_call"]
|
||||
assert [call["namespace"] for call in function_calls] == ["get_weather"]
|
||||
|
||||
|
||||
def test_shared_converter_leaves_tool_search_items_and_namespace_out_by_default():
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
)
|
||||
|
||||
tool_call = {
|
||||
"id": "call-ns",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": "{}"},
|
||||
"provider_specific_fields": {"namespace": "weather"},
|
||||
}
|
||||
search_items = [
|
||||
{"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["weather"]}},
|
||||
{"type": "tool_search_output", "id": "tso_1", "tools": [{"type": "function", "name": "get_weather"}]},
|
||||
]
|
||||
|
||||
input_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
|
||||
[
|
||||
{"role": "user", "content": "weather in Paris"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [tool_call],
|
||||
"provider_specific_fields": {"tool_search_items": search_items},
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call-ns", "content": "sunny"},
|
||||
]
|
||||
)
|
||||
|
||||
assert input_items == [
|
||||
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "weather in Paris"}]},
|
||||
{"type": "function_call", "call_id": "call-ns", "name": "get_weather", "arguments": "{}"},
|
||||
{"type": "function_call_output", "call_id": "call-ns", "output": [{"type": "input_text", "text": "sunny"}]},
|
||||
]
|
||||
|
||||
|
||||
def test_streamed_tool_search_turn_replays_with_its_reasoning_item():
|
||||
from unittest.mock import Mock
|
||||
|
||||
import litellm
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
OpenAiResponsesToChatCompletionStreamIterator,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
||||
|
||||
function_call = {
|
||||
"type": "function_call",
|
||||
"id": "fc_1",
|
||||
"name": "get_weather",
|
||||
"call_id": "call-ts",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"namespace": "get_weather",
|
||||
"status": "completed",
|
||||
}
|
||||
tool_search_call = {
|
||||
"type": "tool_search_call",
|
||||
"id": "tsc_1",
|
||||
"arguments": {"paths": ["get_weather"]},
|
||||
"call_id": None,
|
||||
"created_by": None,
|
||||
"execution": "server",
|
||||
"status": "completed",
|
||||
}
|
||||
tool_search_output = {
|
||||
"type": "tool_search_output",
|
||||
"id": "tso_1",
|
||||
"call_id": None,
|
||||
"execution": "server",
|
||||
"status": "completed",
|
||||
"tools": [{"type": "function", "name": "get_weather", "parameters": {}, "defer_loading": True}],
|
||||
}
|
||||
reasoning = {"type": "reasoning", "id": "rs_1", "encrypted_content": "enc", "summary": []}
|
||||
later_reasoning = {"type": "reasoning", "id": "rs_2", "encrypted_content": "enc_2", "summary": []}
|
||||
iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
|
||||
chunks = [
|
||||
iterator.chunk_parser({"type": ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "item": function_call}),
|
||||
iterator.chunk_parser(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"status": "completed",
|
||||
"output": [reasoning, tool_search_call, tool_search_output, later_reasoning, function_call],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
}
|
||||
),
|
||||
]
|
||||
user = {"role": "user", "content": "weather in Paris"}
|
||||
message = litellm.stream_chunk_builder(chunks, messages=[user]).choices[0].message
|
||||
|
||||
request = LiteLLMResponsesTransformationHandler().transform_request(
|
||||
model="gpt-5.5",
|
||||
messages=[user, message.model_dump(), {"role": "tool", "tool_call_id": "call-ts", "content": "sunny"}],
|
||||
optional_params={"tools": [{"type": "tool_search"}]},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
|
||||
assert [(item["type"], item.get("id")) for item in request["input"][1:-1]] == [
|
||||
("reasoning", "rs_1"),
|
||||
("tool_search_call", "tsc_1"),
|
||||
("tool_search_output", "tso_1"),
|
||||
("reasoning", "rs_2"),
|
||||
("function_call", None),
|
||||
]
|
||||
assert request["input"][2] == {key: value for key, value in tool_search_call.items() if value is not None}
|
||||
assert request["input"][5]["namespace"] == "get_weather"
|
||||
|
||||
|
||||
def test_transform_response_with_reasoning_and_output():
|
||||
"""Test transform_response handles ResponsesAPIResponse with reasoning items and output messages."""
|
||||
from unittest.mock import Mock
|
||||
|
|
@ -3857,6 +4038,207 @@ def test_transform_response_incomplete_reasoning_only_returns_empty_length_choic
|
|||
assert result.usage.completion_tokens_details.reasoning_tokens == 16
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "incomplete_reason", "finish_reason"),
|
||||
[("completed", None, "stop"), ("incomplete", "max_output_tokens", "length")],
|
||||
)
|
||||
def test_search_only_turn_keeps_tool_search_items_for_the_next_request(status, incomplete_reason, finish_reason):
|
||||
search_items = [
|
||||
{"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}, "status": "completed"},
|
||||
{"type": "tool_search_output", "id": "tso_1", "status": "completed", "tools": [{"name": "get_weather"}]},
|
||||
]
|
||||
raw_response = _make_incomplete_responses_api_response(
|
||||
incomplete_reason,
|
||||
[{"type": "reasoning", "id": "rs_raw", "summary": []}, _make_reasoning_only_output_item(), *search_items],
|
||||
status=status,
|
||||
)
|
||||
|
||||
choice = _call_transform_response(LiteLLMResponsesTransformationHandler(), raw_response).choices[0]
|
||||
|
||||
stored = choice.message.provider_specific_fields["tool_search_items"]
|
||||
assert [item["id"] for item in stored] == ["rs_raw", "rs_incomplete", "tsc_1", "tso_1"]
|
||||
assert stored[2:] == search_items
|
||||
assert [item["id"] for item in choice.message.reasoning_items] == ["rs_raw", "rs_incomplete"]
|
||||
assert choice.finish_reason == finish_reason
|
||||
replayed, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
|
||||
[choice.message.model_dump()], replay_tool_search=True
|
||||
)
|
||||
assert [item["id"] for item in replayed] == ["rs_raw", "rs_incomplete", "tsc_1", "tso_1"]
|
||||
|
||||
|
||||
_CLIENT_SEARCH_CALL: Final = {
|
||||
"type": "tool_search_call",
|
||||
"id": "tsc_client",
|
||||
"arguments": {"query": "weather"},
|
||||
"call_id": "call_search",
|
||||
"execution": "client",
|
||||
"status": "completed",
|
||||
}
|
||||
|
||||
|
||||
def _non_streamed_client_search_turn() -> "ModelResponse":
|
||||
raw_response = _make_incomplete_responses_api_response(None, [_CLIENT_SEARCH_CALL], status="completed")
|
||||
return _call_transform_response(LiteLLMResponsesTransformationHandler(), raw_response)
|
||||
|
||||
|
||||
def _streamed_client_search_turn() -> "ModelResponse":
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
OpenAiResponsesToChatCompletionStreamIterator,
|
||||
)
|
||||
|
||||
iterator: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
|
||||
completed: Final = iterator.chunk_parser(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_client_search",
|
||||
"status": "completed",
|
||||
"output": [_CLIENT_SEARCH_CALL],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
}
|
||||
)
|
||||
built: Final = litellm.stream_chunk_builder([completed], messages=[{"role": "user", "content": "weather in Paris"}])
|
||||
assert isinstance(built, litellm.ModelResponse)
|
||||
return built
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"client_search_turn",
|
||||
[_non_streamed_client_search_turn, _streamed_client_search_turn],
|
||||
ids=["non-streaming", "streaming"],
|
||||
)
|
||||
def test_client_search_turn_returns_its_call_and_replays_with_the_client_output(client_search_turn):
|
||||
choice: Final = client_search_turn().choices[0]
|
||||
client_output: Final = {
|
||||
"type": "tool_search_output",
|
||||
"call_id": "call_search",
|
||||
"execution": "client",
|
||||
"status": "completed",
|
||||
"tools": [{"type": "function", "name": "get_weather", "parameters": {}, "defer_loading": True}],
|
||||
}
|
||||
|
||||
assert choice.finish_reason == "stop"
|
||||
assert choice.message.provider_specific_fields["tool_search_items"] == [_CLIENT_SEARCH_CALL]
|
||||
answered: Final = {
|
||||
**choice.message.model_dump(),
|
||||
"provider_specific_fields": {"tool_search_items": [_CLIENT_SEARCH_CALL, client_output]},
|
||||
}
|
||||
request: Final = LiteLLMResponsesTransformationHandler().transform_request(
|
||||
model="gpt-5.5",
|
||||
messages=[{"role": "user", "content": "weather in Paris"}, answered],
|
||||
optional_params={"tools": [{"type": "tool_search", "execution": "client"}]},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
assert request["input"][1:3] == [_CLIENT_SEARCH_CALL, client_output]
|
||||
|
||||
|
||||
def test_reasoning_after_a_tool_search_replays_after_the_search_it_followed():
|
||||
from unittest.mock import Mock
|
||||
|
||||
from openai.types.responses import ResponseFunctionToolCall
|
||||
|
||||
search_items = [
|
||||
{"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}, "status": "completed"},
|
||||
{"type": "tool_search_output", "id": "tso_1", "status": "completed", "tools": [{"name": "get_weather"}]},
|
||||
]
|
||||
function_calls = [
|
||||
ResponseFunctionToolCall(
|
||||
type="function_call", id=f"fc_{city}", call_id=f"call_{city}", name="get_weather", arguments="{}"
|
||||
)
|
||||
for city in ("paris", "berlin")
|
||||
]
|
||||
raw_response = _make_incomplete_responses_api_response(
|
||||
None,
|
||||
[
|
||||
{"type": "reasoning", "id": "rs_search", "summary": [], "encrypted_content": "enc_1"},
|
||||
*search_items,
|
||||
{"type": "reasoning", "id": "rs_call", "summary": [], "encrypted_content": "enc_2"},
|
||||
*function_calls,
|
||||
],
|
||||
status="completed",
|
||||
)
|
||||
message = _call_transform_response(LiteLLMResponsesTransformationHandler(), raw_response).choices[0].message
|
||||
|
||||
request = LiteLLMResponsesTransformationHandler().transform_request(
|
||||
model="gpt-5.5",
|
||||
messages=[
|
||||
{"role": "user", "content": "weather in Paris and Berlin"},
|
||||
message.model_dump(),
|
||||
{"role": "tool", "tool_call_id": "call_paris", "content": "sunny"},
|
||||
{"role": "tool", "tool_call_id": "call_berlin", "content": "rainy"},
|
||||
],
|
||||
optional_params={"tools": [{"type": "tool_search"}]},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
|
||||
assert [(item["type"], item.get("id", item.get("call_id"))) for item in request["input"][1:]] == [
|
||||
("reasoning", "rs_search"),
|
||||
("tool_search_call", "tsc_1"),
|
||||
("tool_search_output", "tso_1"),
|
||||
("reasoning", "rs_call"),
|
||||
("function_call", "call_paris"),
|
||||
("function_call", "call_berlin"),
|
||||
("function_call_output", "call_paris"),
|
||||
("function_call_output", "call_berlin"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tools", "replayed_types"),
|
||||
[
|
||||
(
|
||||
[{"type": "tool_search"}],
|
||||
["message", "tool_search_call", "tool_search_output", "function_call", "function_call_output"],
|
||||
),
|
||||
(
|
||||
[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}],
|
||||
["message", "function_call", "function_call_output"],
|
||||
),
|
||||
],
|
||||
ids=["tool_search_requested", "tool_search_not_requested"],
|
||||
)
|
||||
def test_tool_search_items_replay_only_to_a_request_that_carries_tool_search(tools, replayed_types):
|
||||
from unittest.mock import Mock
|
||||
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "weather in Paris"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
|
||||
}
|
||||
],
|
||||
"provider_specific_fields": {
|
||||
"tool_search_items": [
|
||||
{"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}},
|
||||
{"type": "tool_search_output", "id": "tso_1", "tools": [{"type": "function", "name": "get_weather"}]},
|
||||
]
|
||||
},
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
|
||||
]
|
||||
|
||||
request: Final = LiteLLMResponsesTransformationHandler().transform_request(
|
||||
model="gpt-5-mini",
|
||||
messages=messages,
|
||||
optional_params={"tools": tools},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
|
||||
assert [item["type"] for item in request["input"]] == replayed_types
|
||||
|
||||
|
||||
def test_transform_response_incomplete_content_filter_maps_finish_reason():
|
||||
handler = LiteLLMResponsesTransformationHandler()
|
||||
raw_response = _make_incomplete_responses_api_response("content_filter", [_make_reasoning_only_output_item()])
|
||||
|
|
|
|||
|
|
@ -1485,6 +1485,39 @@ def test_convert_to_model_response_object_preserves_provider_specific_fields_fro
|
|||
== web_search_results
|
||||
)
|
||||
|
||||
def test_convert_to_model_response_object_keeps_reasoning_items_from_a_cached_response():
|
||||
reasoning_items: Final = [
|
||||
{"id": "rs_a", "type": "reasoning", "summary": [], "encrypted_content": "enc-a"},
|
||||
{"id": "rs_b", "type": "reasoning", "summary": [], "encrypted_content": "enc-b"},
|
||||
]
|
||||
result: Final = convert_to_model_response_object(
|
||||
model_response_object=ModelResponse(),
|
||||
response_object={
|
||||
"id": "chatcmpl-cached",
|
||||
"object": "chat.completion",
|
||||
"created": 1728933352,
|
||||
"model": "gpt-5.6",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": None, "reasoning_items": reasoning_items},
|
||||
"finish_reason": "tool_calls",
|
||||
}
|
||||
],
|
||||
},
|
||||
stream=False,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
hidden_params=None,
|
||||
_response_headers=None,
|
||||
convert_tool_call_to_json_mode=False,
|
||||
)
|
||||
assert isinstance(result, ModelResponse)
|
||||
choice: Final = result.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
assert choice.message.reasoning_items == reasoning_items
|
||||
|
||||
|
||||
def test_convert_to_model_response_object_provider_specific_fields_merges_extra_keys():
|
||||
"""
|
||||
Test that provider_specific_fields from the response are merged with
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ replay yields multiple word-shaped chunks instead of a single one-shot
|
|||
content frame, restoring per-token cadence on cache hits.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
|
@ -218,6 +220,37 @@ async def test_async_tool_calls_and_function_call_only_on_first_chunk():
|
|||
assert c.choices[0].delta.function_call is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_reasoning_items_and_provider_specific_fields_only_on_first_chunk():
|
||||
reasoning_items: Final = [{"id": "rs_1", "type": "reasoning", "summary": [], "encrypted_content": "enc-1"}]
|
||||
provider_specific_fields: Final = {
|
||||
"tool_search_items": [{"id": "tsc_1", "type": "tool_search_call", "arguments": {"query": "weather"}}]
|
||||
}
|
||||
payload: Final = {
|
||||
**_async_payload(),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "cached content slices here",
|
||||
"reasoning_items": reasoning_items,
|
||||
"provider_specific_fields": provider_specific_fields,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
chunks: Final = await _collect_async(payload)
|
||||
assert len(chunks) > 1
|
||||
first_delta: Final = chunks[0].choices[0].delta
|
||||
assert getattr(first_delta, "reasoning_items", None) == reasoning_items
|
||||
assert getattr(first_delta, "provider_specific_fields", None) == provider_specific_fields
|
||||
for c in chunks[1:]:
|
||||
assert getattr(c.choices[0].delta, "reasoning_items", None) is None
|
||||
assert getattr(c.choices[0].delta, "provider_specific_fields", None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_logprobs_only_on_first_chunk():
|
||||
payload = _async_payload(content="cached token logprobs")
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import functools
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Final
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -31,6 +31,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
salvage_concatenated_tool_arguments,
|
||||
split_concatenated_json_objects,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
strip_litellm_response_only_fields,
|
||||
system_messages_first,
|
||||
update_messages_with_model_file_ids,
|
||||
with_prompt_cache_breakpoint,
|
||||
|
|
@ -2334,3 +2335,30 @@ def test_native_redacted_thinking_data_only_unwraps_the_converse_base64_layer(da
|
|||
@pytest.mark.parametrize("data", [_NATIVE_REDACTED_DATA, _CONVERSE_REDACTED_DATA], ids=["native", "legacy-wrapped"])
|
||||
def test_converse_redacted_content_wraps_claude_native_data_exactly_once(data: str) -> None:
|
||||
assert converse_redacted_content(data) == _CONVERSE_REDACTED_DATA
|
||||
|
||||
|
||||
def test_strip_litellm_response_only_fields_drops_nested_tool_call_provider_fields() -> None:
|
||||
message: Final = {
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"provider_specific_fields": {"namespace": "ns"},
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"provider_specific_fields": {"namespace": "ns"},
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
stripped: Final = strip_litellm_response_only_fields(message)
|
||||
|
||||
tool_call: Final = cast("dict[str, object]", stripped["tool_calls"][0])
|
||||
assert "provider_specific_fields" not in tool_call
|
||||
function: Final = cast("dict[str, object]", tool_call["function"])
|
||||
assert "provider_specific_fields" not in function
|
||||
assert function["name"] == "get_weather"
|
||||
assert function["arguments"] == '{"city": "Paris"}'
|
||||
|
|
|
|||
|
|
@ -2138,3 +2138,45 @@ def test_stream_chunk_builder_joins_openai_audio_output_pieces() -> None:
|
|||
"expires_at": 1767225600,
|
||||
"transcript": "Yes.",
|
||||
}
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies():
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
reasoning_items: Final[list[dict[str, object]]] = [
|
||||
{"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "enc_1"}
|
||||
]
|
||||
chunks: Final = [
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-rs",
|
||||
created=1,
|
||||
model="gpt-5.5",
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(content="Hello", role="assistant"),
|
||||
)
|
||||
],
|
||||
),
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-rs",
|
||||
created=1,
|
||||
model="gpt-5.5",
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(content=" world", role="assistant", reasoning_items=reasoning_items),
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
response: Final = stream_chunk_builder(chunks=list(chunks))
|
||||
|
||||
assert response is not None
|
||||
assert response.choices[0].message.content == "Hello world"
|
||||
assert response.choices[0].message.reasoning_items == reasoning_items
|
||||
|
|
|
|||
|
|
@ -494,6 +494,46 @@ def test_transform_messages_helper_strips_thinking_blocks_but_keeps_reasoning_co
|
|||
assert out[1]["content"] == "I can help."
|
||||
|
||||
|
||||
def test_tool_search_reply_sent_to_fireworks_drops_litellm_response_fields():
|
||||
tool_search_reply = {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}],
|
||||
"provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]},
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": "weather"},
|
||||
}
|
||||
],
|
||||
}
|
||||
request = FireworksAIConfig().transform_request(
|
||||
model="accounts/fireworks/models/gpt-oss-120b",
|
||||
messages=[
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
tool_search_reply,
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][1] == {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Regression tests for legacy / OpenAPI $ref defs in tool parameters.
|
||||
#
|
||||
|
|
|
|||
|
|
@ -483,3 +483,43 @@ class TestGroqReasoning:
|
|||
parsed_chunk.choices[0].delta.tool_calls[0]["function"]["name"]
|
||||
== "test_function"
|
||||
)
|
||||
|
||||
|
||||
def test_tool_search_reply_sent_to_groq_drops_litellm_response_fields():
|
||||
tool_search_reply = {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}],
|
||||
"provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]},
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": "weather"},
|
||||
}
|
||||
],
|
||||
}
|
||||
request = GroqChatConfig().transform_request(
|
||||
model="openai/gpt-oss-120b",
|
||||
messages=[
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
tool_search_reply,
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][1] == {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
|
|
|||
|
|
@ -830,6 +830,47 @@ class TestMistralStripsOutputOnlyFields:
|
|||
assert assistant_message["content"] == "Follow-up"
|
||||
assert assistant_message["role"] == "assistant"
|
||||
|
||||
def test_tool_search_reply_drops_litellm_response_fields(self):
|
||||
tool_search_reply = {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}],
|
||||
"provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]},
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": "weather"},
|
||||
}
|
||||
],
|
||||
}
|
||||
messages = cast(
|
||||
List[AllMessageValues],
|
||||
[
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
tool_search_reply,
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
|
||||
],
|
||||
)
|
||||
|
||||
result = cast(
|
||||
List[AllMessageValues],
|
||||
MistralConfig().transform_messages(messages=messages, model="mistral-medium-3-5"),
|
||||
)
|
||||
|
||||
assert result[1] == {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
def test_non_assistant_messages_are_untouched(self):
|
||||
messages = cast(
|
||||
List[AllMessageValues],
|
||||
|
|
|
|||
|
|
@ -94,6 +94,22 @@ class RecordingMaskingGuardrail(MockGuardrail):
|
|||
return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
|
||||
|
||||
|
||||
class ToolCallRecordingGuardrail(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.seen_tool_calls: list[object] = []
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.seen_tool_calls.extend(inputs.get("tool_calls", []))
|
||||
return inputs
|
||||
|
||||
|
||||
class LastTextDroppingGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
|
|
@ -1173,6 +1189,39 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
|
|||
assert exc_info.value.status_code == 400
|
||||
assert "Response blocked by guardrail" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_namespaced_function_call_reaches_output_guardrails_in_the_unchanged_tool_call_shape(self):
|
||||
guardrail: Final = ToolCallRecordingGuardrail(guardrail_name="recorder")
|
||||
response: Final = ResponsesAPIResponse(
|
||||
id="resp_namespaced",
|
||||
created_at=1764901066,
|
||||
model="gpt-5.5",
|
||||
object="response",
|
||||
status="completed",
|
||||
output=[
|
||||
ResponseFunctionToolCall(
|
||||
arguments='{"city":"Paris"}',
|
||||
call_id="call_namespaced",
|
||||
name="get_weather",
|
||||
namespace="weather_tools",
|
||||
type="function_call",
|
||||
id="fc_namespaced",
|
||||
status="completed",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
await OpenAIResponsesHandler().process_output_response(response, guardrail)
|
||||
|
||||
assert guardrail.seen_tool_calls == [
|
||||
{
|
||||
"id": "call_namespaced",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
|
||||
"index": 0,
|
||||
}
|
||||
]
|
||||
|
||||
def test_extract_mixed_content_with_text_and_tool_calls(self):
|
||||
"""Test extracting both text and tool calls from response"""
|
||||
handler = OpenAIResponsesHandler()
|
||||
|
|
|
|||
|
|
@ -488,6 +488,46 @@ def test_transform_request_keeps_reasoning_content_strips_internal_fields():
|
|||
_assert_internal_fields_stripped_reasoning_kept(request["messages"])
|
||||
|
||||
|
||||
def test_tool_search_reply_sent_to_together_drops_litellm_response_fields():
|
||||
tool_search_reply = {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}],
|
||||
"provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]},
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
"provider_specific_fields": {"namespace": "weather"},
|
||||
}
|
||||
],
|
||||
}
|
||||
request = TogetherAIChatConfig().transform_request(
|
||||
model="openai/gpt-oss-120b",
|
||||
messages=[
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
tool_search_reply,
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "together_ai"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][1] == {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
async def test_async_transform_request_keeps_reasoning_content_strips_internal_fields():
|
||||
request = await TogetherAIChatConfig().async_transform_request(
|
||||
model=REASONING_MODEL,
|
||||
|
|
|
|||
|
|
@ -750,6 +750,78 @@ class TestToolPermissionGuardrail:
|
|||
assert new_data["function_call"] == "none"
|
||||
assert new_data["tool_choice"] == "none"
|
||||
|
||||
@staticmethod
|
||||
def _tool_search_turn(*loaded_names: str) -> dict:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"provider_specific_fields": {
|
||||
"tool_search_items": [
|
||||
{"type": "tool_search_call", "call_id": "ts_1", "execution": "client"},
|
||||
{
|
||||
"type": "tool_search_output",
|
||||
"call_id": "ts_1",
|
||||
"tools": [{"type": "function", "name": name} for name in loaded_names],
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
def _replayed_run_shell_request(self) -> dict:
|
||||
return {
|
||||
"tools": [{"type": "function", "function": {"name": "get_weather"}}, {"type": "tool_search"}],
|
||||
"messages": [
|
||||
{"role": "user", "content": "check the weather, then clean up the disk"},
|
||||
self._tool_search_turn("get_weather"),
|
||||
self._tool_search_turn("run_shell"),
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _deny_run_shell(on_disallowed_action: str) -> ToolPermissionGuardrail:
|
||||
return ToolPermissionGuardrail(
|
||||
guardrail_name="test-tool-permission",
|
||||
rules=[{"id": "deny_shell", "tool_name": r"^run_shell$", "decision": "deny"}],
|
||||
default_action="allow",
|
||||
on_disallowed_action=on_disallowed_action,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_blocks_tool_loaded_by_replayed_tool_search(self):
|
||||
guardrail = self._deny_run_shell("block")
|
||||
|
||||
with patch.object(guardrail, "should_run_guardrail", return_value=True):
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
cache=DualCache(default_in_memory_ttl=1),
|
||||
data=self._replayed_run_shell_request(),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert excinfo.value.status_code == 400
|
||||
assert "run_shell" in json.dumps(excinfo.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_rewrite_drops_replayed_tool_search_loading_denied_tool(self):
|
||||
guardrail = self._deny_run_shell("rewrite")
|
||||
data = self._replayed_run_shell_request()
|
||||
|
||||
with patch.object(guardrail, "should_run_guardrail", return_value=True):
|
||||
new_data = await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
cache=DualCache(default_in_memory_ttl=1),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert isinstance(new_data, dict)
|
||||
user_turn, weather_search, shell_search = new_data["messages"]
|
||||
assert user_turn == {"role": "user", "content": "check the weather, then clean up the disk"}
|
||||
assert weather_search == self._tool_search_turn("get_weather")
|
||||
assert shell_search == {"role": "assistant", "content": None, "provider_specific_fields": {}}
|
||||
assert [tool.get("function", {}).get("name") for tool in new_data["tools"]] == ["get_weather", None]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_streaming_iterator_hook_plain_text_yields_chunks(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -186,3 +186,39 @@ async def test_response_tool_calls_as_objects(guardrail):
|
|||
await guardrail.apply_guardrail(
|
||||
inputs=inputs, request_data={}, input_type="response"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blocked_tool_loaded_by_a_replayed_tool_search_raises_next_to_allowed_tools(guardrail):
|
||||
request_data = {
|
||||
"tools": [{"type": "function", "function": {"name": "get_weather"}}, {"type": "tool_search"}],
|
||||
"messages": [
|
||||
{"role": "user", "content": "clean up the disk"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"provider_specific_fields": {
|
||||
"tool_search_items": [
|
||||
{"type": "tool_search_call", "call_id": "ts_1", "execution": "client"},
|
||||
{
|
||||
"type": "tool_search_output",
|
||||
"call_id": "ts_1",
|
||||
"tools": [{"type": "function", "name": "run_shell", "parameters": {"type": "object"}}],
|
||||
},
|
||||
]
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
with patch(
|
||||
"litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
|
||||
return_value=_registry_mock({"get_weather": "untrusted", "run_shell": "blocked"}),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=_tool_request_inputs(["get_weather"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail["blocked_tools"] == ["run_shell"]
|
||||
|
|
|
|||
|
|
@ -227,6 +227,89 @@ class TestCheckToolsAllowlist:
|
|||
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
|
||||
assert "restricted_tool" in str(exc_info.value.message)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"loaded_tool",
|
||||
[
|
||||
{"type": "function", "name": "run_shell", "parameters": {"type": "object"}},
|
||||
{
|
||||
"type": "namespace",
|
||||
"name": "ops",
|
||||
"tools": [{"type": "function", "name": "run_shell", "parameters": {"type": "object"}}],
|
||||
},
|
||||
],
|
||||
ids=["function", "namespace_member"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_disallowed_tool_loaded_by_a_replayed_tool_search_raises_on_chat_route(self, loaded_tool):
|
||||
token = _token(metadata={"allowed_tools": ["get_weather"]})
|
||||
body = {
|
||||
"tools": [{"type": "function", "function": {"name": "get_weather"}}, {"type": "tool_search"}],
|
||||
"messages": [
|
||||
{"role": "user", "content": "clean up the disk"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"provider_specific_fields": {
|
||||
"tool_search_items": [
|
||||
{"type": "tool_search_call", "call_id": "ts_1", "execution": "client"},
|
||||
{
|
||||
"type": "tool_search_output",
|
||||
"call_id": "ts_1",
|
||||
"tools": [{"type": "function", "name": "get_weather"}, loaded_tool],
|
||||
},
|
||||
]
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await check_tools_allowlist(
|
||||
request_body=body,
|
||||
valid_token=token,
|
||||
team_object=None,
|
||||
route="/v1/chat/completions",
|
||||
)
|
||||
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
|
||||
assert "run_shell" in str(exc_info.value.message)
|
||||
assert "get_weather" not in str(exc_info.value.message)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"search_tools",
|
||||
[{"extra_body": {"tools": [{"type": "tool_search"}]}}, {}],
|
||||
ids=["extra_body_tool_search", "deployment_adds_tool_search"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_disallowed_tool_loaded_by_a_replayed_tool_search_raises_without_tool_search_in_tools(
|
||||
self, search_tools
|
||||
):
|
||||
token = _token(metadata={"allowed_tools": ["get_weather"]})
|
||||
body = {
|
||||
"tools": [{"type": "function", "function": {"name": "get_weather"}}],
|
||||
**search_tools,
|
||||
"messages": [
|
||||
{"role": "user", "content": "clean up the disk"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"provider_specific_fields": {
|
||||
"tool_search_items": [
|
||||
{"type": "tool_search_call", "call_id": "ts_1", "execution": "client"},
|
||||
{"type": "tool_search_output", "call_id": "ts_1", "tools": [{"name": "run_shell"}]},
|
||||
]
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await check_tools_allowlist(
|
||||
request_body=body,
|
||||
valid_token=token,
|
||||
team_object=None,
|
||||
route="/v1/chat/completions",
|
||||
)
|
||||
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
|
||||
assert "run_shell" in str(exc_info.value.message)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_allowlist_used_when_key_empty(self):
|
||||
token = _token(
|
||||
|
|
|
|||
|
|
@ -361,6 +361,66 @@ def test_responses_bridge_replays_reasoning_items_across_rounds(respx_mock: resp
|
|||
assert reasoning_index < function_call_index, f"reasoning item must precede function_call: {item_types}"
|
||||
|
||||
|
||||
def test_responses_bridge_replays_tool_search_items_across_rounds(respx_mock: respx.MockRouter) -> None:
|
||||
round_one: Final = _responses_payload(
|
||||
"resp_1",
|
||||
[
|
||||
{"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "enc_1"},
|
||||
{"type": "tool_search_call", "id": "tsc_1", "execution": "server", "arguments": {"paths": ["weather"]}},
|
||||
{
|
||||
"type": "tool_search_output",
|
||||
"id": "tso_1",
|
||||
"execution": "server",
|
||||
"tools": [{"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "get_weather"}]}],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"id": "fc_1",
|
||||
"call_id": "call_1",
|
||||
"name": "get_weather",
|
||||
"namespace": "weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"status": "completed",
|
||||
},
|
||||
],
|
||||
)
|
||||
round_two: Final = _responses_payload(
|
||||
"resp_2",
|
||||
[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Paris is 72F", "annotations": []}],
|
||||
}
|
||||
],
|
||||
)
|
||||
respx_mock.post("https://api.openai.com/v1/responses").mock(
|
||||
side_effect=[httpx.Response(200, json=round_one), httpx.Response(200, json=round_two)]
|
||||
)
|
||||
|
||||
litellm.run_tool_loop(
|
||||
model="openai/responses/gpt-5.5",
|
||||
messages=[{"role": "user", "content": "weather in Paris?"}],
|
||||
tools=({"type": "tool_search"}, *WEATHER_TOOLS),
|
||||
execute_tool=_tool_result,
|
||||
api_key="sk-test",
|
||||
)
|
||||
|
||||
second_input: Final = _request_bodies(respx_mock)[1]["input"]
|
||||
replayed: Final = [
|
||||
(item["type"], item.get("id"), item.get("namespace")) for item in second_input if item["type"] != "message"
|
||||
]
|
||||
assert replayed == [
|
||||
("reasoning", "rs_1", None),
|
||||
("tool_search_call", "tsc_1", None),
|
||||
("tool_search_output", "tso_1", None),
|
||||
("function_call", None, "weather"),
|
||||
("function_call_output", None, None),
|
||||
]
|
||||
|
||||
|
||||
async def test_arun_tool_loop_two_rounds(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue