PolyGateway/tests/unit/test_openai_completion.py

"""协议聚合器的分片、防御与快照契约。"""

import copy
import json

import pytest

from polygateway.errors import TransientError
from polygateway.transports._openai_completion import (
    ChatCompletionAccumulator,
    parse_chat_completion,
)


def chunk(delta, **extra):
    """构造单 choice 协议帧。"""
    return (
        "data: "
        + json.dumps({"id": "upstream-1", "choices": [{"index": 0, "delta": delta, **extra}]})
        + "\n\n"
    )


def tool_frames():
    """两个交错分片工具，保留空白及 Unicode。"""
    return [
        chunk({"role": "assistant", "content": None}),
        chunk(
            {
                "tool_calls": [
                    {
                        "index": 1,
                        "id": "b",
                        "type": "function",
                        "function": {"name": "echo", "arguments": '{"v":'},
                    },
                    {
                        "index": 0,
                        "id": "a",
                        "type": "function",
                        "function": {"name": "echo", "arguments": "{"},
                    },
                ]
            }
        ),
        chunk(
            {
                "tool_calls": [
                    {"index": 0, "function": {"arguments": '"v": "中" }'}},
                    {"index": 1, "function": {"arguments": "2}"}},
                ]
            },
            finish_reason="tool_calls",
        ),
    ]


def accumulate():
    """复用双工具网关帧样本。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    for frame in tool_frames():
        state.add_chunk(json.loads(frame.removeprefix("data: ")))
    return state


@pytest.mark.parametrize("role", [None, "user", "tool"])
def test_role_is_required_everywhere(role):
    """缺角色不能被协议成功掩盖。"""
    state = accumulate()
    state.message.pop("role")
    if role:
        state.message["role"] = role
    with pytest.raises(TransientError, match="role"):
        state.finish()
    with pytest.raises(TransientError, match="role"):
        parse_chat_completion({"choices": [{"message": state.message}]})


@pytest.mark.parametrize(
    "change",
    [
        lambda t: t.pop("id"),
        lambda t: t.update(type="custom"),
        lambda t: t["function"].pop("name"),
        lambda t: t["function"].pop("arguments"),
        lambda t: t["function"].update(arguments=12),
    ],
)
def test_incomplete_tool_is_not_success(change):
    """工具必需字段不能靠正文兜底。"""
    state = accumulate()
    state.message["content"] = "text"
    change(state.tools[0])
    with pytest.raises(TransientError):
        state.finish()


@pytest.mark.parametrize(
    "delta",
    [
        [],
        {"audio": {}},
        {"tool_calls": {}},
        {"content": []},
        {"tool_calls": [{"index": True}]},
        {"tool_calls": [{"index": -1}]},
        {"tool_calls": [{"index": 0, "extra": 1}]},
        {"tool_calls": [{"index": 0, "function": {"arguments": None}}]},
    ],
)
def test_invalid_delta_is_domain_error(delta):
    """畸形输入不泄漏裸类型错误。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    with pytest.raises(TransientError):
        state.add_chunk({"choices": [{"index": 0, "delta": delta}]})


def test_identity_repeat_does_not_refresh_progress():
    """身份只初次生效算进展；同值重复不能续命。"""
    state = accumulate()
    repeated = {
        "choices": [
            {"index": 0, "delta": {"tool_calls": [{"index": 0, "id": "a", "type": "function"}]}}
        ]
    }
    assert not state.add_chunk(repeated).progress
    repeated["choices"][0]["delta"]["tool_calls"][0]["id"] = "other"
    with pytest.raises(TransientError, match="冲突"):
        state.add_chunk(repeated)


@pytest.mark.parametrize("delta", [{"role": "assistant"}, {}, {"content": ""}])
def test_metadata_not_progress(delta):
    """空帧和角色不算 token。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    assert not state.add_chunk({"choices": [{"index": 0, "delta": delta}]}).progress


def test_nonstream_snapshot_and_length_preserved():
    """非流式扩展字段与参数原文独立保存，不修复 length 产物。"""
    message = accumulate().finish().message
    message["extension"] = {"nested": [1]}
    message["tool_calls"][0]["function"]["arguments"] = '{"unfinished":'
    original = copy.deepcopy(message)
    result = parse_chat_completion(
        {
            "id": "id",
            "choices": [{"message": message, "finish_reason": "length"}],
            "usage": {"extra": 3},
        }
    )
    message["extension"]["nested"].append(2)
    assert result.message == original
    assert result.finish_reason == "length" and result.usage == {"extra": 3}


@pytest.mark.parametrize("choices", [[], [{"index": 1, "message": {}}], [{}, {}]])
def test_nonstream_single_choice(choices):
    """不会静默取第一个 choice。"""
    with pytest.raises(TransientError):
        parse_chat_completion({"choices": choices})


def test_explicit_empty_tools_remains_present():
    """显式空工具列表与缺省不同，不删除已接收字段。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    state.add_chunk(
        {
            "choices": [
                {"index": 0, "delta": {"role": "assistant", "content": "ok", "tool_calls": []}}
            ]
        }
    )
    assert state.finish().message["tool_calls"] == []


def test_fragmented_name_and_arguments_with_usage():
    """名称/参数分片逐字拼接，usage 扩展原样快照。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    state.add_chunk(
        {
            "id": "response",
            "choices": [
                {
                    "index": 0,
                    "delta": {
                        "role": "assistant",
                        "tool_calls": [
                            {
                                "index": 3,
                                "id": "call",
                                "type": "function",
                                "function": {"name": "ec", "arguments": " "},
                            }
                        ],
                    },
                }
            ],
        }
    )
    progress = state.add_chunk(
        {
            "choices": [
                {
                    "index": 0,
                    "delta": {
                        "content": "text",
                        "reasoning_content": "reason",
                        "tool_calls": [
                            {"index": 3, "function": {"name": "ho", "arguments": "{} "}}
                        ],
                    },
                    "finish_reason": "length",
                }
            ]
        }
    )
    assert (
        progress.content == "text" and progress.reasoning_content == "reason" and progress.progress
    )
    usage = {"prompt_tokens": 7, "completion_tokens": 2, "extra": {"nested": 9}}
    assert not state.add_chunk({"choices": [], "usage": usage}).progress
    result = state.finish()
    usage["extra"]["nested"] = 10
    assert result.message["tool_calls"][0]["function"] == {"name": "echo", "arguments": " {} "}
    assert result.usage["extra"]["nested"] == 9
    assert result.finish_reason == "length"