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_content_usage_and_tool_event():
    parser = ChatCompletionStreamParser()

    items = []
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {"content": "hello"},
                        "finish_reason": None,
                    }
                ]
            }
        )
    )
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 0,
                                    "id": "call_1",
                                    "type": "function",
                                    "function": {
                                        "name": "handoff_note",
                                        "arguments": '{"message":"',
                                    },
                                }
                            ]
                        },
                        "finish_reason": None,
                    }
                ]
            }
        )
    )
    items.extend(
        parser.feed(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 0,
                                    "function": {"arguments": 'from model"}'},
                                }
                            ]
                        },
                        "finish_reason": "tool_calls",
                    }
                ],
                "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"]
    events = [item.event for item in items if item.kind == "event"]
    assert events == [
        ToolCallEvent(
            id="call_1",
            name="handoff_note",
            arguments={"message": "from model"},
            raw_arguments='{"message":"from model"}',
        )
    ]
    usage = [item.usage for item in items if item.kind == "usage"][0]
    assert usage.total_tokens == 13
    assert usage.cached_tokens == 4