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",
    [
        [],
        {"tool_calls": {}},
        {"content": []},
        {"role": "user"},
        {"tool_calls": [{"index": True}]},
        {"tool_calls": [{"index": -1}]},
        {"tool_calls": [{"index": None}]},
        {"tool_calls": [{"index": 0, "function": {"arguments": 12}}]},
    ],
)
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"


# —— 1.3.10(#30): null 优先、未知字段无损合并、finish/usage 位置 ——


def _feed(state, *chunks):
    """依次喂入已解码帧,返回每帧的 ChatDelta。"""
    return [state.add_chunk(c) for c in chunks]


def _frame(delta, finish=None, **top):
    """单 choice 已解码帧。"""
    choice = {"index": 0, "delta": delta}
    if finish is not None:
        choice["finish_reason"] = finish
    return {"choices": [choice], **top}


def test_null_known_and_unknown_fields_are_absent():
    """refusal/tool_calls/role/arguments 为 null 不判失败,键按缺省保留。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    _feed(
        state,
        _frame({"role": None, "refusal": None, "tool_calls": None, "content": None}),
        _frame({"role": "assistant", "content": "hi"}, finish="stop"),
    )
    message = state.finish().message
    assert message["refusal"] is None and message["content"] == "hi"
    assert message["role"] == "assistant"


def test_role_only_null_fails_at_finish():
    """全程只有 role:null 即缺 role,照旧失败。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    _feed(state, _frame({"role": None, "content": "x"}, finish="stop"))
    with pytest.raises(TransientError, match="role"):
        state.finish()


def test_null_arguments_and_null_tool_elements():
    """arguments:null 视缺省;tool_calls 中 null 元素跳过。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    tool = {"index": 0, "id": "a", "type": "function", "function": {"name": "f"}}
    _feed(
        state,
        _frame({"role": "assistant", "tool_calls": [None, tool]}),
        _frame({"tool_calls": [{"index": 0, "function": {"arguments": None}}]}),
        _frame({"tool_calls": [{"index": 0, "function": {"arguments": "{}"}}]}, "tool_calls"),
    )
    assert state.finish().message["tool_calls"][0]["function"] == {"name": "f", "arguments": "{}"}


def test_null_choices_and_delta_are_absent():
    """chunk 层 choices:null 与 choice 层 delta:null 视缺省。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    assert not state.add_chunk({"choices": None}).progress
    assert not state.add_chunk({"choices": [{"index": 0, "delta": None}]}).progress


def test_reasoning_alias_is_progress_and_preserved():
    """只有 reasoning 字段: 计活性、交别名片段、保真原键保留。"""
    for mode in ("openai", "text"):
        state = ChatCompletionAccumulator(response_mode=mode)
        delta = state.add_chunk(_frame({"role": "assistant", "reasoning": "think"}))
        assert delta.progress and delta.reasoning == "think" and delta.reasoning_content == ""
    state.add_chunk(_frame({"reasoning": " more"}))
    openai = ChatCompletionAccumulator(response_mode="openai")
    _feed(openai, _frame({"role": "assistant", "reasoning": "a"}), _frame({"reasoning": "b"}))
    assert openai.raw_message()["reasoning"] == "ab"


def _detail(text, **extra):
    return {"type": "reasoning.text", "index": 0, "format": "x-v1", "text": text, **extra}


def test_reasoning_details_merge_by_index_without_label_concat():
    """OpenRouter reasoning_details 分片: index 合并,format/id/type 不拼接。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    _feed(
        state,
        _frame({"role": "assistant", "reasoning_details": [_detail("he", id="r1")]}),
        _frame({"reasoning_details": [_detail("llo", id="r1")]}),
        _frame({"content": "ok"}, finish="stop"),
    )
    details = state.finish().message["reasoning_details"]
    assert details == [
        {"type": "reasoning.text", "index": 0, "format": "x-v1", "text": "hello", "id": "r1"}
    ]
    assert not state.conflicts


def test_tool_level_nested_extension_preserved():
    """工具层未知嵌套对象(如 thought signature)原样保留以便回传。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    extra = {"google": {"thought_signature": "sig"}}
    tool = {"index": 0, "id": "a", "type": "function", "extra_content": extra}
    _feed(
        state,
        _frame({"role": "assistant", "tool_calls": [{**tool, "function": {"name": "f"}}]}),
        _frame({"tool_calls": [{"index": 0, "function": {"arguments": "{}", "strict": True}}]}),
    )
    snapshot = state.finish().message["tool_calls"][0]
    assert snapshot["extra_content"] == extra and snapshot["function"]["strict"] is True


@pytest.mark.parametrize(
    ("first", "second"),
    [(1, 2), (True, False), ("s", {"o": 1}), ([{"index": 0, "t": "a"}], ["plain"])],
)
def test_conflict_drops_key_and_bans_path(first, second):
    """数值/布尔/形态冲突: 删键、登记路径、后续同路径丢弃、调用成功。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    _feed(
        state,
        _frame({"role": "assistant", "content": "a", "ext": first}),
        _frame({"ext": second}),
        _frame({"ext": first}, finish="stop"),
    )
    assert "ext" not in state.finish().message
    assert state.conflicts == frozenset({"delta.ext"})


def test_plain_arrays_append_and_same_scalar_idempotent():
    """普通数组追加;同值标量幂等;追加数组中的字符串算活性。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    deltas = _feed(
        state,
        _frame({"role": "assistant", "notes": ["a"], "n": 1}),
        _frame({"notes": ["b", None], "n": 1}),
    )
    message = state.raw_message()
    assert message["notes"] == ["a", "b"] and message["n"] == 1
    assert all(d.progress for d in deltas[1:])


def test_one_progress_per_frame_with_many_fragments():
    """同帧多个字段只产生一次进展事件(ChatDelta 单值)。"""
    state = ChatCompletionAccumulator(response_mode="text")
    delta = state.add_chunk(_frame({"content": "a", "reasoning_content": "b", "x": "c"}))
    assert delta.progress is True


def test_label_only_unknown_is_not_progress():
    """未知层仅标签键(type/format/id/index)不刷新活性。"""
    state = ChatCompletionAccumulator(response_mode="openai")
    assert not state.add_chunk(_frame({"meta": {"type": "x", "format": "y"}})).progress


@pytest.mark.parametrize(
    ("frames", "after"),
    [
        ([_frame({"content": "a"}, usage={"x": 1}), _frame({}, finish="stop")], False),
        ([_frame({"content": "a"}, finish="stop", usage={"x": 1})], True),
        ([_frame({"content": "a"}, finish="stop"), {"choices": [], "usage": {"x": 1}}], True),
        ([_frame({"content": "a"}, finish="stop")], False),
    ],
)
def test_usage_after_finish_tracking(frames, after):
    """usage 必须在 finish 同帧或之后到达才算流走到自然末尾。"""
    for mode in ("openai", "text"):
        state = ChatCompletionAccumulator(response_mode=mode)
        _feed(state, *frames)
        assert state.finish_seen and state.finish_reason == "stop"
        assert state.usage_after_finish is after


def test_text_mode_ignores_other_choice_index():
    """文本模式只取 index 缺省或为 0 的 choice,不混入 index=1。"""
    state = ChatCompletionAccumulator(response_mode="text")
    delta = state.add_chunk(
        {
            "choices": [
                {"index": 1, "delta": {"content": "B"}, "finish_reason": "length"},
                {"index": 0, "delta": {"content": "A"}},
            ]
        }
    )
    assert delta.content == "A" and state.finish_reason is None


def test_text_mode_finish_non_string_ignored():
    """文本模式 finish 非字符串记 None,不判失败。"""
    state = ChatCompletionAccumulator(response_mode="text")
    state.add_chunk(_frame({"content": "a"}, finish=5))
    assert state.finish_reason is None and not state.finish_seen


def test_nonstream_null_tool_elements_stripped():
    """非流式 tool_calls 中 null 元素剔除(回传 [null] 会被上游拒),null 列表原样保留。"""
    tool = {"id": "a", "type": "function", "function": {"name": "f", "arguments": "{}"}}
    body = {
        "choices": [{"message": {"role": "assistant", "tool_calls": [None, tool], "refusal": None}}]
    }
    assert parse_chat_completion(body).message["tool_calls"] == [tool]
    body["choices"][0]["message"]["tool_calls"] = None
    assert parse_chat_completion(body).message["tool_calls"] is None