mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
af1828dd32
当 Provider A 生成的 thinking 块被发送到 Provider B 时，签名验证会失败。 本次更新实现了自动整流机制，在遇到签名错误时自动清洗 thinking 块后重试。 主要更改： - 新增 ThinkingRectifier 整流器，移除 thinking 块和 signature 字段 - 新增 ThinkingSignatureException 异常类型 - ErrorClassifier 新增 Thinking 错误模式检测 - FallbackOrchestrator 支持整流后在当前候选重试 - Handler 层传递 request_body_ref 容器支持请求体动态修改 - Usage API 新增 has_rectified 字段标识整流过的请求 - 新增 THINKING_RECTIFIER_ENABLED 配置项控制功能开关 其他改进： - CacheAwareScheduler 支持 exact/convertible 候选分组排序 - StreamProcessor 预读阶段新增格式转换试验 - ProviderAPIKey.api_formats 改为可空（None 表示支持所有格式） - Dockerfile 修复 entrypoint.sh 换行符问题 Closes #115 Co-Authored-By: FredericMN <FredericMN@users.noreply.github.com>
379 lines
12 KiB
Python
379 lines
12 KiB
Python
"""
|
||
ThinkingRectifier 单元测试
|
||
|
||
测试 Thinking 整流器的核心功能：
|
||
- 移除 thinking 和 redacted_thinking 块
|
||
- 移除非 thinking 块上的 signature 字段
|
||
- 条件删除顶层 thinking 参数
|
||
"""
|
||
|
||
import copy
|
||
|
||
import pytest
|
||
|
||
from src.services.message.thinking_rectifier import ThinkingRectifier
|
||
|
||
|
||
class TestRectifyBasic:
|
||
"""测试基本整流功能"""
|
||
|
||
def test_empty_request_body(self) -> None:
|
||
"""空请求体应返回原值"""
|
||
result, modified = ThinkingRectifier.rectify({})
|
||
assert result == {}
|
||
assert modified is False
|
||
|
||
def test_none_request_body(self) -> None:
|
||
"""None 请求体应返回原值"""
|
||
result, modified = ThinkingRectifier.rectify(None) # type: ignore
|
||
assert result is None
|
||
assert modified is False
|
||
|
||
def test_no_messages(self) -> None:
|
||
"""无 messages 字段应返回原值"""
|
||
body = {"model": "claude-3-opus"}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
assert result == body
|
||
assert modified is False
|
||
|
||
def test_empty_messages(self) -> None:
|
||
"""空 messages 列表应返回原值"""
|
||
body = {"model": "claude-3-opus", "messages": []}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
assert result["messages"] == []
|
||
assert modified is False
|
||
|
||
|
||
class TestRemoveThinkingBlocks:
|
||
"""测试移除 thinking 块"""
|
||
|
||
def test_remove_thinking_block(self) -> None:
|
||
"""应移除 thinking 块"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "...", "signature": "abc"},
|
||
{"type": "text", "text": "Hello"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
content = result["messages"][0]["content"]
|
||
assert len(content) == 1
|
||
assert content[0]["type"] == "text"
|
||
|
||
def test_remove_redacted_thinking_block(self) -> None:
|
||
"""应移除 redacted_thinking 块"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "redacted_thinking", "data": "..."},
|
||
{"type": "text", "text": "Hello"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
content = result["messages"][0]["content"]
|
||
assert len(content) == 1
|
||
assert content[0]["type"] == "text"
|
||
|
||
def test_remove_multiple_thinking_blocks(self) -> None:
|
||
"""应移除多个 thinking 块"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "thought 1"},
|
||
{"type": "text", "text": "response 1"},
|
||
],
|
||
},
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "redacted_thinking", "data": "..."},
|
||
{"type": "text", "text": "response 2"},
|
||
],
|
||
},
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
assert len(result["messages"][0]["content"]) == 1
|
||
assert len(result["messages"][1]["content"]) == 1
|
||
|
||
|
||
class TestRemoveSignatureField:
|
||
"""测试移除 signature 字段"""
|
||
|
||
def test_remove_signature_from_text_block(self) -> None:
|
||
"""应从 text 块移除 signature 字段"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "text", "text": "Hello", "signature": "should_remove"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
block = result["messages"][0]["content"][0]
|
||
assert "signature" not in block
|
||
assert block["text"] == "Hello"
|
||
|
||
def test_remove_signature_from_tool_use_block(self) -> None:
|
||
"""应从 tool_use 块移除 signature 字段"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{
|
||
"type": "tool_use",
|
||
"id": "tool_1",
|
||
"name": "search",
|
||
"input": {},
|
||
"signature": "should_remove",
|
||
},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
block = result["messages"][0]["content"][0]
|
||
assert "signature" not in block
|
||
assert block["id"] == "tool_1"
|
||
|
||
|
||
class TestTopLevelThinkingParam:
|
||
"""测试顶层 thinking 参数处理"""
|
||
|
||
def test_remove_thinking_param_when_tool_use_without_thinking_prefix(self) -> None:
|
||
"""整流后有 tool_use 但无 thinking 前缀时应移除 thinking 参数"""
|
||
body = {
|
||
"thinking": {"type": "enabled", "budget_tokens": 10000},
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "tool_use", "id": "t1", "name": "search", "input": {}},
|
||
],
|
||
}
|
||
],
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
assert "thinking" not in result
|
||
# tool_use 应保留
|
||
assert result["messages"][0]["content"][0]["type"] == "tool_use"
|
||
|
||
def test_keep_thinking_param_when_no_tool_use(self) -> None:
|
||
"""整流后无 tool_use 时应保留 thinking 参数"""
|
||
body = {
|
||
"thinking": {"type": "enabled", "budget_tokens": 10000},
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "text", "text": "Hello"},
|
||
],
|
||
}
|
||
],
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
# thinking 参数应保留（只移除了 thinking 块）
|
||
assert "thinking" in result
|
||
assert result["thinking"]["type"] == "enabled"
|
||
|
||
def test_keep_thinking_param_when_disabled(self) -> None:
|
||
"""thinking 参数未启用时应保留"""
|
||
body = {
|
||
"thinking": {"type": "disabled"},
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "tool_use", "id": "t1", "name": "search", "input": {}},
|
||
],
|
||
}
|
||
],
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
# 无 thinking 块需要移除，thinking 参数也不需要移除
|
||
assert modified is False
|
||
assert result["thinking"]["type"] == "disabled"
|
||
|
||
|
||
class TestEdgeCases:
|
||
"""测试边界情况"""
|
||
|
||
def test_non_dict_message_preserved(self) -> None:
|
||
"""非 dict 消息应保留"""
|
||
body = {
|
||
"messages": [
|
||
"string message", # 非 dict
|
||
{"role": "user", "content": "Hello"},
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert result["messages"][0] == "string message"
|
||
assert modified is False
|
||
|
||
def test_string_content_preserved(self) -> None:
|
||
"""字符串 content 应保留"""
|
||
body = {
|
||
"messages": [
|
||
{"role": "user", "content": "Hello"},
|
||
{"role": "assistant", "content": "Hi there"},
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert result["messages"][1]["content"] == "Hi there"
|
||
assert modified is False
|
||
|
||
def test_non_dict_block_in_content_preserved(self) -> None:
|
||
"""content 中的非 dict 块应保留"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
"string block", # 非 dict
|
||
{"type": "text", "text": "Hello"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert result["messages"][0]["content"][0] == "string block"
|
||
assert modified is False
|
||
|
||
def test_empty_content_after_rectify_logs_warning(self) -> None:
|
||
"""整流后 assistant 消息 content 为空时应记录警告（不抛异常）"""
|
||
body = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "only thinking"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
# 应该不抛异常
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
# content 应为空列表
|
||
assert result["messages"][0]["content"] == []
|
||
|
||
def test_deep_copy_does_not_modify_original(self) -> None:
|
||
"""深拷贝应保护原始数据不被修改"""
|
||
original = {
|
||
"messages": [
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "text", "text": "Hello"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
original_copy = copy.deepcopy(original)
|
||
|
||
ThinkingRectifier.rectify(original)
|
||
|
||
# 原始数据不应被修改
|
||
assert original == original_copy
|
||
|
||
|
||
class TestLastAssistantCheck:
|
||
"""测试只检查最后一条 assistant 消息的逻辑"""
|
||
|
||
def test_only_last_assistant_matters_for_thinking_removal(self) -> None:
|
||
"""只有最后一条 assistant 消息决定是否移除 thinking 参数"""
|
||
body = {
|
||
"thinking": {"type": "enabled", "budget_tokens": 10000},
|
||
"messages": [
|
||
# 第一条 assistant 有 tool_use
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "tool_use", "id": "t1", "name": "search", "input": {}},
|
||
],
|
||
},
|
||
{"role": "user", "content": "result"},
|
||
# 最后一条 assistant 无 tool_use
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "text", "text": "Final answer"},
|
||
],
|
||
},
|
||
],
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
# thinking 参数应保留（最后一条 assistant 无 tool_use）
|
||
assert "thinking" in result
|
||
|
||
def test_last_assistant_with_tool_use_removes_thinking(self) -> None:
|
||
"""最后一条 assistant 有 tool_use 时应移除 thinking 参数"""
|
||
body = {
|
||
"thinking": {"type": "enabled", "budget_tokens": 10000},
|
||
"messages": [
|
||
# 第一条 assistant 无 tool_use
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "text", "text": "thinking..."},
|
||
],
|
||
},
|
||
{"role": "user", "content": "continue"},
|
||
# 最后一条 assistant 有 tool_use
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "thinking", "thinking": "..."},
|
||
{"type": "tool_use", "id": "t1", "name": "search", "input": {}},
|
||
],
|
||
},
|
||
],
|
||
}
|
||
result, modified = ThinkingRectifier.rectify(body)
|
||
|
||
assert modified is True
|
||
# thinking 参数应被移除
|
||
assert "thinking" not in result
|