mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-07 08:26:06 +00:00
5b801c0d3e
* refactor(file_io): update file I/O operations and truncation logic * refactor(memory): update file-based memory compaction logic
319 lines
10 KiB
Python
319 lines
10 KiB
Python
"""Tests for Summarizer."""
|
||
|
||
import asyncio
|
||
import datetime
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
from agentscope.message import Msg
|
||
from agentscope.tool import Toolkit
|
||
from test_utils import (
|
||
get_dash_chat_model,
|
||
get_formatter,
|
||
get_token_counter,
|
||
)
|
||
from reme.core.utils import get_logger
|
||
from reme.memory.file_based.components import Summarizer
|
||
from reme.memory.file_based.tools import FileIO
|
||
|
||
|
||
logger = get_logger()
|
||
|
||
|
||
# ANSI 颜色码
|
||
class Colors:
|
||
"""ANSI color codes for terminal output."""
|
||
|
||
GREEN = "\033[92m"
|
||
RED = "\033[91m"
|
||
YELLOW = "\033[93m"
|
||
BLUE = "\033[94m"
|
||
CYAN = "\033[96m"
|
||
BOLD = "\033[1m"
|
||
RESET = "\033[0m"
|
||
|
||
|
||
def print_pass(test_name: str):
|
||
"""打印测试通过信息"""
|
||
print(f"{Colors.GREEN}{Colors.BOLD}✓ {test_name} PASSED{Colors.RESET}")
|
||
|
||
|
||
def print_fail(test_name: str, error: str):
|
||
"""打印测试失败信息"""
|
||
print(f"{Colors.RED}{Colors.BOLD}✗ {test_name} FAILED: {error}{Colors.RESET}")
|
||
|
||
|
||
def print_error(test_name: str, error: str):
|
||
"""打印测试错误信息"""
|
||
print(f"{Colors.YELLOW}{Colors.BOLD}⚠ {test_name} ERROR: {error}{Colors.RESET}")
|
||
|
||
|
||
def print_test_header(test_name: str):
|
||
"""打印测试标题"""
|
||
print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
|
||
print(f"{Colors.BLUE}{Colors.BOLD}Running: {test_name}{Colors.RESET}")
|
||
print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
|
||
|
||
|
||
def create_user_msg(content: str) -> Msg:
|
||
"""Create a user message."""
|
||
return Msg(name="user", role="user", content=content)
|
||
|
||
|
||
def create_assistant_msg(content: str) -> Msg:
|
||
"""Create an assistant message."""
|
||
return Msg(name="assistant", role="assistant", content=content)
|
||
|
||
|
||
def create_tool_use_msg(tool_name: str, tool_input: dict) -> Msg:
|
||
"""Create a message with tool_use content block."""
|
||
return Msg(
|
||
name="assistant",
|
||
role="assistant",
|
||
content=[
|
||
{
|
||
"type": "tool_use",
|
||
"id": "call_123",
|
||
"name": tool_name,
|
||
"input": tool_input,
|
||
},
|
||
],
|
||
)
|
||
|
||
|
||
def create_tool_result_msg(tool_name: str, output: str) -> Msg:
|
||
"""Create a message with tool_result content block."""
|
||
return Msg(
|
||
name="tool",
|
||
role="user",
|
||
content=[
|
||
{
|
||
"type": "tool_result",
|
||
"id": "call_123",
|
||
"name": tool_name,
|
||
"output": output,
|
||
},
|
||
],
|
||
)
|
||
|
||
|
||
def create_toolkit(working_dir: str) -> Toolkit:
|
||
"""Create a default Toolkit with FileIO tools for testing."""
|
||
toolkit = Toolkit()
|
||
file_io = FileIO(working_dir=working_dir)
|
||
toolkit.register_tool_function(file_io.read_file)
|
||
toolkit.register_tool_function(file_io.write_file)
|
||
toolkit.register_tool_function(file_io.edit_file)
|
||
return toolkit
|
||
|
||
|
||
def create_summarizer(working_dir: str = None, memory_dir: str = "memory"):
|
||
"""Create a Summarizer instance for testing."""
|
||
if working_dir is None:
|
||
working_dir = tempfile.mkdtemp()
|
||
|
||
# 确保 memory_dir 存在
|
||
memory_path = Path(working_dir) / memory_dir
|
||
memory_path.mkdir(parents=True, exist_ok=True)
|
||
|
||
return (
|
||
Summarizer(
|
||
working_dir=working_dir,
|
||
memory_dir=memory_dir,
|
||
memory_compact_threshold=4000,
|
||
as_token_counter=get_token_counter(),
|
||
toolkit=create_toolkit(working_dir),
|
||
as_llm=get_dash_chat_model(),
|
||
as_llm_formatter=get_formatter(),
|
||
),
|
||
working_dir,
|
||
)
|
||
|
||
|
||
def test_empty_messages():
|
||
"""Test that empty messages return empty string."""
|
||
summarizer, _ = create_summarizer()
|
||
result = asyncio.run(summarizer.call(messages=[]))
|
||
assert result == "", f"Expected empty string, got: {result}"
|
||
print_pass("test_empty_messages")
|
||
|
||
|
||
def test_short_conversation():
|
||
"""Test summarization of a short conversation."""
|
||
summarizer, working_dir = create_summarizer()
|
||
messages = [
|
||
create_user_msg("Hello, I need help with Python."),
|
||
create_assistant_msg("Sure, I'd be happy to help. What do you need?"),
|
||
create_user_msg("How do I read a file?"),
|
||
create_assistant_msg("You can use open() function: with open('file.txt', 'r') as f: content = f.read()"),
|
||
]
|
||
|
||
logger.info(f"Input messages count: {len(messages)}")
|
||
logger.info(f"Working directory: {working_dir}")
|
||
|
||
result = asyncio.run(summarizer.call(messages=messages))
|
||
|
||
logger.info(f"Result type: {type(result)}")
|
||
logger.info(f"Result: {result}")
|
||
|
||
assert result, "Result should not be empty"
|
||
assert isinstance(result, str), f"Result should be string, got: {type(result)}"
|
||
print_pass("test_short_conversation")
|
||
|
||
|
||
def test_conversation_with_tool_calls():
|
||
"""Test summarization of conversation with tool calls."""
|
||
summarizer, working_dir = create_summarizer()
|
||
messages = [
|
||
create_user_msg("Help me debug this Python script."),
|
||
create_assistant_msg("Let me check the script first."),
|
||
create_tool_use_msg("read_file", {"path": "/app/main.py"}),
|
||
create_tool_result_msg("read_file", "def process():\n data = load_data()\n return data"),
|
||
create_assistant_msg("I found the issue! The script needs error handling."),
|
||
create_user_msg("How should I fix it?"),
|
||
create_assistant_msg("Add try-except block around the load_data() call."),
|
||
]
|
||
|
||
logger.info(f"Working directory: {working_dir}")
|
||
result = asyncio.run(summarizer.call(messages=messages))
|
||
|
||
assert result, "Result should not be empty"
|
||
assert isinstance(result, str), f"Result should be string, got: {type(result)}"
|
||
print_pass("test_conversation_with_tool_calls")
|
||
|
||
|
||
def test_consecutive_summaries():
|
||
"""Test consecutive summaries in the same directory.
|
||
|
||
This test verifies that:
|
||
1. First summary creates the memory file
|
||
2. Second summary reads and updates the existing file
|
||
"""
|
||
# 使用固定的临时目录
|
||
working_dir = tempfile.mkdtemp()
|
||
memory_dir = "memory"
|
||
memory_path = Path(working_dir) / memory_dir
|
||
memory_path.mkdir(parents=True, exist_ok=True)
|
||
|
||
logger.info(f"Working directory: {working_dir}")
|
||
logger.info(f"Memory path: {memory_path}")
|
||
|
||
# 创建 Summarizer 实例
|
||
summarizer = Summarizer(
|
||
working_dir=working_dir,
|
||
memory_dir=memory_dir,
|
||
memory_compact_threshold=4000,
|
||
as_token_counter=get_token_counter(),
|
||
toolkit=create_toolkit(working_dir),
|
||
as_llm=get_dash_chat_model(),
|
||
as_llm_formatter=get_formatter(),
|
||
)
|
||
|
||
# 第一轮对话
|
||
messages_round1 = [
|
||
create_user_msg("My name is Alice and I'm learning Python."),
|
||
create_assistant_msg("Nice to meet you, Alice! Python is a great language to learn."),
|
||
create_user_msg("I prefer using VS Code as my editor."),
|
||
create_assistant_msg("VS Code is excellent for Python development with great extensions."),
|
||
]
|
||
|
||
logger.info("=" * 40)
|
||
logger.info("Round 1: First summary (creating new file)")
|
||
logger.info("=" * 40)
|
||
result1 = asyncio.run(summarizer.call(messages=messages_round1))
|
||
logger.info(f"Round 1 Result:\n{result1}")
|
||
|
||
# 检查文件是否被创建
|
||
today = datetime.datetime.now().strftime("%Y-%m-%d")
|
||
expected_file = memory_path / f"{today}.md"
|
||
logger.info(f"Expected file: {expected_file}")
|
||
|
||
# 列出目录内容
|
||
files_after_round1 = list(memory_path.iterdir())
|
||
logger.info(f"Files after round 1: {files_after_round1}")
|
||
|
||
assert expected_file.exists(), f"Memory file should be created at {expected_file}"
|
||
|
||
# 读取第一轮写入的内容
|
||
content_after_round1 = expected_file.read_text()
|
||
logger.info(f"Content after round 1:\n{content_after_round1}")
|
||
|
||
# 第二轮对话
|
||
messages_round2 = [
|
||
create_user_msg("I also like using Docker for my projects."),
|
||
create_assistant_msg("Docker is great for containerization and deployment."),
|
||
create_user_msg("My favorite framework is FastAPI."),
|
||
create_assistant_msg("FastAPI is excellent for building modern APIs with Python."),
|
||
]
|
||
|
||
logger.info("=" * 40)
|
||
logger.info("Round 2: Second summary (reading and updating existing file)")
|
||
logger.info("=" * 40)
|
||
result2 = asyncio.run(summarizer.call(messages=messages_round2))
|
||
logger.info(f"Round 2 Result:\n{result2}")
|
||
|
||
# 读取第二轮写入后的内容
|
||
content_after_round2 = expected_file.read_text()
|
||
logger.info(f"Content after round 2:\n{content_after_round2}")
|
||
|
||
# 验证
|
||
assert result1, "Round 1 result should not be empty"
|
||
assert result2, "Round 2 result should not be empty"
|
||
|
||
# 验证第二轮内容包含新信息（Docker 或 FastAPI）
|
||
# 注意：具体内容取决于 LLM 的响应
|
||
assert len(content_after_round2) > 0, "Content after round 2 should not be empty"
|
||
|
||
logger.info("=" * 40)
|
||
logger.info("Consecutive summaries test completed successfully!")
|
||
logger.info("=" * 40)
|
||
|
||
print_pass("test_consecutive_summaries")
|
||
|
||
|
||
def run_all_tests():
|
||
"""Run all tests."""
|
||
tests = [
|
||
test_consecutive_summaries,
|
||
test_empty_messages,
|
||
test_short_conversation,
|
||
test_conversation_with_tool_calls,
|
||
]
|
||
|
||
passed = 0
|
||
failed = 0
|
||
|
||
for test in tests:
|
||
try:
|
||
print_test_header(test.__name__)
|
||
test()
|
||
passed += 1
|
||
except AssertionError as e:
|
||
print_fail(test.__name__, str(e))
|
||
failed += 1
|
||
except Exception as e:
|
||
import traceback
|
||
|
||
print_error(test.__name__, str(e))
|
||
traceback.print_exc()
|
||
failed += 1
|
||
|
||
# 打印最终统计结果
|
||
print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
|
||
print(f"{Colors.BOLD}Test Results Summary{Colors.RESET}")
|
||
print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
|
||
print(f"{Colors.GREEN}{Colors.BOLD}✓ Passed: {passed}{Colors.RESET}")
|
||
if failed > 0:
|
||
print(f"{Colors.RED}{Colors.BOLD}✗ Failed: {failed}{Colors.RESET}")
|
||
else:
|
||
print(f"{Colors.GREEN}✗ Failed: {failed}{Colors.RESET}")
|
||
print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
|
||
|
||
if failed == 0:
|
||
print(f"\n{Colors.GREEN}{Colors.BOLD}🎉 All tests passed!{Colors.RESET}")
|
||
else:
|
||
print(f"\n{Colors.RED}{Colors.BOLD}💥 Some tests failed!{Colors.RESET}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_all_tests()
|