agent-lab/tests/test_openai_stream_parser.py

from agent_lab.domain.events import ToolCallEvent
from agent_lab.infrastructure.openai_compatible import ChatCompletionStreamParser


def test_parser_emits_visible_content_usage_and_text_protocol_events():
    parser = ChatCompletionStreamParser()

    items = []
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {"content": "hello"},
                        "finish_reason": None,
                    }
                ]
            }
        )
    )
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {"content": "\n<agent_"},
                        "finish_reason": None,
                    }
                ]
            }
        )
    )
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {
                            "content": "events>\nhandoff_note\nmock_search\n</agent_events>"
                        },
                        "finish_reason": "stop",
                    }
                ],
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 3,
                    "total_tokens": 13,
                    "prompt_tokens_details": {"cached_tokens": 4},
                },
            }
        )
    )

    assert [item.content for item in items if item.kind == "message_delta"] == [
        "hello",
        "\n",
    ]
    events = [item.event for item in items if item.kind == "event"]
    assert events == [
        ToolCallEvent(
            id="event_1",
            name="handoff_note",
            arguments={},
            raw_arguments="{}",
        ),
        ToolCallEvent(
            id="event_2",
            name="mock_search",
            arguments={},
            raw_arguments="{}",
        )
    ]
    usage = [item.usage for item in items if item.kind == "usage"][0]
    assert usage.total_tokens == 13
    assert usage.cached_tokens == 4


def test_parser_emits_provider_tool_call_events_for_event_agent():
    parser = ChatCompletionStreamParser()

    items = []
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 0,
                                    "id": "call_1",
                                    "type": "function",
                                    "function": {
                                        "name": "mock_search",
                                        "arguments": '{"query":"',
                                    },
                                }
                            ]
                        },
                        "finish_reason": None,
                    }
                ]
            }
        )
    )
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 0,
                                    "function": {"arguments": 'latency docs"}'},
                                }
                            ]
                        },
                        "finish_reason": "tool_calls",
                    }
                ]
            }
        )
    )

    assert [item.event for item in items if item.kind == "event"] == [
        ToolCallEvent(
            id="call_1",
            name="mock_search",
            arguments={"query": "latency docs"},
            raw_arguments='{"query":"latency docs"}',
        )
    ]