agent-framework/python/packages/lab/tau2/tests/test_tau2_utils.py

# Copyright (c) Microsoft. All rights reserved.

"""Tests for tau2 utils module."""

import pytest

try:
    from litellm import completion as _litellm_completion  # noqa: F401
except Exception:
    pytest.skip("LiteLLM import surface required by tau2 is unavailable.", allow_module_level=True)

from agent_framework import Content, FunctionTool, Message
from agent_framework_lab_tau2._tau2_utils import (
    convert_agent_framework_messages_to_tau2_messages,
    convert_tau2_tool_to_function_tool,
)
from pydantic import BaseModel
from tau2.data_model.message import AssistantMessage, SystemMessage, ToolCall, ToolMessage, UserMessage


class _DummyToolInput(BaseModel):
    param: str


class _DummyToolResult(BaseModel):
    output: str


class _DummyTau2Tool:
    def __init__(self, name: str, description: str) -> None:
        self.name = name
        self._description = description
        self.params = _DummyToolInput

    def _get_description(self) -> str:
        return self._description

    def __call__(self, **kwargs: str) -> _DummyToolResult:
        return _DummyToolResult(output=kwargs["param"])


def test_convert_tau2_tool_to_function_tool_basic():
    """Test basic conversion from tau2 tool to FunctionTool."""
    tau2_tool = _DummyTau2Tool(name="lookup_booking", description="Lookup booking by id.")

    # Convert the tool
    tool = convert_tau2_tool_to_function_tool(tau2_tool)

    # Verify the conversion
    assert isinstance(tool, FunctionTool)
    assert tool.name == tau2_tool.name
    assert tool.description == tau2_tool._get_description()
    assert tool.input_model == tau2_tool.params

    result = tool.func(param="ABC123")
    assert isinstance(result, _DummyToolResult)
    assert result.output == "ABC123"
    assert callable(tool.func)


def test_convert_tau2_tool_to_function_tool_multiple_tools():
    """Test conversion with multiple tau2 tools."""
    tools = [
        _DummyTau2Tool(name="lookup_booking", description="Lookup booking by id."),
        _DummyTau2Tool(name="cancel_booking", description="Cancel an existing booking."),
        _DummyTau2Tool(name="check_policy", description="Get policy details."),
    ]

    # Convert multiple tools
    function_tools = [convert_tau2_tool_to_function_tool(tool) for tool in tools]

    # Verify all conversions
    for tool, tau2_tool in zip(function_tools, tools, strict=False):
        assert isinstance(tool, FunctionTool)
        assert tool.name == tau2_tool.name
        assert tool.description == tau2_tool._get_description()
        assert tool.input_model == tau2_tool.params
        assert callable(tool.func)


def test_convert_agent_framework_messages_to_tau2_messages_system():
    """Test converting system message."""
    messages = [Message(role="system", contents=[Content.from_text(text="System instruction")])]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], SystemMessage)
    assert tau2_messages[0].role == "system"
    assert tau2_messages[0].content == "System instruction"


def test_convert_agent_framework_messages_to_tau2_messages_user():
    """Test converting user message."""
    messages = [Message(role="user", contents=[Content.from_text(text="Hello assistant")])]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], UserMessage)
    assert tau2_messages[0].role == "user"
    assert tau2_messages[0].content == "Hello assistant"
    assert tau2_messages[0].tool_calls is None


def test_convert_agent_framework_messages_to_tau2_messages_assistant():
    """Test converting assistant message."""
    messages = [Message(role="assistant", contents=[Content.from_text(text="Hello user")])]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], AssistantMessage)
    assert tau2_messages[0].role == "assistant"
    assert tau2_messages[0].content == "Hello user"
    assert tau2_messages[0].tool_calls is None


def test_convert_agent_framework_messages_to_tau2_messages_with_function_call():
    """Test converting message with function call."""
    function_call = Content.from_function_call(call_id="call_123", name="test_function", arguments={"param": "value"})

    messages = [Message(role="assistant", contents=[Content.from_text(text="I'll call a function"), function_call])]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], AssistantMessage)
    assert tau2_messages[0].content == "I'll call a function"
    assert tau2_messages[0].tool_calls is not None
    assert len(tau2_messages[0].tool_calls) == 1

    tool_call = tau2_messages[0].tool_calls[0]
    assert isinstance(tool_call, ToolCall)
    assert tool_call.id == "call_123"
    assert tool_call.name == "test_function"
    assert tool_call.arguments == {"param": "value"}
    assert tool_call.requestor == "assistant"


def test_convert_agent_framework_messages_to_tau2_messages_with_function_result():
    """Test converting message with function result."""
    function_result = Content.from_function_result(call_id="call_123", result={"success": True, "data": "result data"})

    messages = [Message(role="tool", contents=[function_result])]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], ToolMessage)
    assert tau2_messages[0].id == "call_123"
    assert tau2_messages[0].role == "tool"
    assert tau2_messages[0].content is not None
    assert '{"success": true, "data": "result data"}' in tau2_messages[0].content
    assert tau2_messages[0].requestor == "assistant"
    assert tau2_messages[0].error is False


def test_convert_agent_framework_messages_to_tau2_messages_with_error():
    """Test converting function result with error."""
    function_result = Content.from_function_result(
        call_id="call_456", result="Error occurred", exception=Exception("Test error")
    )

    messages = [Message(role="tool", contents=[function_result])]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], ToolMessage)
    assert tau2_messages[0].error is True


def test_convert_agent_framework_messages_to_tau2_messages_multiple_text_contents():
    """Test converting message with multiple text contents."""
    messages = [
        Message(role="user", contents=[Content.from_text(text="First part"), Content.from_text(text="Second part")])
    ]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 1
    assert isinstance(tau2_messages[0], UserMessage)
    assert tau2_messages[0].content == "First part Second part"


def test_convert_agent_framework_messages_to_tau2_messages_complex_scenario():
    """Test converting complex scenario with multiple message types."""
    function_call = Content.from_function_call(call_id="call_789", name="complex_tool", arguments='{"key": "value"}')

    function_result = Content.from_function_result(call_id="call_789", result={"output": "tool result"})

    messages = [
        Message(role="system", contents=[Content.from_text(text="System prompt")]),
        Message(role="user", contents=[Content.from_text(text="User request")]),
        Message(role="assistant", contents=[Content.from_text(text="I'll help you"), function_call]),
        Message(role="tool", contents=[function_result]),
        Message(role="assistant", contents=[Content.from_text(text="Based on the result...")]),
    ]

    tau2_messages = convert_agent_framework_messages_to_tau2_messages(messages)

    assert len(tau2_messages) == 5
    assert isinstance(tau2_messages[0], SystemMessage)
    assert isinstance(tau2_messages[1], UserMessage)
    assert isinstance(tau2_messages[2], AssistantMessage)
    assert isinstance(tau2_messages[3], ToolMessage)
    assert isinstance(tau2_messages[4], AssistantMessage)

    # Check the assistant message with tool call
    assert tau2_messages[2].tool_calls is not None
    assert len(tau2_messages[2].tool_calls) == 1
    assert tau2_messages[2].tool_calls[0].name == "complex_tool"