agentic-mobile-control/tests/test_tool_calling_client.py

from __future__ import annotations

import base64
from typing import Any

import pytest

from runtime.planner_config import PlannerConfig
from runtime.tool_calling_client import (
    AnthropicToolCallingClient,
    OpenAIToolCallingClient,
    ToolCallDecision,
    ToolCallUnavailable,
    build_client,
)
from runtime.tool_specs import FINISH_TASK_SPEC, TAP_SPEC
from tests.fakes import PNG_10X20


class FakeMessages:
    def __init__(self, *, response: object | None = None, error: Exception | None = None) -> None:
        self.response = response
        self.error = error
        self.calls: list[dict[str, Any]] = []

    def create(self, **kwargs: Any) -> object:
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        return self.response


class FakeTransport:
    def __init__(self, messages: FakeMessages) -> None:
        self.messages = messages


class FakeCompletions:
    def __init__(self, *, response: object | None = None, error: Exception | None = None) -> None:
        self.response = response
        self.error = error
        self.calls: list[dict[str, Any]] = []

    def create(self, **kwargs: Any) -> object:
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        return self.response


class FakeChat:
    def __init__(self, completions: FakeCompletions) -> None:
        self.completions = completions


class FakeOpenAITransport:
    def __init__(self, completions: FakeCompletions) -> None:
        self.chat = FakeChat(completions)


# --- Anthropic ---------------------------------------------------------


def test_anthropic_tool_calling_client_sends_forced_single_tool_call_request() -> None:
    messages = FakeMessages(
        response={"content": [{"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}]}
    )
    client = AnthropicToolCallingClient(model="test-model", transport=FakeTransport(messages))

    decision = client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=None,
        tools=[TAP_SPEC, FINISH_TASK_SPEC],
        timeout=2.5,
    )

    assert decision == ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    assert len(messages.calls) == 1
    call = messages.calls[0]
    assert call["model"] == "test-model"
    assert call["timeout"] == 2.5
    assert call["tool_choice"] == {"type": "any", "disable_parallel_tool_use": True}
    assert call["tools"] == [
        {"name": "tap", "description": TAP_SPEC.description, "input_schema": TAP_SPEC.parameters},
        {
            "name": "finish_task",
            "description": FINISH_TASK_SPEC.description,
            "input_schema": FINISH_TASK_SPEC.parameters,
        },
    ]
    assert call["system"][0]["text"] == "system"
    assert call["system"][0]["cache_control"] == {"type": "ephemeral"}
    assert call["messages"] == [{"role": "user", "content": [{"type": "text", "text": "user"}]}]


def test_anthropic_tool_calling_client_includes_image_block_when_screenshot_present() -> None:
    messages = FakeMessages(
        response={
            "content": [
                {"type": "tool_use", "name": "finish_task", "input": {"success": True, "reason": "done"}}
            ]
        }
    )
    client = AnthropicToolCallingClient(model="test-model", transport=FakeTransport(messages))

    client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=PNG_10X20,
        tools=[FINISH_TASK_SPEC],
        timeout=1,
    )

    content = messages.calls[0]["messages"][0]["content"]
    assert content[0] == {
        "type": "image",
        "source": {
            "type": "base64",
            "media_type": "image/png",
            "data": base64.b64encode(PNG_10X20).decode("ascii"),
        },
    }
    assert content[1] == {"type": "text", "text": "user"}


def test_anthropic_tool_calling_client_wraps_transport_errors() -> None:
    messages = FakeMessages(error=TimeoutError("timed out"))
    client = AnthropicToolCallingClient(model="test-model", transport=FakeTransport(messages))

    with pytest.raises(ToolCallUnavailable):
        client.decide(system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1)


@pytest.mark.parametrize(
    "response",
    [
        {"content": []},
        {"content": [{"type": "text", "text": "no tool call"}]},
        {"content": [{"type": "tool_use", "name": "tap", "input": "not-a-dict"}]},
    ],
)
def test_anthropic_tool_calling_client_wraps_malformed_responses(response: object) -> None:
    messages = FakeMessages(response=response)
    client = AnthropicToolCallingClient(model="test-model", transport=FakeTransport(messages))

    with pytest.raises(ToolCallUnavailable):
        client.decide(system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1)


# --- OpenAI --------------------------------------------------------------


def test_openai_tool_calling_client_sends_forced_single_tool_call_request() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {"message": {"tool_calls": [{"function": {"name": "tap", "arguments": '{"x": 1, "y": 2}'}}]}}
            ]
        }
    )
    client = OpenAIToolCallingClient(model="test-model", transport=FakeOpenAITransport(completions))

    decision = client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=None,
        tools=[TAP_SPEC, FINISH_TASK_SPEC],
        timeout=2.5,
    )

    assert decision == ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    assert len(completions.calls) == 1
    call = completions.calls[0]
    assert call["model"] == "test-model"
    assert call["timeout"] == 2.5
    assert call["max_completion_tokens"] == 1024
    assert "max_tokens" not in call
    assert call["tool_choice"] == "required"
    assert call["parallel_tool_calls"] is False
    assert call["tools"] == [
        {
            "type": "function",
            "function": {
                "name": "tap",
                "description": TAP_SPEC.description,
                "parameters": TAP_SPEC.parameters,
            },
        },
        {
            "type": "function",
            "function": {
                "name": "finish_task",
                "description": FINISH_TASK_SPEC.description,
                "parameters": FINISH_TASK_SPEC.parameters,
            },
        },
    ]
    assert call["messages"] == [
        {"role": "system", "content": "system"},
        {"role": "user", "content": "user"},
    ]


def test_openai_tool_calling_client_includes_image_block_when_screenshot_present() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "finish_task",
                                    "arguments": '{"success": true, "reason": "done"}',
                                }
                            }
                        ]
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(model="test-model", transport=FakeOpenAITransport(completions))

    client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=PNG_10X20,
        tools=[FINISH_TASK_SPEC],
        timeout=1,
    )

    user_message = completions.calls[0]["messages"][1]
    assert user_message["role"] == "user"
    assert user_message["content"][0] == {"type": "text", "text": "user"}
    encoded = base64.b64encode(PNG_10X20).decode("ascii")
    assert user_message["content"][1] == {
        "type": "image_url",
        "image_url": {"url": f"data:image/png;base64,{encoded}"},
    }


def test_openai_tool_calling_client_accepts_arguments_already_as_dict() -> None:
    completions = FakeCompletions(
        response={
            "choices": [{"message": {"tool_calls": [{"function": {"name": "tap", "arguments": {"x": 1, "y": 2}}}]}}]
        }
    )
    client = OpenAIToolCallingClient(model="test-model", transport=FakeOpenAITransport(completions))

    decision = client.decide(system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1)

    assert decision == ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})


def test_openai_tool_calling_client_wraps_transport_errors() -> None:
    completions = FakeCompletions(error=TimeoutError("timed out"))
    client = OpenAIToolCallingClient(model="test-model", transport=FakeOpenAITransport(completions))

    with pytest.raises(ToolCallUnavailable):
        client.decide(system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1)


@pytest.mark.parametrize(
    "response",
    [
        {"choices": []},
        {"choices": [{"message": {"tool_calls": []}}]},
        {"choices": [{"message": {"tool_calls": [{"function": {"name": "tap", "arguments": "not-json"}}]}}]},
    ],
)
def test_openai_tool_calling_client_wraps_malformed_responses(response: object) -> None:
    completions = FakeCompletions(response=response)
    client = OpenAIToolCallingClient(model="test-model", transport=FakeOpenAITransport(completions))

    with pytest.raises(ToolCallUnavailable):
        client.decide(system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1)


# --- build_client ----------------------------------------------------------


def test_build_client_selects_provider_and_resolves_default_model() -> None:
    anthropic_client = build_client(PlannerConfig(provider="anthropic", model=""))
    assert isinstance(anthropic_client, AnthropicToolCallingClient)
    assert anthropic_client.model == "claude-sonnet-5"

    openai_client = build_client(PlannerConfig(provider="openai", model=""))
    assert isinstance(openai_client, OpenAIToolCallingClient)
    assert openai_client.model == "gpt-5.6"


def test_build_client_honors_explicit_model_override() -> None:
    client = build_client(PlannerConfig(provider="openai", model="gpt-5.6-custom"))
    assert client.model == "gpt-5.6-custom"