diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 397523288cf..47cd1db9f39 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -102,6 +102,7 @@ class CodexFallbackModel(LiteLLMBaseModel): experimental_supported_tools: tuple[()] = () context_window: int | None base_instructions: str + supports_search_tool: bool = False class CodexStockUpgrade(LiteLLMBaseModel): @@ -134,6 +135,7 @@ class CodexStockModel(LiteLLMBaseModel): upgrade: CodexStockUpgrade | None = None service_tiers: tuple[CodexServiceTier, ...] = () default_service_tier: str | None = None + supports_search_tool: bool = False supported_reasoning_levels: tuple[CodexReasoningLevel, ...] = () context_window: int | None = None @@ -160,6 +162,7 @@ class CodexCatalogRow: reasoning_effort: str | None = None display_name: str | None = None service_tiers: tuple[object, ...] = () + search_tool_support: tuple[bool | None, ...] = () @cache @@ -252,6 +255,12 @@ def configured_service_tiers( return offered +def offered_search_tool(per_deployment: Sequence[bool | None]) -> bool | None: + if False in per_deployment: + return False + return True if per_deployment and None not in per_deployment else None + + def _codex_can_drive(row: CodexCatalogRow) -> bool: return "*" not in row.id and (row.mode is None or row.mode in CODEX_CHAT_MODES) @@ -306,6 +315,7 @@ def _entry( {tier.id: tier for tier in (stock_model.service_tiers if stock_model is not None else ())} ) configured: Final = configured_service_tiers(row.service_tiers, row.id, stock_tiers) + search_tool: Final = offered_search_tool(row.search_tool_support) if stock_model is None: return CodexFallbackModel( slug=row.id, @@ -314,6 +324,7 @@ def _entry( service_tiers=configured if configured is not None else (), context_window=row.max_input_tokens, base_instructions=instructions, + supports_search_tool=search_tool is True, ) upgrade: Final = ( stock_model.upgrade if stock_model.upgrade is not None and stock_model.upgrade.model in served else None @@ -327,6 +338,7 @@ def _entry( "visibility": "list", "supported_in_api": True, "upgrade": upgrade, + **({} if search_tool is None else {"supports_search_tool": search_tool}), **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), } if configured is None: @@ -433,6 +445,7 @@ def _catalog_row( upstream_model=llm_router.get_routable_upstream_model(lookup_id, team_id) if llm_router is not None else None, display_name=llm_router.get_configured_display_name(label_id) if llm_router is not None else None, service_tiers=llm_router.get_configured_service_tiers(lookup_id, team_id) if llm_router is not None else (), + search_tool_support=llm_router.get_search_tool_support(lookup_id, team_id) if llm_router is not None else (), ) diff --git a/litellm/responses/litellm_completion_transformation/client_tool_search.py b/litellm/responses/litellm_completion_transformation/client_tool_search.py new file mode 100644 index 00000000000..030bb014b7a --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/client_tool_search.py @@ -0,0 +1,268 @@ +import json +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final, cast # noqa: TID251 # loaded tool specs are forwarded as the client sent them + +from openai.types.responses.response_input_param import ( + FunctionCallOutput, + ResponseFunctionToolCallParam, + ResponseInputItemParam, +) +from pydantic import ConfigDict, TypeAdapter, ValidationError + +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.openai import ( + ALL_RESPONSES_API_TOOL_PARAMS, + ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, + ResponseInputParam, +) + +from .custom_tools import function_chat_tool +from .tool_search import CLIENT_TOOL_SEARCH_FUNCTION_NAME, has_client_tool_search + +__all__ = [ + "ClientToolSearchBridge", + "bridge_client_tool_search", + "client_tool_search_chat_tool", + "client_tool_search_loaded_tools", +] + +_DEFAULT_SEARCH_DESCRIPTION: Final = ( + "Search the deferred tools by keyword. The matching tools are loaded and become callable on the next turn." +) +_DEFAULT_SEARCH_PARAMETERS: Final = MappingProxyType( + { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query for deferred tools."}, + "limit": {"type": "integer", "description": "Maximum number of tools to return. Defaults to 8."}, + }, + "required": ["query"], + "additionalProperties": False, + } +) + + +class _ClientToolSearchItem(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + execution: str | None = None + call_id: str | None = None + arguments: object = None + tools: tuple[Mapping[str, object], ...] = () + + +class _ToolSpec(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + name: str = "" + description: str = "" + defer_loading: object = None + tools: tuple[Mapping[str, object], ...] = () + + +class _ClientToolSearchTool(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + description: str = "" + parameters: Mapping[str, object] | None = None + + +@dataclass(frozen=True, slots=True) +class _LoadedTool: + namespace: str | None + namespace_description: str + spec: Mapping[str, object] + + @property + def callable_name(self) -> str: + name: Final = str(self.spec.get("name") or "") + return name if self.namespace is None else f"{self.namespace}__{name}" + + +@dataclass(frozen=True, slots=True) +class ClientToolSearchBridge: + input: str | ResponseInputParam + tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] + bridged: bool + + +def client_tool_search_chat_tool(tool: Mapping[str, object]) -> ChatCompletionToolParam: + spec: Final = _ClientToolSearchTool.model_validate(tool) + function: Final = ChatCompletionToolParamFunctionChunk( + name=CLIENT_TOOL_SEARCH_FUNCTION_NAME, + description=spec.description or _DEFAULT_SEARCH_DESCRIPTION, + parameters=dict(spec.parameters or _DEFAULT_SEARCH_PARAMETERS), + ) + return function_chat_tool(function, None, False) + + +def _client_tool_search_item(item: object) -> _ClientToolSearchItem | None: + try: + parsed: Final = _ClientToolSearchItem.model_validate(item) + except ValidationError: + return None + if parsed.type not in ("tool_search_call", "tool_search_output") or parsed.execution != "client": + return None + return parsed + + +_INPUT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def _searched_tools(item: object) -> tuple[Mapping[str, object], ...]: + parsed: Final = _client_tool_search_item(item) + if parsed is None or parsed.type != "tool_search_output" or not parsed.call_id: + return () + return parsed.tools + + +def client_tool_search_loaded_tools(input: object) -> tuple[Mapping[str, object], ...]: + if not isinstance(input, (list, tuple)): + return () + return tuple(chain.from_iterable(map(_searched_tools, _INPUT_ITEMS.validate_python(input)))) + + +def _tool_spec(tool: object) -> _ToolSpec | None: + try: + return _ToolSpec.model_validate(tool) + except ValidationError: + return None + + +def _as_tool_param(spec: Mapping[str, object]) -> ALL_RESPONSES_API_TOOL_PARAMS: + return cast("ALL_RESPONSES_API_TOOL_PARAMS", spec) # cast-ok: client tool specs are forwarded as sent + + +def _without_deferral(spec: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in spec.items() if key != "defer_loading"} + + +def _loaded_tools(tool: Mapping[str, object]) -> tuple[_LoadedTool, ...]: + spec: Final = _tool_spec(tool) + if spec is None: + return () + if spec.type != "namespace": + return (_LoadedTool(namespace=None, namespace_description="", spec=_without_deferral(tool)),) + return tuple( + _LoadedTool(namespace=spec.name, namespace_description=spec.description, spec=_without_deferral(member)) + for member in spec.tools + ) + + +def _all_loaded_tools(tools: Iterable[Mapping[str, object]]) -> tuple[_LoadedTool, ...]: + return tuple(chain.from_iterable(map(_loaded_tools, tools))) + + +def _callable_names(tool: object) -> tuple[str, ...]: + spec: Final = _tool_spec(tool) + if spec is None: + return () + if spec.type != "namespace": + return (spec.name,) + return tuple(f"{spec.name}__{member.get('name') or ''}" for member in spec.tools) + + +def _first_occurrences(loaded: Sequence[_LoadedTool], taken: frozenset[str]) -> tuple[_LoadedTool, ...]: + first_index: Final = {tool.callable_name: index for index, tool in reversed(tuple(enumerate(loaded)))} + return tuple( + tool + for index, tool in enumerate(loaded) + if first_index[tool.callable_name] == index and tool.callable_name not in taken + ) + + +def _regrouped(loaded: Sequence[_LoadedTool]) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: + namespaces: Final = {tool.namespace: tool.namespace_description for tool in reversed(loaded) if tool.namespace} + ordered_namespaces: Final = tuple(dict.fromkeys(tool.namespace for tool in loaded if tool.namespace)) + top_level: Final = tuple(tool.spec for tool in loaded if tool.namespace is None) + grouped: Final = tuple( + { + "type": "namespace", + "name": namespace, + "description": namespaces[namespace], + "tools": tuple(tool.spec for tool in loaded if tool.namespace == namespace), + } + for namespace in ordered_namespaces + ) + return tuple(_as_tool_param(tool) for tool in (*top_level, *grouped)) + + +def _is_deferred(tool: object) -> bool: + spec: Final = _tool_spec(tool) + return spec is not None and spec.defer_loading is True + + +def _without_deferred(tool: ALL_RESPONSES_API_TOOL_PARAMS) -> ALL_RESPONSES_API_TOOL_PARAMS | None: + spec: Final = _tool_spec(tool) + if spec is None: + return tool + if spec.type != "namespace": + return None if spec.defer_loading is True else tool + eager_members: Final = tuple(member for member in spec.tools if not _is_deferred(member)) + if len(eager_members) == len(spec.tools): + return tool + if not eager_members: + return None + return _as_tool_param({**tool, "tools": eager_members}) + + +def _search_output_text(loaded: Sequence[_LoadedTool]) -> str: + if not loaded: + return "No tools matched the search." + return "Loaded these tools, which are now callable: " + ", ".join(tool.callable_name for tool in loaded) + + +def _search_arguments_json(arguments: object) -> str: + return arguments if isinstance(arguments, str) else json.dumps(arguments or {}) + + +def _bridged_item(item: ResponseInputItemParam) -> ResponseInputItemParam | None: + parsed: Final = _client_tool_search_item(item) + if parsed is None: + return item + if not parsed.call_id: + return None + if parsed.type == "tool_search_call": + return ResponseFunctionToolCallParam( + type="function_call", + call_id=parsed.call_id, + name=CLIENT_TOOL_SEARCH_FUNCTION_NAME, + arguments=_search_arguments_json(parsed.arguments), + ) + return FunctionCallOutput( + type="function_call_output", call_id=parsed.call_id, output=_search_output_text(_all_loaded_tools(parsed.tools)) + ) + + +def bridge_client_tool_search( + input: str | ResponseInputParam, + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, + loaded_earlier: Sequence[Mapping[str, object]] = (), +) -> ClientToolSearchBridge: + existing: Final = tuple(tools or ()) + searches_on_client: Final = has_client_tool_search(existing) + items: Final = () if isinstance(input, str) else tuple(input) + search_items: Final = tuple(parsed for item in items if (parsed := _client_tool_search_item(item)) is not None) + if not searches_on_client and not search_items: + return ClientToolSearchBridge(input=input, tools=existing, bridged=False) + eager: Final = ( + tuple(kept for tool in existing if (kept := _without_deferred(tool)) is not None) + if searches_on_client + else existing + ) + loaded: Final = _all_loaded_tools((*loaded_earlier, *client_tool_search_loaded_tools(input))) + taken: Final = frozenset(chain.from_iterable(map(_callable_names, eager))) + bridged_input: Final = ( + input if isinstance(input, str) else [bridged for item in input if (bridged := _bridged_item(item)) is not None] + ) + return ClientToolSearchBridge( + input=bridged_input, + tools=(*eager, *_regrouped(_first_occurrences(loaded, taken))), + bridged=True, + ) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index ea3bda1b8b5..141e97d902d 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -2,17 +2,21 @@ Handler for transforming responses api requests to litellm.completion requests """ -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass from typing import Final from pydantic import TypeAdapter import litellm -from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools +from litellm.responses.litellm_completion_transformation.client_tool_search import bridge_client_tool_search from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary +from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, StoredTurns from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) +from litellm.responses.litellm_completion_transformation.tool_search import has_client_tool_search from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, with_compaction_items_as_params, @@ -28,6 +32,28 @@ from litellm.types.utils import ModelResponse _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +@dataclass(frozen=True, slots=True) +class _BridgedRequest: + input: str | ResponseInputParam + responses_api_request: ResponsesAPIOptionalRequestParams + + +def _bridged_request( + hoisted: HoistedAdditionalTools, + responses_api_request: ResponsesAPIOptionalRequestParams, + loaded_earlier: Sequence[Mapping[str, object]], +) -> _BridgedRequest: + searched: Final = bridge_client_tool_search(hoisted.input, hoisted.tools, loaded_earlier) + return _BridgedRequest( + input=searched.input, + responses_api_request=( + {**responses_api_request, "tools": list(searched.tools)} + if hoisted.hoisted or searched.bridged + else responses_api_request + ), + ) + + class LiteLLMCompletionTransformationHandler: def response_api_handler( self, @@ -48,10 +74,22 @@ class LiteLLMCompletionTransformationHandler: | Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator] ): hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools")) - bridged_input: Final = hoisted.input - bridged_request: Final[ResponsesAPIOptionalRequestParams] = ( - {**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request - ) + previous_response_id: Final = responses_api_request.get("previous_response_id") + if _is_async and previous_response_id and has_client_tool_search(hoisted.tools): + return self._async_response_api_handler_after_earlier_searches( + previous_response_client_ids, + previous_response_id, + hoisted, + model=model, + responses_api_request=responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + ) + bridged: Final = _bridged_request(hoisted, responses_api_request, ()) + bridged_input: Final = bridged.input + bridged_request: Final = bridged.responses_api_request litellm_completion_request: Final[dict] = ( LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, @@ -73,13 +111,8 @@ class LiteLLMCompletionTransformationHandler: **kwargs, ) - completion_args: Final = {} - completion_args.update(kwargs) - completion_args.update(litellm_completion_request) - completion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = litellm.completion( - **completion_args, + **{**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}, ) if isinstance(litellm_completion_response, ModelResponse): @@ -104,9 +137,53 @@ class LiteLLMCompletionTransformationHandler: ) raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") + async def _async_response_api_handler_after_earlier_searches( + self, + previous_response_client_ids: tuple[str, ...], + previous_response_id: str, + hoisted: HoistedAdditionalTools, + /, + *, + model: str, + responses_api_request: ResponsesAPIOptionalRequestParams, + custom_llm_provider: str | None, + stream: bool | None, + extra_headers: Mapping[str, object] | None, + request_kwargs: Mapping[str, object], + ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: + stored_turns: Final = await ResponsesSessionHandler.get_stored_turns_for_previous_response_id( + previous_response_id, previous_response_client_ids + ) + bridged: Final = _bridged_request( + hoisted, + responses_api_request, + ResponsesSessionHandler.client_tool_search_tools_loaded_before( + stored_turns, previous_response_id, previous_response_client_ids + ), + ) + return await self.async_response_api_handler( + previous_response_client_ids, + stored_turns, + litellm_completion_request=( + LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=bridged.input, + responses_api_request=bridged.responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + **request_kwargs, + ) + ), + request_input=bridged.input, + responses_api_request=bridged.responses_api_request, + **request_kwargs, + ) + async def async_response_api_handler( self, previous_response_client_ids: tuple[str, ...] = (), + stored_turns: StoredTurns | None = None, /, *, litellm_completion_request: dict, @@ -122,16 +199,12 @@ class LiteLLMCompletionTransformationHandler: instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), previous_response_client_ids=previous_response_client_ids, + stored_turns=stored_turns, input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)), ) - acompletion_args: Final = {} - acompletion_args.update(kwargs) - acompletion_args.update(litellm_completion_request) - acompletion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( - **acompletion_args, + **{**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}, ) if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 1ab6293b432..ce2396355a2 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -9,6 +9,10 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.responses.litellm_completion_transformation.client_tool_search import ( + bridge_client_tool_search, + client_tool_search_loaded_tools, +) from litellm.responses.litellm_completion_transformation.compaction_items import ( holds_compaction_summary, input_items, @@ -88,6 +92,12 @@ class _ReplayStart: at_output: bool +@dataclass(frozen=True, slots=True) +class StoredTurns: + spend_logs: tuple["SpendLogsPayload", ...] + proxy_server_requests: tuple[Mapping[str, object] | None, ...] + + def _compaction_point( index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None ) -> _ReplayStart | None: @@ -190,6 +200,7 @@ class ResponsesSessionHandler: previous_response_id: str, item_cache_control: bool = False, client_response_ids: tuple[str, ...] = (), + stored_turns: StoredTurns | None = None, ) -> ChatCompletionSession: """ Return the chat completion message history for a previous response id @@ -200,23 +211,21 @@ class ResponsesSessionHandler: ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") - all_spend_logs: Final = [ - spend_log - for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + turns: Final = ( + stored_turns + if stored_turns is not None + else await ResponsesSessionHandler.get_stored_turns_for_previous_response_id( previous_response_id, client_response_ids ) - if _is_a_conversation_turn(spend_log) - ] + ) + all_spend_logs: Final = turns.spend_logs verbose_proxy_logger.debug("found %s spend logs for this response id", len(all_spend_logs)) litellm_session_id: str | None = None if len(all_spend_logs) > 0: litellm_session_id = all_spend_logs[0].get("session_id") - proxy_server_requests: Final = [ - await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) - for spend_log in all_spend_logs - ] + proxy_server_requests: Final = turns.proxy_server_requests replayable_turns: Final = _turns_through_referenced_response( all_spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) ) @@ -252,6 +261,45 @@ class ResponsesSessionHandler: instructions=_latest_replayable_instructions(proxy_server_requests), ) + @staticmethod + async def get_stored_turns_for_previous_response_id( + previous_response_id: str, + client_response_ids: tuple[str, ...] = (), + ) -> StoredTurns: + spend_logs: Final = tuple( + spend_log + for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + previous_response_id, client_response_ids + ) + if _is_a_conversation_turn(spend_log) + ) + return StoredTurns( + spend_logs=spend_logs, + proxy_server_requests=tuple( + [ + await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) + for spend_log in spend_logs + ] + ), + ) + + @staticmethod + def client_tool_search_tools_loaded_before( + stored_turns: StoredTurns, + previous_response_id: str, + client_response_ids: tuple[str, ...] = (), + ) -> tuple[Mapping[str, object], ...]: + turns: Final = _turns_through_referenced_response( + stored_turns.spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) + ) + return tuple( + chain.from_iterable( + client_tool_search_loaded_tools(request.get("input")) + for request in stored_turns.proxy_server_requests[:turns] + if request is not None + ) + ) + @staticmethod def _history_of_turns( spend_logs: Sequence["SpendLogsPayload"], @@ -310,7 +358,7 @@ class ResponsesSessionHandler: if response_input_param: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=response_input_param, + input=bridge_client_tool_search(response_input_param, None).input, responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, @@ -325,7 +373,7 @@ class ResponsesSessionHandler: if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=responses_api_output, + input=bridge_client_tool_search(responses_api_output, None).input, responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 1e8ca85fdbb..9f0058a5958 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1,7 +1,8 @@ import time import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from itertools import filterfalse +from types import MappingProxyType from typing import Any, Final, cast import litellm @@ -18,9 +19,9 @@ from litellm.responses.litellm_completion_transformation.custom_tools import ( ) from litellm.responses.litellm_completion_transformation.reasoning_items import mint_reasoning_item_id from litellm.responses.litellm_completion_transformation.tool_search import ( + ToolSearchExecution, build_tool_search_call, - has_server_tool_search, - is_anthropic_tool_search_use, + tool_search_routing, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -192,16 +193,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools) self._web_search_calls: dict[str, object] = {} # mutable-ok: latest call by provider id self._queued_web_search_call_ids: set[str] = set() # mutable-ok: emitted call ids - self._tool_search_requested: bool = has_server_tool_search(request_tools) - self._tool_search_call_ids: frozenset[str] = frozenset() + self._tool_search = tool_search_routing(request_tools) + self._tool_search_executions: Mapping[str, ToolSearchExecution] = MappingProxyType({}) - def _is_tool_search_call(self, call_id: str, fn_name: str) -> bool: - if call_id in self._tool_search_call_ids: - return True - if not self._tool_search_requested or not is_anthropic_tool_search_use(call_id, fn_name): - return False - self._tool_search_call_ids = self._tool_search_call_ids | {call_id} - return True + def _tool_search_execution(self, call_id: str, fn_name: str) -> ToolSearchExecution | None: + known: Final = self._tool_search_executions.get(call_id) + if known is not None: + return known + execution: Final = self._tool_search.execution_of(call_id, fn_name) + if execution is None: + return None + self._tool_search_executions = MappingProxyType({**self._tool_search_executions, call_id: execution}) + return execution def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: self._sequence_number += 1 @@ -216,18 +219,22 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: self._pending_tool_events.append(self._next_output_item_added(output_index, item)) - def _queue_tool_search_call_added(self, call_id: str, output_index: int, args_delta: str) -> None: + def _queue_tool_search_call_added( + self, call_id: str, output_index: int, args_delta: str, execution: ToolSearchExecution + ) -> None: if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" - item: Final = build_tool_search_call(call_id, "", "in_progress") + item: Final = build_tool_search_call(call_id, "", "in_progress", execution) self._tool_item_id_by_call_id[call_id] = item.id self._queue_output_item_added(output_index, BaseLiteLLMOpenAIResponseObject(**item.model_dump())) self._tool_args_by_call_id[call_id] += args_delta - def _queue_tool_search_call_done(self, call_id: str, output_index: int, final_args: str) -> None: - self._queue_tool_search_call_added(call_id, output_index, "") + def _queue_tool_search_call_done( + self, call_id: str, output_index: int, final_args: str, execution: ToolSearchExecution + ) -> None: + self._queue_tool_search_call_added(call_id, output_index, "", execution) item: Final = build_tool_search_call( - call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed" + call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed", execution ) self._sequence_number += 1 self._pending_tool_events.append( @@ -436,8 +443,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_held_annotations() output_index = self._get_or_assign_tool_output_index(call_id) - if self._is_tool_search_call(call_id, fn_name): - self._queue_tool_search_call_added(call_id, output_index, fn_args_delta) + if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None: + self._queue_tool_search_call_added(call_id, output_index, fn_args_delta, search_execution) continue self._started_tool_call = True @@ -505,8 +512,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_web_search_events(call_id, web_search_call) self._queued_web_search_call_ids.add(call_id) continue - if self._is_tool_search_call(call_id, fn_name): - self._queue_tool_search_call_done(call_id, output_index, fn_args) + if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None: + self._queue_tool_search_call_done(call_id, output_index, fn_args, search_execution) continue # Track if this is a new tool call that wasn't streamed diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index 780a23e02b2..5a1b27f9092 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -1,4 +1,5 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence, Set +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -18,25 +19,31 @@ if TYPE_CHECKING: __all__ = [ "ANTHROPIC_TOOL_SEARCH_TOOL_NAME", "ANTHROPIC_TOOL_SEARCH_TOOL_TYPE", + "CLIENT_TOOL_SEARCH_FUNCTION_NAME", "CompletionRoute", "ToolSearchCallStatus", + "ToolSearchExecution", + "ToolSearchRouting", "answers_tool_search", "anthropic_tool_search_chat_tool", "build_tool_search_call", "chat_config_builds_anthropic_request", "chat_messages_without_tool_search_calls", "chat_route_rejects_unanswered_tool_calls", + "has_client_tool_search", "has_server_tool_search", "holds_tool_search_use", "input_without_unanswered_anthropic_tool_search_calls", "is_anthropic_tool_search_chat_tool", "is_anthropic_tool_search_use", + "is_client_tool_search_tool", "is_openai_tool_search_tool", "is_server_tool_search_tool", "provider_builds_anthropic_request", "provider_chat_config", "runs_tool_search_on_the_server", "tool_search_call_ids", + "tool_search_routing", "tools_with_anthropic_tool_search_mapped", "tools_without_anthropic_tool_search", "unanswered_tool_search_call_ids", @@ -46,7 +53,9 @@ __all__ = [ ANTHROPIC_TOOL_SEARCH_TOOL_TYPE: Final = "tool_search_tool_regex_20251119" ANTHROPIC_TOOL_SEARCH_TOOL_NAME: Final = "tool_search_tool_regex" +CLIENT_TOOL_SEARCH_FUNCTION_NAME: Final = "tool_search" ToolSearchCallStatus: TypeAlias = Literal["in_progress", "completed", "incomplete"] +ToolSearchExecution: TypeAlias = Literal["server", "client"] _ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object]) _MESSAGE_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -79,6 +88,35 @@ def has_server_tool_search(tools: Sequence[object] | None) -> bool: return any(is_server_tool_search_tool(tool) for tool in tools or ()) +def is_client_tool_search_tool(tool: object) -> bool: + try: + parsed: Final = _ToolSearchFields.model_validate(tool) + except ValidationError: + return False + return parsed.type == "tool_search" and parsed.execution == "client" + + +def has_client_tool_search(tools: Sequence[object] | None) -> bool: + return any(is_client_tool_search_tool(tool) for tool in tools or ()) + + +@dataclass(frozen=True, slots=True) +class ToolSearchRouting: + server: bool + client: bool + + def execution_of(self, tool_id: str, tool_name: str) -> ToolSearchExecution | None: + if self.server and is_anthropic_tool_search_use(tool_id, tool_name): + return "server" + if self.client and tool_name == CLIENT_TOOL_SEARCH_FUNCTION_NAME: + return "client" + return None + + +def tool_search_routing(tools: Sequence[object] | None) -> ToolSearchRouting: + return ToolSearchRouting(server=has_server_tool_search(tools), client=has_client_tool_search(tools)) + + def anthropic_tool_search_chat_tool() -> AnthropicToolSearchToolRegex: return AnthropicToolSearchToolRegex(type=ANTHROPIC_TOOL_SEARCH_TOOL_TYPE, name=ANTHROPIC_TOOL_SEARCH_TOOL_NAME) @@ -346,13 +384,18 @@ def _parsed_arguments(arguments_json: str) -> Mapping[str, object] | None: return None -def build_tool_search_call(tool_id: str, arguments_json: str, status: ToolSearchCallStatus) -> OutputToolSearchCall: +def build_tool_search_call( + tool_id: str, + arguments_json: str, + status: ToolSearchCallStatus, + execution: ToolSearchExecution = "server", +) -> OutputToolSearchCall: arguments: Final = _parsed_arguments(arguments_json) return OutputToolSearchCall( id=openai_shaped_tool_call_item_id("tool_search_call", tool_id), call_id=tool_id, arguments={} if arguments is None else arguments, - execution="server", + execution=execution, status="incomplete" if arguments is None else status, type="tool_search_call", ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 8b076ae922e..ce8841c9190 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -64,6 +64,7 @@ from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, + StoredTurns, ) from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex from litellm.types.llms.base import CachedTokensDetails @@ -124,6 +125,7 @@ from litellm.types.utils import ( Usage, ) +from .client_tool_search import client_tool_search_chat_tool from .compaction_items import ( carries_only_compaction, compaction_output_item, @@ -151,16 +153,17 @@ from .custom_tools import ( from .reasoning_items import decode_thinking_blocks, encode_thinking_blocks, mint_reasoning_item_id from .tool_search import ( CompletionRoute, + ToolSearchRouting, answers_tool_search, anthropic_tool_search_chat_tool, build_tool_search_call, has_server_tool_search, holds_tool_search_use, is_anthropic_tool_search_chat_tool, - is_anthropic_tool_search_use, is_server_tool_search_tool, provider_builds_anthropic_request, tool_search_call_ids, + tool_search_routing, without_tool_search_calls, without_tool_search_uses, ) @@ -1342,6 +1345,7 @@ class LiteLLMCompletionResponsesConfig: request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), input_replaces_history: bool = False, + stored_turns: StoredTurns | None = None, ) -> dict[str, object]: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1357,6 +1361,7 @@ class LiteLLMCompletionResponsesConfig: ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id, item_cache_control=keeps_item_marks, client_response_ids=previous_response_client_ids, + stored_turns=stored_turns, ) chat_completion_session: Final = ( await lookup(previous_response_id) @@ -2818,9 +2823,11 @@ class LiteLLMCompletionResponsesConfig: if tool_type == "custom": converted: Final = convert_custom_tool_to_function_tool(tool) return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) - if tool_type == "tool_search" and validated_tool_search_execution(tool.get("execution")) != "client": + if tool_type == "tool_search" and validated_tool_search_execution(tool.get("execution")) == "client": + return ResponsesToolChatForm(chat_tools=(client_tool_search_chat_tool(tool),), web_search_options=None) + if tool_type == "tool_search": return ResponsesToolChatForm(chat_tools=(anthropic_tool_search_chat_tool(),), web_search_options=None) - if tool_type in ("computer_use", "image_generation", "local_shell", "shell", "tool_search"): + if tool_type in ("computer_use", "image_generation", "local_shell", "shell"): verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", @@ -2960,7 +2967,7 @@ class LiteLLMCompletionResponsesConfig: request_tools: Final = responses_api_request.get("tools") if responses_api_request is not None else None custom_tool_names: Final = extract_custom_tool_names(request_tools) namespace_tool_names: Final = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools) - tool_search_requested: Final = has_server_tool_search(request_tools) + tool_search: Final = tool_search_routing(request_tools) web_search_calls: Final = LiteLLMCompletionResponsesConfig._web_search_calls_by_call_id( chat_completion_response @@ -2977,7 +2984,7 @@ class LiteLLMCompletionResponsesConfig: if ( server_tool_call := LiteLLMCompletionResponsesConfig._server_tool_call_item( - tool_id, tool_name, tool_arguments, web_search_calls, tool_search_requested + tool_id, tool_name, tool_arguments, web_search_calls, tool_search ) ) is not None: responses_tools.append(server_tool_call) @@ -3042,14 +3049,15 @@ class LiteLLMCompletionResponsesConfig: tool_name: str, tool_arguments: str, web_search_calls: Mapping[str, ResponseFunctionWebSearch], - tool_search_requested: bool, + tool_search: ToolSearchRouting, ) -> ResponseFunctionWebSearch | OutputToolSearchCall | None: web_search_call: Final = web_search_calls.get(tool_id) if web_search_call is not None: return web_search_call - if tool_search_requested and is_anthropic_tool_search_use(tool_id, tool_name): - return build_tool_search_call(tool_id, tool_arguments, "completed") - return None + execution: Final = tool_search.execution_of(tool_id, tool_name) + if execution is None: + return None + return build_tool_search_call(tool_id, tool_arguments, "completed", execution) @staticmethod def _web_search_calls_by_call_id( diff --git a/litellm/router.py b/litellm/router.py index e01f6474ded..02ad6275b82 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -723,6 +723,29 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: + from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request + + configured: Final = _configured_model_info(deployment).get("supports_search_tool") + if isinstance(configured, bool): + return configured + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + model: Final = params.get("model") + provider: Final = params.get("custom_llm_provider") + if not isinstance(model, str): + return None + try: + resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( + model=model, custom_llm_provider=provider if isinstance(provider, str) else None + ) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=resolved_model, provider=LlmProviders(resolved_provider) + ) + except (litellm.BadRequestError, ValueError): + return None + return True if chat_config_builds_anthropic_request(chat_config) else None + + def _every_deployment_refused( deployments: Sequence[Mapping[str, object]], refused_deployment_ids: tuple[str, ...] ) -> bool: @@ -12234,6 +12257,15 @@ class Router: for deployment in self._routable_deployments(model_name, team_id) ) + def get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: + """ + Whether each deployment a request from team_id to model_name can route + to (see _routable_deployments) answers Codex's client `tool_search`: + its model_info.supports_search_tool when that is a bool, else True when + its route builds an Anthropic Messages request, else None (unknown). + """ + return tuple(_search_tool_support(deployment) for deployment in self._routable_deployments(model_name, team_id)) + def get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: """ Return the `litellm_params.model` of the first deployment a request diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index f936d23188c..c7781819ae4 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -107,6 +107,7 @@ ignored_function_names = [ "_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py "_request_fallback_list", # Tested through every mid-stream retry test in test_router.py "_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call + "_search_tool_support", # Tested through the test_get_search_tool_support_* tests "_mid_stream_retry_trigger", # Tested through the retry policy mid-stream budget test in test_router.py "_anthropic_messages_group_retry_policy", # Tested through the retry budget precedence test in test_router.py "_anthropic_messages_resolved_retry_policy", # Tested through the malformed retry policy tests in test_router.py diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json index 428fd9b4011..52e5bb43625 100644 --- a/tests/code_coverage_tests/router_surface_allowlist.json +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -102,6 +102,7 @@ "get_routable_upstream_model", "get_router_deployment_budget_limiter", "get_routing_group", + "get_search_tool_support", "get_settings", "get_valid_args", "guardrail_list", diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index c3c5330c440..0e450b08ebe 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -155,3 +155,4 @@ - {id: llm.responses.anthropic.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /v1/responses streams Responses events"} - {id: llm.responses.anthropic.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /v1/responses: cost header and spend row agree"} - {id: llm.responses.anthropic.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic function_call_output round trip over /v1/responses"} +- {id: llm.responses.anthropic.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: tool_search, streaming: nonstream, assertions: [works], source: "llm_translation/test_responses_codex_tool_search_e2e.py", rationale: "Codex client tool_search over /v1/responses finds and loads a deferred tool on Claude"} diff --git a/tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py b/tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py new file mode 100644 index 00000000000..79fd61e247a --- /dev/null +++ b/tests/e2e/llm_translation/test_responses_codex_tool_search_e2e.py @@ -0,0 +1,205 @@ +"""Live e2e: Codex's client tool_search over /v1/responses on Claude. + +Codex holds MCP tools back behind its client-executed `tool_search` tool when the model catalog says +`supports_search_tool`. Over the completion bridge the proxy offers that search to Claude as a plain function, +hands Claude's call back as a `tool_search_call` Codex runs itself, and loads the tools Codex returns in +`tool_search_output` so Claude can call them on the next turn +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from openai.types.responses import ( + FunctionToolParam, + Response, + ResponseFunctionToolCall, + ResponseInputParam, + ResponseToolSearchCall, + ToolParam, + ToolSearchToolParam, +) +from proxy_client import ProxyClient +from pydantic import BaseModel, ConfigDict +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +ANTHROPIC_BACKEND: Final = "anthropic/claude-haiku-5-5" +DEFERRED_TOOL_COUNT: Final = 50 +CALENDAR_TOOL_NAME: Final = "mcp__codex_apps__calendar_create_event" + +CODEX_TOOL_SEARCH: Final[ToolSearchToolParam] = { + "type": "tool_search", + "execution": "client", + "description": ( + "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next " + "model call. Some of the tools may not have been provided to you upfront, and you should use this tool " + "(`tool_search`) to search for the required tools." + ), + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query for deferred tools."}, + "limit": {"type": "number", "description": "Maximum number of tools to return. Defaults to 8."}, + }, + "required": ["query"], + "additionalProperties": False, + }, +} + +SHELL_TOOL: Final[FunctionToolParam] = { + "type": "function", + "name": "shell", + "description": "Runs a shell command.", + "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}, "required": ["cmd"]}, + "strict": False, +} + +CALENDAR_TOOL: Final[FunctionToolParam] = { + "type": "function", + "name": CALENDAR_TOOL_NAME, + "description": "Create a calendar event.", + "parameters": { + "type": "object", + "properties": {"title": {"type": "string"}}, + "required": ["title"], + "additionalProperties": False, + }, + "strict": False, + "defer_loading": True, +} + +PROMPT: Final = "Create a calendar event titled Standup. Find the calendar tool with tool_search first, then call it." + + +class _CalendarArguments(BaseModel): + model_config = ConfigDict(frozen=True) + + title: str + + +def _filler_tool(index: int) -> FunctionToolParam: + return { + "type": "function", + "name": f"mcp__codex_apps__records_{index:02d}", + "description": f"List, inspect, update and archive records of family {index} in the connected workspace.", + "parameters": { + "type": "object", + "properties": { + "title": {"type": "string", "description": "Title of the record to act on"}, + "notes": {"type": "string", "description": "Free-form notes and context for the operation"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "priority": {"type": "integer", "minimum": 0, "maximum": 5}, + }, + "required": ["title"], + "additionalProperties": False, + }, + "strict": False, + "defer_loading": True, + } + + +DEFERRED_TOOLS: Final[tuple[FunctionToolParam, ...]] = ( + *(_filler_tool(index) for index in range(DEFERRED_TOOL_COUNT - 1)), + CALENDAR_TOOL, +) + + +def _eager(tool: FunctionToolParam) -> ToolParam: + return FunctionToolParam( + type="function", + name=tool["name"], + description=tool.get("description"), + parameters=tool.get("parameters"), + strict=False, + ) + + +def _input_tokens(response: Response) -> int: + assert response.usage is not None, f"response carried no usage: {response.id}" + return response.usage.input_tokens + + +def _search_call(response: Response) -> ResponseToolSearchCall: + calls: Final = tuple(item for item in response.output if isinstance(item, ResponseToolSearchCall)) + assert len(calls) == 1, f"expected one tool_search_call, got {response.output!r}" + return calls[0] + + +class TestResponsesCodexToolSearchOverAnthropic: + @pytest.mark.covers("llm.responses.anthropic.tool_search.nonstream.works") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.ANTHROPIC,), + models=(ANTHROPIC_BACKEND,), + capabilities=(Capability.TOOL_SEARCH, Capability.FUNCTION_CALLING), + mode=Mode.NONSTREAM, + ) + ) + def test_claude_finds_a_deferred_tool_through_codex_tool_search_and_calls_it( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = f"e2e-codex-tool-search-{unique_marker()}" + model_id: Final = proxy.create_model( + model, LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY") + ) + resources.defer(lambda: proxy.delete_model(model_id)) + client: Final = sdk.openai(resources.key()) + searchable_tools: Final[list[ToolParam]] = [SHELL_TOOL, CODEX_TOOL_SEARCH] + + eager: Final = client.responses.create( + model=model, + input=PROMPT, + tools=[SHELL_TOOL, *(_eager(tool) for tool in DEFERRED_TOOLS)], + extra_body=NO_PROXY_CACHE, + ) + first: Final = client.responses.create( + model=model, input=PROMPT, tools=searchable_tools, extra_body=NO_PROXY_CACHE + ) + + search: Final = _search_call(first) + assert search.execution == "client", f"tool_search ran on the wrong side: {search!r}" + assert search.call_id, f"tool_search_call had no call_id: {search!r}" + assert not any(isinstance(item, ResponseFunctionToolCall) for item in first.output), first.output + assert _input_tokens(first) * 4 < _input_tokens(eager), ( + f"deferred first turn used {_input_tokens(first)} input tokens against {_input_tokens(eager)} eager" + ) + + follow_up: Final[ResponseInputParam] = [ + {"role": "user", "content": PROMPT}, + { + "type": "tool_search_call", + "call_id": search.call_id, + "execution": "client", + "arguments": search.arguments, + }, + { + "type": "tool_search_output", + "call_id": search.call_id, + "status": "completed", + "execution": "client", + "tools": [CALENDAR_TOOL], + }, + ] + second: Final = client.responses.create( + model=model, input=follow_up, tools=searchable_tools, extra_body=NO_PROXY_CACHE + ) + + calendar_call: Final = next( + ( + item + for item in second.output + if isinstance(item, ResponseFunctionToolCall) and item.name == CALENDAR_TOOL_NAME + ), + None, + ) + assert calendar_call is not None, f"Claude did not call the loaded tool: {second.output!r}" + assert "standup" in _CalendarArguments.model_validate_json(calendar_call.arguments).title.lower() diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 3716162692d..6334fc1482e 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -780,7 +780,7 @@ class TestCodexModelSync: "availability_nux": None, "default_reasoning_level": None, } - assert "supports_search_tool" not in entries["m"] + assert entries["m"]["supports_search_tool"] is False def test_alias_gets_no_upgrade_nudge_away_from_its_own_name(self, tmp_path): listing = self._listing(self._row("gpt-5.6-terra"), self._row("api-54")) @@ -939,7 +939,7 @@ class TestCodexModelSync: assert result == ModelSyncArgs(("-c", f"model_catalog_json={json.dumps(str(catalog))}")) entry = self._entries(tmp_path)["api-terra"] assert entry["slug"] == "api-terra" - assert "supports_search_tool" not in entry + assert entry["supports_search_tool"] is False def test_stock_catalog_is_decoded_as_utf8_regardless_of_locale(self, tmp_path): description = "Modelo equilibrado para el trabajo diario, con acentos y ñ." diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index b43e43df7c9..75b3ee70aab 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -507,8 +507,11 @@ def test_catalog_rows_read_the_router_by_each_entry_lookup_id(): upstream_model="openai/gpt-5.5", display_name="Team 5.5", service_tiers=(["ultrafast"], None), + search_tool_support=(None, None), + ), + CodexCatalogRow( + id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,), search_tool_support=(None,) ), - CodexCatalogRow(id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,)), ) @@ -530,7 +533,12 @@ def test_catalog_rows_read_an_alias_off_its_target_under_the_alias_name(): (row,) = codex_catalog_rows(listing, (("gpt-6", "gpt-6"),), router) assert row == CodexCatalogRow( - id="gpt-6", mode="chat", max_input_tokens=9, upstream_model="openai/gpt-6-astra", service_tiers=(["ultrafast"],) + id="gpt-6", + mode="chat", + max_input_tokens=9, + upstream_model="openai/gpt-6-astra", + service_tiers=(["ultrafast"],), + search_tool_support=(None,), ) (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6", "gpt-6"),), router))["models"] assert (entry["slug"], entry["display_name"], [tier["id"] for tier in entry["service_tiers"]]) == ( @@ -611,3 +619,51 @@ def test_catalog_rows_without_a_router_carry_only_the_listing(): listing = ({"id": "plain", "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"},) assert codex_catalog_rows(listing, (("plain", "plain"),), None) == (CodexCatalogRow(id="plain", mode="chat"),) + + +def test_codex_tool_search_is_offered_off_the_route_and_model_info_and_only_when_every_deployment_answers_it(): + claude = "anthropic/claude-sonnet-5-5" + converse = "bedrock/us.anthropic.claude-sonnet-5-5" + router = litellm.Router( + model_list=[ + {"model_name": "claude", "litellm_params": {"model": claude}}, + { + "model_name": "claude-off", + "litellm_params": {"model": claude}, + "model_info": {"supports_search_tool": False}, + }, + {"model_name": "converse", "litellm_params": {"model": converse}}, + { + "model_name": "converse-on", + "litellm_params": {"model": converse}, + "model_info": {"supports_search_tool": True}, + }, + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5"}}, + { + "model_name": "gpt-5.5-off", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"supports_search_tool": False}, + }, + {"model_name": "mixed", "litellm_params": {"model": claude}}, + {"model_name": "mixed", "litellm_params": {"model": converse}}, + ] + ) + names = ("claude", "claude-off", "converse", "converse-on", "gpt-5.5", "gpt-5.5-off", "mixed") + listing = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "x", "mode": "chat", "max_input_tokens": 1000000} + for name in names + ) + + models = json.loads(codex_model_list_body(listing, tuple((name, name) for name in names), router))["models"] + + stock_support = bundled_codex_models()["gpt-5.5"].model_dump()["supports_search_tool"] + assert {entry["slug"]: entry["supports_search_tool"] for entry in models} == { + "claude": True, + "claude-off": False, + "converse": False, + "converse-on": True, + "gpt-5.5": stock_support, + "gpt-5.5-off": False, + "mixed": False, + } + assert models[0]["context_window"] == 1000000 diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index 36f2a8d5acd..5c6ee0ffe5d 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -27,9 +27,7 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem from .conftest import normalize # type: ignore[import-not-found] -def _stub_model_info_response( - model_id: str = "gpt-4", provider: str = "openai" -) -> dict: +def _stub_model_info_response(model_id: str = "gpt-4", provider: str = "openai") -> dict: return { "id": model_id, "object": "model", @@ -54,6 +52,7 @@ def patched_models(monkeypatch): router.get_routable_upstream_model = MagicMock(return_value="gpt-4") router.get_configured_display_name = MagicMock(return_value=None) router.get_configured_service_tiers = MagicMock(return_value=()) + router.get_search_tool_support = MagicMock(return_value=()) monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) @@ -70,9 +69,7 @@ def patched_models(monkeypatch): def _fake_create_model_info_response(model_id, provider="openai", **kwargs): return _stub_model_info_response(model_id=model_id, provider=provider) - monkeypatch.setattr( - proxy_utils, "create_model_info_response", _fake_create_model_info_response - ) + monkeypatch.setattr(proxy_utils, "create_model_info_response", _fake_create_model_info_response) monkeypatch.setattr(proxy_utils, "validate_model_access", lambda **kwargs: None) @@ -111,9 +108,7 @@ def test_get_models_happy_path(client, auth_as, patched_models, path): @pytest.mark.parametrize("path", ["/v1/models", "/models"]) -def test_get_models_anthropic_format_when_header_present( - client, auth_as, patched_models, path -): +def test_get_models_anthropic_format_when_header_present(client, auth_as, patched_models, path): """Pins: ``GET /v1/models`` returns the Anthropic-native models shape when the caller sends an ``anthropic-version`` header (Claude Code gateway discovery), while the default OpenAI shape is unchanged without it.""" @@ -133,9 +128,7 @@ def test_get_models_anthropic_format_when_header_present( @pytest.mark.parametrize("path", ["/v1/models", "/models"]) -def test_anthropic_format_exposes_token_limits( - client, auth_as, patched_models, monkeypatch, path -): +def test_anthropic_format_exposes_token_limits(client, auth_as, patched_models, monkeypatch, path): """Claude Code sizes requests off the listing, so the Anthropic-native entries carry the same token limits the OpenAI listing resolves, with the output budget named max_tokens as the Messages API names it.""" @@ -150,9 +143,7 @@ def test_anthropic_format_exposes_token_limits( "max_output_tokens": 64000, } - monkeypatch.setattr( - proxy_utils, "create_model_info_response", _create_model_info_response - ) + monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) with auth_as(): response = client.get(path, headers={"anthropic-version": "2023-06-01"}) @@ -227,9 +218,7 @@ def test_anthropic_format_uses_configured_display_name(client, auth_as, patched_ @pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) -def test_anthropic_display_name_resolved_via_internal_team_key( - client, auth_as, patched_models, monkeypatch, params -): +def test_anthropic_display_name_resolved_via_internal_team_key(client, auth_as, patched_models, monkeypatch, params): """For a team-scoped row the configured display name must be looked up by the internal routing key while the entry itself is keyed by the public name, so the clean name lands on the id the client actually sees.""" @@ -263,14 +252,10 @@ def test_anthropic_display_name_resolved_via_internal_team_key( "get_available_models_for_user", _fake_get_available_models_for_user, ) - monkeypatch.setattr( - model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] - ) + monkeypatch.setattr(model_checks, "get_complete_model_list", lambda **kwargs: [internal_name]) with auth_as(): - response = client.get( - "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} - ) + response = client.get("/v1/models", params=params, headers={"anthropic-version": "2023-06-01"}) assert response.status_code == 200 (entry,) = response.json()["data"] @@ -311,9 +296,7 @@ def test_get_model_by_id_not_found(client, auth_as, patched_models, path): @pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) -def test_anthropic_format_returns_public_team_model_name( - client, auth_as, patched_models, monkeypatch, params -): +def test_anthropic_format_returns_public_team_model_name(client, auth_as, patched_models, monkeypatch, params): """Regression: the Anthropic-native listing must go through the same team name translation as the OpenAI listing, so a caller never sees the internal ``model_name_{team_id}_{uuid}`` routing key.""" @@ -344,14 +327,10 @@ def test_anthropic_format_returns_public_team_model_name( "get_available_models_for_user", _fake_get_available_models_for_user, ) - monkeypatch.setattr( - model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] - ) + monkeypatch.setattr(model_checks, "get_complete_model_list", lambda **kwargs: [internal_name]) with auth_as(): - response = client.get( - "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} - ) + response = client.get("/v1/models", params=params, headers={"anthropic-version": "2023-06-01"}) assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"] @@ -382,7 +361,9 @@ def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( patched_models.model_group_alias = {} patched_models.has_model_id.return_value = False - patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: ( + frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + ) monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) with auth_as(): @@ -402,7 +383,9 @@ def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( def test_anthropic_format_keeps_served_ids_for_other_anthropic_clients(client, auth_as, patched_models, path): """An Anthropic SDK asking for the vendor shape gets the served ids: the view is Claude Code's alone.""" with auth_as(): - response = client.get(path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"}) + response = client.get( + path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"} + ) assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4", "claude-sonnet"] diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 36df738c0bb..ef198b487b9 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -17,6 +17,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.litellm_completion_transformation.session_handler import StoredTurns from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -55,6 +56,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): request_kwargs: Mapping[str, object] | None = None, previous_response_client_ids: tuple[str, ...] = (), input_replaces_history: bool = False, + stored_turns: StoredTurns | None = None, ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request @@ -118,8 +120,10 @@ async def test_aresponses_forwards_timeout_to_acompletion(): class _FakeSpendLogsDB: def __init__(self, spend_logs: Sequence[Mapping[str, object]]) -> None: self._spend_logs = spend_logs + self.queries = 0 async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + self.queries += 1 return self._spend_logs @@ -1010,9 +1014,7 @@ class _RecordingCustomLLM(CustomLLM): def completion(self, *args: object, **kwargs: object) -> ModelResponse: self.sent_messages.extend(cast(list[object], kwargs["messages"])) - return litellm.completion( - model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog." - ) + return litellm.completion(model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog.") def test_responses_over_a_custom_provider_send_a_compaction_item_as_summary_text( @@ -1198,3 +1200,372 @@ def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_th ], } ] + + +_CODEX_TOOL_SEARCH: Final = { + "type": "tool_search", + "execution": "client", + "description": "Search the deferred app tools.", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}, "limit": {"type": "integer"}}, + "required": ["query"], + "additionalProperties": False, + }, +} +_SHELL_TOOL: Final = { + "type": "function", + "name": "shell", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}, +} +_CALENDAR_PARAMETERS: Final = { + "type": "object", + "properties": {"title": {"type": "string"}}, + "required": ["title"], + "additionalProperties": False, +} +_CODEX_SEARCH_CALL: Final = { + "type": "tool_search_call", + "call_id": "search-1", + "execution": "client", + "arguments": {"query": "calendar create", "limit": 1}, +} +_CODEX_SEARCH_OUTPUT: Final = { + "type": "tool_search_output", + "call_id": "search-1", + "status": "completed", + "execution": "client", + "tools": [ + { + "type": "function", + "name": "mcp__codex_apps__calendar_create_event", + "description": "Create a calendar event.", + "defer_loading": True, + "parameters": _CALENDAR_PARAMETERS, + } + ], +} + + +def _deferred_namespace(name: str, *members: str) -> Mapping[str, object]: + return { + "type": "namespace", + "name": name, + "description": f"{name} tools", + "tools": [ + {"type": "function", "name": member, "defer_loading": True, "parameters": _CALENDAR_PARAMETERS} + for member in members + ], + } + + +def _claude_tool_use(tool_id: str, name: str, tool_input: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_tool_use", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "tool_use", "id": tool_id, "name": name, "input": tool_input}], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +def _sent(route: respx.Route) -> Mapping[str, object]: + return json.loads(route.calls.last.request.content) + + +def test_responses_over_anthropic_offer_codex_tool_search_as_a_function_and_hold_back_deferred_tools( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "ok"}) + ) + + litellm.responses( + model=_MODEL, + input="Create a calendar event", + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL, _deferred_namespace("mcp__vercel__", "list_projects")], + api_key="mock_api_key", + ) + + sent_tools: Final = _sent(route)["tools"] + assert [tool["name"] for tool in sent_tools] == ["tool_search", "shell"] + assert sent_tools[0]["description"] == _CODEX_TOOL_SEARCH["description"] + assert sent_tools[0]["input_schema"] == _CODEX_TOOL_SEARCH["parameters"] + assert "defer_loading" not in json.dumps(sent_tools) + + +@pytest.mark.parametrize("tool_input", [{"query": "calendar create", "limit": 1}, {"query": "calendar create"}]) +def test_claude_calling_tool_search_comes_back_as_a_client_tool_search_call( + respx_mock: respx.MockRouter, tool_input: Mapping[str, object] +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_search", "tool_search", tool_input) + ) + + response: Final = litellm.responses( + model=_MODEL, input="Create a calendar event", tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], api_key="mock_api_key" + ) + + (item,) = response.output + assert (item.type, item.execution, item.call_id, item.arguments) == ( + "tool_search_call", + "client", + "toolu_search", + tool_input, + ) + + +def test_claude_calling_a_function_named_tool_search_stays_a_function_call_without_codex_tool_search( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_search", "tool_search", {"query": "calendar"}) + ) + own_tool_search: Final = {**_SHELL_TOOL, "name": "tool_search"} + + response: Final = litellm.responses(model=_MODEL, input="Search", tools=[own_tool_search], api_key="mock_api_key") + + (item,) = response.output + assert (item.type, item.name, json.loads(item.arguments)) == ("function_call", "tool_search", {"query": "calendar"}) + + +def test_streamed_claude_tool_search_is_one_client_tool_search_call_item(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_search_stream", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_search", "name": "tool_search", "input": {}}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "calendar create",'}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": ' "limit": 1}'}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input="Create a calendar event", + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + stream=True, + api_key="mock_api_key", + ) + ) + + done: Final = [event.item for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] + assert [(item.type, item.execution, item.call_id, item.arguments) for item in done] == [ + ("tool_search_call", "client", "toolu_search", {"query": "calendar create", "limit": 1}) + ] + assert not any(event.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA for event in events) + assert [item.type for item in events[-1].response.output] == ["tool_search_call"] + + +def test_codex_tool_search_output_loads_the_found_tool_and_answers_the_search_on_anthropic( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_event", "mcp__codex_apps__calendar_create_event", {"title": "Standup"}) + ) + + response: Final = litellm.responses( + model=_MODEL, + input=[ + {"role": "user", "content": "Create a calendar event called Standup"}, + _CODEX_SEARCH_CALL, + _CODEX_SEARCH_OUTPUT, + ], + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + api_key="mock_api_key", + ) + + sent: Final = _sent(route) + assert [(tool["name"], tool["input_schema"]) for tool in sent["tools"][1:]] == [ + ("shell", _SHELL_TOOL["parameters"]), + ("mcp__codex_apps__calendar_create_event", _CALENDAR_PARAMETERS), + ] + assert "defer_loading" not in json.dumps(sent["tools"]) + (assistant, answer) = sent["messages"][1:] + assert [block for block in assistant["content"] if block["type"] == "tool_use"] == [ + {"type": "tool_use", "id": "search-1", "name": "tool_search", "input": {"query": "calendar create", "limit": 1}} + ] + (result,) = answer["content"] + assert result["type"] == "tool_result" and result["tool_use_id"] == "search-1" + assert "mcp__codex_apps__calendar_create_event" in json.dumps(result["content"]) + (item,) = response.output + assert (item.type, item.name, json.loads(item.arguments)) == ( + "function_call", + "mcp__codex_apps__calendar_create_event", + {"title": "Standup"}, + ) + + +def test_a_namespace_codex_loads_through_tool_search_is_callable_and_its_calls_keep_the_namespace( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_claude_tool_use("toolu_projects", "mcp__vercel____list_projects", {"title": "x"}) + ) + loaded: Final = _deferred_namespace("mcp__vercel__", "list_projects", "get_project") + + response: Final = litellm.responses( + model=_MODEL, + input=[ + {"role": "user", "content": "List my Vercel projects"}, + {**_CODEX_SEARCH_CALL, "arguments": {"query": "vercel projects"}}, + {**_CODEX_SEARCH_OUTPUT, "tools": [loaded]}, + ], + tools=[_CODEX_TOOL_SEARCH, loaded], + api_key="mock_api_key", + ) + + assert [tool["name"] for tool in _sent(route)["tools"]] == [ + "tool_search", + "mcp__vercel____list_projects", + "mcp__vercel____get_project", + ] + (item,) = response.output + assert (item.type, item.namespace, item.name) == ("function_call", "mcp__vercel__", "list_projects") + + +def test_codex_tool_search_output_loading_the_same_tool_twice_sends_it_once(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "ok"}) + ) + + litellm.responses( + model=_MODEL, + input=[ + {"role": "user", "content": "Create two events"}, + _CODEX_SEARCH_CALL, + _CODEX_SEARCH_OUTPUT, + {**_CODEX_SEARCH_CALL, "call_id": "search-2"}, + {**_CODEX_SEARCH_OUTPUT, "call_id": "search-2"}, + ], + tools=[_CODEX_TOOL_SEARCH], + api_key="mock_api_key", + ) + + assert [tool["name"] for tool in _sent(route)["tools"]] == ["tool_search", "mcp__codex_apps__calendar_create_event"] + + +def _stored_tool_call_turn( + request_id: str, request_input: object, tool_call_id: str, name: str, arguments: str +) -> Mapping[str, object]: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "codex-session", + "proxy_server_request": {"model": _MODEL, "input": request_input, "tools": [_CODEX_TOOL_SEARCH, _SHELL_TOOL]}, + "response": { + "id": request_id, + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": tool_call_id, "type": "function", "function": {"name": name, "arguments": arguments}} + ], + }, + } + ], + }, + } + + +_SEARCH_TURN: Final = _stored_tool_call_turn( + "chatcmpl-search-turn", + "Create a calendar event titled Standup.", + "search-1", + "tool_search", + '{"query": "calendar"}', +) +_SEARCH_TURN_STORED_AS_RESPONSE: Final = { + **_SEARCH_TURN, + "response": {"id": "chatcmpl-search-turn", "object": "response", "output": [_CODEX_SEARCH_CALL]}, +} +_CALENDAR_TURN: Final = _stored_tool_call_turn( + "chatcmpl-calendar-turn", + [_CODEX_SEARCH_OUTPUT], + "toolu_standup", + "mcp__codex_apps__calendar_create_event", + '{"title": "Standup"}', +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "search_turn", [_SEARCH_TURN, _SEARCH_TURN_STORED_AS_RESPONSE], ids=["chat_completion", "responses_output"] +) +async def test_previous_response_id_keeps_the_tools_an_earlier_codex_tool_search_loaded( + search_turn: Mapping[str, object], +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + prisma_client: Final = _FakePrismaClient([search_turn, _CALENDAR_TURN]) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma_client): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-calendar-turn", + input=[ + {"type": "function_call_output", "call_id": "toolu_standup", "output": "Created Standup."}, + {"role": "user", "content": "Now create another event titled Retro."}, + ], + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + sent: Final = json.loads(anthropic.body) + assert [tool["name"] for tool in sent["tools"]] == [ + "tool_search", + "shell", + "mcp__codex_apps__calendar_create_event", + ] + assert [ + [(block["type"], block.get("id") or block.get("tool_use_id")) for block in message["content"]] + for message in sent["messages"][1:] + ] == [ + [("tool_use", "search-1")], + [("tool_result", "search-1")], + [("tool_use", "toolu_standup")], + [("tool_result", "toolu_standup"), ("text", None)], + ] + assert prisma_client.db.queries == 1 diff --git a/tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py b/tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py new file mode 100644 index 00000000000..dc975b3d112 --- /dev/null +++ b/tests/unit/responses/litellm_completion_transformation/test_client_tool_search.py @@ -0,0 +1,68 @@ +from typing import Final + +from litellm.responses.litellm_completion_transformation.client_tool_search import bridge_client_tool_search +from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam + +_CODEX_TOOL_SEARCH: Final = {"type": "tool_search", "execution": "client"} +_SHELL: Final = {"type": "function", "name": "shell", "parameters": {"type": "object", "properties": {}}} +_SEARCH_CALL: Final = { + "type": "tool_search_call", + "call_id": "search-1", + "execution": "client", + "arguments": {"query": "shell"}, +} + + +def _search_output(*tools: dict[str, object]) -> dict[str, object]: + return { + "type": "tool_search_output", + "call_id": "search-1", + "status": "completed", + "execution": "client", + "tools": list(tools), + } + + +def test_requests_without_client_tool_search_pass_through_untouched() -> None: + request_input: Final[ResponseInputParam] = [ + {"type": "message", "role": "user", "content": "hi"}, + {"type": "tool_search_call", "call_id": "srv-1", "execution": "server", "arguments": {"query": "x"}}, + ] + tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = [_SHELL, {"type": "tool_search", "execution": "server"}] + + bridged: Final = bridge_client_tool_search(request_input, tools) + + assert bridged.bridged is False + assert bridged.input is request_input + assert bridged.tools == tuple(tools) + + +def test_a_search_that_found_nothing_answers_the_call_and_loads_no_tools() -> None: + bridged: Final = bridge_client_tool_search([_SEARCH_CALL, _search_output()], [_CODEX_TOOL_SEARCH, _SHELL]) + + assert bridged.tools == (_CODEX_TOOL_SEARCH, _SHELL) + assert bridged.input[1] == { + "type": "function_call_output", + "call_id": "search-1", + "output": "No tools matched the search.", + } + + +def test_a_loaded_tool_already_offered_eagerly_is_not_sent_twice() -> None: + loaded_shell: Final = {**_SHELL, "description": "loaded copy", "defer_loading": True} + + bridged: Final = bridge_client_tool_search( + [_SEARCH_CALL, _search_output(loaded_shell)], [_CODEX_TOOL_SEARCH, _SHELL] + ) + + assert bridged.tools == (_CODEX_TOOL_SEARCH, _SHELL) + + +def test_deferred_request_tools_are_held_back_only_when_the_client_searches() -> None: + deferred: Final = {**_SHELL, "name": "deploy", "defer_loading": True} + + with_search: Final = bridge_client_tool_search("hi", [_CODEX_TOOL_SEARCH, _SHELL, deferred]) + without_search: Final = bridge_client_tool_search("hi", [_SHELL, deferred]) + + assert with_search.tools == (_CODEX_TOOL_SEARCH, _SHELL) + assert without_search.tools == (_SHELL, deferred) diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 475628569a0..4250c78c669 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2,6 +2,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from copy import deepcopy +from types import MappingProxyType from typing import Final, Literal, cast import pytest @@ -25,6 +26,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, ) from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks +from litellm.responses.litellm_completion_transformation.tool_search import tool_search_routing from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, ChatCompletionSession, @@ -1288,12 +1290,7 @@ class TestFunctionCallTransformation: def test_parallel_tool_calls_dropped_when_no_chat_tools_remain(self) -> None: transform: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request - codex_tool_search: Final = { - "type": "tool_search", - "execution": "client", - "description": "Searches over deferred tool metadata with BM25.", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, - } + hosted_tool_search: Final = {"type": "tool_search", "execution": "server"} function_tool: Final = { "type": "function", "name": "get_goal", @@ -1311,7 +1308,7 @@ class TestFunctionCallTransformation: hosted_only_result: Final = transform( model="azure/gpt-5.4-mini", input="Reply with just the word pong.", - responses_api_request={"tools": [codex_tool_search], "parallel_tool_calls": True}, + responses_api_request={"tools": [hosted_tool_search], "parallel_tool_calls": True}, custom_llm_provider="azure", ) function_tools_result: Final = transform( @@ -1770,7 +1767,7 @@ class TestToolTransformation: assert len(result_tools) == 0 assert web_search_options is None - def test_transform_codex_tools_drops_hosted_tool_search(self) -> None: + def test_transform_codex_tools_offers_client_tool_search_as_a_function(self) -> None: codex_tools: Final = [ { "type": "function", @@ -1823,16 +1820,22 @@ class TestToolTransformation: }, {"type": "web_search", "external_web_access": False, "search_content_types": ["text", "image"]}, ] - function_and_custom_count: Final = sum(1 for tool in codex_tools if tool["type"] in ("function", "custom")) + codex_search: Final = codex_tools[3] ( result_tools, web_search_options, ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=codex_tools) - assert not any(tool.get("type") == "tool_search" for tool in result_tools) assert all(tool.get("type") == "function" for tool in result_tools) - assert len(result_tools) == function_and_custom_count + assert [tool["function"]["name"] for tool in result_tools] == [ + "exec_command", + "write_stdin", + "apply_patch", + "tool_search", + ] + assert result_tools[3]["function"]["description"] == codex_search["description"] + assert result_tools[3]["function"]["parameters"] == codex_search["parameters"] assert web_search_options is not None @pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) @@ -1877,9 +1880,6 @@ class TestToolTransformation: assert [tool["type"] for tool in result["tools"]] == ["function", "function"] assert all("defer_loading" not in tool for tool in result["tools"]) - @pytest.mark.parametrize( - "search_tools", [[], [{"type": "tool_search", "execution": "client"}]], ids=["none", "client"] - ) @pytest.mark.parametrize( ("model", "provider"), [ @@ -1890,7 +1890,7 @@ class TestToolTransformation: ], ) def test_claude_routes_keep_the_callers_deferral_when_no_server_search_tool_is_sent( - self, model: str, provider: str, search_tools: list[dict[str, object]] + self, model: str, provider: str ) -> None: function_tools: Final = [ tool for tool in _deferred_tools_with_server_tool_search() if tool["type"] != "tool_search" @@ -1899,7 +1899,7 @@ class TestToolTransformation: result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, input="What is the weather in Paris right now?", - responses_api_request={"tools": [*search_tools, *function_tools]}, + responses_api_request={"tools": function_tools}, custom_llm_provider=provider, ) @@ -4361,8 +4361,8 @@ class TestEnsureOutputItemContentPartAdded: iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(None) iterator._web_search_calls = {} iterator._queued_web_search_call_ids = set() - iterator._tool_search_requested = False - iterator._tool_search_call_ids = set() + iterator._tool_search = tool_search_routing(None) + iterator._tool_search_executions = MappingProxyType({}) return iterator def _make_text_chunk(self): @@ -6754,7 +6754,6 @@ def test_summary_replay_drops_the_history_its_compaction_item_replaces() -> None ] - @pytest.mark.parametrize( ("call", "output"), [ @@ -6905,7 +6904,14 @@ def test_compaction_replay_invents_no_call_for_an_output_whose_call_was_never_st [ pytest.param( _COMPACTION_REQUESTED, - [("user", ()), ("assistant", ("call_0",)), ("tool", "Sunny."), ("assistant", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + [ + ("user", ()), + ("assistant", ("call_0",)), + ("tool", "Sunny."), + ("assistant", ()), + ("assistant", ("call_0",)), + ("tool", "Rain."), + ], id="native_compaction_block", ), pytest.param( @@ -7047,7 +7053,9 @@ async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chun assert incomplete_event.type == "response.incomplete" assert incomplete_event.response.status == "incomplete" assert incomplete_event.response.usage is not None - assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text) + assert incomplete_event.response.usage.output_tokens == litellm.token_counter( + model="gpt-4o-mini", text=relayed_text + ) assert incomplete_event.response.usage.input_tokens > 0 diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d59fb69dcaf..b55b3c5efc2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -59,6 +59,7 @@ from litellm.router import ( _is_retriable_anthropic_status, _responses_stream_holds_event, _without_line_breaks, + _search_tool_support, Span, ) from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request @@ -13517,6 +13518,80 @@ def test_get_configured_service_tiers_reads_only_the_deployments_a_request_can_r assert router.get_configured_service_tiers("gpt-6-astra") == (["ultrafast"],) +@pytest.mark.parametrize( + "litellm_params, model_info, expected", + [ + ({"model": "anthropic/claude-fable-5-1"}, None, True), + ({"model": "bedrock/invoke/us.anthropic.claude-fable-5-1"}, None, True), + ({"model": "openai/gpt-6-astra"}, None, None), + ({"model": "gemini/gemini-3.8-flash"}, None, None), + ({"model": "sonar-pro", "api_base": "https://api.perplexity.ai"}, None, None), + ({"model": "gemini/gemini-3.8-flash"}, {"supports_search_tool": True}, True), + ({"model": "anthropic/claude-fable-5-1"}, {"supports_search_tool": False}, False), + ], + ids=[ + "anthropic-defaults-on", + "bedrock-invoke-anthropic-defaults-on", + "openai-passthrough-unset", + "other-bridged-provider-unset", + "provider-inferred-from-api-base-unset", + "model-info-opt-in-wins", + "model-info-opt-out-wins", + ], +) +def test_get_search_tool_support_defaults_from_the_provider_and_lets_model_info_override( + litellm_params: dict[str, str], model_info: dict[str, bool] | None, expected: bool | None +): + deployment: Final = {"model_name": "codex-model", "litellm_params": litellm_params} + router: Final = litellm.Router(model_list=[{**deployment, "model_info": model_info} if model_info else deployment]) + + assert router.get_search_tool_support("codex-model") == (expected,) + + +def test_get_search_tool_support_reports_each_routable_deployment_of_the_group(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "mixed", "litellm_params": {"model": "anthropic/claude-fable-5-1"}}, + {"model_name": "mixed", "litellm_params": {"model": "openai/gpt-6-astra"}}, + { + "model_name": "mixed", + "litellm_params": {"model": "anthropic/claude-fable-5-1", "api_base": "https://paused.example"}, + "model_info": {"blocked": True, "supports_search_tool": False}, + }, + ] + ) + + assert router.get_search_tool_support("mixed") == (True, None) + assert router.get_search_tool_support("not-a-real-model") == () + + +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=False) +@pytest.mark.parametrize("provider", ["chatgpt", "github_copilot"]) +@pytest.mark.parametrize("explicit_provider", [False, True]) +async def test_search_tool_catalog_lookup_needs_no_credentials( + provider: str, explicit_provider: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "chatgpt")) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path / "copilot")) + (tmp_path / "chatgpt").mkdir() + (tmp_path / "chatgpt" / "auth.json").write_text( + json.dumps({"access_token": "expired-access", "refresh_token": "test-refresh", "expires_at": 1}) + ) + (tmp_path / "copilot").mkdir() + (tmp_path / "copilot" / "api-key.json").write_text(json.dumps({"token": "expired-key", "expires_at": 1})) + (tmp_path / "copilot" / "access-token").write_text("test-access") + respx_mock.route().respond(401, json={"error": "invalid_grant"}) + params: Final = ( + {"model": "catalog-model", "custom_llm_provider": provider} + if explicit_provider + else {"model": f"{provider}/catalog-model"} + ) + + assert _search_tool_support({"model_name": "catalog-model", "litellm_params": params}) is None + assert len(respx_mock.calls) == 0 + + @pytest.mark.parametrize( "alias_value, expected_group", [