tkmind_go/packages/exchange/tests/test_utils.py

from typing import Literal
import pytest
from exchange import utils
from unittest.mock import patch
from exchange.message import Message
from exchange.content import Text, ToolResult
from exchange.providers.utils import messages_to_openai_spec, encode_image


def test_encode_image():
    image_path = "tests/test_image.png"
    encoded_image = encode_image(image_path)

    # Adjust this string based on the actual initial part of your base64-encoded image.
    expected_start = "iVBORw0KGgo"
    assert encoded_image.startswith(expected_start)


def test_create_object_id() -> None:
    prefix = "test"
    object_id = utils.create_object_id(prefix)
    assert object_id.startswith(prefix + "_")
    assert len(object_id) == len(prefix) + 1 + 24  # prefix + _ + 24 chars


def test_compact() -> None:
    content = "This   is \n\n   a test"
    compacted = utils.compact(content)
    assert compacted == "This is a test"


def test_parse_docstring() -> None:
    def dummy_func(a, b, c):
        """
        This function does something.

        Args:
            a (int): The first parameter.
            b (str): The second parameter.
            c (list): The third parameter.
        """
        pass

    description, parameters = utils.parse_docstring(dummy_func)
    assert description == "This function does something."
    assert parameters == [
        {"name": "a", "annotation": "int", "description": "The first parameter."},
        {"name": "b", "annotation": "str", "description": "The second parameter."},
        {"name": "c", "annotation": "list", "description": "The third parameter."},
    ]


def test_parse_docstring_no_description() -> None:
    def dummy_func(a, b, c):
        """
        Args:
            a (int): The first parameter.
            b (str): The second parameter.
            c (list): The third parameter.
        """
        pass

    with pytest.raises(ValueError) as e:
        utils.parse_docstring(dummy_func)

    assert "Attempted to load from a function" in str(e.value)


def test_parse_docstring_with_optional_params() -> None:
    from typing import Optional, List

    def dummy_func(a: int, b: List[int], c: Literal["foo", "bar"] = "foo", d: Optional[str] = None) -> None:
        """This function does something.

        Args:
            a (int): The first required parameter.
            b (List[int]): The second parameter.
            c (Literal["foo", "bar"], optional): A parameter with a literal default value. Defaults to "foo".
            d (Optional[str], optional): Optional fourth parameter. Defaults to None.
        """
        pass

    description, parameters = utils.parse_docstring(dummy_func)
    assert description == "This function does something."
    assert parameters == [
        {"name": "a", "annotation": "int", "description": "The first required parameter."},
        {"name": "b", "annotation": "List[int]", "description": "The second parameter."},
        {
            "name": "c",
            "annotation": 'Literal["foo", "bar"]',
            "description": 'A parameter with a literal default value. Defaults to "foo".',
        },
        {"name": "d", "annotation": "Optional[str]", "description": "Optional fourth parameter. Defaults to None."},
    ]


def test_json_schema() -> None:
    def dummy_func(a: int, b: str, c: list) -> None:
        pass

    schema = utils.json_schema(dummy_func)

    assert schema == {
        "type": "object",
        "properties": {
            "a": {"type": "integer"},
            "b": {"type": "string"},
            "c": {"type": "string"},
        },
        "required": ["a", "b", "c"],
    }


def test_json_schema_with_optional_params() -> None:
    from typing import Optional, List

    def dummy_func(
        a: int,
        b: Literal["foo", "bar"] = "foo",
        c: Optional[List[int]] = None,
        d: Optional[str] = None,
    ) -> None:
        pass

    schema = utils.json_schema(dummy_func)

    assert schema == {
        "type": "object",
        "properties": {
            "a": {"type": "integer"},
            "b": {"enum": ["foo", "bar"], "default": "foo"},
            "c": {"type": "array", "items": {"type": "integer"}, "default": None},
            "d": {"type": "string", "default": None},
        },
        "required": ["a"],
    }


def test_load_plugins() -> None:
    class DummyEntryPoint:
        def __init__(self, name, plugin):
            self.name = name
            self.plugin = plugin

        def load(self):
            return self.plugin

    with patch("exchange.utils.entry_points") as entry_points_mock:
        entry_points_mock.return_value = [
            DummyEntryPoint("plugin1", object()),
            DummyEntryPoint("plugin2", object()),
        ]

        plugins = utils.load_plugins("dummy_group")

        assert "plugin1" in plugins
        assert "plugin2" in plugins
        assert len(plugins) == 2


def test_messages_to_openai_spec():
    # Use provided test image
    png_path = "tests/test_image.png"

    # Create a list of messages as input
    messages = [
        Message(role="user", content=[Text(text="Hello, Assistant!")]),
        Message(role="assistant", content=[Text(text="Here is a text with tool usage")]),
        Message(
            role="tool",
            content=[ToolResult(tool_use_id="1", output=f'"image:{png_path}')],
        ),
    ]

    # Call the function
    output = messages_to_openai_spec(messages)

    assert "This tool result included an image that is uploaded in the next message." in str(output)
    assert "{'role': 'user', 'content': [{'type': 'image_url'" in str(output)