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