mlflow--mlflow/tests/pyfunc/test_responses_agent_validation.py

import pytest
from pydantic import ValidationError

from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
    responses_to_cc,
    to_chat_completions_input,
)
from mlflow.types.responses_helpers import FunctionCallOutput, Message


def test_responses_request_validation():
    with pytest.raises(ValueError, match="content.0.text"):
        ResponsesAgentRequest(**{
            "input": [
                {
                    "type": "message",
                    "id": "1",
                    "status": "completed",
                    "role": "assistant",
                    "content": [
                        {
                            "type": "output_text",
                        }
                    ],
                }
            ],
        })

    with pytest.raises(ValueError, match="role"):
        ResponsesAgentRequest(**{
            "input": [
                {
                    "type": "message",
                    "id": "1",
                    "status": "completed",
                    "role": "asdf",
                    "content": [
                        {
                            "type": "output_text",
                            "text": "asdf",
                        }
                    ],
                }
            ],
        })


def test_message_content_validation():
    # Test that None content is rejected (by Pydantic validation)
    with pytest.raises(ValidationError, match="Input should be a valid"):
        Message(role="assistant", content=None, type="message")

    # Test that empty string content is allowed
    message_empty_str = Message(role="assistant", content="", type="message")
    assert message_empty_str.content == ""

    # Test that empty list content is allowed
    message_empty_list = Message(role="assistant", content=[], type="message")
    assert message_empty_list.content == []


def test_responses_response_validation():
    with pytest.raises(ValueError, match="output.0.content.0.text"):
        ResponsesAgentResponse(**{
            "output": [
                {
                    "type": "message",
                    "id": "1",
                    "status": "completed",
                    "role": "assistant",
                    "content": [
                        {
                            "type": "output_text",
                        }
                    ],
                }
            ],
        })


def test_responses_stream_event_validation():
    with pytest.raises(ValueError, match="content must not be an empty"):
        ResponsesAgentStreamEvent(**{
            "type": "response.output_item.done",
            "output_index": 0,
            "item": {
                "type": "message",
                "status": "in_progress",
                "role": "assistant",
                "content": [],
                "id": "1",
            },
        })

    with pytest.raises(ValueError, match="Invalid status"):
        ResponsesAgentStreamEvent(**{
            "type": "response.output_item.done",
            "output_index": 0,
            "item": {
                "type": "message",
                "status": "asdf",
                "role": "assistant",
                "content": [
                    {
                        "type": "output_text",
                        "text": "asdf",
                    }
                ],
                "id": "1",
            },
        })

    with pytest.raises(ValueError, match="item.content.0.annotations.0.url"):
        ResponsesAgentStreamEvent(
            **{
                "type": "response.output_item.done",
                "output_index": 1,
                "item": {
                    "type": "message",
                    "id": "msg_67ed73ed2c288191b0f0f445e21c66540fbd8030171e9b0c",
                    "status": "completed",
                    "role": "assistant",
                    "content": [
                        {
                            "type": "output_text",
                            "text": "On T",
                            "annotations": [
                                {
                                    "type": "url_citation",
                                    "start_index": 359,
                                    "end_index": 492,
                                    "title": "NBA roundup:",
                                },
                            ],
                        }
                    ],
                },
            },
        )
    with pytest.raises(ValueError, match="delta"):
        ResponsesAgentStreamEvent(
            **{
                "type": "response.output_text.delta",
                "item_id": "msg_67eda402cba48191a1c35b84af04fc3c0a4363ad71e9395a",
                "output_index": 0,
                "content_index": 0,
            },
        )

    with pytest.raises(ValueError, match="annotation.url"):
        ResponsesAgentStreamEvent(
            **{
                "type": "response.output_text.annotation.added",
                "item_id": "msg_67ed73ed2c288191b0f0f445e21c66540fbd8030171e9b0c",
                "output_index": 1,
                "content_index": 0,
                "annotation_index": 0,
                "annotation": {
                    "type": "url_citation",
                    "start_index": 359,
                    "end_index": 492,
                    "title": "NBA roundup: Wolves overcome Nikola",
                },
            },
        )


@pytest.mark.parametrize(
    "output",
    [
        "hello",
        [{"type": "input_text", "text": "result"}],
        [
            {
                "type": "input_text",
                "text": '{"content":{"queryAttachments":[]},"status":"COMPLETED"}',
            }
        ],
    ],
)
def test_function_call_output_accepts_string_and_list(output):
    FunctionCallOutput(call_id="c", output=output)
    ResponsesAgentStreamEvent(
        type="response.output_item.done",
        item={"type": "function_call_output", "call_id": "c", "output": output},
    )


@pytest.mark.parametrize(
    ("output", "expected"),
    [
        ("hello", "hello"),
        ([{"key": "value"}], '[{"key": "value"}]'),
        ({"a": 1}, '{"a": 1}'),
        (12345, "12345"),
    ],
)
def test_responses_to_cc_stringifies_function_call_output(output, expected):
    result = responses_to_cc({"type": "function_call_output", "call_id": "c", "output": output})
    assert result[0]["content"] == expected


def test_responses_to_cc_fallback_to_str_on_non_serializable():
    class NonSerializable:
        pass

    result = responses_to_cc({
        "type": "function_call_output",
        "call_id": "c",
        "output": [NonSerializable()],
    })
    assert isinstance(result[0]["content"], str)


def test_function_call_output_round_trip():
    raw_item = {
        "call_id": "toolu_bdrk_017fvUyTS6oaCDYg6GVL3X7j",
        "output": [{"type": "input_text", "text": '{"status":"COMPLETED"}'}],
        "type": "function_call_output",
    }
    event = ResponsesAgentStreamEvent(type="response.output_item.done", item=raw_item)
    response_items = [event.item]
    dumped_items = [
        item.model_dump() if hasattr(item, "model_dump") else item for item in response_items
    ]
    cc_messages = to_chat_completions_input(dumped_items)
    assert cc_messages[0]["role"] == "tool"
    assert isinstance(cc_messages[0]["content"], str)
    assert "input_text" in cc_messages[0]["content"]