feat(mem-agent): introduce version 4 memory agents and tools

This commit is contained in:
jinli.yl 2026-01-18 15:50:46 +08:00
parent e6ad682ede
commit 08b771b6c4
31 changed files with 2143 additions and 74 deletions

View file

@ -0,0 +1,327 @@
"""
Compute Question Answering statistics from eval_reme_simple_v4.py results.
This script processes the output from eval_reme_simple_v4.py and computes
comprehensive QA metrics.
Usage:
python bench/halumem/compute_qa_stats_v4.py --results_file bench_results/reme_simple_v4/eval_results.jsonl
"""
import json
import os
from pathlib import Path
from typing import Any
from loguru import logger
def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
"""Compute question answering metrics."""
total = len(qa_records)
if total == 0:
return {
"correct_qa_ratio(all)": 0,
"hallucination_qa_ratio(all)": 0,
"omission_qa_ratio(all)": 0,
"correct_qa_ratio(valid)": 0,
"hallucination_qa_ratio(valid)": 0,
"omission_qa_ratio(valid)": 0,
"qa_valid_num": 0,
"qa_num": 0
}
correct = 0
hallucination = 0
omission = 0
valid = 0
for qa in qa_records:
result_type = qa.get("result_type", "")
if result_type in ["Correct", "Hallucination", "Omission"]:
valid += 1
if result_type == "Correct":
correct += 1
elif result_type == "Hallucination":
hallucination += 1
elif result_type == "Omission":
omission += 1
metrics = {
"correct_qa_ratio(all)": correct / total,
"hallucination_qa_ratio(all)": hallucination / total,
"omission_qa_ratio(all)": omission / total,
"qa_valid_num": valid,
"qa_num": total
}
if valid > 0:
metrics.update({
"correct_qa_ratio(valid)": correct / valid,
"hallucination_qa_ratio(valid)": hallucination / valid,
"omission_qa_ratio(valid)": omission / valid
})
else:
metrics.update({
"correct_qa_ratio(valid)": 0,
"hallucination_qa_ratio(valid)": 0,
"omission_qa_ratio(valid)": 0
})
return metrics
def compute_time_metrics(results_file: str) -> dict[str, float]:
"""Compute timing metrics from evaluation results."""
add_duration = 0
search_duration = 0
with open(results_file, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
user_data = json.loads(line)
for session in user_data.get("sessions", []):
add_duration += session.get("add_dialogue_duration_ms", 0)
eval_results = session.get("evaluation_results", {})
for qa in eval_results.get("question_answering_records", []):
search_duration += qa.get("search_duration_ms", 0)
# Convert to minutes
return {
"add_dialogue_duration_time": add_duration / 1000 / 60,
"search_memory_duration_time": search_duration / 1000 / 60,
"total_duration_time": (add_duration + search_duration) / 1000 / 60
}
def load_from_tmp_dir(tmp_dir: str) -> tuple[str, list[dict]]:
"""Load data from tmp directory and generate eval_results.jsonl file."""
tmp_path = Path(tmp_dir)
parent_dir = tmp_path.parent
eval_results_file = parent_dir / "eval_results.jsonl"
print(f"\n📁 Loading data from tmp directory: {tmp_dir}")
print(f"📝 Will generate: {eval_results_file}")
# Collect all user directories
user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()]
print(f" Found {len(user_dirs)} user directories")
users_data = []
for user_dir in user_dirs:
user_name = user_dir.name
# Load all session files for this user (sorted by session number)
session_files = sorted(
[f for f in user_dir.iterdir()
if f.name.startswith("session_") and f.suffix == ".json"],
key=lambda f: int(f.stem.split("_")[1]) # Sort by session number
)
if not session_files:
print(f" ⚠️ No session files found for user: {user_name}")
continue
# Load first session to get user metadata
with open(session_files[0], "r", encoding="utf-8") as f:
first_session = json.load(f)
user_data = {
"uuid": first_session["uuid"],
"user_name": first_session["user_name"],
"sessions": []
}
# Load all sessions
for session_file in session_files:
with open(session_file, "r", encoding="utf-8") as f:
session_data = json.load(f)
# Remove redundant user metadata
session_data.pop("uuid", None)
session_data.pop("user_name", None)
user_data["sessions"].append(session_data)
users_data.append(user_data)
print(f" ✓ Loaded user {user_name}: {len(session_files)} sessions")
# Write to eval_results.jsonl
with open(eval_results_file, "w", encoding="utf-8") as f:
for user_data in users_data:
f.write(json.dumps(user_data, ensure_ascii=False) + "\n")
print(f" ✅ Generated: {eval_results_file}")
return str(eval_results_file), users_data
def main(input_path: str):
"""Main function to compute statistics from eval results."""
if not os.path.exists(input_path):
logger.error(f"Input path not found: {input_path}")
return
print("\n" + "=" * 80)
print("COMPUTING QUESTION ANSWERING STATISTICS - REME V4")
print("=" * 80)
# Determine if input is a directory (tmp) or file (eval_results.jsonl)
if os.path.isdir(input_path):
results_file, users_data = load_from_tmp_dir(input_path)
else:
results_file = input_path
users_data = None
print(f"\n📁 Using existing results file: {results_file}")
# Collect all QA records with metadata
qa_records = []
qa_records_with_metadata = [] # Store records with user/session/question info
user_count = 0
session_count = 0
with open(results_file, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
user_data = json.loads(line)
user_count += 1
user_name = user_data.get("user_name", "Unknown")
valid_session_idx = 0 # Track the index of valid (non-skipped) sessions
for original_idx, session in enumerate(user_data.get("sessions", [])):
if session.get("is_generated_qa_session"):
continue
session_count += 1
eval_results = session.get("evaluation_results", {})
session_qa_records = eval_results.get("question_answering_records", [])
# Add records with metadata
for qa_idx, qa in enumerate(session_qa_records):
qa_records.append(qa)
qa_records_with_metadata.append({
"user_name": user_name,
"session_idx": valid_session_idx,
"original_session_idx": original_idx,
"question_idx": qa_idx,
"qa_record": qa
})
valid_session_idx += 1
print(f"\n📊 Data loaded:")
print(f" Users: {user_count}")
print(f" Sessions: {session_count}")
print(f" QA Records: {len(qa_records)}")
# Compute metrics
print("\n🔄 Computing metrics...")
qa_metrics = compute_qa_metrics(qa_records)
time_metrics = compute_time_metrics(results_file)
final_results = {
"overall_score": {
"question_answering": qa_metrics,
"time_consuming": time_metrics
},
"question_answering_records": qa_records
}
# Save final report
output_dir = Path(results_file).parent
report_file = output_dir / "reme_eval_stat_result.json"
with open(report_file, "w", encoding="utf-8") as f:
json.dump(final_results, f, ensure_ascii=False, indent=4)
print(f"\n✅ Statistics saved to: {report_file}")
# Print summary
print("\n" + "=" * 80)
print("EVALUATION SUMMARY - REME V4")
print("=" * 80)
print("\n📊 Question Answering:")
print(f" Correct (all): {qa_metrics['correct_qa_ratio(all)']:.4f}")
print(f" Hallucination (all): {qa_metrics['hallucination_qa_ratio(all)']:.4f}")
print(f" Omission (all): {qa_metrics['omission_qa_ratio(all)']:.4f}")
print(f" Correct (valid): {qa_metrics['correct_qa_ratio(valid)']:.4f}")
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
print(f"\n⏱️ Time Metrics:")
print(f" Memory Addition: {time_metrics['add_dialogue_duration_time']:.2f} min")
print(f" Memory Search: {time_metrics['search_memory_duration_time']:.2f} min")
print(f" Total: {time_metrics['total_duration_time']:.2f} min")
# Print non-Correct QA records
print("\n" + "=" * 80)
print("NON-CORRECT QA RECORDS")
print("=" * 80)
non_correct_records = [
record for record in qa_records_with_metadata
if record["qa_record"].get("result_type") not in ["Correct", ""]
]
if non_correct_records:
print(f"\nFound {len(non_correct_records)} non-correct records:\n")
for record in non_correct_records:
user_name = record["user_name"]
session_idx = record["session_idx"]
original_idx = record["original_session_idx"]
question_idx = record["question_idx"]
qa = record["qa_record"]
result_type = qa.get("result_type", "Unknown")
question = qa.get("question", "N/A")
answer = qa.get("answer", "N/A")
print(f"👤 User: {user_name}")
print(f"📅 Session: {original_idx} (valid session index: {session_idx})")
print(f"❓ Question #{question_idx}")
print(f"🏷️ Result Type: {result_type}")
print(f"💬 Question: {question}")
print(f"💡 Answer: {answer}")
print("-" * 80)
else:
print("\n✅ All QA records are Correct!")
print("\n" + "=" * 80)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description="Compute QA statistics from eval_reme_simple_v4.py results"
)
parser.add_argument(
"--results_file",
type=str,
required=False,
help="Path to eval_results.jsonl file (e.g., bench_results/reme_simple_v4/eval_results.jsonl)"
)
parser.add_argument(
"--tmp_dir",
type=str,
required=False,
help="Path to tmp directory (e.g., bench_results/reme_simple_v4/tmp)"
)
args = parser.parse_args()
# Determine input path
if args.tmp_dir:
input_path = args.tmp_dir
elif args.results_file:
input_path = args.results_file
else:
parser.error("Either --results_file or --tmp_dir must be provided")
main(input_path=input_path)

View file

@ -9,7 +9,7 @@ are available in the tmp directory. It will:
4. Aggregate results and compute metrics
Usage:
python bench/halumem/compute_stats_from_tmp.py --tmp_dir bench_results/reme/tmp
python bench/halumem/compute_stats_from_tmp.py --tmp_dir bench_results/reme_simple_v4/tmp
"""
import asyncio

View file

@ -17,6 +17,7 @@ import asyncio
import json
import os
import re
import shutil
import time
from dataclasses import dataclass
from datetime import datetime, timezone
@ -497,6 +498,13 @@ class HaluMemEvaluatorV3:
# Clear existing data
await self.reme.vector_store.delete_all()
# Clear meta_memory directory
meta_memory_path = Path(f"meta_memory/{self.reme.vector_store.collection_name}")
if meta_memory_path.exists():
shutil.rmtree(meta_memory_path)
logger.info(f"Cleared meta_memory directory: {meta_memory_path}")
meta_memory_path.mkdir(parents=True, exist_ok=True)
# Load user data
all_users = self.data_loader.load_jsonl(self.config.data_path)
users_to_process = all_users[:self.config.user_num]

View file

@ -0,0 +1,671 @@
"""
HaluMem Benchmark Evaluator for ReMe - Question Answering
A modular evaluation pipeline that:
1. Loads HaluMem benchmark data
2. Processes user sessions through ReMe (summarization + retrieval)
3. Evaluates question answering performance
4. Generates comprehensive metrics
Usage:
python bench/halumem/eval_reme_simple_v4.py \
--data_path /Users/yuli/workspace/HaluMem/data/HaluMem-Medium.jsonl \
--top_k 20 --user_num 100 --max_concurrency 20
"""
import asyncio
import json
import os
import re
import shutil
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from loguru import logger
from eval_tools import evaluation_for_question2
from reme_ai.core.enumeration import MemoryType
from reme_ai.core.schema import MemoryNode
from reme_ai.reme import ReMe
# ==================== Configuration ====================
@dataclass
class EvalConfig:
"""Evaluation configuration parameters."""
data_path: str
top_k: int = 20
user_num: int = 1
max_concurrency: int = 2
batch_size: int = 20
output_dir: str = "bench_results/reme_simple_v4"
# ==================== Utilities ====================
class DataLoader:
"""Handles loading and parsing of HaluMem data."""
@staticmethod
def load_jsonl(file_path: str) -> list[dict]:
"""Load all entries from a JSONL file."""
with open(file_path, "r", encoding="utf-8") as f:
return [json.loads(line.strip()) for line in f if line.strip()]
@staticmethod
def extract_user_name(persona_info: str) -> str:
"""Extract user name from persona info string."""
match = re.search(r"Name:\s*(.*?); Gender:", persona_info)
if not match:
raise ValueError(f"No name found in persona_info: {persona_info}")
return match.group(1).strip()
@staticmethod
def format_dialogue_messages(dialogue: list[dict]) -> list[dict]:
"""Format dialogue into ReMe message format with conversation_time (user messages only)."""
return [
{
"role": turn["role"],
"content": turn["content"],
"time_created": datetime.strptime(
turn["timestamp"], "%b %d, %Y, %H:%M:%S"
)
.replace(tzinfo=timezone.utc)
.strftime("%Y-%m-%d %H:%M:%S"),
}
for turn in dialogue
if turn["role"] == "user" # Only include user messages
]
@staticmethod
def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str:
"""Format dialogue into string for evaluation."""
formatted_turns = []
for turn in dialogue:
timestamp = datetime.strptime(
turn["timestamp"], "%b %d, %Y, %H:%M:%S"
).replace(tzinfo=timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
# Use user_name if role is 'user' and user_name is provided
role = user_name if turn['role'] == 'user' and user_name else turn['role']
formatted_turns.append(
f"Role: {role}\n"
f"Content: {turn['content']}\n"
f"Time: {timestamp}"
)
return "\n\n".join(formatted_turns)
class FileManager:
"""Manages file I/O operations."""
def __init__(self, base_dir: str):
self.base_dir = Path(base_dir)
self.tmp_dir = self.base_dir / "tmp"
self.tmp_dir.mkdir(parents=True, exist_ok=True)
def get_user_dir(self, user_name: str) -> Path:
"""Get the directory path for a user."""
user_dir = self.tmp_dir / user_name
user_dir.mkdir(parents=True, exist_ok=True)
return user_dir
def get_session_file(self, user_name: str, session_id: int) -> Path:
"""Get the file path for a specific session."""
return self.get_user_dir(user_name) / f"session_{session_id}.json"
def save_session(self, user_name: str, session_id: int, data: dict):
"""Save session data to file."""
file_path = self.get_session_file(user_name, session_id)
with open(file_path, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2)
logger.info(f"✅ Saved session {session_id} to {file_path}")
def load_session(self, user_name: str, session_id: int) -> dict | None:
"""Load session data from file."""
file_path = self.get_session_file(user_name, session_id)
if not file_path.exists():
return None
with open(file_path, "r", encoding="utf-8") as f:
return json.load(f)
def user_has_cache(self, user_name: str) -> bool:
"""Check if user has cached results."""
user_dir = self.get_user_dir(user_name)
return any(f.name.startswith("session_") and f.suffix == ".json"
for f in user_dir.iterdir())
def combine_results(self, output_file: str):
"""Combine all user session files into a single JSONL file."""
with open(output_file, "w", encoding="utf-8") as f_out:
for user_dir in self.tmp_dir.iterdir():
if not user_dir.is_dir():
continue
session_files = sorted([
f for f in user_dir.iterdir()
if f.name.startswith("session_") and f.suffix == ".json"
])
if not session_files:
continue
# Load first session to get user metadata
with open(session_files[0], "r", encoding="utf-8") as f_in:
first_session = json.load(f_in)
user_data = {
"uuid": first_session["uuid"],
"user_name": first_session["user_name"],
"sessions": []
}
# Load all sessions
for session_file in session_files:
with open(session_file, "r", encoding="utf-8") as f_in:
session_data = json.load(f_in)
# Remove redundant user metadata
session_data.pop("uuid", None)
session_data.pop("user_name", None)
user_data["sessions"].append(session_data)
f_out.write(json.dumps(user_data, ensure_ascii=False) + "\n")
# ==================== Memory Operations ====================
class MemoryProcessor:
"""Handles ReMe memory operations."""
def __init__(self, reme: ReMe):
self.reme = reme
async def add_memories(
self,
user_id: str,
messages: list[dict],
batch_size: int = 10000
) -> tuple[list[str], list[list[dict]], float]:
"""
Add memories in batches using ReMe and return extracted memory contents.
Returns:
tuple: (extracted_memories, agent_messages, total_duration_ms)
"""
added_memories: list[MemoryNode] = []
deleted_memories: list[str] = []
all_agent_messages: list = []
total_duration_ms = 0
for i in range(0, len(messages), batch_size):
batch = messages[i:i + batch_size]
start = time.time()
memory_nodes, agent_messages, success = await self.reme.summary_v4(
messages=batch,
user_id=user_id
)
duration_ms = (time.time() - start) * 1000
total_duration_ms += duration_ms
# Save agent messages for this batch
if agent_messages:
all_agent_messages.extend(agent_messages)
if memory_nodes:
for node in memory_nodes:
if isinstance(node, MemoryNode) and node.memory_type == MemoryType.HISTORY:
continue
if isinstance(node, MemoryNode):
added_memories.append(node)
if isinstance(node, str):
deleted_memories.append(node)
extracted_memories = deleted_memories
extracted_memories += ["[delete]" + n.format_memory() for n in added_memories if n.memory_id in deleted_memories]
extracted_memories += ["[add]" + n.format_memory() for n in added_memories if n.memory_id not in deleted_memories]
return extracted_memories, all_agent_messages, total_duration_ms
async def search_memory(
self,
query: str,
user_id: str,
top_k: int = 20
) -> tuple[str, list, float]:
"""
Search memory using ReMe and return response.
Returns:
tuple: (response, agent_messages, duration_ms)
"""
start = time.time()
response, agent_messages, success = await self.reme.retrieve_v4(
query=query,
user_id=user_id,
top_k=top_k
)
duration_ms = (time.time() - start) * 1000
return response, agent_messages, duration_ms
# ==================== Evaluation ====================
class QuestionAnsweringEvaluator:
"""Evaluates question answering performance."""
def __init__(self, memory_processor: MemoryProcessor, top_k: int):
self.memory_processor = memory_processor
self.top_k = top_k
async def evaluate_questions(
self,
questions: list[dict],
user_name: str,
uuid: str,
session_id: int,
formatted_dialogue: str
) -> list[dict]:
"""Evaluate all questions for a session."""
results = []
for qa in questions:
response, agent_messages, duration_ms = await self.memory_processor.search_memory(
query=qa["question"],
user_id=user_name,
top_k=self.top_k
)
# Evaluate response
evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]])
eval_result = await evaluation_for_question2(
qa["question"],
qa["answer"],
evidence_text,
response,
formatted_dialogue
)
# Build result record
qa_result = {
**qa,
"uuid": uuid,
"session_id": session_id,
"system_response": response,
"retrieve_messages": [m.model_dump() for m in agent_messages],
"search_duration_ms": duration_ms,
"result_type": eval_result.get("evaluation_result"),
"question_answering_reasoning": eval_result.get("reasoning", "")
}
results.append(qa_result)
return results
class MetricsAggregator:
"""Aggregates evaluation metrics."""
@staticmethod
def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
"""Compute question answering metrics."""
total = len(qa_records)
if total == 0:
return {
"correct_qa_ratio(all)": 0,
"hallucination_qa_ratio(all)": 0,
"omission_qa_ratio(all)": 0,
"correct_qa_ratio(valid)": 0,
"hallucination_qa_ratio(valid)": 0,
"omission_qa_ratio(valid)": 0,
"qa_valid_num": 0,
"qa_num": 0
}
correct = 0
hallucination = 0
omission = 0
valid = 0
for qa in qa_records:
result_type = qa.get("result_type", "")
if result_type in ["Correct", "Hallucination", "Omission"]:
valid += 1
if result_type == "Correct":
correct += 1
elif result_type == "Hallucination":
hallucination += 1
elif result_type == "Omission":
omission += 1
metrics = {
"correct_qa_ratio(all)": correct / total,
"hallucination_qa_ratio(all)": hallucination / total,
"omission_qa_ratio(all)": omission / total,
"qa_valid_num": valid,
"qa_num": total
}
if valid > 0:
metrics.update({
"correct_qa_ratio(valid)": correct / valid,
"hallucination_qa_ratio(valid)": hallucination / valid,
"omission_qa_ratio(valid)": omission / valid
})
else:
metrics.update({
"correct_qa_ratio(valid)": 0,
"hallucination_qa_ratio(valid)": 0,
"omission_qa_ratio(valid)": 0
})
return metrics
@staticmethod
def compute_time_metrics(eval_results_file: str) -> dict[str, float]:
"""Compute timing metrics from evaluation results."""
add_duration = 0
search_duration = 0
with open(eval_results_file, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
user_data = json.loads(line)
for session in user_data["sessions"]:
add_duration += session.get("add_dialogue_duration_ms", 0)
eval_results = session.get("evaluation_results", {})
for qa in eval_results.get("question_answering_records", []):
search_duration += qa.get("search_duration_ms", 0)
# Convert to minutes
return {
"add_dialogue_duration_time": add_duration / 1000 / 60,
"search_memory_duration_time": search_duration / 1000 / 60,
"total_duration_time": (add_duration + search_duration) / 1000 / 60
}
# ==================== Main Pipeline ====================
class HaluMemEvaluatorV4:
def __init__(self, config: EvalConfig):
self.config = config
self.reme = ReMe()
self.file_manager = FileManager(config.output_dir)
self.memory_processor = MemoryProcessor(self.reme)
self.qa_evaluator = QuestionAnsweringEvaluator(
self.memory_processor,
config.top_k
)
self.data_loader = DataLoader()
async def process_session(
self,
session: dict,
session_id: int,
user_name: str,
uuid: str
) -> dict:
"""Process a single session using ReMe."""
session_data = {
"uuid": uuid,
"user_name": user_name,
"session_id": session_id,
"memory_points": session["memory_points"]
}
# Skip generated QA sessions
if session.get("is_generated_qa_session", False):
session_data["is_generated_qa_session"] = True
return session_data
dialogue = session["dialogue"]
formatted_messages = self.data_loader.format_dialogue_messages(dialogue)
extracted_memories, agent_messages, duration_ms = await self.memory_processor.add_memories(
user_id=user_name,
messages=formatted_messages,
batch_size=self.config.batch_size
)
session_data.update({
"dialogue": dialogue,
"extracted_memories": extracted_memories,
"summary_messages": [m.model_dump() for m in agent_messages],
"add_dialogue_duration_ms": duration_ms
})
# Evaluate questions if present
if "questions" in session:
formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name)
qa_results = await self.qa_evaluator.evaluate_questions(
questions=session["questions"],
user_name=user_name,
uuid=uuid,
session_id=session_id,
formatted_dialogue=formatted_dialogue
)
session_data["evaluation_results"] = {
"question_answering_records": qa_results
}
return session_data
async def process_user(self, user_data: dict) -> dict:
"""Process all sessions for a user."""
user_name = self.data_loader.extract_user_name(user_data["persona_info"])
uuid = user_data["uuid"]
logger.info(f"Processing user: {user_name}")
for idx, session in enumerate(user_data["sessions"]):
logger.info(f" Session {idx + 1}/{len(user_data['sessions'])}")
session_data = await self.process_session(
session=session,
session_id=idx,
user_name=user_name,
uuid=uuid
)
self.file_manager.save_session(user_name, idx, session_data)
return {"uuid": uuid, "user_name": user_name, "status": "ok"}
async def run_evaluation(self):
"""Run the complete evaluation pipeline using ReMe."""
start_time = time.time()
# Clear existing data
await self.reme.vector_store.delete_all()
# Clear meta_memory directory
meta_memory_path = Path(f"meta_memory/{self.reme.vector_store.collection_name}")
if meta_memory_path.exists():
shutil.rmtree(meta_memory_path)
logger.info(f"Cleared meta_memory directory: {meta_memory_path}")
meta_memory_path.mkdir(parents=True, exist_ok=True)
# Load user data
all_users = self.data_loader.load_jsonl(self.config.data_path)
users_to_process = all_users[:self.config.user_num]
print("\n" + "=" * 80)
print("HALUMEM EVALUATION - REME - QUESTION ANSWERING")
print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}")
print("=" * 80 + "\n")
# Process users with concurrency control
semaphore = asyncio.Semaphore(self.config.max_concurrency)
async def process_with_cache_check(idx: int, user_data: dict):
async with semaphore:
user_name = self.data_loader.extract_user_name(user_data["persona_info"])
# Check cache
if self.file_manager.user_has_cache(user_name):
print(f"⚡ [{idx}/{len(users_to_process)}] Skipping {user_name} (cached)")
return {"user_name": user_name, "status": "cached"}
print(f"🔄 [{idx}/{len(users_to_process)}] Processing {user_name}...")
result = await self.process_user(user_data)
print(f"✅ [{idx}/{len(users_to_process)}] Completed {user_name}")
return result
tasks = [
process_with_cache_check(idx, user)
for idx, user in enumerate(users_to_process, 1)
]
await asyncio.gather(*tasks)
# Combine results
output_file = os.path.join(self.config.output_dir, "eval_results.jsonl")
self.file_manager.combine_results(output_file)
elapsed = time.time() - start_time
print(f"\n✅ Processing completed in {elapsed:.2f}s")
print(f"📁 Results: {output_file}\n")
# Aggregate metrics
await self.aggregate_and_report(output_file)
async def aggregate_and_report(self, results_file: str):
"""Aggregate results and generate final report."""
print("=" * 80)
print("AGGREGATING METRICS")
print("=" * 80 + "\n")
# Collect all QA records
qa_records = []
with open(results_file, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
user_data = json.loads(line)
for session in user_data["sessions"]:
if session.get("is_generated_qa_session"):
continue
eval_results = session.get("evaluation_results", {})
qa_records.extend(
eval_results.get("question_answering_records", [])
)
# Compute metrics
qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records)
time_metrics = MetricsAggregator.compute_time_metrics(results_file)
final_results = {
"overall_score": {
"question_answering": qa_metrics,
"time_consuming": time_metrics
},
"question_answering_records": qa_records
}
# Save final report
report_file = os.path.join(self.config.output_dir, "eval_statistics.json")
with open(report_file, "w", encoding="utf-8") as f:
json.dump(final_results, f, ensure_ascii=False, indent=4)
print(f"📊 Statistics saved to: {report_file}\n")
# Print summary
self._print_summary(qa_metrics, time_metrics)
def _print_summary(self, qa_metrics: dict, time_metrics: dict):
"""Print evaluation summary."""
print("=" * 80)
print("EVALUATION SUMMARY - REME")
print("=" * 80 + "\n")
print("📊 Question Answering:")
print(f" Correct (all): {qa_metrics['correct_qa_ratio(all)']:.4f}")
print(f" Hallucination (all): {qa_metrics['hallucination_qa_ratio(all)']:.4f}")
print(f" Omission (all): {qa_metrics['omission_qa_ratio(all)']:.4f}")
print(f" Correct (valid): {qa_metrics['correct_qa_ratio(valid)']:.4f}")
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
print(f"\n⏱️ Time Metrics:")
print(f" Memory Addition: {time_metrics['add_dialogue_duration_time']:.2f} min")
print(f" Memory Search: {time_metrics['search_memory_duration_time']:.2f} min")
print(f" Total: {time_metrics['total_duration_time']:.2f} min")
print("\n" + "=" * 80)
# ==================== Entry Point ====================
def main(
data_path: str,
top_k: int = 20,
user_num: int = 1,
max_concurrency: int = 2
):
"""Main entry point for ReMe evaluation."""
config = EvalConfig(
data_path=data_path,
top_k=top_k,
user_num=user_num,
max_concurrency=max_concurrency
)
evaluator = HaluMemEvaluatorV4(config)
asyncio.run(evaluator.run_evaluation())
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description="Evaluate ReMe on HaluMem benchmark (Question Answering)"
)
parser.add_argument(
"--data_path",
type=str,
required=True,
help="Path to HaluMem JSONL file"
)
parser.add_argument(
"--top_k",
type=int,
default=20,
help="Number of memories to retrieve (default: 20)"
)
parser.add_argument(
"--user_num",
type=int,
default=1,
help="Number of users to evaluate (default: 1)"
)
parser.add_argument(
"--max_concurrency",
type=int,
default=2,
help="Maximum concurrent user processing (default: 2)"
)
args = parser.parse_args()
main(
data_path=args.data_path,
top_k=args.top_k,
user_num=args.user_num,
max_concurrency=args.max_concurrency
)

View file

@ -424,10 +424,7 @@ EVALUATION_PROMPT_FOR_QUESTION: |
EVALUATION_PROMPT_FOR_QUESTION2: |
You are an **evaluation expert for AI memory system question answering**.
**Dialogue:**
{dialogue}
Based **only** on the provided **“Question”**, **“Reference Answer”**, and **“Key Memory Points”** (the essential facts needed to derive the reference answer), strictly evaluate the **accuracy** of the **“Memory System Response.”** Classify it as one of **“Correct”**, **“Hallucination”**, or **“Omission.”** Do **not** use any external knowledge or subjective inference. Finally, output your judgment **strictly** in the specified JSON format.
Based **only** on the provided **"Question"**, **"Reference Answer"**, and **"Key Memory Points"** (the essential facts needed to derive the reference answer), strictly evaluate the **accuracy** of the **"Memory System Response."** Classify it as one of **"Correct"**, **"Hallucination"**, or **"Omission."** Do **not** use any external knowledge or subjective inference. Finally, output your judgment **strictly** in the specified JSON format.
# Evaluation Criteria
@ -435,36 +432,45 @@ EVALUATION_PROMPT_FOR_QUESTION2: |
### 1. Correct
* The “Memory System Response” accurately answers the “Question,” and its content is **semantically equivalent** to the “Reference Answer.”
* It contains **no contradictions** with the “Key Memory Points” or “Reference Answer.”
* It introduces **no unsupported details** beyond the “Key Memory Points” that could alter the conclusion.
* The "Memory System Response" accurately answers the "Question," and its content is **semantically equivalent** to the "Reference Answer."
* It contains **no contradictions** with the "Key Memory Points" or "Reference Answer."
* **Extra details not present in the Key Memory Points are allowed and should not be penalized**, as long as they:
- Do not contradict the Key Memory Points or Reference Answer
- Do not change or mislead the core conclusion
- Are reasonable additional context that the memory system may have retained from the conversation
* The memory system may have stored additional information beyond the Key Memory Points. Such extra information should be treated as **supplementary context** rather than hallucination, provided it does not conflict with the core answer.
* Synonyms, paraphrasing, and reasonable summarization are acceptable.
### 2. Hallucination
* The “Memory System Response” includes information or facts that **contradict or are inconsistent** with the “Reference Answer” or the “Key Memory Points.”
* When the “Reference Answer” is labeled as *unknown/uncertain*, yet the response provides a specific verifiable fact or conclusion.
* Extra irrelevant information that does **not change** the conclusion is **not** considered hallucination by itself; however, if it **changes or misleads** the conclusion, or **contradicts** the “Key Memory Points,” it should be judged as a **Hallucination**.
* The "Memory System Response" includes information or facts that **contradict or are inconsistent** with the "Reference Answer" or the "Key Memory Points."
* The response provides information that **directly contradicts** known facts from the Key Memory Points.
* When the "Reference Answer" is labeled as *unknown/uncertain*, yet the response provides a specific verifiable fact or conclusion.
* **Important:** Extra information that is NOT in Key Memory Points is **NOT automatically a hallucination**. Only classify as hallucination if the extra information:
- Directly contradicts the Key Memory Points or Reference Answer
- Changes or misleads the core conclusion in a way that makes the answer incorrect
- Provides a definitive answer when the Reference Answer indicates uncertainty
### 3. Omission
* The response is **incomplete** compared to the “Reference Answer.”
* It explicitly states “dont know,” “cant remember,” or “no related memory,” even though relevant information exists in the “Key Memory Points.”
* The response is **incomplete** compared to the "Reference Answer."
* It explicitly states "don't know," "can't remember," or "no related memory," even though relevant information exists in the "Key Memory Points."
* For multi-element questions, **all elements must be correct and present**; omission of **any** element is considered an **Omission**.
## Priority Rules (Conflict Handling)
* If the response contains **both missing necessary information** and **fabricated/contradictory information**, classify it as **Hallucination**.
* If there is **no fabrication/contradiction** but some necessary information is missing, classify it as **Omission**.
* Only when the meaning is **fully equivalent** to the reference answer should it be classified as **Correct**.
* If the core answer is correct and complete, classify as **Correct** even if there are extra details not in Key Memory Points (as long as they don't contradict or mislead).
## Detailed Guidelines and Tolerance
* Equivalent expressions of numbers, times, and units are acceptable, but the **numerical values themselves must not differ**.
* For multi-element questions, **all elements must be complete and accurate**; missing any element counts as **Omission**.
* If the reference answer is *“unknown / cannot be determined”* and the system provides a definite fact, that is a **Hallucination**.
If the system also answers *“unknown”* (without guessing), it may be **Correct**.
* The evaluation must rely **only** on the *Reference Answer*, *Key Memory Points*, and *System Response* — no external context, world knowledge, or speculative reasoning is allowed.
* If the reference answer is *"unknown / cannot be determined"* and the system provides a definite fact, that is a **Hallucination**.
If the system also answers *"unknown"* (without guessing), it may be **Correct**.
* **Focus on evaluating whether the core answer to the question is correct**, not whether the response is limited to only the Key Memory Points.
* Extra contextual information (e.g., additional preferences, related details) should be viewed as enrichment, not as errors, unless they contradict or mislead.
# Information for Evaluation
@ -487,7 +493,7 @@ EVALUATION_PROMPT_FOR_QUESTION2: |
```json
{{
"reasoning": "Provide a concise and traceable evaluation rationale: first compare the systems response with the Key Memory Points (which were correctly used, which were missing, and whether there was any fabrication/contradiction), then assess its consistency with the Reference Answer, and finally state the classification basis.",
"reasoning": "Provide a concise and traceable evaluation rationale: first verify that the system's response correctly includes all required elements from the Reference Answer, then check if any information contradicts the Key Memory Points or Reference Answer. Extra details not in Key Memory Points should be noted but not penalized unless they contradict or mislead. Finally state the classification basis.",
"evaluation_result": "Correct | Hallucination | Omission"
}}
```

View file

@ -68,7 +68,9 @@ class BaseLLM(ABC):
if tool.name not in tool_dict:
continue
if not tool.check_argument():
# First try sanitizing arguments
if not tool.sanitize_and_check_argument():
logger.error(f"Tool call {tool.name} has invalid JSON arguments after sanitization attempt: {tool.arguments}")
raise ValueError(f"Tool call {tool.name} has invalid JSON arguments: {tool.arguments}")
validated_tools.append(tool.simple_output_dump())

View file

@ -177,6 +177,43 @@ class ToolCall(BaseModel):
return True
except Exception:
return False
def sanitize_and_check_argument(self) -> bool:
"""
Attempt to sanitize and validate arguments JSON.
Common issues from LLM streaming:
- Extra closing brackets: }]}] -> }]
- Missing closing brackets
- Trailing commas
"""
if not self.arguments or not self.arguments.strip():
return False
try:
# First try parsing as-is
_ = json.loads(self.arguments)
return True
except json.JSONDecodeError:
pass
# Try to fix common issues
sanitized = self.arguments.strip()
# Remove trailing extra brackets/braces
# Pattern: if it ends with multiple closing chars, try removing extras
while len(sanitized) > 1:
try:
json.loads(sanitized)
self.arguments = sanitized # Update with sanitized version
return True
except json.JSONDecodeError:
# Try removing last character
if sanitized[-1] in ']}':
sanitized = sanitized[:-1].rstrip()
else:
break
return False
def simple_output_dump(self) -> dict:
"""Convert ToolCall to output format dictionary for API responses."""

View file

@ -152,7 +152,7 @@ class BaseMemoryAgent(BaseOp, metaclass=ABCMeta):
tool_call_id=op.tool_call.id,
)
tool_result_messages.append(tool_message)
logger.info(f"[{self.__class__.__name__}] step{step + 1}.{j} join tool_result={tool_result[:500]}...\n\n")
logger.info(f"[{self.__class__.__name__}] step{step + 1}.{j} join tool_result={tool_result[:2000]}...\n\n")
return tool_result_messages
async def react(self, messages: list[Message]):
@ -209,3 +209,8 @@ class BaseMemoryAgent(BaseOp, metaclass=ABCMeta):
def author(self) -> str:
"""Returns the LLM model name as the author identifier."""
return self.llm.model_name
@property
def history_node(self):
"""Returns the history node."""
return self.context.get("history_node", None)

View file

@ -17,6 +17,10 @@ system_prompt: |
### Step 1: Extract Conversation Memories
Use `AddMemory` to extract key personal facts from the conversation.
- Extract: preferences, habits, status, personal details, decisions, conclusions
- **Format**: Use third-person perspective to record what **{memory_target}** said, did, or expressed at specific times
- **Consolidation**: Merge related information under the same topic into ONE memory entry
- Group similar facts (e.g., multiple food preferences → one food preference entry)
- Avoid creating separate entries for closely related information
- Keep entries concise and distinct (no duplicates, no omissions)
- Record `conversation_time` for each memory (format: 2020-01-01 00:00:00; use 0000-00-00 00:00:00 if unavailable)

View file

@ -5,13 +5,13 @@ tool: |
NEVER hallucinate or fabricate information not present in retrieved memories.
system_prompt: |
You are a memory retrieval agent. Search for relevant memories to answer the user's question following this strategy:
You are a memory agent. Search for relevant memories to answer the user's question following this strategy:
## Available Meta Memories
Format: "- <memory_type>(<memory_target>): <description>"
{meta_memory_info}
## User Context
## User's Question
{context}
## Three-Step Retrieval Strategy
@ -19,7 +19,6 @@ system_prompt: |
**STEP 1: Read User Profile (REQUIRED FIRST)**
- Use `read_user_profile` with memory_type and memory_target from available meta memories
- Check if the user profile directly answers the question
- If sufficient information found, provide the answer and STOP
**STEP 2: Vector Search (If Step 1 insufficient)**
- Use `retrieve_memory` with memory_type, memory_target, and query
@ -44,10 +43,8 @@ system_prompt: |
- Try multiple history_id entries if needed
## Response Rules
- Answer ONLY based on retrieved information - NEVER guess or fabricate
- If nothing found after all three steps: State clearly "I don't know. I cannot find relevant information to answer this question."
- If nothing found after all three steps: State clearly "I don't know. "
- Be persistent: try multiple angles in each step before moving to the next
- Once you find sufficient information, provide a direct answer
user_message: |
Retrieve relevant memories and answer the question using the three-step strategy.
Answer the question using the three-step strategy.

View file

@ -0,0 +1,11 @@
from .reme_summarizer_v4 import ReMeSummarizerV4
from .reme_retriever_v4 import ReMeRetrieverV4
from .personal_summarizer_v4 import PersonalSummarizerV4
from .personal_retriever_v4 import PersonalRetrieverV4
__all__ = [
"ReMeSummarizerV4",
"ReMeRetrieverV4",
"PersonalSummarizerV4",
"PersonalRetrieverV4",
]

View file

@ -0,0 +1,46 @@
from ..base_memory_agent import BaseMemoryAgent
from ...core.enumeration import Role, MemoryType
from ...core.schema import Message
from ...core.utils import format_messages
from ...mem_tool.v4 import ReadUserProfile
class PersonalRetrieverV4(BaseMemoryAgent):
memory_type: MemoryType = MemoryType.PERSONAL
async def build_messages(self) -> list[Message]:
if self.context.get("query"):
context = self.context.query
elif self.context.get("messages"):
context = format_messages(self.context.messages)
else:
raise ValueError("input must have either `query` or `messages`")
read_profile_tool = ReadUserProfile()
await read_profile_tool.call(memory_type=self.memory_type.value, memory_target=self.memory_target)
messages = [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=read_profile_tool.output,
context=context,
)),
Message(
role=Role.USER,
content=self.get_prompt("user_message"),
),
]
return messages
async def _acting_step(self, assistant_message: Message, step: int, **kwargs) -> list[Message]:
return await super()._acting_step(
assistant_message,
step,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
**kwargs,
)

View file

@ -0,0 +1,37 @@
tool: |
Retrieve relevant personal memories to answer user questions through vector search and history reading.
system_prompt: |
You are a memory agent managing **{memory_type}** memories about **{memory_target}**.
## User Profile
{user_profile}
## Question
{context}
## Retrieval Strategy
**Tool 1: Vector Search (`retrieve_memory`)
- Try at least 3-5 different queries before moving to next tool:
* Direct question
* Reformulated phrasings
* Entity-focused queries
* Different keyword combinations
- If no results: retry with different time ranges or remove time constraints: [start, end] in YYYYMMDD format
* Example: [20200101, 20200102] for 20200101 <= time <= 20200102
* Single-sided: [0, 20200102] or [20200101, 99999999]
**Tool 2: Read Context (`read_history`) - ONLY AFTER Tool 1**
- Use this ONLY after completing multiple retrieve_memory attempts
- Use history_id from retrieved memories to read original conversations
- Prioritize most relevant or recent entries
- Read multiple if needed for complete context
**Response**
- **CRITICAL: Answer ONLY based on retrieved memories. Do NOT hallucinate or infer information not present in the search results.**
- Try multiple angles before giving up
- State "nothing found after thorough search" if nothing found after thorough search
user_message: |
Answer the question using the retrieval strategy.

View file

@ -0,0 +1,127 @@
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ...core.enumeration import Role, MemoryType
from ...core.schema import Message, MemoryNode
class PersonalSummarizerV4(BaseMemoryAgent):
memory_type: MemoryType = MemoryType.PERSONAL
async def build_messages_phase1(self) -> list[Message]:
"""Build messages for phase 1: AddSummaryMemory"""
history_node: MemoryNode = self.context.history_node
messages = [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt_phase1",
context=history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
)),
Message(
role=Role.USER,
content=self.get_prompt("user_message_phase1"),
),
]
return messages
async def build_messages_phase2(self, user_profile: str) -> list[Message]:
"""Build messages for phase 2: UpdateUserProfile"""
history_node: MemoryNode = self.context.history_node
messages = [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt_phase2",
context=history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=user_profile,
)),
Message(
role=Role.USER,
content=self.get_prompt("user_message_phase2"),
),
]
return messages
async def _acting_step(self, assistant_message: Message, step: int, **kwargs) -> list[Message]:
return await super()._acting_step(
assistant_message,
step,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
history_node=self.history_node,
author=self.author,
**kwargs,
)
async def execute(self):
"""Execute in two phases: 1) AddSummaryMemory, 2) UpdateUserProfile"""
# Log available tools
for i, tool in enumerate(self.tools):
logger.info(
f"[{self.__class__.__name__}] step0.{i} "
f"tool_call={tool.tool_call.name}",
)
# Phase 1: AddSummaryMemory
logger.info(f"[{self.__class__.__name__}] Starting Phase 1: AddSummaryMemory")
# Filter tools for phase 1 (only AddSummaryMemory)
original_tools = self.tools.copy()
self.tools = [t for t in self.tools if t.tool_call.name == "add_summary_memory"]
messages_phase1 = await self.build_messages_phase1()
for i, message in enumerate(messages_phase1):
logger.info(
f"[{self.__class__.__name__}] phase1.step0.{i} {message.role} "
f"{message.simple_dump(enable_json_dump=True)}",
)
messages_phase1, success_phase1 = await self.react(messages_phase1)
if not success_phase1:
logger.warning(f"[{self.__class__.__name__}] Phase 1 did not complete successfully")
# Phase 2: Read user profile and UpdateUserProfile
logger.info(f"[{self.__class__.__name__}] Starting Phase 2: UpdateUserProfile")
# Restore original tools and get ReadUserProfile tool
self.tools = original_tools
read_profile_tool = next((t for t in self.tools if t.tool_call.name == "read_user_profile"), None)
user_profile = ""
if read_profile_tool:
# Call ReadUserProfile to load current profile
logger.info(f"[{self.__class__.__name__}] Loading user profile with ReadUserProfile")
await read_profile_tool.call(memory_type=self.memory_type.value, memory_target=self.memory_target)
user_profile = str(read_profile_tool.output)
logger.info(f"[{self.__class__.__name__}] User profile loaded: {user_profile}...")
else:
logger.warning(f"[{self.__class__.__name__}] ReadUserProfile tool not found")
# Filter tools for phase 2 (only UpdateUserProfile)
self.tools = [t for t in self.tools if t.tool_call.name == "update_user_profile"]
messages_phase2 = await self.build_messages_phase2(user_profile)
for i, message in enumerate(messages_phase2):
logger.info(
f"[{self.__class__.__name__}] phase2.step0.{i} {message.role} "
f"{message.simple_dump(enable_json_dump=True)}",
)
messages_phase2, success_phase2 = await self.react(messages_phase2)
# Restore original tools
self.tools = original_tools
# Set final output and messages
self.messages = messages_phase1 + messages_phase2
self.success = success_phase1 and success_phase2
if self.success and messages_phase2:
self.output = messages_phase2[-1].content
else:
self.output = "Memory processing completed with issues."

View file

@ -0,0 +1,45 @@
tool: |
Extract and update personal memories about the user from conversation context.
Identify preferences, habits, background, relationships, and key facts.
system_prompt_phase1: |
You are a memory agent managing **{memory_type}** memories about **{memory_target}**.
## Latest Conversation:
{context}
Message format: `round<index> [<timestamp>] <role/name>: <content>` (timestamp: YYYY-MM-DD HH:MM:SS).
**CRITICAL**: Extract ONLY explicitly stated information. DO NOT infer, assume, or fabricate.
## Task: Extract Memories with `AddSummaryMemory`
Summarize all important information about **{memory_target}**
- Set `conversation_time` (format: 2020-01-01 00:00:00; use 0000-00-00 00:00:00 if unavailable)
user_message_phase1: |
Extract personal memories from the conversation using `AddSummaryMemory`.
# capturing complete contexts with preconditions, causes, and consequences
system_prompt_phase2: |
You are a memory agent managing **{memory_type}** memories about **{memory_target}**.
## Latest Conversation:
{context}
Message format: `round<index> [<timestamp>] <role/name>: <content>` (timestamp: YYYY-MM-DD HH:MM:SS).
**CRITICAL**: Extract ONLY explicitly stated information. DO NOT infer, assume, or fabricate.
## Current User Profile:
{user_profile}
## Task: Update Profile with `UpdateUserProfile`
Synchronize profile/memories with new information from the conversation, including **{memory_target}**' current status:
- `profile_ids_to_delete`: Remove outdated, conflicting, or redundant entries.
- `profiles_to_add`: Add new profiles/memories with `conversation_time`, e.g. `YYYY-MM-DD HH:MM:SS`, {memory_target} did something.
- Maintain profiles that are concise, mutually exclusive, and collectively comprehensive with no information loss.
user_message_phase2: |
Update user profile using `UpdateUserProfile` based on the conversation and current profile.

View file

@ -0,0 +1,53 @@
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ...core.enumeration import Role
from ...core.schema import Message
from ...core.utils import format_messages
class ReMeRetrieverV4(BaseMemoryAgent):
def __init__(self, meta_memories: list[dict] | None = None, **kwargs):
super().__init__(**kwargs)
self.meta_memories: list[dict] = meta_memories or []
async def _read_meta_memories(self) -> str:
from ...mem_tool import ReadMetaMemory
meta_memory_info = ReadMetaMemory().format_memory_metadata(self.meta_memories)
logger.info(f"meta_memory_info={meta_memory_info}")
return meta_memory_info
async def build_messages(self) -> list[Message]:
if self.context.get("query"):
user_query = self.context.query
elif self.context.get("messages"):
user_query = format_messages(self.context.messages)
else:
raise ValueError("Input must have either `query` or `messages`")
messages = [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
meta_memory_info=await self._read_meta_memories(),
user_query=user_query,
),
),
Message(
role=Role.USER,
content=self.get_prompt("user_message"),
),
]
return messages
async def _acting_step(self, assistant_message: Message, step: int, **kwargs) -> list[Message]:
return await super()._acting_step(
assistant_message,
step,
query=self.context.get("query", ""),
messages=self.context.get("messages", []),
**kwargs,
)

View file

@ -0,0 +1,25 @@
tool: |
Retrieve information from specialized memory agents to answer user queries.
system_prompt: |
You are a Memory Retrieval Orchestrator responsible for querying specialized agents to answer user questions.
# User Query
{user_query}
## Available Memory Agents
Each line indicates a specialized Memory Agent that stores and retrieves memories within a specific dimension (memory_type + memory_target).
Format: "- <memory_type>(<memory_target>): <description>"
{meta_memory_info}
## Your Task
1. Use the `hands_off` tool to retrieve information from relevant agents
- Specify `memory_type` and `memory_target` for each query
- The `memory_type` and `memory_target` must **exactly match** existing entries in the "Available Memory Agents" listed above
- Do NOT query agents that don't exist above
- You can query multiple agents if needed
2. Answer the user query STRICTLY based on the `hands_off` results
3. If the retrieved information is insufficient to answer the query, respond: "nothing found after thorough search."
user_message: |
Please retrieve relevant information from the existing agents and provide an answer based on the results.

View file

@ -0,0 +1,63 @@
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ...core.enumeration import Role, MemoryType
from ...core.schema import Message, MemoryNode
from ...core.utils import format_messages
class ReMeSummarizerV4(BaseMemoryAgent):
def __init__(self, meta_memories: list[dict] | None = None, **kwargs):
super().__init__(**kwargs)
self.meta_memories: list[dict] = meta_memories or []
async def _read_meta_memories(self) -> str:
from ...mem_tool import ReadMetaMemory
meta_memory_info = ReadMetaMemory().format_memory_metadata(self.meta_memories)
logger.info(f"meta_memory_info={meta_memory_info}")
return meta_memory_info
async def build_messages(self) -> list[Message]:
self.context.messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages]
history_content = self.description + "\n" + format_messages(self.context.messages)
self.context.history_node = history_node = MemoryNode(
memory_type=MemoryType.HISTORY,
memory_target="",
when_to_use=history_content[:100],
content=history_content,
ref_memory_id="",
author=self.author,
metadata={},
)
logger.info(f"Adding summary node: {history_node.model_dump_json(indent=2, exclude_none=True)}")
await self.vector_store.delete(history_node.memory_id)
await self.vector_store.insert([history_node.to_vector_node()])
messages = [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
meta_memory_info=await self._read_meta_memories(),
context=history_node.content,
),
),
Message(
role=Role.USER,
content=self.get_prompt("user_message"),
),
]
return messages
async def _acting_step(self, assistant_message: Message, step: int, **kwargs) -> list[Message]:
return await super()._acting_step(
assistant_message,
step,
messages=self.context.messages,
history_node=self.context.history_node,
author=self.author,
**kwargs,
)

View file

@ -0,0 +1,26 @@
tool: |
Orchestrate memory updates across specialized memory agents.
system_prompt: |
You are a Memory Orchestrator responsible for routing memory tasks to specialized agents based on the context.
# Context
{context}
## Available Memory Agents
Each line indicates a specialized Memory Agent dedicated to deep summarization and updating of memories within a specific dimension (memory_type + memory_target).
Format: "- <memory_type>(<memory_target>): <description>"
{meta_memory_info}
## Your Task
Use the `hands_off` tool to distribute memory tasks to specialized agents:
1. Analyze the context and identify which memory dimensions require updates
2. Specify `memory_type` and `memory_target` for each task
- The `memory_type` and `memory_target` must **exactly match** existing entries in the "Available Memory Agents" listed above
- Do NOT create new agents or use memory_type/memory_target combinations that don't exist above
3. Multiple tasks can be specified to enable parallel processing by specialized agents
Note: If the context contains no memorable information (e.g., simple greetings), return `<NO_MEMORY_NEEDED>`.
user_message: |
Please analyze the context and route memory tasks to the appropriate existing agents.

View file

@ -80,11 +80,21 @@ class BaseMemoryTool(BaseOp, metaclass=ABCMeta):
"""Get the reference memory ID from context."""
return self.context.get("ref_memory_id", "")
@property
def description(self) -> str:
"""Get the description from context."""
return self.context.get("description", "")
@property
def messages_formated(self) -> str:
"""Get the formated messages from context."""
return self.context.get("messages_formated", "")
@property
def history_node(self) -> MemoryNode:
"""Get the history node from context."""
return self.context.get("history_node")
@property
def retrieved_nodes(self) -> list[MemoryNode]:
"""Get the retrieved nodes from context."""
@ -115,8 +125,4 @@ class BaseMemoryTool(BaseOp, metaclass=ABCMeta):
author=author or self.author,
metadata=metadata or {},
)
# logger.opt(depth=1).info(
# f"[{self.__class__.__name__}] build node={node.model_dump_json(indent=2, exclude_none=True)}",
# )
return node

View file

@ -28,7 +28,7 @@ class AddMemory(BaseMemoryTool):
"description": "memory content",
},
"conversation_time": {
"type": "object",
"type": "string",
"description": "conversation time, e.g. '2020-01-01 00:00:00'",
}
},

View file

@ -49,7 +49,7 @@ class ReadUserProfile(BaseMemoryTool):
# Convert to MemoryNode objects and sort by conversation_time (oldest first)
memory_nodes = [MemoryNode(**node_data) for node_data in cached_data]
memory_nodes.sort(
key=lambda node: node.metadata.get("conversation_time", "")
key=lambda n: n.metadata.get("conversation_time", "")
)
memory_formated = []

View file

@ -1,42 +1,43 @@
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.context import C
from ...core.schema.memory_node import MemoryNode
@C.register_op()
class UpdateUserProfile(BaseMemoryTool):
def __init__(self, **kwargs):
kwargs["enable_multiple"] = True
super().__init__(**kwargs)
def _build_tool_description(self) -> str:
return "Update user profile."
def _build_multiple_parameters(self) -> dict:
return {
"type": "object",
"properties": {
"profile_ids_to_delete": {
"type": "array",
"description": self.get_prompt("profile_ids_to_delete"),
"description": "profile_ids_to_delete",
"items": {"type": "string"},
},
"profiles_to_add": {
"type": "array",
"description": self.get_prompt("profiles_to_add"),
"description": "profiles_to_add",
"items": {
"type": "object",
"properties": {
"profile_content": {
"type": "string",
"description": self.get_prompt("profile_content"),
"description": "profile_content",
},
"timestamp": {
"conversation_time": {
"type": "string",
"description": self.get_prompt("timestamp"),
"description": "conversation_time, e.g. '2020-01-01 00:00:00'",
},
},
"required": ["profile_content", "timestamp"],
"required": ["profile_content", "conversation_time"],
},
},
},
@ -44,22 +45,16 @@ class UpdateUserProfile(BaseMemoryTool):
}
async def execute(self):
memory_type = "personal"
memory_target = self.memory_target
assert memory_target, "memory_target is not configured."
cache_key = f"{memory_type}_{memory_target}"
profile_ids_to_delete = self.context.get("profile_ids_to_delete", [])
profile_ids_to_delete = [m for m in profile_ids_to_delete if m]
profile_ids_to_delete = list(dict.fromkeys(profile_ids_to_delete))
profiles_to_add = self.context.get("profiles_to_add", [])
if not profile_ids_to_delete and not profiles_to_add:
self.output = "No memories to remove or add. Operation has been done."
return
cache_key = f"{self.memory_type}_{self.memory_target}"
cached_data = self.meta_memory.load(cache_key, auto_clean=False)
existing_memory_nodes = []
if cached_data:
@ -70,9 +65,7 @@ class UpdateUserProfile(BaseMemoryTool):
if profile_ids_to_delete:
profile_ids_set = set(profile_ids_to_delete)
existing_memory_nodes = [
node for node in existing_memory_nodes if node.memory_id not in profile_ids_set
]
existing_memory_nodes = [node for node in existing_memory_nodes if node.memory_id not in profile_ids_set]
removed_count = len(profile_ids_to_delete)
logger.info(f"Removed {removed_count} memories from user profile.")
@ -80,24 +73,13 @@ class UpdateUserProfile(BaseMemoryTool):
if profiles_to_add:
for mem in profiles_to_add:
profile_content = mem.get("profile_content", "")
timestamp = mem.get("timestamp", "")
if not profile_content:
logger.warning("Skipping memory with empty content")
continue
memory_node = self._build_memory_node(
conversation_time = mem.get("conversation_time", "")
new_memory_nodes.append(self._build_memory_node(
memory_content=profile_content,
when_to_use="",
metadata={"timestamp": timestamp}
)
memory_node.memory_type = MemoryNode.MemoryType.PERSONAL
memory_node.memory_target = memory_target
new_memory_nodes.append(memory_node)
added_count = len(new_memory_nodes)
logger.info(f"Added {added_count} new memories to user profile.")
metadata={"conversation_time": conversation_time}
))
logger.info(f"Added {len(new_memory_nodes)} new memories to user profile.")
updated_memory_nodes = existing_memory_nodes + new_memory_nodes
@ -114,5 +96,4 @@ class UpdateUserProfile(BaseMemoryTool):
self.output = f"Successfully {' and '.join(operations)} in user profile."
else:
self.output = "Operation has been done."
logger.info(self.output)

View file

@ -0,0 +1,15 @@
from .add_summary_memory import AddSummaryMemory
from .hands_off import HandsOff
from .read_history import ReadHistory
from .read_user_profile import ReadUserProfile
from .retrieve_memory import RetrieveMemory
from .update_user_profile import UpdateUserProfile
__all__ = [
"AddSummaryMemory",
"HandsOff",
"ReadHistory",
"ReadUserProfile",
"RetrieveMemory",
"UpdateUserProfile",
]

View file

@ -0,0 +1,63 @@
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.schema import MemoryNode
class AddSummaryMemory(BaseMemoryTool):
def __init__(self, **kwargs):
kwargs["enable_multiple"] = False
super().__init__(**kwargs)
def _build_tool_description(self) -> str:
return "Add a summary memory to the vector store for future retrieval."
@staticmethod
def _build_item_schema() -> tuple[dict, list[str]]:
properties = {
"conversation_time": {"type": "string", "description": "conversation time, e.g. '2020-01-01 00:00:00'"},
"summary_memory": {"type": "string", "description": "summary_memory"},
}
return properties, ["conversation_time", "summary_memory"]
def _build_parameters(self) -> dict:
properties, required = self._build_item_schema()
return {
"type": "object",
"properties": properties,
"required": required,
}
async def execute(self):
summary_memory = self.context.get("summary_memory", "")
conversation_time = self.context.get("conversation_time", "")
if not summary_memory:
self.output = "No summary_memory provided for addition."
return
metadata: dict = {"conversation_time": conversation_time}
try:
metadata["time_int"] = int(conversation_time.split(" ")[0].replace("-", ""))
except Exception:
pass
memory_node = MemoryNode(
memory_type=self.memory_type,
memory_target=self.memory_target,
when_to_use="",
content=summary_memory,
ref_memory_id=self.history_node.memory_id,
author=self.author,
metadata=metadata,
)
vector_node = memory_node.to_vector_node()
vector_id = vector_node.vector_id
await self.vector_store.delete(vector_ids=[vector_id])
await self.vector_store.insert(nodes=[vector_node])
self.memory_nodes.append(memory_node)
self.output = f"Successfully added summary memory to vector_store."
logger.info(self.output)

View file

@ -0,0 +1,112 @@
from typing import TYPE_CHECKING
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.enumeration import MemoryType
from ...core.schema import Message
if TYPE_CHECKING:
from ...mem_agent import BaseMemoryAgent
class HandsOff(BaseMemoryTool):
def __init__(self, memory_agents: list["BaseMemoryAgent"], **kwargs):
kwargs["enable_multiple"] = True
kwargs["sub_ops"] = memory_agents or []
super().__init__(**kwargs)
from ...mem_agent import BaseMemoryAgent
self.sub_ops: list[BaseMemoryAgent] = [a for a in self.sub_ops if isinstance(a, BaseMemoryAgent)]
self.messages: list[Message] = []
@property
def memory_agent_dict(self) -> dict[MemoryType, "BaseMemoryAgent"]:
return {a.memory_type: a for a in self.sub_ops}
def _build_tool_description(self) -> str:
return "Distribute memory tasks to appropriate memory agents."
def _build_multiple_parameters(self) -> dict:
return {
"type": "object",
"properties": {
"memory_tasks": {
"type": "array",
"description": "List of memory tasks to distribute to specific agents",
"items": {
"type": "object",
"properties": {
"memory_type": {
"type": "string",
"description": "Type of memory to handle",
"enum": [k.value for k in self.memory_agent_dict],
},
"memory_target": {
"type": "string",
"description": "Target or context for the memory operation",
},
},
"required": ["memory_type", "memory_target"],
},
},
},
"required": ["memory_tasks"],
}
async def execute(self):
tasks = []
seen = set()
for task in self.context.get("memory_tasks", []):
memory_type = MemoryType(task.get("memory_type", ""))
memory_target = task.get("memory_target", "")
# Deduplicate tasks with same memory_type and memory_target
task_key = (memory_type, memory_target)
if task_key in seen:
logger.info(f"Skipping duplicate task: memory_type={memory_type.value}, memory_target={memory_target}")
continue
seen.add(task_key)
tasks.append({
"memory_type": memory_type,
"memory_target": memory_target,
})
if not tasks:
self.output = "No valid memory tasks to execute."
return
agent_list = []
for i, task in enumerate(tasks):
memory_type: MemoryType = task["memory_type"]
memory_target: str = task["memory_target"]
agent = self.memory_agent_dict[memory_type].copy()
agent_list.append([agent, memory_type, memory_target])
logger.info(f"Task {i}: Submitting {memory_type.value} agent for target={memory_target}")
self.submit_async_task(
agent.call,
memory_type=memory_type,
memory_target=memory_target,
query=self.context.get("query", ""),
messages=self.context.get("messages", []),
description=self.context.get("description"),
history_node=self.context.get("history_node"),
)
await self.join_async_tasks()
results = []
for i, (agent, memory_type, memory_target) in enumerate(agent_list):
if agent.memory_nodes:
self.memory_nodes.extend(agent.memory_nodes)
if agent.messages:
self.messages.extend(agent.messages)
results.append(f"{memory_type.value} {memory_target} agent result: {agent.output}")
self.output = "\n".join(results)
logger.info(f"Completed {len(results)} hands-off task(s):\n{self.output}")

View file

@ -0,0 +1,38 @@
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.schema import MemoryNode
class ReadHistory(BaseMemoryTool):
def __init__(self, **kwargs):
kwargs["enable_multiple"] = False
super().__init__(**kwargs)
def _build_tool_description(self) -> str:
return "Read original history dialogue."
def _build_parameters(self) -> dict:
return {
"type": "object",
"properties": {
"history_id": {
"type": "string",
"description": "history_id",
},
},
"required": ["history_id"],
}
async def execute(self):
history_id = self.context.get("history_id", "")
nodes = await self.vector_store.get(vector_ids=[history_id])
if not nodes:
self.output = f"No history: {history_id}"
logger.warning(self.output)
return
memory = MemoryNode.from_vector_node(nodes[0])
self.output = memory.content
logger.info(f"Successfully read history memory: {history_id}")

View file

@ -0,0 +1,62 @@
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.schema.memory_node import MemoryNode
class ReadUserProfile(BaseMemoryTool):
def __init__(self, add_memory_type_target: bool = False, **kwargs):
kwargs["enable_multiple"] = False
super().__init__(**kwargs)
self.add_memory_type_target = add_memory_type_target
def _build_tool_description(self) -> str:
return "Read user profile."
def _build_parameters(self) -> dict:
if self.add_memory_type_target:
return {
"type": "object",
"properties": {
"memory_type": {
"type": "string",
"description": "memory_type",
},
"memory_target": {
"type": "string",
"description": "memory_target",
},
},
"required": ["memory_type", "memory_target"],
}
else:
return {
"type": "object",
"properties": {},
"required": [],
}
async def execute(self):
cache_key = f"{self.memory_type}_{self.memory_target}".replace(" ", "_").lower()
cached_data = self.meta_memory.load(cache_key, auto_clean=False)
if not cached_data:
self.output = ""
logger.info(f"empty cached_data={cache_key}")
return
memory_nodes = [MemoryNode(**node_data) for node_data in cached_data]
memory_nodes.sort(key=lambda n: n.metadata.get("conversation_time", ""))
memory_formated = []
for node in memory_nodes:
node_formated = f"profile_id={node.memory_id} profile_content={node.content}"
if "conversation_time" in node.metadata and node.metadata["conversation_time"]:
node_formated += f" conversation_time={node.metadata['conversation_time']}"
if node.ref_memory_id:
node_formated += f" history_id={node.ref_memory_id}"
memory_formated.append(node_formated.strip())
self.output = "\n".join(memory_formated)
logger.info(f"Read {len(memory_formated)} nodes from cache key: {cache_key}")

View file

@ -0,0 +1,103 @@
import json
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.schema import MemoryNode
from ...core.utils import deduplicate_memories
class RetrieveMemory(BaseMemoryTool):
def __init__(self, top_k: int = 20, **kwargs):
super().__init__(**kwargs)
self.top_k: int = top_k
def _build_tool_description(self) -> str:
return "Retrieve memories using vector similarity search."
def _build_multiple_parameters(self) -> dict:
return {
"type": "object",
"properties": {
"query_items": {
"type": "array",
"description": "query_items",
"items": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "query",
},
"time_range": {
"type": "string",
"description": "time_range(optional), e.g. [20200101, 20200101]",
},
},
"required": ["query"],
},
},
},
"required": ["query_items"],
}
async def execute(self):
query_items: list[dict] = self.context.get("query_items", [])
memory_nodes: list[MemoryNode] = []
for query_item in query_items:
query = query_item.get("query")
time_range = query_item.get("time_range", "")
filter_dict: dict = {
"memory_type": self.memory_type.value,
"memory_target": self.memory_target,
}
if time_range:
# Handle different time_range formats
if isinstance(time_range, str):
try:
time_range = json.loads(time_range)
except json.JSONDecodeError:
# If it's a plain string like "20250907", treat it as a single date
time_range = time_range
# Convert to list format [start, end]
if isinstance(time_range, (list, tuple)):
if len(time_range) == 1:
# Single element list, use it for both start and end
filter_dict["time_int"] = [int(time_range[0]), int(time_range[0])]
else:
# Two element list/tuple
filter_dict["time_int"] = [int(time_range[0]), int(time_range[1])]
else:
# Single value (int or string), use it for both start and end
filter_dict["time_int"] = [int(time_range), int(time_range)]
logger.info(f"memory_type={self.memory_type} memory_target={self.memory_target} query={query} "
f"filter_dict={filter_dict}")
nodes = await self.vector_store.search(query=query, limit=self.top_k, filters=filter_dict)
memory_nodes.extend([MemoryNode.from_vector_node(n) for n in nodes])
memory_nodes = deduplicate_memories(memory_nodes)
retrieved_memory_ids = {node.memory_id for node in self.retrieved_nodes if node.memory_id}
new_memory_nodes = [node for node in memory_nodes if node.memory_id not in retrieved_memory_ids]
self.retrieved_nodes.extend(new_memory_nodes)
self.memory_nodes = new_memory_nodes
if not new_memory_nodes:
self.output = "No new memory_nodes found matching the query (duplicates removed)."
else:
output = []
for node in new_memory_nodes:
line = ""
if "conversation_time" in node.metadata and node.metadata["conversation_time"]:
line += f"conversation_time={node.metadata['conversation_time']} "
line += node.content.strip() + " "
if node.ref_memory_id:
line += f"history_id={node.ref_memory_id} "
output.append(line.strip())
self.output = "\n".join(output)
logger.info(f"Retrieved {len(memory_nodes)} memory_nodes, {len(new_memory_nodes)} new after deduplication")

View file

@ -0,0 +1,105 @@
from loguru import logger
from ..base_memory_tool import BaseMemoryTool
from ...core.schema.memory_node import MemoryNode
from ...core.utils import deduplicate_memories
class UpdateUserProfile(BaseMemoryTool):
def __init__(self, **kwargs):
kwargs["enable_multiple"] = True
super().__init__(**kwargs)
def _build_tool_description(self) -> str:
return "Update user profile."
def _build_multiple_parameters(self) -> dict:
return {
"type": "object",
"properties": {
"profile_ids_to_delete": {
"type": "array",
"description": "profile_ids_to_delete",
"items": {
"type": "string"
},
},
"profiles_to_add": {
"type": "array",
"description": "profiles_to_add",
"items": {
"type": "object",
"properties": {
"conversation_time": {
"type": "string",
"description": "conversation_time, e.g. '2020-01-01 00:00:00'",
},
"profile_content": {
"type": "string",
"description": "profile_content",
},
},
"required": ["profile_content", "conversation_time"],
},
},
},
"required": ["profile_ids_to_delete", "profiles_to_add"],
}
async def execute(self):
profile_ids_to_delete = self.context.get("profile_ids_to_delete", [])
profile_ids_to_delete = [m for m in profile_ids_to_delete if m]
profile_ids_to_delete = list(dict.fromkeys(profile_ids_to_delete))
profiles_to_add = self.context.get("profiles_to_add", [])
if not profile_ids_to_delete and not profiles_to_add:
self.output = "No profiles to remove or add. Operation has been done."
return
cache_key = f"{self.memory_type}_{self.memory_target}".replace(" ", "_").lower()
cached_data = self.meta_memory.load(cache_key, auto_clean=False)
if cached_data:
existing_memory_nodes = [MemoryNode(**node_data) for node_data in cached_data]
else:
existing_memory_nodes = []
removed_count = 0
if profile_ids_to_delete:
original_count = len(existing_memory_nodes)
existing_memory_nodes = [n for n in existing_memory_nodes if n.memory_id not in profile_ids_to_delete]
removed_count = original_count - len(existing_memory_nodes)
logger.info(f"Removed {removed_count} profiles.")
added_count = 0
new_memory_nodes = []
if profiles_to_add:
for mem in profiles_to_add:
memory_node = MemoryNode(
memory_type=self.memory_type,
memory_target=self.memory_target,
when_to_use="",
content=mem.get("profile_content", ""),
ref_memory_id=self.ref_memory_id,
author=self.author,
metadata={"conversation_time": mem.get("conversation_time", "")},
)
new_memory_nodes.append(memory_node)
added_count = len(new_memory_nodes)
logger.info(f"Added {added_count} new profiles.")
updated_memory_nodes = deduplicate_memories(existing_memory_nodes + new_memory_nodes)
nodes_data = [node.model_dump(exclude_none=True) for node in updated_memory_nodes]
self.meta_memory.save(cache_key, nodes_data)
operations = []
if removed_count > 0:
operations.append(f"removed {removed_count} old profiles")
if added_count > 0:
operations.append(f"added {added_count} new profiles")
if operations:
self.output = f"Successfully {' and '.join(operations)} in user profile."
else:
self.output = "Operation has been done."
logger.info(self.output)

View file

@ -18,6 +18,12 @@ from .mem_agent.v3 import (
ReMeRetrieverV3,
ReMeSummarizerV3,
)
from .mem_agent.v4 import (
PersonalSummarizerV4,
PersonalRetrieverV4,
ReMeRetrieverV4,
ReMeSummarizerV4,
)
from .mem_tool import (
HandsOffTool,
ReadHistoryMemory,
@ -42,6 +48,14 @@ from .mem_tool.v3 import (
SummaryAndHandsOff as SummaryAndHandsOffV3,
UpdateUserProfile,
)
from .mem_tool.v4 import (
AddSummaryMemory as AddSummaryMemoryV4,
HandsOff as HandsOffV4,
ReadHistory as ReadHistoryV4,
ReadUserProfile as ReadUserProfileV4,
RetrieveMemory as RetrieveMemoryV4,
UpdateUserProfile as UpdateUserProfileV4,
)
@singleton
@ -348,9 +362,9 @@ class ReMe(Application):
personal_summarizer_v3 = PersonalSummarizerV3(
tools=[
AddMemoryV3(),
ReadUserProfile(add_memory_type_target=False),
UpdateUserProfile(),
AddMemoryV3(enable_thinking_params=True),
ReadUserProfile(enable_thinking_params=True, add_memory_type_target=False),
UpdateUserProfile(enable_thinking_params=True),
],
)
@ -394,9 +408,9 @@ class ReMe(Application):
reme_retriever_v3 = ReMeRetrieverV3(
meta_memories=meta_memories,
tools=[
ReadUserProfile(add_memory_type_target=True),
RetrieveMemory(top_k=top_k),
ReadHistoryV3(),
ReadUserProfile(enable_thinking_params=True, add_memory_type_target=True),
RetrieveMemory(enable_thinking_params=True, top_k=top_k),
ReadHistoryV3(enable_thinking_params=True),
],
)
@ -409,3 +423,83 @@ class ReMe(Application):
else:
raise NotImplementedError
async def summary_v4(
self,
messages: list[dict],
description: str = "",
user_id: str = "",
assistant_id: str = "",
enable_thinking_params: bool = False,
**kwargs,
):
"""Summarizes messages using V4 workflow with simplified memory management."""
if user_id:
meta_memories = [
{
"memory_type": "personal",
"memory_target": user_id,
},
]
messages = self._prepare_messages(messages, user_id, assistant_id)
personal_summarizer_v4 = PersonalSummarizerV4(
tools=[
AddSummaryMemoryV4(enable_thinking_params=enable_thinking_params),
ReadUserProfileV4(enable_thinking_params=enable_thinking_params),
UpdateUserProfileV4(enable_thinking_params=enable_thinking_params),
],
)
reme_summarizer_v4 = ReMeSummarizerV4(
meta_memories=meta_memories,
tools=[HandsOffV4(memory_agents=[personal_summarizer_v4])],
)
await reme_summarizer_v4.call(messages=messages, description=description, **kwargs)
return reme_summarizer_v4.memory_nodes, reme_summarizer_v4.messages, reme_summarizer_v4.success
else:
raise NotImplementedError
async def retrieve_v4(
self,
query: str = "",
messages: list[dict] | None = None,
description: str = "",
user_id: str = "",
assistant_id: str = "",
top_k: int = 20,
enable_thinking_params: bool = False,
**kwargs,
):
"""Retrieves relevant memories using V4 workflow with enhanced retrieval."""
if user_id:
messages = self._prepare_messages(messages, user_id, assistant_id)
meta_memories = [
{
"memory_type": "personal",
"memory_target": user_id,
},
]
personal_retriever_v4 = PersonalRetrieverV4(
tools=[
RetrieveMemoryV4(enable_thinking_params=enable_thinking_params, top_k=top_k),
ReadHistoryV4(enable_thinking_params=enable_thinking_params),
],
)
reme_retriever_v4 = ReMeRetrieverV4(
meta_memories=meta_memories,
tools=[HandsOffV4(memory_agents=[personal_retriever_v4])],
)
await reme_retriever_v4.call(query=query, messages=messages, description=description, **kwargs)
return reme_retriever_v4.output, reme_retriever_v4.messages, reme_retriever_v4.success
else:
raise NotImplementedError