2aa6e9240f
- 新增 context_rewrite 工具：允许 AI 在单轮内压缩工具调用上下文 - 清理 persona_query_max 和 task_query_max 死配置（commit 9be60ba 引入） - 同步全栈：后端/前端/文档/测试 18 个文件 - 测试：70/70 通过
219 lines
10 KiB
Python
219 lines
10 KiB
Python
"""context_rewrite 全面测试 - 单元 + 集成"""
|
||
|
||
import pytest
|
||
import json
|
||
import tempfile
|
||
import os
|
||
from core.tools.memory_tools import TOOLS, MEMORY_TOOLS
|
||
from core.tool_executor import execute_tool, execute_context_rewrite
|
||
from core.tool_limiter import ToolLimiter, ToolLimits
|
||
from core import EmbeddedGraphDB
|
||
|
||
|
||
# ========== 工具定义测试 ==========
|
||
|
||
class TestToolDefinition:
|
||
def test_context_rewrite_in_tools(self):
|
||
tool_names = [t["function"]["name"] for t in TOOLS]
|
||
assert "context_rewrite" in tool_names
|
||
|
||
def test_context_rewrite_in_memory_tools(self):
|
||
tool_names = [t["function"]["name"] for t in MEMORY_TOOLS]
|
||
assert "context_rewrite" in tool_names
|
||
|
||
def test_context_rewrite_has_required_params(self):
|
||
tool_def = None
|
||
for t in MEMORY_TOOLS:
|
||
if t["function"]["name"] == "context_rewrite":
|
||
tool_def = t
|
||
break
|
||
assert tool_def is not None
|
||
assert "summary" in tool_def["function"]["parameters"]["required"]
|
||
|
||
def test_context_rewrite_description_not_empty(self):
|
||
for t in MEMORY_TOOLS:
|
||
if t["function"]["name"] == "context_rewrite":
|
||
assert len(t["function"]["description"]) > 100
|
||
break
|
||
|
||
|
||
# ========== 执行器测试 ==========
|
||
|
||
class TestContextRewriteExecutor:
|
||
@pytest.fixture
|
||
def db(self):
|
||
with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f:
|
||
db_path = f.name
|
||
db = EmbeddedGraphDB(db_path)
|
||
yield db
|
||
db.close()
|
||
os.unlink(db_path)
|
||
|
||
def test_valid_summary(self, db):
|
||
args = {"summary": "[工具调用总结: 本次总结了 2 次工具调用 | 调用工具: memory_recall, memory_recall]\n\n- 查询人设图：未找到"}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "success"
|
||
assert result["message"] == "上下文已压缩"
|
||
assert "memory_recall" in result["summary"]
|
||
|
||
def test_missing_marker(self, db):
|
||
args = {"summary": "查询人设图：未找到"}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "error"
|
||
assert "必须包含" in result["message"]
|
||
|
||
def test_empty_summary(self, db):
|
||
args = {"summary": ""}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "error"
|
||
|
||
def test_marker_only(self, db):
|
||
args = {"summary": "[工具调用总结"}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "success"
|
||
|
||
def test_via_execute_tool(self, db):
|
||
args = {"summary": "[工具调用总结: 本次总结了 1 次工具调用 | 调用工具: memory_recall]\n\n- 查询记忆：找到 3 个实体"}
|
||
result_str = execute_tool(db, "context_rewrite", args)
|
||
result = json.loads(result_str)
|
||
assert result["status"] == "success"
|
||
|
||
def test_unicode_content(self, db):
|
||
args = {"summary": "[工具调用总结: 本次总结了 3 次工具调用 | 调用工具: memory_recall, memory_commit, task_create]\n\n- 查询：找到实体\"用户\"\n- 写入：{\"subject\": \"用户\", \"relation\": \"喜欢\"}\n- 任务：Task_测试"}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "success"
|
||
assert "用户" in result["summary"]
|
||
|
||
def test_newlines_preserved(self, db):
|
||
summary = "[工具调用总结: 本次总结了 2 次工具调用 | 调用工具: memory_recall, memory_recall]\n\n- 查询1：结果1\n- 查询2：结果2"
|
||
args = {"summary": summary}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["summary"] == summary
|
||
|
||
def test_long_summary(self, db):
|
||
summary = "[工具调用总结: 本次总结了 5 次工具调用 | 调用工具: memory_recall, memory_recall, memory_commit, task_create, task_set_state]\n\n" + "详细结果\n" * 50
|
||
args = {"summary": summary}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "success"
|
||
assert len(result["summary"]) == len(summary)
|
||
|
||
def test_special_json_chars(self, db):
|
||
args = {"summary": '[工具调用总结: 本次总结了 1 次工具调用 | 调用工具: memory_commit]\n\n- 写入：{"subject": "测试", "relation": "包含\"引号"}'}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "success"
|
||
|
||
def test_missing_summary_key(self, db):
|
||
args = {}
|
||
result = execute_context_rewrite(db, args)
|
||
assert result["status"] == "error"
|
||
|
||
|
||
# ========== 工具限流器测试 ==========
|
||
|
||
class TestContextRewriteLimiter:
|
||
def test_classified_as_memory_query(self):
|
||
limiter = ToolLimiter(ToolLimits(memory_query_max=1))
|
||
category, operation = limiter._classify_tool("context_rewrite", {})
|
||
assert category == "memory"
|
||
assert operation == "query"
|
||
|
||
def test_counts_toward_memory_query_limit(self):
|
||
limiter = ToolLimiter(ToolLimits(memory_query_max=1))
|
||
allowed, _ = limiter.can_call("context_rewrite", {})
|
||
assert allowed
|
||
limiter.record_call("context_rewrite", {})
|
||
allowed, reason = limiter.can_call("context_rewrite", {})
|
||
assert not allowed
|
||
assert "一般记忆查询次数已达上限" in reason
|
||
|
||
def test_does_not_affect_memory_update(self):
|
||
limiter = ToolLimiter(ToolLimits(memory_update_max=1))
|
||
limiter.record_call("context_rewrite", {})
|
||
allowed, _ = limiter.can_call("memory_commit", {})
|
||
assert allowed
|
||
|
||
def test_reset_clears_count(self):
|
||
limiter = ToolLimiter(ToolLimits(memory_query_max=1))
|
||
limiter.record_call("context_rewrite", {})
|
||
limiter.reset()
|
||
allowed, _ = limiter.can_call("context_rewrite", {})
|
||
assert allowed
|
||
|
||
|
||
# ========== 集成测试：messages_history 压缩流程 ==========
|
||
|
||
class TestMessagesHistoryCompression:
|
||
def test_compression_preserves_user_message(self):
|
||
messages_history = [
|
||
{"role": "user", "content": "我们之前聊过成语接龙吗？"},
|
||
{"role": "assistant", "content": None, "tool_calls": [{"id": "tc1", "type": "function", "function": {"name": "memory_recall", "arguments": '{"query_intent": "人设"}'}}]},
|
||
{"role": "tool", "tool_call_id": "tc1", "content": "===== 记忆检索结果 =====\n\n(未找到相关记忆)\n=============================="},
|
||
{"role": "assistant", "content": None, "tool_calls": [{"id": "tc2", "type": "function", "function": {"name": "memory_recall", "arguments": '{"query_intent": "工作记忆"}'}}]},
|
||
{"role": "tool", "tool_call_id": "tc2", "content": "===== 记忆检索结果 =====\n\n实体 (2 个):\n - Task_成语接龙 (类型: unknown, 提及: 1次)\n=============================="},
|
||
]
|
||
|
||
summary = "[工具调用总结: 本次总结了 2 次工具调用 | 调用工具: memory_recall, memory_recall]\n\n- 查询人设图：未找到人设\n- 查询工作记忆链：发现 Task_成语接龙，状态已暂停"
|
||
|
||
user_msg = messages_history[0]
|
||
messages_history[:] = [
|
||
user_msg,
|
||
{"role": "assistant", "content": summary}
|
||
]
|
||
|
||
assert len(messages_history) == 2
|
||
assert messages_history[0]["role"] == "user"
|
||
assert messages_history[0]["content"] == "我们之前聊过成语接龙吗？"
|
||
assert messages_history[1]["role"] == "assistant"
|
||
assert "memory_recall" in messages_history[1]["content"]
|
||
assert "成语接龙" in messages_history[1]["content"]
|
||
|
||
def test_compression_removes_json_noise(self):
|
||
messages_history = [
|
||
{"role": "user", "content": "查询用户信息"},
|
||
{"role": "assistant", "content": None, "tool_calls": [{"id": "tc1", "type": "function", "function": {"name": "memory_recall", "arguments": '{"query_intent": "用户"}'}}]},
|
||
{"role": "tool", "tool_call_id": "tc1", "content": json.dumps({"entities": [{"name": "用户", "type": "person", "mention_count": 5}], "relations": [{"source": "用户", "target": "Python", "type": "喜欢"}]})},
|
||
]
|
||
|
||
summary = "[工具调用总结: 本次总结了 1 次工具调用 | 调用工具: memory_recall]\n\n- 查询用户：找到用户实体，提及5次，喜欢Python"
|
||
|
||
user_msg = messages_history[0]
|
||
messages_history[:] = [user_msg, {"role": "assistant", "content": summary}]
|
||
|
||
for msg in messages_history[1:]:
|
||
assert "entities" not in msg.get("content", "")
|
||
assert "relations" not in msg.get("content", "")
|
||
|
||
def test_compression_retains_tool_meta_cognition(self):
|
||
messages_history = [
|
||
{"role": "user", "content": "查询"},
|
||
{"role": "assistant", "content": None, "tool_calls": [{"id": "tc1", "type": "function", "function": {"name": "memory_recall", "arguments": "{}"}}]},
|
||
{"role": "tool", "tool_call_id": "tc1", "content": "结果"},
|
||
]
|
||
|
||
summary = "[工具调用总结: 本次总结了 1 次工具调用 | 调用工具: memory_recall]\n\n- 查询记忆：无结果"
|
||
|
||
user_msg = messages_history[0]
|
||
messages_history[:] = [user_msg, {"role": "assistant", "content": summary}]
|
||
|
||
content = messages_history[1]["content"]
|
||
assert "工具调用总结" in content
|
||
assert "memory_recall" in content
|
||
assert "1 次" in content
|
||
|
||
def test_multiple_compressions_in_sequence(self):
|
||
messages_history = [
|
||
{"role": "user", "content": "多轮查询"},
|
||
]
|
||
|
||
for i in range(3):
|
||
messages_history.append({"role": "assistant", "content": None, "tool_calls": [{"id": f"tc{i}", "type": "function", "function": {"name": "memory_recall", "arguments": "{}"}}]})
|
||
messages_history.append({"role": "tool", "tool_call_id": f"tc{i}", "content": f"结果{i}"})
|
||
|
||
summary = f"[工具调用总结: 本次总结了 {i+1} 次工具调用 | 调用工具: memory_recall]\n\n- 第{i+1}轮查询：结果{i}"
|
||
user_msg = messages_history[0]
|
||
messages_history[:] = [user_msg, {"role": "assistant", "content": summary}]
|
||
|
||
assert len(messages_history) == 2
|
||
assert messages_history[0]["content"] == "多轮查询"
|
||
assert "第3轮查询" in messages_history[1]["content"]
|