OpenJarvis/tests/engine/test_openai_compat_tools.py

"""Tests for tool_calls extraction in _OpenAICompatibleEngine and OllamaEngine."""

from __future__ import annotations

import json

import httpx

from openjarvis.core.types import Message, Role
from openjarvis.engine._openai_compat import _OpenAICompatibleEngine
from openjarvis.engine.ollama import OllamaEngine

# ---------------------------------------------------------------------------
# _OpenAICompatibleEngine tests
# ---------------------------------------------------------------------------


class _TestEngine(_OpenAICompatibleEngine):
    engine_id = "test"
    _default_host = "http://localhost:9999"


class TestOpenAICompatToolCalls:
    def test_no_tool_calls(self, respx_mock):
        """When no tool_calls in response, result has no tool_calls key."""
        respx_mock.post("http://localhost:9999/v1/chat/completions").mock(
            return_value=httpx.Response(200, json={
                "choices": [
                    {"message": {"content": "Hi"}, "finish_reason": "stop"},
                ],
                "usage": {
                    "prompt_tokens": 5,
                    "completion_tokens": 2,
                    "total_tokens": 7,
                },
                "model": "test",
            })
        )
        engine = _TestEngine()
        result = engine.generate(
            [Message(role=Role.USER, content="Hello")],
            model="test",
        )
        assert "tool_calls" not in result
        assert result["content"] == "Hi"

    def test_with_tool_calls(self, respx_mock):
        """Extract tool_calls from OpenAI-format response."""
        respx_mock.post("http://localhost:9999/v1/chat/completions").mock(
            return_value=httpx.Response(200, json={
                "choices": [{
                    "message": {
                        "content": None,
                        "tool_calls": [{
                            "id": "call_abc",
                            "type": "function",
                            "function": {
                                "name": "calculator",
                                "arguments": '{"expression":"2+2"}',
                            },
                        }],
                    },
                    "finish_reason": "tool_calls",
                }],
                "usage": {
                    "prompt_tokens": 5,
                    "completion_tokens": 10,
                    "total_tokens": 15,
                },
                "model": "test",
            })
        )
        engine = _TestEngine()
        result = engine.generate(
            [Message(role=Role.USER, content="Calculate 2+2")],
            model="test",
        )
        assert "tool_calls" in result
        assert len(result["tool_calls"]) == 1
        tc = result["tool_calls"][0]
        assert tc["id"] == "call_abc"
        assert tc["name"] == "calculator"
        assert tc["arguments"] == '{"expression":"2+2"}'
        assert result["content"] == ""  # None → ""

    def test_tools_kwarg_passthrough(self, respx_mock):
        """Tools kwarg is spread into the payload via **kwargs."""
        captured = {}

        def capture(request):
            captured["body"] = json.loads(request.content)
            return httpx.Response(200, json={
                "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}],
                "usage": {},
                "model": "test",
            })

        respx_mock.post("http://localhost:9999/v1/chat/completions").mock(
            side_effect=capture,
        )
        engine = _TestEngine()
        engine.generate(
            [Message(role=Role.USER, content="Hello")],
            model="test",
            tools=[{"type": "function", "function": {"name": "calc"}}],
        )
        assert "tools" in captured["body"]

    def test_multiple_tool_calls(self, respx_mock):
        respx_mock.post("http://localhost:9999/v1/chat/completions").mock(
            return_value=httpx.Response(200, json={
                "choices": [{
                    "message": {
                        "content": "",
                        "tool_calls": [
                            {
                                "id": "c1", "type": "function",
                                "function": {"name": "a", "arguments": "{}"},
                            },
                            {
                                "id": "c2", "type": "function",
                                "function": {"name": "b", "arguments": "{}"},
                            },
                        ],
                    },
                    "finish_reason": "tool_calls",
                }],
                "usage": {},
                "model": "test",
            })
        )
        engine = _TestEngine()
        result = engine.generate(
            [Message(role=Role.USER, content="Use tools")],
            model="test",
        )
        assert len(result["tool_calls"]) == 2


# ---------------------------------------------------------------------------
# OllamaEngine tests
# ---------------------------------------------------------------------------


class TestOllamaToolCalls:
    def test_no_tool_calls(self, respx_mock):
        respx_mock.post("http://localhost:11434/api/chat").mock(
            return_value=httpx.Response(200, json={
                "message": {"content": "Hi"},
                "model": "test",
                "prompt_eval_count": 5,
                "eval_count": 2,
            })
        )
        engine = OllamaEngine()
        result = engine.generate(
            [Message(role=Role.USER, content="Hello")],
            model="test",
        )
        assert "tool_calls" not in result

    def test_with_tool_calls(self, respx_mock):
        respx_mock.post("http://localhost:11434/api/chat").mock(
            return_value=httpx.Response(200, json={
                "message": {
                    "content": "",
                    "tool_calls": [{
                        "function": {
                            "name": "calculator",
                            "arguments": '{"expression":"3*3"}',
                        },
                    }],
                },
                "model": "test",
                "prompt_eval_count": 5,
                "eval_count": 3,
            })
        )
        engine = OllamaEngine()
        result = engine.generate(
            [Message(role=Role.USER, content="3*3")],
            model="test",
        )
        assert "tool_calls" in result
        assert result["tool_calls"][0]["name"] == "calculator"

    def test_tools_in_payload(self, respx_mock):
        captured = {}

        def capture(request):
            captured["body"] = json.loads(request.content)
            return httpx.Response(200, json={
                "message": {"content": "ok"},
                "model": "test",
            })

        respx_mock.post("http://localhost:11434/api/chat").mock(side_effect=capture)
        engine = OllamaEngine()
        engine.generate(
            [Message(role=Role.USER, content="Hello")],
            model="test",
            tools=[{"type": "function", "function": {"name": "calc"}}],
        )
        assert "tools" in captured["body"]

    def test_no_tools_no_tools_key(self, respx_mock):
        captured = {}

        def capture(request):
            captured["body"] = json.loads(request.content)
            return httpx.Response(200, json={
                "message": {"content": "ok"},
                "model": "test",
            })

        respx_mock.post("http://localhost:11434/api/chat").mock(side_effect=capture)
        engine = OllamaEngine()
        engine.generate(
            [Message(role=Role.USER, content="Hello")],
            model="test",
        )
        assert "tools" not in captured["body"]