litellm/tests/unit/test_tool_loop.py

import json
from typing import Final

import httpx
import pytest
import respx

import litellm
from litellm.tool_loop import ToolLoopMaxRoundsExceeded
from litellm.types.llms.openai import ChatCompletionToolMessage
from litellm.types.utils import ChatCompletionMessageToolCall

OPENAI_CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions"
WEATHER_TOOLS: Final = (
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the weather for a city",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string"}},
                "required": ["city"],
            },
        },
    },
)


def _openai_response(content: str | None, tool_calls: list | None = None, **message_fields: object) -> dict:
    return {
        "id": "chatcmpl-tool-loop",
        "object": "chat.completion",
        "created": 1739462947,
        "model": "gpt-5-mini",
        "choices": [
            {
                "index": 0,
                "finish_reason": "tool_calls" if tool_calls else "stop",
                "message": {
                    "role": "assistant",
                    "content": content,
                    "tool_calls": tool_calls,
                    **message_fields,
                },
            }
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
    }


def _tool_call(call_id: str, name: str, arguments: dict) -> dict:
    return {
        "id": call_id,
        "type": "function",
        "function": {"name": name, "arguments": json.dumps(arguments)},
    }


def _tool_result(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
    return ChatCompletionToolMessage(role="tool", content='{"temp": "72F"}', tool_call_id=tool_call.id or "")


def _request_bodies(respx_mock: respx.MockRouter) -> list[dict]:
    return [json.loads(call.request.content) for call in respx_mock.calls]


def test_final_answer_without_tool_calls_returns_content(respx_mock: respx.MockRouter) -> None:
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        return_value=httpx.Response(200, json=_openai_response("done"))
    )
    executor_called: Final = []

    def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executor_called.append(tc)
        return _tool_result(tc)

    answer: Final = litellm.run_tool_loop(
        model="openai/gpt-5-mini",
        messages=[{"role": "user", "content": "hi"}],
        tools=WEATHER_TOOLS,
        execute_tool=executor,
        api_key="sk-test",
    )

    assert answer == "done"
    assert executor_called == []
    assert route.call_count == 1


def test_two_rounds_appends_assistant_and_tool_messages_in_order(respx_mock: respx.MockRouter) -> None:
    tool_calls: Final = [
        _tool_call("call_1", "get_weather", {"city": "Paris"}),
        _tool_call("call_2", "get_weather", {"city": "Tokyo"}),
    ]
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        side_effect=[
            httpx.Response(200, json=_openai_response(None, tool_calls)),
            httpx.Response(200, json=_openai_response("Paris 72F, Tokyo 60F")),
        ]
    )
    executed: Final = []

    def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executed.append(tc)
        return _tool_result(tc)

    messages: Final = [{"role": "user", "content": "weather in Paris and Tokyo?"}]
    messages_snapshot: Final = [dict(message) for message in messages]

    answer: Final = litellm.run_tool_loop(
        model="openai/gpt-5-mini",
        messages=messages,
        tools=WEATHER_TOOLS,
        execute_tool=executor,
        api_key="sk-test",
    )

    assert answer == "Paris 72F, Tokyo 60F"
    assert route.call_count == 2
    assert [tc.id for tc in executed] == ["call_1", "call_2"]
    assert [tc.function.name for tc in executed] == ["get_weather", "get_weather"]
    assert [tc.function.arguments for tc in executed] == [
        '{"city": "Paris"}',
        '{"city": "Tokyo"}',
    ]

    second_body: Final = _request_bodies(respx_mock)[1]
    assert second_body["messages"] == [
        {"role": "user", "content": "weather in Paris and Tokyo?"},
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": "call_1",
                    "type": "function",
                    "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
                },
                {
                    "id": "call_2",
                    "type": "function",
                    "function": {"name": "get_weather", "arguments": '{"city": "Tokyo"}'},
                },
            ],
        },
        {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_1"},
        {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_2"},
    ]

    assert len(messages) == len(messages_snapshot)
    assert messages == messages_snapshot


def test_response_format_and_tools_forwarded_every_round(respx_mock: respx.MockRouter) -> None:
    respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        side_effect=[
            httpx.Response(200, json=_openai_response(None, [_tool_call("call_1", "get_weather", {"city": "Paris"})])),
            httpx.Response(200, json=_openai_response('{"summary": "sunny"}')),
        ]
    )
    response_format: Final = {
        "type": "json_schema",
        "json_schema": {
            "name": "weather_report",
            "schema": {
                "type": "object",
                "properties": {"summary": {"type": "string"}},
                "required": ["summary"],
            },
        },
    }

    litellm.run_tool_loop(
        model="openai/gpt-5-mini",
        messages=[{"role": "user", "content": "weather?"}],
        tools=WEATHER_TOOLS,
        execute_tool=_tool_result,
        response_format=response_format,
        api_key="sk-test",
    )

    bodies: Final = _request_bodies(respx_mock)
    assert len(bodies) == 2
    for body in bodies:
        assert body["response_format"] == response_format
        assert body["tools"] == list(WEATHER_TOOLS)


def test_max_rounds_exceeded_raises_without_executing_last_round(respx_mock: respx.MockRouter) -> None:
    tool_call: Final = _tool_call("call_1", "get_weather", {"city": "Paris"})
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        return_value=httpx.Response(200, json=_openai_response(None, [tool_call]))
    )
    executed: Final = []

    def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executed.append(tc)
        return _tool_result(tc)

    with pytest.raises(ToolLoopMaxRoundsExceeded) as exc_info:
        litellm.run_tool_loop(
            model="openai/gpt-5-mini",
            messages=[{"role": "user", "content": "weather?"}],
            tools=WEATHER_TOOLS,
            execute_tool=executor,
            max_rounds=2,
            api_key="sk-test",
        )

    assert exc_info.value.max_rounds == 2
    assert route.call_count == 2
    assert [tc.id for tc in executed] == ["call_1"]


def test_max_rounds_below_one_rejected_before_any_request(respx_mock: respx.MockRouter) -> None:
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        return_value=httpx.Response(200, json=_openai_response("done"))
    )

    with pytest.raises(ValueError, match="max_rounds must be >= 1"):
        litellm.run_tool_loop(
            model="openai/gpt-5-mini",
            messages=[{"role": "user", "content": "hi"}],
            tools=WEATHER_TOOLS,
            execute_tool=_tool_result,
            max_rounds=0,
            api_key="sk-test",
        )

    assert route.call_count == 0


def test_stream_rejected_before_any_request(respx_mock: respx.MockRouter) -> None:
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        return_value=httpx.Response(200, json=_openai_response("done"))
    )

    with pytest.raises(ValueError, match="stream=True is not supported"):
        litellm.run_tool_loop(
            model="openai/gpt-5-mini",
            messages=[{"role": "user", "content": "hi"}],
            tools=WEATHER_TOOLS,
            execute_tool=_tool_result,
            stream=True,
            api_key="sk-test",
        )

    assert route.call_count == 0


def test_custom_tool_call_raises_type_error_without_executing(respx_mock: respx.MockRouter) -> None:
    custom_response: Final = _openai_response(
        None,
        [{"id": "call_custom", "type": "custom", "custom": {"name": "apply_patch", "input": "*** patch"}}],
    )
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        return_value=httpx.Response(200, json=custom_response)
    )
    executor_called: Final = []

    def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executor_called.append(tc)
        return _tool_result(tc)

    with pytest.raises(TypeError, match="custom tool call call_custom"):
        litellm.run_tool_loop(
            model="openai/gpt-5-mini",
            messages=[{"role": "user", "content": "hi"}],
            tools=WEATHER_TOOLS,
            execute_tool=executor,
            api_key="sk-test",
        )

    assert executor_called == []
    assert route.call_count == 1


def _responses_payload(response_id: str, output: list) -> dict:
    return {
        "id": response_id,
        "object": "response",
        "created_at": 1734366691,
        "status": "completed",
        "model": "gpt-5.5",
        "output": output,
        "parallel_tool_calls": True,
        "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
        "error": None,
        "incomplete_details": None,
        "instructions": None,
        "metadata": None,
        "temperature": None,
        "tool_choice": "auto",
        "tools": [],
        "top_p": None,
        "max_output_tokens": None,
        "previous_response_id": None,
        "reasoning": None,
        "truncation": None,
        "user": None,
    }


def test_responses_bridge_replays_reasoning_items_across_rounds(respx_mock: respx.MockRouter) -> None:
    round_one: Final = _responses_payload(
        "resp_1",
        [
            {"type": "reasoning", "id": "rs_abc123", "summary": [], "encrypted_content": "enc_xyz"},
            {
                "type": "function_call",
                "id": "fc_1",
                "call_id": "call_1",
                "name": "get_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": []}],
            }
        ],
    )
    route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
        side_effect=[httpx.Response(200, json=round_one), httpx.Response(200, json=round_two)]
    )
    executed: Final = []

    def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executed.append(tc)
        return _tool_result(tc)

    answer: Final = litellm.run_tool_loop(
        model="openai/responses/gpt-5.5",
        messages=[{"role": "user", "content": "weather in Paris?"}],
        tools=WEATHER_TOOLS,
        execute_tool=executor,
        api_key="sk-test",
    )

    assert answer == "Paris is 72F"
    assert route.call_count == 2
    assert [tc.id for tc in executed] == ["fc_1"]

    second_input: Final = _request_bodies(respx_mock)[1]["input"]
    item_types: Final = [item.get("type") for item in second_input]
    reasoning_index: Final = next(i for i, item in enumerate(second_input) if item.get("type") == "reasoning")
    function_call_index: Final = next(
        i for i, item in enumerate(second_input) if item.get("type") == "function_call"
    )
    reasoning_item: Final = second_input[reasoning_index]
    assert reasoning_item["id"] == "rs_abc123"
    assert reasoning_item["encrypted_content"] == "enc_xyz"
    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

    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler())
    tool_calls: Final = [
        _tool_call("call_1", "get_weather", {"city": "Paris"}),
        _tool_call("call_2", "get_weather", {"city": "Tokyo"}),
    ]
    route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        side_effect=[
            httpx.Response(200, json=_openai_response(None, tool_calls)),
            httpx.Response(200, json=_openai_response("Paris 72F, Tokyo 60F")),
        ]
    )
    executed: Final = []

    async def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executed.append(tc)
        return _tool_result(tc)

    messages: Final = [{"role": "user", "content": "weather in Paris and Tokyo?"}]
    messages_snapshot: Final = [dict(message) for message in messages]

    answer: Final = await litellm.arun_tool_loop(
        model="openai/gpt-5-mini",
        messages=messages,
        tools=WEATHER_TOOLS,
        execute_tool=executor,
        api_key="sk-test",
    )

    assert answer == "Paris 72F, Tokyo 60F"
    assert route.call_count == 2
    assert [tc.id for tc in executed] == ["call_1", "call_2"]

    second_body: Final = _request_bodies(respx_mock)[1]
    assert second_body["messages"] == [
        {"role": "user", "content": "weather in Paris and Tokyo?"},
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": "call_1",
                    "type": "function",
                    "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
                },
                {
                    "id": "call_2",
                    "type": "function",
                    "function": {"name": "get_weather", "arguments": '{"city": "Tokyo"}'},
                },
            ],
        },
        {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_1"},
        {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_2"},
    ]
    assert messages == messages_snapshot


ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages"
TOOL_SEARCH_CALL: Final = {
    "type": "server_tool_use",
    "id": "srvtoolu_01Search",
    "name": "tool_search_tool_bm25",
    "input": {"query": "weather"},
}
WEATHER_CALL: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}}


def _tool_search_result(tool_names: list[str]) -> dict:
    return {
        "type": "tool_search_tool_result",
        "tool_use_id": "srvtoolu_01Search",
        "content": {
            "type": "tool_search_tool_search_result",
            "tool_references": [{"type": "tool_reference", "tool_name": name} for name in tool_names],
        },
    }


def _anthropic_response(content: list[dict], stop_reason: str) -> dict:
    return {
        "id": "msg_tool_loop",
        "type": "message",
        "role": "assistant",
        "model": "claude-opus-5-5",
        "content": content,
        "stop_reason": stop_reason,
        "stop_sequence": None,
        "usage": {"input_tokens": 10, "output_tokens": 5},
    }


def test_tool_search_call_is_replayed_with_its_result_instead_of_executed(respx_mock: respx.MockRouter) -> None:
    search_turn: Final = [TOOL_SEARCH_CALL, _tool_search_result(["get_weather"]), WEATHER_CALL]
    respx_mock.post(ANTHROPIC_MESSAGES_URL).mock(
        side_effect=[
            httpx.Response(200, json=_anthropic_response(search_turn, "tool_use")),
            httpx.Response(200, json=_anthropic_response([{"type": "text", "text": "Paris is 72F"}], "end_turn")),
        ]
    )
    executed: Final = []

    def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executed.append(tc.function.name)
        return _tool_result(tc)

    answer: Final = litellm.run_tool_loop(
        model="anthropic/claude-opus-5-5",
        messages=[{"role": "user", "content": "weather in Paris?"}],
        tools=WEATHER_TOOLS,
        execute_tool=executor,
        api_key="sk-ant-test",
    )

    assert answer == "Paris is 72F"
    assert executed == ["get_weather"]
    second_messages: Final = _request_bodies(respx_mock)[1]["messages"]
    assert second_messages[1]["content"] == search_turn
    assert [block["tool_use_id"] for block in second_messages[2]["content"]] == ["toolu_01Weather"]


async def test_arun_tool_loop_answers_after_a_search_that_found_nothing(
    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler

    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler())
    route: Final = respx_mock.post(ANTHROPIC_MESSAGES_URL).mock(
        return_value=httpx.Response(
            200,
            json=_anthropic_response(
                [TOOL_SEARCH_CALL, _tool_search_result([]), {"type": "text", "text": "I have no flight tool."}],
                "end_turn",
            ),
        )
    )
    executed: Final = []

    async def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
        executed.append(tc.function.name)
        return _tool_result(tc)

    answer: Final = await litellm.arun_tool_loop(
        model="anthropic/claude-opus-5-5",
        messages=[{"role": "user", "content": "book me a flight"}],
        tools=WEATHER_TOOLS,
        execute_tool=executor,
        api_key="sk-ant-test",
    )

    assert answer == "I have no flight tool."
    assert executed == []
    assert route.call_count == 1


def test_client_tool_turn_is_replayed_without_its_provider_specific_fields(respx_mock: respx.MockRouter) -> None:
    respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
        side_effect=[
            httpx.Response(
                200,
                json=_openai_response(
                    None,
                    [_tool_call("call_1", "get_weather", {"city": "Paris"})],
                    reasoning_details=[{"type": "reasoning.text", "text": "Check the weather tool"}],
                ),
            ),
            httpx.Response(200, json=_openai_response("Paris is 72F")),
        ]
    )

    answer: Final = litellm.run_tool_loop(
        model="openai/gpt-5-mini",
        messages=[{"role": "user", "content": "weather in Paris?"}],
        tools=WEATHER_TOOLS,
        execute_tool=_tool_result,
        api_key="sk-test",
    )

    assert answer == "Paris is 72F"
    replayed_assistant: Final = _request_bodies(respx_mock)[1]["messages"][1]
    assert replayed_assistant["tool_calls"][0]["id"] == "call_1"
    assert "provider_specific_fields" not in replayed_assistant


def test_responses_bridge_replays_commentary_reasoning_and_tool_call_in_order(respx_mock: respx.MockRouter) -> None:
    round_one: Final = _responses_payload(
        "resp_1",
        [
            {
                "type": "message",
                "id": "msg_commentary",
                "status": "completed",
                "role": "assistant",
                "phase": "commentary",
                "content": [{"type": "output_text", "text": "Checking Paris.", "annotations": []}],
            },
            {"type": "reasoning", "id": "rs_abc123", "summary": [], "encrypted_content": "enc_xyz"},
            {
                "type": "function_call",
                "id": "fc_1",
                "call_id": "call_weather_paris",
                "name": "get_weather",
                "arguments": '{"city": "Paris"}',
                "status": "completed",
            },
        ],
    )
    round_two: Final = _responses_payload(
        "resp_2",
        [
            {
                "type": "message",
                "id": "msg_1",
                "status": "completed",
                "role": "assistant",
                "phase": "final_answer",
                "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)]
    )

    answer: Final = litellm.run_tool_loop(
        model="openai/responses/gpt-5.6",
        messages=[{"role": "user", "content": "weather in Paris?"}],
        tools=WEATHER_TOOLS,
        execute_tool=_tool_result,
        api_key="sk-test",
    )

    assert answer == "Paris is 72F"
    assert _request_bodies(respx_mock)[1]["input"][1:] == [
        {
            "type": "message",
            "role": "assistant",
            "content": [{"type": "output_text", "text": "Checking Paris."}],
            "phase": "commentary",
        },
        {"type": "reasoning", "id": "rs_abc123", "summary": [], "encrypted_content": "enc_xyz"},
        {"type": "function_call", "call_id": "call_weather_paris", "name": "get_weather", "arguments": '{"city": "Paris"}'},
        {
            "type": "function_call_output",
            "call_id": "call_weather_paris",
            "output": [{"type": "input_text", "text": '{"temp": "72F"}'}],
        },
    ]