"""协议聚合器的分片、防御与快照契约。""" 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"