341 lines
13 KiB
Python
341 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
ChatMessage / ToolCall / WriteRequest schema 扩展 + _parse_input 双风格规范化测试
|
|
|
|
覆盖：
|
|
- ChatMessage / ToolCall 基础字段与 to_dict / from_dict round-trip
|
|
- WriteRequest 新增 workspace_id / branch 字段
|
|
- _parse_input：纯字符串、旧式 plain messages、OpenAI 风格、Anthropic 风格、
|
|
嵌套 tool_result 内容、多 tool_use+tool_result 配对、非法 JSON args、
|
|
空 content 防御
|
|
- _parse_input 不变量：所有 ChatMessage.content 永远是 string
|
|
"""
|
|
|
|
from typing import List
|
|
|
|
import pytest
|
|
|
|
from hy_memory.client import HyMemoryClient
|
|
from hy_memory.pipelines.base import ChatMessage, ToolCall, WriteRequest
|
|
|
|
|
|
# ================================================================
|
|
# ToolCall / ChatMessage 基础
|
|
# ================================================================
|
|
|
|
class TestToolCall:
|
|
def test_default_construction(self):
|
|
tc = ToolCall()
|
|
assert tc.id == "" and tc.name == "" and tc.arguments == {}
|
|
|
|
def test_to_dict(self):
|
|
tc = ToolCall(id="t1", name="Read", arguments={"path": "a.py"})
|
|
assert tc.to_dict() == {"id": "t1", "name": "Read", "arguments": {"path": "a.py"}}
|
|
|
|
def test_from_dict_round_trip(self):
|
|
d = {"id": "x", "name": "Bash", "arguments": {"command": "ls"}}
|
|
assert ToolCall.from_dict(d).to_dict() == d
|
|
|
|
def test_from_dict_handles_missing_keys(self):
|
|
tc = ToolCall.from_dict({})
|
|
assert tc.id == "" and tc.name == "" and tc.arguments == {}
|
|
|
|
def test_from_dict_handles_none_arguments(self):
|
|
tc = ToolCall.from_dict({"id": "x", "name": "Y", "arguments": None})
|
|
assert tc.arguments == {}
|
|
|
|
|
|
class TestChatMessageBackwardCompat:
|
|
"""已有 chat 路径用法不能改变行为"""
|
|
|
|
def test_old_construction_still_works(self):
|
|
m = ChatMessage(role="user", content="hi")
|
|
assert m.role == "user" and m.content == "hi"
|
|
assert m.tool_calls == []
|
|
assert m.tool_call_id is None
|
|
assert m.tool_name is None
|
|
|
|
def test_old_to_dict_format_unchanged(self):
|
|
# 不带 tool 信息时 to_dict 输出格式与历史完全一致
|
|
m = ChatMessage(role="user", content="hi")
|
|
assert m.to_dict() == {"role": "user", "content": "hi"}
|
|
|
|
def test_old_from_dict_unchanged(self):
|
|
m = ChatMessage.from_dict({"role": "assistant", "content": "ok"})
|
|
assert m.role == "assistant" and m.content == "ok"
|
|
assert not m.has_tool_calls() and not m.is_tool_message()
|
|
|
|
|
|
class TestChatMessageNewFields:
|
|
def test_to_dict_includes_tool_calls(self):
|
|
tc = ToolCall(id="t1", name="Read", arguments={"path": "a.py"})
|
|
m = ChatMessage(role="assistant", content="", tool_calls=[tc])
|
|
d = m.to_dict()
|
|
assert d["role"] == "assistant" and d["content"] == ""
|
|
assert d["tool_calls"] == [tc.to_dict()]
|
|
assert "tool_call_id" not in d # 默认 None 不输出
|
|
assert "tool_name" not in d
|
|
|
|
def test_to_dict_includes_tool_result_fields(self):
|
|
m = ChatMessage(role="tool", content="42 lines", tool_call_id="t1", tool_name="Read")
|
|
d = m.to_dict()
|
|
assert d["tool_call_id"] == "t1"
|
|
assert d["tool_name"] == "Read"
|
|
assert "tool_calls" not in d # 空列表不输出
|
|
|
|
def test_helpers(self):
|
|
# has_tool_calls
|
|
assert not ChatMessage(role="user").has_tool_calls()
|
|
assert ChatMessage(role="assistant",
|
|
tool_calls=[ToolCall(id="x", name="Y")]).has_tool_calls()
|
|
# is_tool_message
|
|
assert ChatMessage(role="tool").is_tool_message()
|
|
assert not ChatMessage(role="user").is_tool_message()
|
|
|
|
def test_from_dict_round_trip_with_tool_calls(self):
|
|
d = {
|
|
"role": "assistant",
|
|
"content": "I'll read it",
|
|
"tool_calls": [{"id": "x", "name": "Bash", "arguments": {"cmd": "ls"}}],
|
|
}
|
|
m = ChatMessage.from_dict(d)
|
|
assert m.has_tool_calls()
|
|
assert m.tool_calls[0].arguments == {"cmd": "ls"}
|
|
|
|
|
|
# ================================================================
|
|
# WriteRequest 扩展
|
|
# ================================================================
|
|
|
|
class TestWriteRequestExtension:
|
|
def test_default_workspace_branch_none(self):
|
|
wr = WriteRequest(content="x", user_id="u1")
|
|
assert wr.workspace_id is None
|
|
assert wr.branch is None
|
|
|
|
def test_set_workspace_and_branch(self):
|
|
wr = WriteRequest(content="x", user_id="u1",
|
|
workspace_id="github.com/x/y", branch="main")
|
|
assert wr.workspace_id == "github.com/x/y"
|
|
assert wr.branch == "main"
|
|
|
|
def test_existing_fields_unaffected(self):
|
|
wr = WriteRequest(content="x", user_id="u1", agent_id="a1",
|
|
session_id="s1")
|
|
assert wr.user_id == "u1" and wr.agent_id == "a1" and wr.session_id == "s1"
|
|
|
|
|
|
# ================================================================
|
|
# _parse_input — 主路径
|
|
# ================================================================
|
|
|
|
class TestParseInputBackwardCompat:
|
|
def test_string_input(self):
|
|
content, msgs, ctype = HyMemoryClient._parse_input("hello")
|
|
assert content == "hello" and msgs == [] and ctype == "text"
|
|
|
|
def test_invalid_input_type(self):
|
|
with pytest.raises(TypeError):
|
|
HyMemoryClient._parse_input(42)
|
|
|
|
def test_plain_messages(self):
|
|
data = [
|
|
{"role": "user", "content": "你好"},
|
|
{"role": "assistant", "content": "我帮你"},
|
|
]
|
|
_, msgs, ctype = HyMemoryClient._parse_input(data)
|
|
assert ctype == "messages"
|
|
assert len(msgs) == 2
|
|
assert msgs[0].role == "user" and msgs[0].content == "你好"
|
|
assert msgs[1].role == "assistant" and msgs[1].content == "我帮你"
|
|
# 不会因为新字段引入额外消息
|
|
assert all(not m.tool_calls for m in msgs)
|
|
|
|
|
|
class TestParseInputOpenAIStyle:
|
|
def test_assistant_with_tool_calls_and_tool_message(self):
|
|
data = [
|
|
{"role": "user", "content": "读 a.py"},
|
|
{"role": "assistant", "content": None,
|
|
"tool_calls": [{"id": "t1", "type": "function",
|
|
"function": {"name": "Read",
|
|
"arguments": '{"path": "a.py"}'}}]},
|
|
{"role": "tool", "tool_call_id": "t1", "name": "Read",
|
|
"content": "file ok\n42 lines"},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert len(msgs) == 3
|
|
# assistant
|
|
assert msgs[1].role == "assistant"
|
|
assert msgs[1].content == "" # null → ""
|
|
assert len(msgs[1].tool_calls) == 1
|
|
tc = msgs[1].tool_calls[0]
|
|
assert tc.id == "t1" and tc.name == "Read"
|
|
# arguments 已被 json.loads 成 dict
|
|
assert tc.arguments == {"path": "a.py"}
|
|
# tool message
|
|
assert msgs[2].role == "tool"
|
|
assert msgs[2].tool_call_id == "t1"
|
|
assert msgs[2].tool_name == "Read"
|
|
assert msgs[2].content == "file ok\n42 lines"
|
|
|
|
def test_invalid_json_arguments_become_raw(self):
|
|
data = [
|
|
{"role": "assistant", "content": None,
|
|
"tool_calls": [{"id": "x",
|
|
"function": {"name": "Bad",
|
|
"arguments": "{not valid"}}]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert msgs[0].tool_calls[0].arguments == {"_raw": "{not valid"}
|
|
|
|
def test_empty_arguments_string(self):
|
|
data = [
|
|
{"role": "assistant", "content": None,
|
|
"tool_calls": [{"id": "x",
|
|
"function": {"name": "F", "arguments": ""}}]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert msgs[0].tool_calls[0].arguments == {}
|
|
|
|
def test_tool_message_with_null_content(self):
|
|
data = [{"role": "tool", "tool_call_id": "t1", "content": None}]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert len(msgs) == 1
|
|
assert msgs[0].role == "tool"
|
|
assert msgs[0].content == ""
|
|
|
|
|
|
class TestParseInputAnthropicStyle:
|
|
def test_text_plus_tool_use_block(self):
|
|
data = [
|
|
{"role": "user", "content": "读 a.py"},
|
|
{"role": "assistant", "content": [
|
|
{"type": "text", "text": "让我读一下"},
|
|
{"type": "tool_use", "id": "t2", "name": "Read",
|
|
"input": {"path": "a.py"}},
|
|
]},
|
|
{"role": "user", "content": [
|
|
{"type": "tool_result", "tool_use_id": "t2",
|
|
"content": "42 lines"},
|
|
]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
# user1 + assistant(text+tool_call) + tool(从 user 拆出来的)
|
|
assert len(msgs) == 3
|
|
assert msgs[0].role == "user" and msgs[0].content == "读 a.py"
|
|
assert msgs[1].role == "assistant"
|
|
assert msgs[1].content == "让我读一下"
|
|
assert len(msgs[1].tool_calls) == 1
|
|
assert msgs[1].tool_calls[0].id == "t2"
|
|
assert msgs[1].tool_calls[0].arguments == {"path": "a.py"}
|
|
# tool_result 拆为独立 role=tool 消息
|
|
assert msgs[2].role == "tool"
|
|
assert msgs[2].tool_call_id == "t2"
|
|
assert msgs[2].content == "42 lines"
|
|
|
|
def test_nested_tool_result_content_list(self):
|
|
data = [
|
|
{"role": "user", "content": [
|
|
{"type": "tool_result", "tool_use_id": "t3",
|
|
"content": [
|
|
{"type": "text", "text": "output line 1"},
|
|
{"type": "text", "text": "line 2"},
|
|
]},
|
|
]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert len(msgs) == 1
|
|
assert msgs[0].role == "tool"
|
|
assert msgs[0].content == "output line 1\nline 2"
|
|
|
|
def test_multiple_tool_use_and_tool_result_pairs(self):
|
|
data = [
|
|
{"role": "assistant", "content": [
|
|
{"type": "tool_use", "id": "a", "name": "Read",
|
|
"input": {"path": "x.py"}},
|
|
{"type": "tool_use", "id": "b", "name": "Read",
|
|
"input": {"path": "y.py"}},
|
|
]},
|
|
{"role": "user", "content": [
|
|
{"type": "tool_result", "tool_use_id": "a", "content": "X"},
|
|
{"type": "tool_result", "tool_use_id": "b", "content": "Y"},
|
|
]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
# assistant(空 text + 2 tool_calls) + 2 个独立 tool 消息
|
|
assert len(msgs) == 3
|
|
assert msgs[0].role == "assistant"
|
|
assert msgs[0].content == ""
|
|
assert len(msgs[0].tool_calls) == 2
|
|
assert msgs[1].role == "tool" and msgs[1].tool_call_id == "a"
|
|
assert msgs[1].content == "X"
|
|
assert msgs[2].role == "tool" and msgs[2].tool_call_id == "b"
|
|
assert msgs[2].content == "Y"
|
|
|
|
def test_assistant_only_text_block(self):
|
|
# 只有 text block 时不应产生空 tool_calls 列表
|
|
data = [
|
|
{"role": "assistant", "content": [
|
|
{"type": "text", "text": "纯文本回复"},
|
|
]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert len(msgs) == 1
|
|
assert msgs[0].role == "assistant"
|
|
assert msgs[0].content == "纯文本回复"
|
|
assert msgs[0].tool_calls == []
|
|
|
|
def test_empty_content_block_list_drops_message(self):
|
|
# content=[] 且无任何 tool_calls 不应保留这条消息
|
|
data = [{"role": "assistant", "content": []}]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert len(msgs) == 0
|
|
|
|
def test_unknown_block_type_silently_ignored(self):
|
|
data = [
|
|
{"role": "assistant", "content": [
|
|
{"type": "image", "source": {"data": "..."}}, # 未知 block
|
|
{"type": "text", "text": "看图"},
|
|
]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
assert len(msgs) == 1
|
|
assert msgs[0].content == "看图"
|
|
|
|
|
|
class TestParseInputInvariants:
|
|
"""所有路径下 ChatMessage.content 必须是 string"""
|
|
|
|
@pytest.mark.parametrize("data", [
|
|
[{"role": "user", "content": "plain"}],
|
|
[{"role": "assistant", "content": None,
|
|
"tool_calls": [{"id": "x", "function":
|
|
{"name": "F", "arguments": "{}"}}]}],
|
|
[{"role": "assistant", "content":
|
|
[{"type": "text", "text": "a"},
|
|
{"type": "tool_use", "id": "i", "name": "N", "input": {}}]}],
|
|
[{"role": "user", "content":
|
|
[{"type": "tool_result", "tool_use_id": "i", "content": "r"}]}],
|
|
[{"role": "tool", "tool_call_id": "i", "content": None}],
|
|
])
|
|
def test_content_always_str(self, data):
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
for m in msgs:
|
|
assert isinstance(m.content, str), \
|
|
f"role={m.role} content type is {type(m.content).__name__}"
|
|
|
|
def test_role_tool_arguments_always_dict(self):
|
|
data = [
|
|
{"role": "assistant", "content": None,
|
|
"tool_calls": [
|
|
{"id": "1", "function": {"name": "A", "arguments": '{"a": 1}'}},
|
|
{"id": "2", "function": {"name": "B", "arguments": "bogus"}},
|
|
{"id": "3", "function": {"name": "C", "arguments": ""}},
|
|
{"id": "4", "function": {"name": "D", "arguments": None}},
|
|
]},
|
|
]
|
|
_, msgs, _ = HyMemoryClient._parse_input(data)
|
|
for tc in msgs[0].tool_calls:
|
|
assert isinstance(tc.arguments, dict)
|