mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
feat(mem-agent): introduce version 4 memory agents and tools
This commit is contained in:
parent
e6ad682ede
commit
08b771b6c4
31 changed files with 2143 additions and 74 deletions
327
bench/halumem/compute_qa_stats_v4.py
Normal file
327
bench/halumem/compute_qa_stats_v4.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
671
bench/halumem/eval_reme_simple_v4.py
Normal file
671
bench/halumem/eval_reme_simple_v4.py
Normal 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
|
||||
)
|
||||
|
|
@ -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 “don’t know,” “can’t 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 system’s 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"
|
||||
}}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
11
reme_ai/mem_agent/v4/__init__.py
Normal file
11
reme_ai/mem_agent/v4/__init__.py
Normal 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",
|
||||
]
|
||||
46
reme_ai/mem_agent/v4/personal_retriever_v4.py
Normal file
46
reme_ai/mem_agent/v4/personal_retriever_v4.py
Normal 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,
|
||||
)
|
||||
37
reme_ai/mem_agent/v4/personal_retriever_v4.yaml
Normal file
37
reme_ai/mem_agent/v4/personal_retriever_v4.yaml
Normal 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.
|
||||
127
reme_ai/mem_agent/v4/personal_summarizer_v4.py
Normal file
127
reme_ai/mem_agent/v4/personal_summarizer_v4.py
Normal 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."
|
||||
45
reme_ai/mem_agent/v4/personal_summarizer_v4.yaml
Normal file
45
reme_ai/mem_agent/v4/personal_summarizer_v4.yaml
Normal 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.
|
||||
53
reme_ai/mem_agent/v4/reme_retriever_v4.py
Normal file
53
reme_ai/mem_agent/v4/reme_retriever_v4.py
Normal 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,
|
||||
)
|
||||
25
reme_ai/mem_agent/v4/reme_retriever_v4.yaml
Normal file
25
reme_ai/mem_agent/v4/reme_retriever_v4.yaml
Normal 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.
|
||||
63
reme_ai/mem_agent/v4/reme_summarizer_v4.py
Normal file
63
reme_ai/mem_agent/v4/reme_summarizer_v4.py
Normal 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,
|
||||
)
|
||||
26
reme_ai/mem_agent/v4/reme_summarizer_v4.yaml
Normal file
26
reme_ai/mem_agent/v4/reme_summarizer_v4.yaml
Normal 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.
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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'",
|
||||
}
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
15
reme_ai/mem_tool/v4/__init__.py
Normal file
15
reme_ai/mem_tool/v4/__init__.py
Normal 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",
|
||||
]
|
||||
63
reme_ai/mem_tool/v4/add_summary_memory.py
Normal file
63
reme_ai/mem_tool/v4/add_summary_memory.py
Normal 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)
|
||||
112
reme_ai/mem_tool/v4/hands_off.py
Normal file
112
reme_ai/mem_tool/v4/hands_off.py
Normal 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}")
|
||||
38
reme_ai/mem_tool/v4/read_history.py
Normal file
38
reme_ai/mem_tool/v4/read_history.py
Normal 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}")
|
||||
62
reme_ai/mem_tool/v4/read_user_profile.py
Normal file
62
reme_ai/mem_tool/v4/read_user_profile.py
Normal 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}")
|
||||
103
reme_ai/mem_tool/v4/retrieve_memory.py
Normal file
103
reme_ai/mem_tool/v4/retrieve_memory.py
Normal 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")
|
||||
105
reme_ai/mem_tool/v4/update_user_profile.py
Normal file
105
reme_ai/mem_tool/v4/update_user_profile.py
Normal 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)
|
||||
106
reme_ai/reme.py
106
reme_ai/reme.py
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue