From 394257c8ef819b554b298a3654bf0c0b464bffea Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:46:55 +0000 Subject: [PATCH] test: add mocked unit tests for LiteLLM response parsing (bug prevention coverage) --- tests/llm/test_litellm_api.py | 236 +++++++++++++++++++++++++++++++++- 1 file changed, 234 insertions(+), 2 deletions(-) diff --git a/tests/llm/test_litellm_api.py b/tests/llm/test_litellm_api.py index f5685a6e1..fb7eb9691 100644 --- a/tests/llm/test_litellm_api.py +++ b/tests/llm/test_litellm_api.py @@ -1,9 +1,11 @@ +import json import os from functools import partial +from types import SimpleNamespace import pytest -from agentlab.llm.litellm_api import LiteLLMModelArgs -from agentlab.llm.response_api import APIPayload, LLMOutput +from agentlab.llm.litellm_api import LiteLLMAPIMessageBuilder, LiteLLMModel, LiteLLMModelArgs +from agentlab.llm.response_api import APIPayload, LLMOutput, ToolCall, ToolCalls chat_api_tools = [ { @@ -161,6 +163,236 @@ def test_force_tool_call(): ), f"Model:{name},Expected all tool calls to be 'get_time', but got {called_fn_names} with force call {force_call_tool}" +class FakeMessage: + """Minimal stand-in for a litellm/OpenAI chat completion message.""" + + def __init__(self, data: dict, tool_calls=None): + self._data = dict(data) + self.tool_calls = tool_calls + + def to_dict(self): + data = dict(self._data) + if self.tool_calls is not None: + data["tool_calls"] = self.tool_calls + return data + + +def make_raw_tool_call(call_id: str, name: str, arguments, as_json: bool = True): + return { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments) if as_json else arguments, + }, + } + + +def make_response(content="", reasoning=None, text=None, tool_calls=None): + data = {"content": content} + if reasoning is not None: + data["reasoning"] = reasoning + if text is not None: + data["text"] = text + message = FakeMessage(data, tool_calls=tool_calls) + return SimpleNamespace(choices=[SimpleNamespace(message=message)]) + + +@pytest.fixture +def model(): + return LiteLLMModel(model_name="openai/gpt-4.1", api_key="fake-key") + + +def test_extract_thinking_wraps_reasoning(model): + response = make_response(content="answer", reasoning="because", text="prefix ") + assert ( + model._extract_thinking_content_from_response(response) + == "because\nprefix answer" + ) + + +def test_extract_thinking_without_reasoning(model): + response = make_response(content="answer") + assert model._extract_thinking_content_from_response(response) == "answer" + + +def test_extract_thinking_custom_wrap_tag(model): + response = make_response(content="", reasoning="because") + assert ( + model._extract_thinking_content_from_response(response, wrap_tag="reason") + == "because\n" + ) + + +def test_extract_tool_calls_parses_arguments(model): + response = make_response( + tool_calls=[ + make_raw_tool_call("call_1", "get_weather", {"location": "Paris"}), + make_raw_tool_call("call_2", "get_time", {"location": "Delhi"}), + ] + ) + tool_calls = model._extract_tool_calls_from_response(response) + assert [call.name for call in tool_calls] == ["get_weather", "get_time"] + assert tool_calls.tool_calls[0].arguments == {"location": "Paris"} + assert tool_calls.raw_calls is response + + +def test_extract_tool_calls_returns_none_without_tool_calls(model): + assert model._extract_tool_calls_from_response(make_response(content="hi")) is None + + +def test_extract_tool_calls_keeps_only_first_when_configured(): + model = LiteLLMModel( + model_name="openai/gpt-4.1", api_key="fake-key", use_only_first_toolcall=True + ) + response = make_response( + tool_calls=[ + make_raw_tool_call("call_1", "get_weather", {"location": "Paris"}), + make_raw_tool_call("call_2", "get_time", {"location": "Delhi"}), + ] + ) + tool_calls = model._extract_tool_calls_from_response(response) + assert len(tool_calls) == 1 + assert tool_calls.tool_calls[0].name == "get_weather" + + +def test_extract_tool_calls_malformed_json_raises(model): + response = make_response( + tool_calls=[make_raw_tool_call("call_1", "get_weather", "{not json", as_json=False)] + ) + with pytest.raises(json.JSONDecodeError): + model._extract_tool_calls_from_response(response) + + +def test_extract_env_actions_single_toolcall(model): + tool_calls = ToolCalls( + tool_calls=[ToolCall(name="get_weather", arguments={"location": "Paris"})] + ) + assert model._extract_env_actions_from_toolcalls(tool_calls) == "get_weather(location='Paris')" + + +def test_extract_env_actions_multiple_toolcalls_are_joined(model): + tool_calls = ToolCalls( + tool_calls=[ + ToolCall(name="get_weather", arguments={"location": "Paris"}), + ToolCall(name="get_time", arguments={}), + ] + ) + assert ( + model._extract_env_actions_from_toolcalls(tool_calls) + == "get_weather(location='Paris')\nget_time()" + ) + + +def test_extract_env_actions_empty_toolcalls(model): + assert model._extract_env_actions_from_toolcalls(ToolCalls()) is None + assert model._extract_env_actions_from_toolcalls(None) is None + + +def test_extract_env_actions_from_text_response_is_unsupported(model): + assert model._extract_env_actions_from_text_response(make_response(content="click(1)")) is None + + +def test_parse_response_with_tool_calls(model): + response = make_response( + content="here you go", + reasoning="thinking", + tool_calls=[make_raw_tool_call("call_1", "get_weather", {"location": "Paris"})], + ) + output = model._parse_response(response) + assert isinstance(output, LLMOutput) + assert output.raw_response is response + assert output.think == "thinking\nhere you go" + assert output.action == "get_weather(location='Paris')" + assert len(output.tool_calls) == 1 + + +def test_parse_response_without_tool_calls(model): + output = model._parse_response(make_response(content="no tools")) + assert output.action is None + assert output.tool_calls is None + assert output.think == "no tools" + + +def test_parse_response_text_action_space_has_no_action(model): + model.action_space_as_tools = False + output = model._parse_response(make_response(content="click(1)")) + assert output.action is None + + +def make_responded_tool_calls(n_calls: int = 2, image_response: bool = False): + raw_calls = [ + make_raw_tool_call(f"call_{i}", f"tool_{i}", {"arg": i}) for i in range(1, n_calls + 1) + ] + response = make_response(tool_calls=raw_calls) + tool_calls = [] + for raw_call in raw_calls: + call = ToolCall(name=raw_call["function"]["name"], arguments={}, raw_call=raw_call) + if image_response: + call.response_image("data:image/png;base64,abc") + else: + call.response_text(f"result of {raw_call['id']}") + tool_calls.append(call) + return ToolCalls(tool_calls=tool_calls, raw_calls=response) + + +def test_handle_tool_call_builds_provider_messages(): + responded = make_responded_tool_calls(n_calls=2) + builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(responded) + output = builder.prepare_message() + + assert output[0] is responded.raw_calls.choices[0].message + assert output[1:] == [ + { + "name": "tool_1", + "role": "tool", + "tool_call_id": "call_1", + "content": "result of call_1", + }, + { + "name": "tool_2", + "role": "tool", + "tool_call_id": "call_2", + "content": "result of call_2", + }, + ] + + +def test_handle_tool_call_truncates_raw_calls_to_first(): + responded = make_responded_tool_calls(n_calls=2) + builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(responded) + output = builder.prepare_message(use_only_first_toolcall=True) + + assert len(output[0].tool_calls) == 1 + assert output[0].tool_calls[0]["id"] == "call_1" + + +def test_handle_tool_call_rejects_image_response(): + responded = make_responded_tool_calls(n_calls=1, image_response=True) + builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(responded) + with pytest.raises(AssertionError, match="Image output is not supported"): + builder.prepare_message() + + +def test_handle_tool_call_without_responded_tool_calls(): + builder = LiteLLMAPIMessageBuilder("tool") + with pytest.raises(ValueError, match="No tool calls found"): + builder.handle_tool_call() + + +def test_format_tools_for_chat_completion(): + formatted = LiteLLMModel.format_tools_for_chat_completion(chat_api_tools) + assert formatted[0] == { + "type": "function", + "function": { + "name": chat_api_tools[0]["name"], + "description": chat_api_tools[0]["description"], + "parameters": chat_api_tools[0]["parameters"], + }, + } + assert LiteLLMModel.format_tools_for_chat_completion(None) is None + + if __name__ == "__main__": test_multi_action_tool_calls() test_force_tool_call()