diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1795bc17..3d6fcf96 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,7 +3,7 @@ repos: rev: v6.0.0 hooks: - id: check-ast - exclude: ^(test/|cookbook/) + exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench) - id: check-yaml - id: check-xml - id: check-toml @@ -14,18 +14,18 @@ repos: rev: v4.0.0 hooks: - id: add-trailing-comma - exclude: ^(test/|cookbook/) + exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench) - repo: https://github.com/psf/black rev: 25.9.0 hooks: - id: black - exclude: ^(test/|cookbook/) + exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench) args: [--line-length=120] - repo: https://github.com/PyCQA/flake8 rev: 7.3.0 hooks: - id: flake8 - exclude: ^(test/|cookbook/) + exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench) args: [ "--extend-ignore=E203", "--max-line-length=120" @@ -44,6 +44,10 @@ repos: | \.demo$ | \.md$ | \.html$ + | reme_ai/core_old/ + | reme_ai/mem_agent/ + | reme_ai/mem_tool/ + | bench ) args: [ --disable=W0511, diff --git a/bench/eval_reme_old.py b/bench/eval_reme_old.py index 51265da7..41fde3b1 100644 --- a/bench/eval_reme_old.py +++ b/bench/eval_reme_old.py @@ -10,8 +10,8 @@ from datetime import datetime, timezone from tqdm import tqdm -from reme_ai.core.enumeration import Role -from reme_ai.core.schema import Message, MemoryNode +from reme_ai.core_old.enumeration import Role +from reme_ai.core_old.schema import Message, MemoryNode from reme_ai.reme import ReMe TEMPLATE_REME = """Memories for user {user_id}: diff --git a/bench/halumem/analyze_dataset_stats.py b/bench/halumem/analyze_dataset_stats.py index d5a213b4..061f96ae 100644 --- a/bench/halumem/analyze_dataset_stats.py +++ b/bench/halumem/analyze_dataset_stats.py @@ -38,29 +38,29 @@ class DatasetStats: total_users: int total_sessions: int total_dialogues: int - + avg_sessions_per_user: float avg_dialogues_per_session: float avg_dialogue_length_per_session: float - + # 详细分布 sessions_per_user_list: list[int] dialogues_per_session_list: list[int] dialogue_lengths_per_session_list: list[int] - + # Content 统计 total_contents: int # 所有对话回合的 content 总数 content_sizes: list[int] # 每个 content 的大小（字符数） min_content_size: int max_content_size: int percentiles: dict[str, float] # 分位点统计（全部） - + # 按 role 分类的 Content 统计 total_user_contents: int total_assistant_contents: int user_percentiles: dict[str, float] # user 角色的分位点 assistant_percentiles: dict[str, float] # assistant 角色的分位点 - + # Session 分割统计 total_chunks_after_split: int # 按 5000 字符分割后的总 chunk 数 chunks_per_user_list: list[int] # 每个用户分割后的 chunk 数量 @@ -69,14 +69,14 @@ class DatasetStats: class DatasetAnalyzer: """数据集分析器""" - + def __init__(self, data_path: str): self.data_path = data_path self.user_stats_list: list[UserStats] = [] self.all_content_sizes: list[int] = [] # 收集所有 content 的大小 self.user_content_sizes: list[int] = [] # user 角色的 content 大小 self.assistant_content_sizes: list[int] = [] # assistant 角色的 content 大小 - + @staticmethod def extract_user_name(persona_info: str) -> str: """从 persona_info 中提取用户名""" @@ -84,7 +84,7 @@ class DatasetAnalyzer: if not match: return "Unknown" return match.group(1).strip() - + @staticmethod def calculate_dialogue_length(dialogue: list[dict]) -> int: """计算对话的总长度（字符数）""" @@ -93,7 +93,7 @@ class DatasetAnalyzer: content = turn.get("content", "") total_length += len(content) return total_length - + @staticmethod def split_session_into_chunks(dialogue: list[dict], max_length: int = 5000) -> int: """ @@ -102,23 +102,23 @@ class DatasetAnalyzer: 1. 每次添加 2 个对话回合（user-assistant 对） 2. 如果添加后超过 max_length，就开始新的 chunk 3. 但是每个 chunk 至少包含 2 个对话回合 - + 返回分割后的 chunk 数量 """ if not dialogue: return 0 - + chunks = [] current_chunk = [] current_length = 0 - + # 每次处理 2 个对话回合 i = 0 while i < len(dialogue): # 取 2 个对话回合（如果不足 2 个，取剩余的） pair = dialogue[i:i+2] pair_length = sum(len(turn.get("content", "")) for turn in pair) - + # 如果当前 chunk 为空，直接添加（保证至少 2 个） if not current_chunk: current_chunk.extend(pair) @@ -137,72 +137,72 @@ class DatasetAnalyzer: current_chunk.extend(pair) current_length += pair_length i += len(pair) - + # 添加最后一个 chunk if current_chunk: chunks.append(current_chunk) - + return len(chunks) - + def load_and_analyze(self): """加载并分析数据集""" logger.info(f"Loading data from: {self.data_path}") - + with open(self.data_path, "r", encoding="utf-8") as f: for line_num, line in enumerate(f, 1): if not line.strip(): continue - + try: user_data = json.loads(line) self._analyze_user(user_data) except json.JSONDecodeError as e: logger.error(f"Error parsing line {line_num}: {e}") continue - + logger.info(f"Analyzed {len(self.user_stats_list)} users") - + def _analyze_user(self, user_data: dict): """分析单个用户的数据""" user_name = self.extract_user_name(user_data.get("persona_info", "")) uuid = user_data.get("uuid", "") sessions = user_data.get("sessions", []) - + dialogues_per_session = [] dialogue_lengths_per_session = [] session_time_ranges = [] total_chunks = 0 - + for session in sessions: dialogue = session.get("dialogue", []) num_dialogues = len(dialogue) dialogue_length = self.calculate_dialogue_length(dialogue) - + dialogues_per_session.append(num_dialogues) dialogue_lengths_per_session.append(dialogue_length) - + # 收集 session 的时间范围 start_time = session.get("start_time", None) end_time = session.get("end_time", None) session_time_ranges.append((start_time, end_time)) - + # 计算这个 session 分割后的 chunk 数量 num_chunks = self.split_session_into_chunks(dialogue, max_length=5000) total_chunks += num_chunks - + # 收集每个 content 的大小，并按 role 分类 for turn in dialogue: content = turn.get("content", "") content_size = len(content) role = turn.get("role", "") - + self.all_content_sizes.append(content_size) - + if role == "user": self.user_content_sizes.append(content_size) elif role == "assistant": self.assistant_content_sizes.append(content_size) - + user_stats = UserStats( user_name=user_name, uuid=uuid, @@ -212,25 +212,25 @@ class DatasetAnalyzer: num_chunks_after_split=total_chunks, session_time_ranges=session_time_ranges ) - + self.user_stats_list.append(user_stats) - + def compute_dataset_stats(self) -> DatasetStats: """计算整体数据集统计""" total_users = len(self.user_stats_list) - + sessions_per_user_list = [u.num_sessions for u in self.user_stats_list] total_sessions = sum(sessions_per_user_list) - + dialogues_per_session_list = [] dialogue_lengths_per_session_list = [] - + for user in self.user_stats_list: dialogues_per_session_list.extend(user.dialogues_per_session) dialogue_lengths_per_session_list.extend(user.dialogue_lengths_per_session) - + total_dialogues = sum(dialogues_per_session_list) - + # 计算平均值 avg_sessions_per_user = total_sessions / total_users if total_users > 0 else 0 avg_dialogues_per_session = ( @@ -240,43 +240,43 @@ class DatasetAnalyzer: sum(dialogue_lengths_per_session_list) / len(dialogue_lengths_per_session_list) if dialogue_lengths_per_session_list else 0 ) - + # Content 统计 total_contents = len(self.all_content_sizes) min_content_size = min(self.all_content_sizes) if self.all_content_sizes else 0 max_content_size = max(self.all_content_sizes) if self.all_content_sizes else 0 - + # 计算分位点 (10%, 15%, 20%, ..., 95%) percentile_points = list(range(10, 100, 5)) # 10, 15, 20, ..., 95 - + # 全部 content 的分位点 percentiles = {} if self.all_content_sizes: content_array = np.array(self.all_content_sizes) for p in percentile_points: percentiles[f"p{p}"] = float(np.percentile(content_array, p)) - + # user 角色的分位点 user_percentiles = {} if self.user_content_sizes: user_array = np.array(self.user_content_sizes) for p in percentile_points: user_percentiles[f"p{p}"] = float(np.percentile(user_array, p)) - + # assistant 角色的分位点 assistant_percentiles = {} if self.assistant_content_sizes: assistant_array = np.array(self.assistant_content_sizes) for p in percentile_points: assistant_percentiles[f"p{p}"] = float(np.percentile(assistant_array, p)) - + # Session 分割统计 chunks_per_user_list = [u.num_chunks_after_split for u in self.user_stats_list] total_chunks_after_split = sum(chunks_per_user_list) avg_chunks_per_user = ( total_chunks_after_split / total_users if total_users > 0 else 0 ) - + return DatasetStats( total_users=total_users, total_sessions=total_sessions, @@ -300,16 +300,16 @@ class DatasetAnalyzer: chunks_per_user_list=chunks_per_user_list, avg_chunks_per_user=avg_chunks_per_user ) - + @staticmethod def _print_percentiles(percentiles: dict[str, float]): """打印分位点统计（辅助函数）""" if not percentiles: print(" (无数据)") return - + sorted_percentiles = sorted(percentiles.keys(), key=lambda x: int(x[1:])) - + # 每行显示 5 个分位点，让输出更紧凑 for i in range(0, len(sorted_percentiles), 5): line_items = [] @@ -318,36 +318,36 @@ class DatasetAnalyzer: p_num = percentile_key[1:] # 去掉 'p' 前缀 line_items.append(f"{p_num}%: {percentile_value:.0f}") print(f" {' | '.join(line_items)}") - + def print_summary(self, stats: DatasetStats): """打印统计摘要""" print("\n" + "=" * 80) print("HALUMEM DATASET STATISTICS") print("=" * 80 + "\n") - + print("📊 总体统计:") print(f" 总用户数: {stats.total_users}") print(f" 总 Session 数: {stats.total_sessions}") print(f" 总对话数: {stats.total_dialogues}") - + print(f"\n📈 平均值:") print(f" 每个用户的平均 Session 数: {stats.avg_sessions_per_user:.2f}") print(f" 每个 Session 的平均对话数: {stats.avg_dialogues_per_session:.2f}") print(f" 每个 Session 的平均对话长度（字符）: {stats.avg_dialogue_length_per_session:.2f}") - + print(f"\n📊 分布统计:") if stats.sessions_per_user_list: print(f" 每用户 Session 数 - 最小: {min(stats.sessions_per_user_list)}, " f"最大: {max(stats.sessions_per_user_list)}") - + if stats.dialogues_per_session_list: print(f" 每 Session 对话数 - 最小: {min(stats.dialogues_per_session_list)}, " f"最大: {max(stats.dialogues_per_session_list)}") - + if stats.dialogue_lengths_per_session_list: print(f" 每 Session 对话长度 - 最小: {min(stats.dialogue_lengths_per_session_list)}, " f"最大: {max(stats.dialogue_lengths_per_session_list)}") - + print(f"\n💬 Content 详细统计:") print(f" 总 Content 数量: {stats.total_contents}") print(f" User 消息数: {stats.total_user_contents}") @@ -355,34 +355,34 @@ class DatasetAnalyzer: print(f" Content 大小（字符数）:") print(f" 最小值: {stats.min_content_size}") print(f" 最大值: {stats.max_content_size}") - + if stats.content_sizes: avg_content_size = sum(stats.content_sizes) / len(stats.content_sizes) print(f" 平均值: {avg_content_size:.2f}") - + print(f"\n📈 Content 大小分位点 (全部):") self._print_percentiles(stats.percentiles) - + print(f"\n📈 Content 大小分位点 (User 角色):") self._print_percentiles(stats.user_percentiles) - + print(f"\n📈 Content 大小分位点 (Assistant 角色):") self._print_percentiles(stats.assistant_percentiles) - + print(f"\n✂️ Session 分割统计 (按 5000 字符分割):") print(f" 原始 Session 总数: {stats.total_sessions}") print(f" 分割后 Chunk 总数: {stats.total_chunks_after_split}") print(f" 每个用户平均 Chunk 数: {stats.avg_chunks_per_user:.2f}") print(f" Chunk/Session 比例: {stats.total_chunks_after_split / stats.total_sessions:.2f}x") - + print("\n" + "=" * 80) - + def print_per_user_stats(self): """打印每个用户的详细统计""" print("\n" + "=" * 80) print("PER-USER STATISTICS") print("=" * 80 + "\n") - + for idx, user_stats in enumerate(self.user_stats_list, 1): avg_dialogues = ( sum(user_stats.dialogues_per_session) / len(user_stats.dialogues_per_session) @@ -392,50 +392,50 @@ class DatasetAnalyzer: sum(user_stats.dialogue_lengths_per_session) / len(user_stats.dialogue_lengths_per_session) if user_stats.dialogue_lengths_per_session else 0 ) - + print(f"[{idx}] {user_stats.user_name} (UUID: {user_stats.uuid[:8]}...)") print(f" Session 数: {user_stats.num_sessions}") print(f" 分割后 Chunk 数: {user_stats.num_chunks_after_split}") print(f" 平均每 Session 对话数: {avg_dialogues:.2f}") print(f" 平均每 Session 对话长度: {avg_length:.2f} 字符") print() - + def print_first_user_session_times(self): """打印第一个用户的每个 session 的时间范围""" if not self.user_stats_list: print("\n没有用户数据") return - + first_user = self.user_stats_list[0] - + print("\n" + "=" * 80) print(f"第一个用户的 Session 时间统计") print("=" * 80 + "\n") print(f"用户名: {first_user.user_name}") print(f"UUID: {first_user.uuid}") print(f"总 Session 数: {first_user.num_sessions}\n") - + print("-" * 80) print(f"{'Session #':<12} {'开始时间':<30} {'结束时间':<30}") print("-" * 80) - + for idx, (start_time, end_time) in enumerate(first_user.session_time_ranges, 1): start_str = str(start_time) if start_time is not None else "无" end_str = str(end_time) if end_time is not None else "无" print(f"{idx:<12} {start_str:<30} {end_str:<30}") - + print("=" * 80) - + def print_user_split_summary(self): """打印每个用户的分割统计摘要（表格形式）""" print("\n" + "=" * 80) print("PER-USER SESSION SPLIT SUMMARY (按 5000 字符分割)") print("=" * 80 + "\n") - + # 表头 print(f"{'序号':<6} {'用户名':<25} {'原始Sessions':<15} {'分割后Chunks':<15} {'比例':<10}") print("-" * 80) - + # 每个用户的数据 for idx, user_stats in enumerate(self.user_stats_list, 1): ratio = ( @@ -444,17 +444,17 @@ class DatasetAnalyzer: ) print(f"{idx:<6} {user_stats.user_name[:24]:<25} {user_stats.num_sessions:<15} " f"{user_stats.num_chunks_after_split:<15} {ratio:.2f}x") - + print("-" * 80) - + # 总计 total_sessions = sum(u.num_sessions for u in self.user_stats_list) total_chunks = sum(u.num_chunks_after_split for u in self.user_stats_list) overall_ratio = total_chunks / total_sessions if total_sessions > 0 else 0 - + print(f"{'总计':<6} {'':<25} {total_sessions:<15} {total_chunks:<15} {overall_ratio:.2f}x") print("=" * 80) - + def save_results(self, output_path: str, stats: DatasetStats): """保存统计结果到 JSON 文件""" results = { @@ -512,10 +512,10 @@ class DatasetAnalyzer: for u in self.user_stats_list ] } - + with open(output_path, "w", encoding="utf-8") as f: json.dump(results, f, ensure_ascii=False, indent=2) - + logger.info(f"Results saved to: {output_path}") @@ -525,27 +525,27 @@ def main(data_path: str, output_path: str = None, show_per_user: bool = False): if not Path(data_path).exists(): logger.error(f"File not found: {data_path}") return - + # 创建分析器并执行分析 analyzer = DatasetAnalyzer(data_path) analyzer.load_and_analyze() - + # 计算统计数据 stats = analyzer.compute_dataset_stats() - + # 打印摘要 analyzer.print_summary(stats) - + # 打印第一个用户的 session 时间统计 analyzer.print_first_user_session_times() - + # 打印每个用户的分割统计摘要（始终显示） analyzer.print_user_split_summary() - + # 打印每个用户的详细统计（可选） if show_per_user: analyzer.print_per_user_stats() - + # 保存结果到文件 if output_path: analyzer.save_results(output_path, stats) @@ -557,7 +557,7 @@ def main(data_path: str, output_path: str = None, show_per_user: bool = False): if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="Analyze HaluMem dataset statistics" ) @@ -578,9 +578,9 @@ if __name__ == "__main__": action="store_true", help="Show detailed statistics for each user" ) - + args = parser.parse_args() - + main( data_path=args.data_path, output_path=args.output_path, diff --git a/bench/halumem/analyze_results.py b/bench/halumem/analyze_results.py index 4814001e..630ae5e3 100644 --- a/bench/halumem/analyze_results.py +++ b/bench/halumem/analyze_results.py @@ -13,73 +13,73 @@ from typing import Dict, List, Tuple def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"): """ 分析评估结果目录。 - + Args: tmp_dir: 临时结果目录路径 """ tmp_path = Path(tmp_dir) - + if not tmp_path.exists(): print(f"❌ 目录不存在: {tmp_dir}") return - + # 统计数据 result_counter = Counter() non_correct_results = [] # 存储非Correct结果的详细信息 - + # 遍历所有用户目录 user_dirs = sorted([d for d in tmp_path.iterdir() if d.is_dir()]) - + if not user_dirs: print(f"❌ {tmp_dir} 下没有用户目录") return - + print(f"📁 找到 {len(user_dirs)} 个用户目录\n") print("=" * 80) print("开始分析...") print("=" * 80 + "\n") - + total_sessions = 0 total_questions = 0 - + # 遍历每个用户目录 for user_dir in user_dirs: user_name = user_dir.name - + # 获取该用户的所有session文件 session_files = sorted([ - f for f in user_dir.iterdir() + f for f in user_dir.iterdir() if f.name.startswith("session_") and f.suffix == ".json" ]) - + if not session_files: continue - + # 遍历每个session for session_file in session_files: try: with open(session_file, "r", encoding="utf-8") as f: session_data = json.load(f) - + session_id = session_data.get("session_id", -1) total_sessions += 1 - + # 跳过生成的QA session if session_data.get("is_generated_qa_session", False): continue - + # 获取评估结果 eval_results = session_data.get("evaluation_results", {}) qa_records = eval_results.get("question_answering_records", []) - + # 分析每个问题的结果 for qa_idx, qa_record in enumerate(qa_records): result_type = qa_record.get("result_type", "Unknown") - + # 统计result_type result_counter[result_type] += 1 total_questions += 1 - + # 如果不是Correct，记录详细信息 if result_type != "Correct": non_correct_results.append({ @@ -91,42 +91,42 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"): "answer": qa_record.get("answer", ""), "system_response": qa_record.get("system_response", "") }) - + except Exception as e: print(f"⚠️ 读取文件失败: {session_file}, 错误: {e}") continue - + # 输出统计结果 print("\n" + "=" * 80) print("统计结果") print("=" * 80 + "\n") - + print(f"📊 总用户数: {len(user_dirs)}") print(f"📊 总Session数: {total_sessions}") print(f"📊 总问题数: {total_questions}\n") - + if total_questions == 0: print("❌ 没有找到任何问题数据") return - + # 输出result_type分布 print("=" * 80) print("Result Type 分布") print("=" * 80 + "\n") - + # 按数量降序排列 sorted_results = sorted(result_counter.items(), key=lambda x: x[1], reverse=True) - + for result_type, count in sorted_results: ratio = count / total_questions * 100 print(f" {result_type:20s}: {count:5d} ({ratio:6.2f}%)") - + # 输出非Correct结果的详细信息 if non_correct_results: print("\n" + "=" * 80) print(f"非 Correct 结果详情 (共 {len(non_correct_results)} 条)") print("=" * 80 + "\n") - + for idx, result in enumerate(non_correct_results, 1): print(f"[{idx}] {result['result_type']}") print(f" 用户: {result['user_name']}") @@ -135,10 +135,10 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"): print(f" 正确答案: {result['answer']}") print(f" 系统回答: {result['system_response'][:200]}{'...' if len(result['system_response']) > 200 else ''}") print() - + else: print("\n🎉 所有问题都是 Correct!") - + # 保存详细报告到文件 report_file = Path(tmp_dir).parent / "analysis_report.json" report_data = { @@ -148,16 +148,16 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"): "total_questions": total_questions, "result_type_distribution": dict(result_counter), "result_type_ratio": { - result_type: count / total_questions + result_type: count / total_questions for result_type, count in result_counter.items() } }, "non_correct_results": non_correct_results } - + with open(report_file, "w", encoding="utf-8") as f: json.dump(report_data, f, ensure_ascii=False, indent=2) - + print("=" * 80) print(f"📄 详细报告已保存到: {report_file}") print("=" * 80) @@ -165,7 +165,7 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"): if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="分析 ReMe 评估结果中的 result_type 分布" ) @@ -175,6 +175,6 @@ if __name__ == "__main__": default="bench_results/reme_simple/tmp", help="临时结果目录路径 (默认: bench_results/reme_simple/tmp)" ) - + args = parser.parse_args() analyze_results(args.tmp_dir) diff --git a/bench/halumem/compute_qa_stats_v4.py b/bench/halumem/compute_qa_stats_v4.py index c7bdfab2..afd07e4b 100644 --- a/bench/halumem/compute_qa_stats_v4.py +++ b/bench/halumem/compute_qa_stats_v4.py @@ -26,9 +26,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": 0, "qa_num": 0 } - + correct = hallucination = omission = valid = 0 - + for qa in qa_records: result_type = qa.get("result_type", "") if result_type == "Correct": @@ -40,7 +40,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: elif result_type == "Omission": omission += 1 valid += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -51,26 +51,26 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": valid, "qa_num": total } - + return metrics def compute_time_metrics(results_file: str) -> dict[str, float]: """Compute timing metrics from evaluation results.""" add_duration = 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) - + return { "add_dialogue_duration_time": add_duration / 1000 / 60, "search_memory_duration_time": search_duration / 1000 / 60, @@ -82,73 +82,73 @@ def load_from_tmp_dir(tmp_dir: str) -> str: """Load data from tmp directory and generate eval_results.jsonl file.""" tmp_path = Path(tmp_dir) eval_results_file = tmp_path.parent / "eval_results.jsonl" - + print(f"\n📁 Loading from: {tmp_dir}") print(f"📝 Generating: {eval_results_file}") - + user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()] print(f" Found {len(user_dirs)} users") - + users_data = [] for user_dir in user_dirs: 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]) ) - + if not session_files: continue - + 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": [] } - + for session_file in session_files: with open(session_file, "r", encoding="utf-8") as f: session_data = json.load(f) session_data.pop("uuid", None) session_data.pop("user_name", None) user_data["sessions"].append(session_data) - + users_data.append(user_data) print(f" ✓ {user_dir.name}: {len(session_files)} sessions") - + 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) def main(input_path: str): """Main function to compute statistics from eval results.""" - + if not os.path.exists(input_path): print(f"❌ Error: Path not found: {input_path}") return - + print("\n" + "=" * 80) print("REME V4 - QUESTION ANSWERING STATISTICS") print("=" * 80) - + # Load or generate eval_results.jsonl if os.path.isdir(input_path): results_file = load_from_tmp_dir(input_path) else: results_file = input_path print(f"\n📁 Using: {results_file}") - + # Collect QA records with metadata qa_records = [] qa_with_metadata = [] user_count = session_count = 0 - + with open(results_file, "r", encoding="utf-8") as f: for line in f: if not line.strip(): @@ -156,15 +156,15 @@ def main(input_path: str): user_data = json.loads(line) user_count += 1 user_name = user_data.get("user_name", "Unknown") - + valid_session_idx = 0 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", {}) - + for qa_idx, qa in enumerate(eval_results.get("question_answering_records", [])): qa_records.append(qa) qa_with_metadata.append({ @@ -173,22 +173,22 @@ def main(input_path: str): "question_idx": qa_idx, "qa_record": qa }) - + valid_session_idx += 1 - + print(f"\n📊 Data Summary:") print(f" Users: {user_count}") print(f" Sessions: {session_count}") print(f" QA Records: {len(qa_records)}") - + # Compute metrics qa_metrics = compute_qa_metrics(qa_records) time_metrics = compute_time_metrics(results_file) - + # Save results output_dir = Path(results_file).parent report_file = output_dir / "reme_eval_stat_result.json" - + final_results = { "overall_score": { "question_answering": qa_metrics, @@ -196,12 +196,12 @@ def main(input_path: str): }, "question_answering_records": qa_records } - + with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=4) - + print(f"\n✅ Results saved to: {report_file}") - + # Print metrics print("\n" + "=" * 80) print("📊 QUESTION ANSWERING METRICS") @@ -213,27 +213,27 @@ def main(input_path: str): 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 error records print("\n" + "=" * 80) print("❌ ERROR RECORDS (Non-Correct)") print("=" * 80) - + error_records = [r for r in qa_with_metadata if r["qa_record"].get("result_type") not in ["Correct", ""]] - + if not error_records: print("\n✅ All QA records are correct!") else: print(f"\nFound {len(error_records)} error records:\n") - + for idx, record in enumerate(error_records, 1): qa = record["qa_record"] - + print(f"\n{'━' * 80}") print(f"❌ ERROR #{idx}") print(f"{'━' * 80}") @@ -264,19 +264,19 @@ def main(input_path: str): print("\n".join(lines)) else: print(f" {reason}") - + 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, help="Path to eval_results.jsonl file") parser.add_argument("--tmp_dir", type=str, help="Path to tmp directory") - + args = parser.parse_args() - + if args.tmp_dir: main(input_path=args.tmp_dir) elif args.results_file: diff --git a/bench/halumem/compute_stats_from_tmp.py b/bench/halumem/compute_stats_from_tmp.py index 5ab401ba..2d0201ec 100644 --- a/bench/halumem/compute_stats_from_tmp.py +++ b/bench/halumem/compute_stats_from_tmp.py @@ -358,7 +358,7 @@ async def main_async(tmp_dir: str): # Determine paths parent_dir = os.path.dirname(tmp_dir) frame = "reme" - + output_file_stage1 = os.path.join(parent_dir, f"{frame}_eval_results.jsonl") output_file_stage2 = os.path.join(parent_dir, f"{frame}_eval_stat_result.json") @@ -392,7 +392,7 @@ async def main_async(tmp_dir: str): # Load all users and process user_data_list = list(enumerate(iter_jsonl(output_file_stage1), 1)) - + for idx, user_data in user_data_list: uuid = user_data["uuid"] tmp_file = os.path.join(tmp_dir2, f"{uuid}.json") diff --git a/bench/halumem/eval_baseline_simple.py b/bench/halumem/eval_baseline_simple.py index 6da0bbed..34be3d0a 100644 --- a/bench/halumem/eval_baseline_simple.py +++ b/bench/halumem/eval_baseline_simple.py @@ -44,13 +44,13 @@ class EvalConfig: 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.""" @@ -58,7 +58,7 @@ class DataLoader: 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.""" @@ -74,7 +74,7 @@ class DataLoader: } for turn in dialogue ] - + @staticmethod def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str: """Format dialogue into string for evaluation (only user messages).""" @@ -83,14 +83,14 @@ class DataLoader: # Skip assistant messages - only include user messages if turn['role'] != 'user': continue - + 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 provided role = user_name if user_name else 'user' - + formatted_turns.append( f"Role: {role}\n" f"Content: {turn['content']}\n" @@ -101,29 +101,29 @@ class DataLoader: 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) @@ -131,38 +131,38 @@ class FileManager: 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" + 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() + 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: @@ -171,7 +171,7 @@ class FileManager: 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") @@ -206,10 +206,10 @@ Please respond in JSON format with the following structure: class BaselineQuestionAnsweringEvaluator: """Evaluates question answering performance using direct LLM inference (no memory system).""" - + def __init__(self): pass - + async def answer_question( self, question: str, @@ -217,18 +217,18 @@ class BaselineQuestionAnsweringEvaluator: ) -> tuple[str, str, float]: """ Answer a question using the dialogue history directly. - + Returns: tuple: (answer, reasoning, duration_ms) """ start = time.time() - + # Format prompt prompt = BASELINE_QA_PROMPT.format( dialogue=formatted_dialogue, question=question ) - + # Get answer from LLM try: # model_name = "qwen3-max" @@ -240,10 +240,10 @@ class BaselineQuestionAnsweringEvaluator: logger.error(f"Error getting answer from LLM: {e}") answer = "Error: Failed to get answer" reasoning = str(e) - + duration_ms = (time.time() - start) * 1000 return answer, reasoning, duration_ms - + async def evaluate_questions( self, questions: list[dict], @@ -254,14 +254,14 @@ class BaselineQuestionAnsweringEvaluator: ) -> list[dict]: """Evaluate all questions for a session.""" results = [] - + for qa in questions: # Get answer directly from LLM answer, reasoning, duration_ms = await self.answer_question( question=qa["question"], formatted_dialogue=formatted_dialogue ) - + # Evaluate response evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]]) eval_result = await evaluation_for_question2( @@ -271,7 +271,7 @@ class BaselineQuestionAnsweringEvaluator: answer, formatted_dialogue ) - + # Build result record qa_result = { **qa, @@ -284,13 +284,13 @@ class BaselineQuestionAnsweringEvaluator: "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.""" @@ -306,15 +306,15 @@ class MetricsAggregator: "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": @@ -323,7 +323,7 @@ class MetricsAggregator: hallucination += 1 elif result_type == "Omission": omission += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -331,7 +331,7 @@ class MetricsAggregator: "qa_valid_num": valid, "qa_num": total } - + if valid > 0: metrics.update({ "correct_qa_ratio(valid)": correct / valid, @@ -344,25 +344,25 @@ class MetricsAggregator: "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.""" answer_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"]: eval_results = session.get("evaluation_results", {}) for qa in eval_results.get("question_answering_records", []): answer_duration += qa.get("answer_duration_ms", 0) - + # Convert to minutes return { "answer_duration_time": answer_duration / 1000 / 60, @@ -374,13 +374,13 @@ class MetricsAggregator: class HaluMemBaselineEvaluator: """Main evaluator orchestrating the baseline evaluation pipeline.""" - + def __init__(self, config: EvalConfig): self.config = config self.file_manager = FileManager(config.output_dir) self.qa_evaluator = BaselineQuestionAnsweringEvaluator() self.data_loader = DataLoader() - + async def process_session( self, session: dict, @@ -395,16 +395,16 @@ class HaluMemBaselineEvaluator: "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 - + # Store dialogue dialogue = session["dialogue"] session_data["dialogue"] = dialogue - + # Evaluate questions if present if "questions" in session: formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name) @@ -415,25 +415,25 @@ class HaluMemBaselineEvaluator: 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"] - + total_sessions = len(user_data["sessions"]) logger.info(f"Processing user: {user_name} ({total_sessions} sessions)") - + # Semaphore for concurrency control within user sessions semaphore = asyncio.Semaphore(self.config.max_concurrency) completed_count = [0] # Use list to allow modification in nested async function - + async def process_session_with_log(idx: int, session: dict): async with semaphore: session_data = await self.process_session( @@ -442,65 +442,65 @@ class HaluMemBaselineEvaluator: user_name=user_name, uuid=uuid ) - + self.file_manager.save_session(user_name, idx, session_data) - + # Update and log completion completed_count[0] += 1 print(f"✅ {user_name} complete {completed_count[0]}/{total_sessions}") - + # Process all sessions in parallel tasks = [ process_session_with_log(idx, session) for idx, session in enumerate(user_data["sessions"]) ] await asyncio.gather(*tasks) - + return {"uuid": uuid, "user_name": user_name, "status": "ok"} - + async def run_evaluation(self): """Run the complete evaluation pipeline.""" start_time = time.time() - + # 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 BASELINE EVALUATION - DIRECT QA WITHOUT MEMORY SYSTEM") print(f"Users: {len(users_to_process)} | Session Concurrency: {self.config.max_concurrency}") print("=" * 80 + "\n") - + # Process users sequentially (for loop) for idx, user_data in enumerate(users_to_process, 1): 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)") continue - + print(f"🔄 [{idx}/{len(users_to_process)}] Processing {user_name}...") await self.process_user(user_data) print(f"✅ [{idx}/{len(users_to_process)}] User {user_name} completed\n") - + # 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: @@ -508,20 +508,20 @@ class HaluMemBaselineEvaluator: 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, @@ -529,23 +529,23 @@ class HaluMemBaselineEvaluator: }, "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") 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}") @@ -554,7 +554,7 @@ class HaluMemBaselineEvaluator: 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" Answer Duration: {time_metrics['answer_duration_time']:.2f} min") print(f" Total: {time_metrics['total_duration_time']:.2f} min") @@ -574,14 +574,14 @@ def main( user_num=user_num, max_concurrency=max_concurrency ) - + evaluator = HaluMemBaselineEvaluator(config) asyncio.run(evaluator.run_evaluation()) if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="Evaluate Baseline (Direct QA) on HaluMem benchmark" ) @@ -603,9 +603,9 @@ if __name__ == "__main__": default=2, help="Maximum concurrent user processing (default: 2)" ) - + args = parser.parse_args() - + main( data_path=args.data_path, user_num=args.user_num, diff --git a/bench/halumem/eval_reme.py b/bench/halumem/eval_reme.py index d3044535..3254c332 100644 --- a/bench/halumem/eval_reme.py +++ b/bench/halumem/eval_reme.py @@ -35,8 +35,8 @@ from eval_tools import ( evaluation_for_update_memory, ) from llms import llm_request -from reme_ai.core.enumeration import MemoryType -from reme_ai.core.schema import MemoryNode +from reme_ai.core_old.enumeration import MemoryType +from reme_ai.core_old.schema import MemoryNode from reme_ai.reme import ReMe # Template for formatting memories (from shared YAML config) @@ -685,7 +685,7 @@ async def main_async( user_data_list = user_data_list[:total_users] print(f"Processing {total_users} users with max concurrency {max_concurrency}...") - + # Create semaphore to limit concurrency for Stage 1 semaphore_stage1 = asyncio.Semaphore(max_concurrency) @@ -694,11 +694,11 @@ async def main_async( async with semaphore_stage1: uuid = user_data['uuid'] tmp_file = os.path.join(tmp_dir, f"{uuid}.json") - + if os.path.exists(tmp_file): print(f"⚡ Skipping user {uuid} ({idx}/{total_users}) — cached result found.") return {"uuid": uuid, "status": "cached", "path": tmp_file} - + print(f"[{idx}/{total_users}] Processing user {uuid}...") result = await process_user_stage1(user_data, top_k, save_path) print(f"[{idx}/{total_users}] ✅ Finished {uuid} ({result['status']})") @@ -733,7 +733,7 @@ async def main_async( # Load all users and process sequentially user_data_list = list(enumerate(iter_jsonl(output_file_stage1), 1)) - + for idx, user_data in user_data_list: uuid = user_data["uuid"] tmp_file = os.path.join(tmp_dir2, f"{uuid}.json") diff --git a/bench/halumem/eval_reme_simple.py b/bench/halumem/eval_reme_simple.py index aff0934d..4635c556 100644 --- a/bench/halumem/eval_reme_simple.py +++ b/bench/halumem/eval_reme_simple.py @@ -26,8 +26,8 @@ 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.core_old.enumeration import MemoryType +from reme_ai.core_old.schema import MemoryNode from reme_ai.reme import ReMe @@ -48,13 +48,13 @@ class EvalConfig: 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.""" @@ -62,7 +62,7 @@ class DataLoader: 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.""" @@ -78,7 +78,7 @@ class DataLoader: } for turn in dialogue ] - + @staticmethod def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str: """Format dialogue into string for evaluation.""" @@ -87,10 +87,10 @@ class DataLoader: 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" @@ -101,29 +101,29 @@ class DataLoader: 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) @@ -131,38 +131,38 @@ class FileManager: 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" + 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() + 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: @@ -171,7 +171,7 @@ class FileManager: 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") @@ -179,19 +179,19 @@ class FileManager: class MemoryProcessor: """Handles ReMe memory operations.""" - + def __init__(self, reme: ReMe): self.reme = reme - + async def add_memories( - self, - user_id: str, + self, + user_id: str, messages: list[dict], batch_size: int = 20 ) -> tuple[list[str], list[list[dict]], float]: """ Add memories in batches and return extracted memory contents. - + Returns: tuple: (extracted_memories, agent_messages, total_duration_ms) """ @@ -202,19 +202,19 @@ class MemoryProcessor: 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_v2( - messages=batch, + 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: @@ -230,23 +230,23 @@ class MemoryProcessor: 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, + self, + query: str, + user_id: str, top_k: int = 20 ) -> tuple[str, list, float]: """ Search memory and return response. - + Returns: tuple: (response, agent_messages, duration_ms) """ start = time.time() response, agent_messages, success = await self.reme.retrieve_v2( - query=query, - user_id=user_id, + query=query, + user_id=user_id, top_k=top_k ) duration_ms = (time.time() - start) * 1000 @@ -257,11 +257,11 @@ class MemoryProcessor: 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], @@ -272,7 +272,7 @@ class QuestionAnsweringEvaluator: ) -> list[dict]: """Evaluate all questions for a session.""" results = [] - + for qa in questions: # Search memory for answer response, agent_messages, duration_ms = await self.memory_processor.search_memory( @@ -280,7 +280,7 @@ class QuestionAnsweringEvaluator: 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( @@ -290,7 +290,7 @@ class QuestionAnsweringEvaluator: response, formatted_dialogue ) - + # Build result record qa_result = { **qa, @@ -303,13 +303,13 @@ class QuestionAnsweringEvaluator: "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.""" @@ -325,15 +325,15 @@ class MetricsAggregator: "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": @@ -342,7 +342,7 @@ class MetricsAggregator: hallucination += 1 elif result_type == "Omission": omission += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -350,7 +350,7 @@ class MetricsAggregator: "qa_valid_num": valid, "qa_num": total } - + if valid > 0: metrics.update({ "correct_qa_ratio(valid)": correct / valid, @@ -363,28 +363,28 @@ class MetricsAggregator: "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, @@ -397,18 +397,18 @@ class MetricsAggregator: class HaluMemEvaluator: """Main evaluator orchestrating the entire pipeline.""" - + 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, + self.memory_processor, config.top_k ) self.data_loader = DataLoader() - + async def process_session( self, session: dict, @@ -423,29 +423,29 @@ class HaluMemEvaluator: "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 - + # Format and add dialogue to memory 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) @@ -456,90 +456,90 @@ class HaluMemEvaluator: 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.""" start_time = time.time() - + # Clear existing data await self.reme.vector_store.delete_all() - + # 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 - 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) + 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: @@ -547,20 +547,20 @@ class HaluMemEvaluator: 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, @@ -568,23 +568,23 @@ class HaluMemEvaluator: }, "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") 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}") @@ -593,7 +593,7 @@ class HaluMemEvaluator: 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") @@ -616,14 +616,14 @@ def main( user_num=user_num, max_concurrency=max_concurrency ) - + evaluator = HaluMemEvaluator(config) asyncio.run(evaluator.run_evaluation()) if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="Evaluate ReMe on HaluMem benchmark (Question Answering)" ) @@ -651,9 +651,9 @@ if __name__ == "__main__": default=2, help="Maximum concurrent user processing (default: 2)" ) - + args = parser.parse_args() - + main( data_path=args.data_path, top_k=args.top_k, diff --git a/bench/halumem/eval_reme_simple_v3.py b/bench/halumem/eval_reme_simple_v3.py index 7c0dfd8e..43cee5d9 100644 --- a/bench/halumem/eval_reme_simple_v3.py +++ b/bench/halumem/eval_reme_simple_v3.py @@ -27,8 +27,8 @@ 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.core_old.enumeration import MemoryType +from reme_ai.core_old.schema import MemoryNode from reme_ai.reme import ReMe @@ -49,13 +49,13 @@ class EvalConfig: 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.""" @@ -63,7 +63,7 @@ class DataLoader: 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).""" @@ -80,7 +80,7 @@ class DataLoader: 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.""" @@ -89,10 +89,10 @@ class DataLoader: 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" @@ -103,29 +103,29 @@ class DataLoader: 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) @@ -133,38 +133,38 @@ class FileManager: 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" + 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() + 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: @@ -173,7 +173,7 @@ class FileManager: 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") @@ -181,19 +181,19 @@ class FileManager: class MemoryProcessor: """Handles ReMe V3 memory operations.""" - + def __init__(self, reme: ReMe): self.reme = reme - + async def add_memories( - self, - user_id: str, + self, + user_id: str, messages: list[dict], batch_size: int = 10000 ) -> tuple[list[str], list[list[dict]], float]: """ Add memories in batches using ReMe V3 and return extracted memory contents. - + Returns: tuple: (extracted_memories, agent_messages, total_duration_ms) """ @@ -201,24 +201,24 @@ class MemoryProcessor: 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() - + # Use summary_v3 instead of summary_v2 memory_nodes, agent_messages, success = await self.reme.summary_v3( - messages=batch, + 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: @@ -234,28 +234,28 @@ class MemoryProcessor: 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, + self, + query: str, + user_id: str, top_k: int = 20 ) -> tuple[str, list, float]: """ Search memory using ReMe V3 and return response. - + Returns: tuple: (response, agent_messages, duration_ms) """ start = time.time() - + # Use retrieve_v3 instead of retrieve_v2 response, agent_messages, success = await self.reme.retrieve_v3( - query=query, - user_id=user_id, + query=query, + user_id=user_id, top_k=top_k ) - + duration_ms = (time.time() - start) * 1000 return response, agent_messages, duration_ms @@ -264,11 +264,11 @@ class MemoryProcessor: 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], @@ -279,7 +279,7 @@ class QuestionAnsweringEvaluator: ) -> list[dict]: """Evaluate all questions for a session.""" results = [] - + for qa in questions: # Search memory for answer using V3 response, agent_messages, duration_ms = await self.memory_processor.search_memory( @@ -287,7 +287,7 @@ class QuestionAnsweringEvaluator: 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( @@ -297,7 +297,7 @@ class QuestionAnsweringEvaluator: response, formatted_dialogue ) - + # Build result record qa_result = { **qa, @@ -310,13 +310,13 @@ class QuestionAnsweringEvaluator: "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.""" @@ -332,15 +332,15 @@ class MetricsAggregator: "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": @@ -349,7 +349,7 @@ class MetricsAggregator: hallucination += 1 elif result_type == "Omission": omission += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -357,7 +357,7 @@ class MetricsAggregator: "qa_valid_num": valid, "qa_num": total } - + if valid > 0: metrics.update({ "correct_qa_ratio(valid)": correct / valid, @@ -370,28 +370,28 @@ class MetricsAggregator: "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, @@ -404,18 +404,18 @@ class MetricsAggregator: class HaluMemEvaluatorV3: """Main evaluator orchestrating the entire ReMe V3 pipeline.""" - + 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, + self.memory_processor, config.top_k ) self.data_loader = DataLoader() - + async def process_session( self, session: dict, @@ -430,29 +430,29 @@ class HaluMemEvaluatorV3: "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 - + # Format and add dialogue to memory using V3 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) @@ -463,97 +463,97 @@ class HaluMemEvaluatorV3: 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 V3.""" 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 V3 - 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) + 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: @@ -561,20 +561,20 @@ class HaluMemEvaluatorV3: 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, @@ -582,23 +582,23 @@ class HaluMemEvaluatorV3: }, "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 V3") 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}") @@ -607,7 +607,7 @@ class HaluMemEvaluatorV3: 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") @@ -630,14 +630,14 @@ def main( user_num=user_num, max_concurrency=max_concurrency ) - + evaluator = HaluMemEvaluatorV3(config) asyncio.run(evaluator.run_evaluation()) if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="Evaluate ReMe V3 on HaluMem benchmark (Question Answering)" ) @@ -665,9 +665,9 @@ if __name__ == "__main__": default=2, help="Maximum concurrent user processing (default: 2)" ) - + args = parser.parse_args() - + main( data_path=args.data_path, top_k=args.top_k, diff --git a/bench/halumem/eval_reme_simple_v4.py b/bench/halumem/eval_reme_simple_v4.py index 825697cb..63cb4f37 100644 --- a/bench/halumem/eval_reme_simple_v4.py +++ b/bench/halumem/eval_reme_simple_v4.py @@ -27,8 +27,8 @@ from typing import Any from loguru import logger from eval_tools import evaluation_for_question2, answer_question_with_memories -from reme_ai.core.enumeration import MemoryType -from reme_ai.core.schema import MemoryNode +from reme_ai.core_old.enumeration import MemoryType +from reme_ai.core_old.schema import MemoryNode from reme_ai.reme import ReMe @@ -49,13 +49,13 @@ class EvalConfig: 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.""" @@ -63,7 +63,7 @@ class DataLoader: 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).""" @@ -80,7 +80,7 @@ class DataLoader: 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.""" @@ -89,10 +89,10 @@ class DataLoader: 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" @@ -103,29 +103,29 @@ class DataLoader: 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) @@ -133,38 +133,38 @@ class FileManager: 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" + 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() + 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: @@ -173,7 +173,7 @@ class FileManager: 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") @@ -181,19 +181,19 @@ class FileManager: class MemoryProcessor: """Handles ReMe memory operations.""" - + def __init__(self, reme: ReMe): self.reme = reme - + async def add_memories( - self, - user_id: str, + 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) """ @@ -201,23 +201,23 @@ class MemoryProcessor: 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, + 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: @@ -233,39 +233,39 @@ class MemoryProcessor: 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, + self, + query: str, + user_id: str, top_k: int = 20 ) -> tuple[dict, list, float]: """ Search memory using ReMe and return structured answer with reasoning. - + Returns: tuple: (answer_dict, agent_messages, duration_ms) answer_dict contains: {"reasoning": str, "answer": str, "memories": str} """ start = time.time() - + # Retrieve memories from ReMe memories_response, agent_messages, success = await self.reme.retrieve_v4( - query=query, - user_id=user_id, + query=query, + user_id=user_id, top_k=top_k ) - + # Use LLM to generate structured answer from memories answer_result = await answer_question_with_memories( question=query, memories=memories_response, user_id=user_id ) - + # Add original memories to the result answer_result["memories"] = memories_response - + duration_ms = (time.time() - start) * 1000 return answer_result, agent_messages, duration_ms @@ -274,11 +274,11 @@ class MemoryProcessor: 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], @@ -289,19 +289,19 @@ class QuestionAnsweringEvaluator: ) -> list[dict]: """Evaluate all questions for a session.""" results = [] - + for qa in questions: answer_dict, agent_messages, duration_ms = await self.memory_processor.search_memory( query=qa["question"], user_id=user_name, top_k=self.top_k ) - + # Extract answer and reasoning from the structured response system_answer = answer_dict.get("answer", "") system_reasoning = answer_dict.get("reasoning", "") retrieved_memories = answer_dict.get("memories", "") - + # Evaluate response evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]]) eval_result = await evaluation_for_question2( @@ -311,7 +311,7 @@ class QuestionAnsweringEvaluator: system_answer, formatted_dialogue ) - + # Build result record qa_result = { **qa, @@ -326,13 +326,13 @@ class QuestionAnsweringEvaluator: "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.""" @@ -348,15 +348,15 @@ class MetricsAggregator: "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": @@ -365,7 +365,7 @@ class MetricsAggregator: hallucination += 1 elif result_type == "Omission": omission += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -373,7 +373,7 @@ class MetricsAggregator: "qa_valid_num": valid, "qa_num": total } - + if valid > 0: metrics.update({ "correct_qa_ratio(valid)": correct / valid, @@ -386,28 +386,28 @@ class MetricsAggregator: "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, @@ -426,11 +426,11 @@ class HaluMemEvaluatorV4: self.file_manager = FileManager(config.output_dir) self.memory_processor = MemoryProcessor(self.reme) self.qa_evaluator = QuestionAnsweringEvaluator( - self.memory_processor, + self.memory_processor, config.top_k ) self.data_loader = DataLoader() - + async def process_session( self, session: dict, @@ -445,7 +445,7 @@ class HaluMemEvaluatorV4: "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 @@ -453,20 +453,20 @@ class HaluMemEvaluatorV4: 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) @@ -477,97 +477,97 @@ class HaluMemEvaluatorV4: 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) + 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: @@ -575,20 +575,20 @@ class HaluMemEvaluatorV4: 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, @@ -596,23 +596,23 @@ class HaluMemEvaluatorV4: }, "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}") @@ -621,7 +621,7 @@ class HaluMemEvaluatorV4: 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") @@ -644,14 +644,14 @@ def main( 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)" ) @@ -679,9 +679,9 @@ if __name__ == "__main__": default=2, help="Maximum concurrent user processing (default: 2)" ) - + args = parser.parse_args() - + main( data_path=args.data_path, top_k=args.top_k, diff --git a/bench/halumem/eval_tools.py b/bench/halumem/eval_tools.py index 29868ca8..e0fae45d 100644 --- a/bench/halumem/eval_tools.py +++ b/bench/halumem/eval_tools.py @@ -141,12 +141,12 @@ async def answer_question_with_memories( ): """ Answer a question using retrieved memories with PROMPT_MEMZERO_JSON template. - + Args: question: The question to answer memories: The retrieved memories (formatted as context) user_id: Optional user ID for context formatting - + Returns: dict with 'reasoning' and 'answer' fields """ @@ -158,13 +158,13 @@ async def answer_question_with_memories( ) else: context = f"Memories:\n{memories}" - + # Use PROMPT_MEMZERO_JSON template for structured JSON response prompt = _PROMPTS["PROMPT_MEMZERO_JSON"].format( context=context, question=question ) - + # result = await llm_request_for_json(prompt, model_name="qwen3-max") result = await llm_request_for_json(prompt, model_name="qwen3-30b-a3b-instruct-2507") diff --git a/bench/halumem/halumem.yaml b/bench/halumem/halumem.yaml index 51419c37..56dde2db 100644 --- a/bench/halumem/halumem.yaml +++ b/bench/halumem/halumem.yaml @@ -11,7 +11,7 @@ PROMPT_MEMZERO_JSON: | 1. **Historical Dialogue** (highest priority) - Direct conversation content 2. **Extracted Memories** (medium priority) - Summarized memory points 3. **User Profile** (lowest priority) - General user information - + # Question: {question} @@ -35,7 +35,7 @@ PROMPT_MEMZERO_JSON2: | 1. **Historical Dialogue** (highest priority) - Direct conversation content 2. **Extracted Memories** (medium priority) - Summarized memory points 3. **User Profile** (lowest priority) - General user information - + # Question: {question} @@ -467,7 +467,7 @@ EVALUATION_PROMPT_FOR_QUESTION: | "evaluation_result": "Correct | Hallucination | Omission" }} ``` - + EVALUATION_PROMPT_FOR_QUESTION2: | You are an **evaluation expert for AI memory system question answering**. diff --git a/bench/halumem/llms.py b/bench/halumem/llms.py index 24edc703..5a59b2b2 100644 --- a/bench/halumem/llms.py +++ b/bench/halumem/llms.py @@ -5,8 +5,8 @@ import re from tenacity import retry, stop_after_attempt, wait_random_exponential, before_sleep_log -from reme_ai.core.schema import Message -from reme_ai.core.utils import load_env +from reme_ai.core_old.schema import Message +from reme_ai.core_old.utils import load_env from reme_ai.reme import ReMe logger = logging.getLogger(__name__) @@ -28,12 +28,12 @@ reme = ReMe() ) async def llm_request(prompt, model_name: str = "qwen3-max", **kwargs) -> str: """Make an LLM request using ReMe's LLM with optional model override. - + Args: prompt: The prompt to send to the LLM model_name: Optional model name to override the default model (default: "qwen3-max") **kwargs: Additional arguments to pass to the chat method - + Returns: The assistant's response content """ @@ -61,15 +61,15 @@ async def llm_request(prompt, model_name: str = "qwen3-max", **kwargs) -> str: async def llm_request_for_json(prompt, model_name: str = "qwen-flash", **kwargs): # async def llm_request_for_json(prompt, model_name: str = "qwen3-max", **kwargs): """Make an LLM request expecting JSON response using ReMe's LLM. - + Args: prompt: The prompt to send to the LLM model_name: Optional model name to override the default model (default: "qwen3-max") **kwargs: Additional arguments to pass to the chat method - + Returns: Parsed JSON object from the LLM response - + Raises: ValueError: If no JSON block is found in the model output """ diff --git a/bench/human_in_the_loop/compute_qa_stats.py b/bench/human_in_the_loop/compute_qa_stats.py index 65ac6dc8..69d65e55 100644 --- a/bench/human_in_the_loop/compute_qa_stats.py +++ b/bench/human_in_the_loop/compute_qa_stats.py @@ -22,9 +22,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": 0, "qa_num": 0 } - + correct = hallucination = omission = valid = 0 - + for qa in qa_records: result_type = qa.get("result_type", "") if result_type == "Correct": @@ -36,7 +36,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: elif result_type == "Omission": omission += 1 valid += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -47,21 +47,21 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": valid, "qa_num": total } - + return metrics def compute_time_metrics(users_data: list[dict]) -> dict[str, float]: """Compute timing metrics from evaluation results.""" add_duration = search_duration = 0 - + for user_data in users_data: for session in user_data.get("sessions", []): add_duration += session.get("add_dialogue_duration_ms", 0) eval_results = session.get("session", {}).get("evaluation_results", {}) for qa in eval_results.get("question_answering_records", []): search_duration += qa.get("search_duration_ms", 0) - + return { "add_dialogue_duration_time": add_duration / 1000 / 60, "search_memory_duration_time": search_duration / 1000 / 60, @@ -72,21 +72,21 @@ def compute_time_metrics(users_data: list[dict]) -> dict[str, float]: def load_from_tmp_dir(tmp_dir: str) -> list[dict]: """Load data from tmp directory.""" tmp_path = Path(tmp_dir) - + # Try flat file structure first (conversation_{user}_session_{idx}.json) json_files = [f for f in tmp_path.iterdir() if f.is_file() and f.suffix == ".json"] - + if json_files: # Group files by user users_dict = defaultdict(list) - + for json_file in json_files: with open(json_file, "r", encoding="utf-8") as f: session_data = json.load(f) user_name = session_data.get("user_name") if user_name: users_dict[user_name].append(session_data) - + # Sort sessions by session_idx for each user users_data = [] for user_name, sessions in users_dict.items(): @@ -103,71 +103,71 @@ def load_from_tmp_dir(tmp_dir: str) -> list[dict]: session_copy.pop("user_name", None) user_data["sessions"].append(session_copy) users_data.append(user_data) - + return users_data - + # Fallback to directory structure (user_name/session_{idx}.json) user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()] - + users_data = [] for user_dir in user_dirs: session_files = sorted( [f for f in user_dir.iterdir() if "session_" in f.name and f.suffix == ".json"], key=lambda f: int(f.stem.split("_")[-1]) ) - + if not session_files: continue - + 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": [] } - + for session_file in session_files: with open(session_file, "r", encoding="utf-8") as f: session_data = json.load(f) session_data.pop("uuid", None) session_data.pop("user_name", None) user_data["sessions"].append(session_data) - + users_data.append(user_data) - + return users_data def main(tmp_dir: str): """Main function to compute statistics from tmp directory.""" tmp_path = Path(tmp_dir) - + if not tmp_path.exists() or not tmp_path.is_dir(): print(f"❌ Error: Directory not found: {tmp_dir}") return - + # Load data from tmp directory users_data = load_from_tmp_dir(tmp_dir) - + # Collect QA records with metadata qa_records = [] qa_with_metadata = [] user_count = session_count = 0 - + for user_data in users_data: user_count += 1 user_name = user_data.get("user_name", "Unknown") - + valid_session_idx = 0 for session in user_data.get("sessions", []): if session.get("is_generated_qa_session"): continue - + session_count += 1 eval_results = session.get("session", {}).get("evaluation_results", {}) - + for qa_idx, qa in enumerate(eval_results.get("question_answering_records", [])): qa_records.append(qa) qa_with_metadata.append({ @@ -176,17 +176,17 @@ def main(tmp_dir: str): "question_idx": qa_idx, "qa_record": qa }) - + valid_session_idx += 1 - + # Compute metrics qa_metrics = compute_qa_metrics(qa_records) time_metrics = compute_time_metrics(users_data) - + # Save results output_dir = tmp_path.parent report_file = output_dir / "reme_eval_stat_result.json" - + final_results = { "overall_score": { "question_answering": qa_metrics, @@ -194,10 +194,10 @@ def main(tmp_dir: str): }, "question_answering_records": qa_records } - + with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=4) - + # Print summary print(f"\n📊 Data: {user_count} users, {session_count} sessions, {len(qa_records)} QA records") print(f"\n✅ Metrics:") @@ -205,12 +205,12 @@ def main(tmp_dir: str): print(f" Valid: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}") print(f"\n⏱️ Time: {time_metrics['total_duration_time']:.2f} min (Add: {time_metrics['add_dialogue_duration_time']:.2f} | Search: {time_metrics['search_memory_duration_time']:.2f})") print(f"\n💾 Results saved: {report_file}") - + # Print error records print(f"\n{'='*80}\n❌ ERROR RECORDS ({len([r for r in qa_with_metadata if r['qa_record'].get('result_type') not in ['Correct', '']])} errors)\n{'='*80}") - + error_records = [r for r in qa_with_metadata if r["qa_record"].get("result_type") not in ["Correct", ""]] - + if error_records: for idx, record in enumerate(error_records, 1): qa = record["qa_record"] @@ -218,13 +218,13 @@ def main(tmp_dir: str): print(f" Q: {qa.get('question', 'N/A')}") print(f" Expected: {qa.get('answer', 'N/A')}") print(f" Got: {qa.get('system_response', 'N/A')}") - + print() if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser(description="Compute QA statistics from tmp directory") parser.add_argument( "tmp_dir", @@ -232,6 +232,6 @@ if __name__ == "__main__": default="./data", type=str, help="Path to tmp directory containing user session data (default: ./data)") - + args = parser.parse_args() main(tmp_dir=args.tmp_dir) diff --git a/bench/human_in_the_loop/eval.yaml b/bench/human_in_the_loop/eval.yaml index 41d36943..3e685bc7 100644 --- a/bench/human_in_the_loop/eval.yaml +++ b/bench/human_in_the_loop/eval.yaml @@ -64,7 +64,7 @@ EVALUATION_PROMPT_FOR_QUESTION: | "evaluation_result": "Correct | Hallucination | Omission" }} ``` - + EVALUATION_PROMPT_FOR_QUESTION2: | You are an **evaluation expert for AI memory system question answering**. diff --git a/bench/human_in_the_loop/reevaluate_qa.py b/bench/human_in_the_loop/reevaluate_qa.py index 98b61b26..a013616c 100644 --- a/bench/human_in_the_loop/reevaluate_qa.py +++ b/bench/human_in_the_loop/reevaluate_qa.py @@ -16,8 +16,8 @@ from collections import defaultdict from pathlib import Path from typing import Any -from reme_ai.core.schema import Message -from reme_ai.core.utils import load_env +from reme_ai.core_old.schema import Message +from reme_ai.core_old.utils import load_env from reme_ai.reme import ReMe from tenacity import retry, stop_after_attempt, wait_random_exponential @@ -82,7 +82,7 @@ async def evaluate_qa_record( prompt_version: str = "v1" ) -> dict: """Evaluate a single QA record using LLM with specified prompt version. - + Args: question: The question to evaluate reference_answer: The reference answer @@ -90,9 +90,9 @@ async def evaluate_qa_record( response: System response to evaluate dialogue: Dialogue context (optional) model_name: LLM model name - prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION, + prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION, "v2" for EVALUATION_PROMPT_FOR_QUESTION2 - + Returns: dict with evaluation_result and reasoning """ @@ -101,7 +101,7 @@ async def evaluate_qa_record( prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION2"] else: prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION"] - + # Format prompt prompt = prompt_template.format( question=question, @@ -110,7 +110,7 @@ async def evaluate_qa_record( response=response, dialogue=dialogue or "N/A" ) - + result = await llm_request_for_json(prompt, model_name=model_name) return result @@ -118,14 +118,14 @@ async def evaluate_qa_record( def load_from_data_dir(data_dir: str) -> list[dict]: """Load data from data directory (same as compute_qa_stats.py).""" data_path = Path(data_dir) - + # Try flat file structure first (conversation_{user}_session_{idx}.json) json_files = [f for f in data_path.iterdir() if f.is_file() and f.suffix == ".json"] - + if json_files: # Group files by user users_dict = defaultdict(list) - + for json_file in json_files: with open(json_file, "r", encoding="utf-8") as f: session_data = json.load(f) @@ -135,18 +135,18 @@ def load_from_data_dir(data_dir: str) -> list[dict]: "file": json_file, "data": session_data }) - + # Sort sessions by session_idx for each user users_data = [] for user_name, sessions in users_dict.items(): sessions_sorted = sorted( - sessions, + sessions, key=lambda s: s["data"].get("session_idx", 0) ) users_data.extend(sessions_sorted) - + return users_data - + return [] @@ -155,7 +155,7 @@ def format_dialogue_context(session_data: dict) -> str: dialogue = session_data.get("session", {}).get("dialogue", []) if not dialogue: return "N/A" - + formatted_turns = [] for turn in dialogue: role = turn.get("role", "unknown") @@ -175,34 +175,34 @@ async def reevaluate_session( parallel: bool = True ) -> dict: """Re-evaluate all QA records in a session using multiple models and prompts. - + Args: session_file: Path to session file session_data: Session data dict models: List of model names to use for evaluation prompt_versions: List of prompt versions ("v1", "v2") - parallel: If True, use asyncio.gather for parallel execution; + parallel: If True, use asyncio.gather for parallel execution; if False, execute sequentially - + Returns: Updated session data with evaluation results for each model+prompt combination - + Note: Request rate limiting is handled by base_llm.py's request_interval mechanism. """ eval_results = session_data.get("session", {}).get("evaluation_results", {}) qa_records = eval_results.get("question_answering_records", []) - + if not qa_records: print(f" ⏭️ No QA records found") return session_data - + total_evals = len(models) * len(prompt_versions) * len(qa_records) print(f" 🔍 Re-evaluating {len(qa_records)} QA records with {len(models)} models × {len(prompt_versions)} prompts = {total_evals} evaluations...") - + # Format dialogue context once dialogue_context = format_dialogue_context(session_data) - + async def evaluate_single_combination( idx: int, qa: dict, @@ -210,19 +210,19 @@ async def reevaluate_session( prompt_version: str ) -> tuple[int, str, str, dict]: """Evaluate a single QA record with specific model and prompt. - + Note: Rate limiting is handled by BaseLLM's request_interval mechanism. """ question = qa.get("question", "") reference_answer = qa.get("answer", "") - + # Get key memory points from evidence evidence = qa.get("evidence", []) key_memory_points = "\n".join([e.get("memory_content", "") for e in evidence]) - + # Get system response system_response = qa.get("system_response", "") - + try: # Call LLM for evaluation eval_result = await evaluate_qa_record( @@ -234,28 +234,28 @@ async def reevaluate_session( model_name=model_name, prompt_version=prompt_version ) - + result = { "result_type": eval_result.get("evaluation_result", "Invalid"), "reasoning": eval_result.get("reasoning", "") } - + return idx, model_name, prompt_version, result - + except Exception as e: print(f" ❌ QA[{idx+1}] {model_name}/{prompt_version}: Error: {e}") return idx, model_name, prompt_version, { "result_type": "Error", "reasoning": f"Evaluation error: {str(e)}" } - + # Create all evaluation tasks (all combinations of models, prompts, and QA records) tasks = [] for idx, qa in enumerate(qa_records): for model_name in models: for prompt_version in prompt_versions: tasks.append(evaluate_single_combination(idx, qa, model_name, prompt_version)) - + # Execute evaluations based on parallel mode if parallel: print(f" ⚡ Starting {len(tasks)} parallel evaluations (rate limited by LLM layer)...") @@ -268,18 +268,18 @@ async def reevaluate_session( results.append(result) if i % 10 == 0 or i == len(tasks): print(f" ⏳ Progress: {i}/{len(tasks)} evaluations completed") - + # Organize results by QA index, then by model and prompt # Structure: qa_records[idx]["evaluations"][model][prompt_version] = {result_type, reasoning} for idx, qa in enumerate(qa_records): if "evaluations" not in qa: qa["evaluations"] = {} - + # Initialize evaluations structure for model_name in models: if model_name not in qa["evaluations"]: qa["evaluations"][model_name] = {} - + # Fill in results completed_count = 0 for qa_idx, model_name, prompt_version, result in results: @@ -287,7 +287,7 @@ async def reevaluate_session( completed_count += 1 if completed_count % 10 == 0 or completed_count == len(results): print(f" ✅ Completed {completed_count}/{len(results)} evaluations") - + # Set default result_type to first model's v1 result for compatibility if models and prompt_versions: default_model = models[0] @@ -296,21 +296,21 @@ async def reevaluate_session( default_eval = qa["evaluations"].get(default_model, {}).get(default_prompt, {}) qa["result_type"] = default_eval.get("result_type", "Invalid") qa["question_answering_reasoning"] = default_eval.get("reasoning", "") - + # Update session data if "session" not in session_data: session_data["session"] = {} if "evaluation_results" not in session_data["session"]: session_data["session"]["evaluation_results"] = {} - + session_data["session"]["evaluation_results"]["question_answering_records"] = qa_records - + # Save updated session data with open(session_file, "w", encoding="utf-8") as f: json.dump(session_data, f, ensure_ascii=False, indent=2) - + print(f" 💾 Updated session saved with all evaluations") - + return session_data @@ -328,9 +328,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": 0, "qa_num": 0 } - + correct = hallucination = omission = valid = 0 - + for qa in qa_records: result_type = qa.get("result_type", "") if result_type == "Correct": @@ -342,7 +342,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: elif result_type == "Omission": omission += 1 valid += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -353,7 +353,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": valid, "qa_num": total } - + return metrics @@ -364,28 +364,28 @@ async def main( parallel: bool = True ): """Main function to re-evaluate QA records from data directory with multiple models and prompts. - + Args: data_dir: Path to data directory models: List of model names (e.g., ["qwen3-max", "qwen-flash"]) prompt_versions: List of prompt versions (e.g., ["v1", "v2"]) parallel: If True, use parallel execution; if False, use sequential execution - + Note: Request rate limiting is automatically handled by base_llm.py's request_interval mechanism. """ data_path = Path(data_dir) - + if not data_path.exists() or not data_path.is_dir(): print(f"❌ Error: Directory not found: {data_dir}") return - + # Default values if models is None: models = ["qwen3-max"] if prompt_versions is None: prompt_versions = ["v1"] - + print("=" * 80) print("RE-EVALUATING QA RECORDS WITH MULTIPLE MODELS & PROMPTS") print(f"Models: {', '.join(models)}") @@ -393,27 +393,27 @@ async def main( print(f"Execution mode: {'Parallel' if parallel else 'Sequential'}") print("Note: Request rate limiting handled by LLM layer (base_llm.py)") print("=" * 80 + "\n") - + # Load data from directory sessions = load_from_data_dir(data_dir) - + if not sessions: print(f"❌ No session files found in {data_dir}") return - + print(f"📂 Found {len(sessions)} session files\n") - + # Process each session all_qa_records = [] - + for idx, session_info in enumerate(sessions, 1): session_file = session_info["file"] session_data = session_info["data"] user_name = session_data.get("user_name", "Unknown") session_idx = session_data.get("session_idx", 0) - + print(f"[{idx}/{len(sessions)}] {user_name} - Session {session_idx}") - + updated_session = await reevaluate_session( session_file=session_file, session_data=session_data, @@ -421,25 +421,25 @@ async def main( prompt_versions=prompt_versions, parallel=parallel ) - + # Collect QA records for metrics eval_results = updated_session.get("session", {}).get("evaluation_results", {}) qa_records = eval_results.get("question_answering_records", []) all_qa_records.extend(qa_records) - + print() - + # Compute and display metrics for each model+prompt combination print("=" * 80) print("UPDATED METRICS (BY MODEL & PROMPT)") print("=" * 80 + "\n") - + for model_name in models: for prompt_version in prompt_versions: prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2" print(f"\n📊 {model_name} / {prompt_name}:") print("─" * 80) - + # Extract QA records for this model+prompt combination model_qa_records = [] for qa in all_qa_records: @@ -452,10 +452,10 @@ async def main( "question_answering_reasoning": eval_data.get("reasoning", "") } model_qa_records.append(qa_copy) - + if model_qa_records: metrics = compute_qa_metrics(model_qa_records) - + print(f" Correct (all): {metrics['correct_qa_ratio(all)']:.4f}") print(f" Hallucination (all): {metrics['hallucination_qa_ratio(all)']:.4f}") print(f" Omission (all): {metrics['omission_qa_ratio(all)']:.4f}") @@ -463,17 +463,17 @@ async def main( print(f" Hallucination (valid): {metrics['hallucination_qa_ratio(valid)']:.4f}") print(f" Omission (valid): {metrics['omission_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {metrics['qa_valid_num']}/{metrics['qa_num']}") - + # Save detailed results with all evaluations report_file = data_path.parent / "reme_eval_stat_result_detailed.json" - + # Create summary for each model+prompt combination evaluation_summary = {} for model_name in models: evaluation_summary[model_name] = {} for prompt_version in prompt_versions: prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2" - + # Extract QA records for this combination model_qa_records = [] for qa in all_qa_records: @@ -485,28 +485,28 @@ async def main( "question_answering_reasoning": eval_data.get("reasoning", "") } model_qa_records.append(qa_copy) - + metrics = compute_qa_metrics(model_qa_records) evaluation_summary[model_name][prompt_name] = { "metrics": metrics, "qa_records": model_qa_records } - + final_results = { "evaluation_summary": evaluation_summary, "all_qa_records_with_evaluations": all_qa_records } - + with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=2) - + print(f"\n💾 Detailed results saved: {report_file}") print("\n" + "=" * 80) if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="Re-evaluate QA records from data directory using multiple models and prompts in parallel. " "Request rate limiting is automatically handled by base_llm.py's request_interval mechanism." @@ -539,9 +539,9 @@ if __name__ == "__main__": action="store_true", help="Use sequential execution instead of parallel (default: parallel)" ) - + args = parser.parse_args() - + asyncio.run(main( data_dir=args.data_dir, models=args.models, diff --git a/bench/human_in_the_loop2/compute_qa_stats.py b/bench/human_in_the_loop2/compute_qa_stats.py index 65ac6dc8..69d65e55 100644 --- a/bench/human_in_the_loop2/compute_qa_stats.py +++ b/bench/human_in_the_loop2/compute_qa_stats.py @@ -22,9 +22,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": 0, "qa_num": 0 } - + correct = hallucination = omission = valid = 0 - + for qa in qa_records: result_type = qa.get("result_type", "") if result_type == "Correct": @@ -36,7 +36,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: elif result_type == "Omission": omission += 1 valid += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -47,21 +47,21 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": valid, "qa_num": total } - + return metrics def compute_time_metrics(users_data: list[dict]) -> dict[str, float]: """Compute timing metrics from evaluation results.""" add_duration = search_duration = 0 - + for user_data in users_data: for session in user_data.get("sessions", []): add_duration += session.get("add_dialogue_duration_ms", 0) eval_results = session.get("session", {}).get("evaluation_results", {}) for qa in eval_results.get("question_answering_records", []): search_duration += qa.get("search_duration_ms", 0) - + return { "add_dialogue_duration_time": add_duration / 1000 / 60, "search_memory_duration_time": search_duration / 1000 / 60, @@ -72,21 +72,21 @@ def compute_time_metrics(users_data: list[dict]) -> dict[str, float]: def load_from_tmp_dir(tmp_dir: str) -> list[dict]: """Load data from tmp directory.""" tmp_path = Path(tmp_dir) - + # Try flat file structure first (conversation_{user}_session_{idx}.json) json_files = [f for f in tmp_path.iterdir() if f.is_file() and f.suffix == ".json"] - + if json_files: # Group files by user users_dict = defaultdict(list) - + for json_file in json_files: with open(json_file, "r", encoding="utf-8") as f: session_data = json.load(f) user_name = session_data.get("user_name") if user_name: users_dict[user_name].append(session_data) - + # Sort sessions by session_idx for each user users_data = [] for user_name, sessions in users_dict.items(): @@ -103,71 +103,71 @@ def load_from_tmp_dir(tmp_dir: str) -> list[dict]: session_copy.pop("user_name", None) user_data["sessions"].append(session_copy) users_data.append(user_data) - + return users_data - + # Fallback to directory structure (user_name/session_{idx}.json) user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()] - + users_data = [] for user_dir in user_dirs: session_files = sorted( [f for f in user_dir.iterdir() if "session_" in f.name and f.suffix == ".json"], key=lambda f: int(f.stem.split("_")[-1]) ) - + if not session_files: continue - + 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": [] } - + for session_file in session_files: with open(session_file, "r", encoding="utf-8") as f: session_data = json.load(f) session_data.pop("uuid", None) session_data.pop("user_name", None) user_data["sessions"].append(session_data) - + users_data.append(user_data) - + return users_data def main(tmp_dir: str): """Main function to compute statistics from tmp directory.""" tmp_path = Path(tmp_dir) - + if not tmp_path.exists() or not tmp_path.is_dir(): print(f"❌ Error: Directory not found: {tmp_dir}") return - + # Load data from tmp directory users_data = load_from_tmp_dir(tmp_dir) - + # Collect QA records with metadata qa_records = [] qa_with_metadata = [] user_count = session_count = 0 - + for user_data in users_data: user_count += 1 user_name = user_data.get("user_name", "Unknown") - + valid_session_idx = 0 for session in user_data.get("sessions", []): if session.get("is_generated_qa_session"): continue - + session_count += 1 eval_results = session.get("session", {}).get("evaluation_results", {}) - + for qa_idx, qa in enumerate(eval_results.get("question_answering_records", [])): qa_records.append(qa) qa_with_metadata.append({ @@ -176,17 +176,17 @@ def main(tmp_dir: str): "question_idx": qa_idx, "qa_record": qa }) - + valid_session_idx += 1 - + # Compute metrics qa_metrics = compute_qa_metrics(qa_records) time_metrics = compute_time_metrics(users_data) - + # Save results output_dir = tmp_path.parent report_file = output_dir / "reme_eval_stat_result.json" - + final_results = { "overall_score": { "question_answering": qa_metrics, @@ -194,10 +194,10 @@ def main(tmp_dir: str): }, "question_answering_records": qa_records } - + with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=4) - + # Print summary print(f"\n📊 Data: {user_count} users, {session_count} sessions, {len(qa_records)} QA records") print(f"\n✅ Metrics:") @@ -205,12 +205,12 @@ def main(tmp_dir: str): print(f" Valid: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}") print(f"\n⏱️ Time: {time_metrics['total_duration_time']:.2f} min (Add: {time_metrics['add_dialogue_duration_time']:.2f} | Search: {time_metrics['search_memory_duration_time']:.2f})") print(f"\n💾 Results saved: {report_file}") - + # Print error records print(f"\n{'='*80}\n❌ ERROR RECORDS ({len([r for r in qa_with_metadata if r['qa_record'].get('result_type') not in ['Correct', '']])} errors)\n{'='*80}") - + error_records = [r for r in qa_with_metadata if r["qa_record"].get("result_type") not in ["Correct", ""]] - + if error_records: for idx, record in enumerate(error_records, 1): qa = record["qa_record"] @@ -218,13 +218,13 @@ def main(tmp_dir: str): print(f" Q: {qa.get('question', 'N/A')}") print(f" Expected: {qa.get('answer', 'N/A')}") print(f" Got: {qa.get('system_response', 'N/A')}") - + print() if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser(description="Compute QA statistics from tmp directory") parser.add_argument( "tmp_dir", @@ -232,6 +232,6 @@ if __name__ == "__main__": default="./data", type=str, help="Path to tmp directory containing user session data (default: ./data)") - + args = parser.parse_args() main(tmp_dir=args.tmp_dir) diff --git a/bench/human_in_the_loop2/eval.yaml b/bench/human_in_the_loop2/eval.yaml index 41d36943..3e685bc7 100644 --- a/bench/human_in_the_loop2/eval.yaml +++ b/bench/human_in_the_loop2/eval.yaml @@ -64,7 +64,7 @@ EVALUATION_PROMPT_FOR_QUESTION: | "evaluation_result": "Correct | Hallucination | Omission" }} ``` - + EVALUATION_PROMPT_FOR_QUESTION2: | You are an **evaluation expert for AI memory system question answering**. diff --git a/bench/human_in_the_loop2/reevaluate_qa.py b/bench/human_in_the_loop2/reevaluate_qa.py index a65c4d1c..a5e6d0fd 100644 --- a/bench/human_in_the_loop2/reevaluate_qa.py +++ b/bench/human_in_the_loop2/reevaluate_qa.py @@ -16,8 +16,8 @@ from collections import defaultdict from pathlib import Path from typing import Any -from reme_ai.core.schema import Message -from reme_ai.core.utils import load_env +from reme_ai.core_old.schema import Message +from reme_ai.core_old.utils import load_env from reme_ai.reme import ReMe from tenacity import retry, stop_after_attempt, wait_random_exponential @@ -82,7 +82,7 @@ async def evaluate_qa_record( prompt_version: str = "v1" ) -> dict: """Evaluate a single QA record using LLM with specified prompt version. - + Args: question: The question to evaluate reference_answer: The reference answer @@ -90,9 +90,9 @@ async def evaluate_qa_record( response: System response to evaluate dialogue: Dialogue context (optional) model_name: LLM model name - prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION, + prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION, "v2" for EVALUATION_PROMPT_FOR_QUESTION2 - + Returns: dict with evaluation_result and reasoning """ @@ -101,7 +101,7 @@ async def evaluate_qa_record( prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION2"] else: prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION"] - + # Format prompt prompt = prompt_template.format( question=question, @@ -110,7 +110,7 @@ async def evaluate_qa_record( response=response, dialogue=dialogue or "N/A" ) - + result = await llm_request_for_json(prompt, model_name=model_name) return result @@ -118,14 +118,14 @@ async def evaluate_qa_record( def load_from_data_dir(data_dir: str) -> list[dict]: """Load data from data directory (same as compute_qa_stats.py).""" data_path = Path(data_dir) - + # Try flat file structure first (conversation_{user}_session_{idx}.json) json_files = [f for f in data_path.iterdir() if f.is_file() and f.suffix == ".json"] - + if json_files: # Group files by user users_dict = defaultdict(list) - + for json_file in json_files: with open(json_file, "r", encoding="utf-8") as f: session_data = json.load(f) @@ -135,18 +135,18 @@ def load_from_data_dir(data_dir: str) -> list[dict]: "file": json_file, "data": session_data }) - + # Sort sessions by session_idx for each user users_data = [] for user_name, sessions in users_dict.items(): sessions_sorted = sorted( - sessions, + sessions, key=lambda s: s["data"].get("session_idx", 0) ) users_data.extend(sessions_sorted) - + return users_data - + return [] @@ -155,7 +155,7 @@ def format_dialogue_context(session_data: dict) -> str: dialogue = session_data.get("session", {}).get("dialogue", []) if not dialogue: return "N/A" - + formatted_turns = [] for turn in dialogue: role = turn.get("role", "unknown") @@ -175,34 +175,34 @@ async def reevaluate_session( parallel: bool = True ) -> dict: """Re-evaluate all QA records in a session using multiple models and prompts. - + Args: session_file: Path to session file session_data: Session data dict models: List of model names to use for evaluation prompt_versions: List of prompt versions ("v1", "v2") - parallel: If True, use asyncio.gather for parallel execution; + parallel: If True, use asyncio.gather for parallel execution; if False, execute sequentially - + Returns: Updated session data with evaluation results for each model+prompt combination - + Note: Request rate limiting is handled by base_llm.py's request_interval mechanism. """ eval_results = session_data.get("session", {}).get("evaluation_results", {}) qa_records = eval_results.get("question_answering_records", []) - + if not qa_records: print(f" ⏭️ No QA records found") return session_data - + total_evals = len(models) * len(prompt_versions) * len(qa_records) print(f" 🔍 Re-evaluating {len(qa_records)} QA records with {len(models)} models × {len(prompt_versions)} prompts = {total_evals} evaluations...") - + # Format dialogue context once dialogue_context = format_dialogue_context(session_data) - + async def evaluate_single_combination( idx: int, qa: dict, @@ -210,19 +210,19 @@ async def reevaluate_session( prompt_version: str ) -> tuple[int, str, str, dict]: """Evaluate a single QA record with specific model and prompt. - + Note: Rate limiting is handled by BaseLLM's request_interval mechanism. """ question = qa.get("question", "") reference_answer = qa.get("answer", "") - + # Get key memory points from evidence evidence = qa.get("evidence", []) key_memory_points = "\n".join([e.get("memory_content", "") for e in evidence]) - + # Get system response system_response = qa.get("system_response", "") - + try: # Call LLM for evaluation eval_result = await evaluate_qa_record( @@ -234,28 +234,28 @@ async def reevaluate_session( model_name=model_name, prompt_version=prompt_version ) - + result = { "result_type": eval_result.get("evaluation_result", "Invalid"), "reasoning": eval_result.get("reasoning", "") } - + return idx, model_name, prompt_version, result - + except Exception as e: print(f" ❌ QA[{idx+1}] {model_name}/{prompt_version}: Error: {e}") return idx, model_name, prompt_version, { "result_type": "Error", "reasoning": f"Evaluation error: {str(e)}" } - + # Create all evaluation tasks (all combinations of models, prompts, and QA records) tasks = [] for idx, qa in enumerate(qa_records): for model_name in models: for prompt_version in prompt_versions: tasks.append(evaluate_single_combination(idx, qa, model_name, prompt_version)) - + # Execute evaluations based on parallel mode if parallel: print(f" ⚡ Starting {len(tasks)} parallel evaluations (rate limited by LLM layer)...") @@ -268,18 +268,18 @@ async def reevaluate_session( results.append(result) if i % 10 == 0 or i == len(tasks): print(f" ⏳ Progress: {i}/{len(tasks)} evaluations completed") - + # Organize results by QA index, then by model and prompt # Structure: qa_records[idx]["evaluations"][model][prompt_version] = {result_type, reasoning} for idx, qa in enumerate(qa_records): if "evaluations" not in qa: qa["evaluations"] = {} - + # Initialize evaluations structure for model_name in models: if model_name not in qa["evaluations"]: qa["evaluations"][model_name] = {} - + # Fill in results completed_count = 0 for qa_idx, model_name, prompt_version, result in results: @@ -287,7 +287,7 @@ async def reevaluate_session( completed_count += 1 if completed_count % 10 == 0 or completed_count == len(results): print(f" ✅ Completed {completed_count}/{len(results)} evaluations") - + # Set default result_type to first model's v1 result for compatibility if models and prompt_versions: default_model = models[0] @@ -296,21 +296,21 @@ async def reevaluate_session( default_eval = qa["evaluations"].get(default_model, {}).get(default_prompt, {}) qa["result_type"] = default_eval.get("result_type", "Invalid") qa["question_answering_reasoning"] = default_eval.get("reasoning", "") - + # Update session data if "session" not in session_data: session_data["session"] = {} if "evaluation_results" not in session_data["session"]: session_data["session"]["evaluation_results"] = {} - + session_data["session"]["evaluation_results"]["question_answering_records"] = qa_records - + # Save updated session data with open(session_file, "w", encoding="utf-8") as f: json.dump(session_data, f, ensure_ascii=False, indent=2) - + print(f" 💾 Updated session saved with all evaluations") - + return session_data @@ -328,9 +328,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": 0, "qa_num": 0 } - + correct = hallucination = omission = valid = 0 - + for qa in qa_records: result_type = qa.get("result_type", "") if result_type == "Correct": @@ -342,7 +342,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: elif result_type == "Omission": omission += 1 valid += 1 - + metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, @@ -353,7 +353,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: "qa_valid_num": valid, "qa_num": total } - + return metrics @@ -364,28 +364,28 @@ async def main( parallel: bool = True ): """Main function to re-evaluate QA records from data directory with multiple models and prompts. - + Args: data_dir: Path to data directory models: List of model names (e.g., ["qwen3-max", "qwen-flash"]) prompt_versions: List of prompt versions (e.g., ["v1", "v2"]) parallel: If True, use parallel execution; if False, use sequential execution - + Note: Request rate limiting is automatically handled by base_llm.py's request_interval mechanism. """ data_path = Path(data_dir) - + if not data_path.exists() or not data_path.is_dir(): print(f"❌ Error: Directory not found: {data_dir}") return - + # Default values if models is None: models = ["qwen3-max"] if prompt_versions is None: prompt_versions = ["v1"] - + print("=" * 80) print("RE-EVALUATING QA RECORDS WITH MULTIPLE MODELS & PROMPTS") print(f"Models: {', '.join(models)}") @@ -393,27 +393,27 @@ async def main( print(f"Execution mode: {'Parallel' if parallel else 'Sequential'}") print("Note: Request rate limiting handled by LLM layer (base_llm.py)") print("=" * 80 + "\n") - + # Load data from directory sessions = load_from_data_dir(data_dir) - + if not sessions: print(f"❌ No session files found in {data_dir}") return - + print(f"📂 Found {len(sessions)} session files\n") - + # Process each session all_qa_records = [] - + for idx, session_info in enumerate(sessions, 1): session_file = session_info["file"] session_data = session_info["data"] user_name = session_data.get("user_name", "Unknown") session_idx = session_data.get("session_idx", 0) - + print(f"[{idx}/{len(sessions)}] {user_name} - Session {session_idx}") - + updated_session = await reevaluate_session( session_file=session_file, session_data=session_data, @@ -421,25 +421,25 @@ async def main( prompt_versions=prompt_versions, parallel=parallel ) - + # Collect QA records for metrics eval_results = updated_session.get("session", {}).get("evaluation_results", {}) qa_records = eval_results.get("question_answering_records", []) all_qa_records.extend(qa_records) - + print() - + # Compute and display metrics for each model+prompt combination print("=" * 80) print("UPDATED METRICS (BY MODEL & PROMPT)") print("=" * 80 + "\n") - + for model_name in models: for prompt_version in prompt_versions: prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2" print(f"\n📊 {model_name} / {prompt_name}:") print("─" * 80) - + # Extract QA records for this model+prompt combination model_qa_records = [] for qa in all_qa_records: @@ -452,10 +452,10 @@ async def main( "question_answering_reasoning": eval_data.get("reasoning", "") } model_qa_records.append(qa_copy) - + if model_qa_records: metrics = compute_qa_metrics(model_qa_records) - + print(f" Correct (all): {metrics['correct_qa_ratio(all)']:.4f}") print(f" Hallucination (all): {metrics['hallucination_qa_ratio(all)']:.4f}") print(f" Omission (all): {metrics['omission_qa_ratio(all)']:.4f}") @@ -463,17 +463,17 @@ async def main( print(f" Hallucination (valid): {metrics['hallucination_qa_ratio(valid)']:.4f}") print(f" Omission (valid): {metrics['omission_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {metrics['qa_valid_num']}/{metrics['qa_num']}") - + # Save detailed results with all evaluations report_file = data_path.parent / "reme_eval_stat_result_detailed.json" - + # Create summary for each model+prompt combination evaluation_summary = {} for model_name in models: evaluation_summary[model_name] = {} for prompt_version in prompt_versions: prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2" - + # Extract QA records for this combination model_qa_records = [] for qa in all_qa_records: @@ -485,28 +485,28 @@ async def main( "question_answering_reasoning": eval_data.get("reasoning", "") } model_qa_records.append(qa_copy) - + metrics = compute_qa_metrics(model_qa_records) evaluation_summary[model_name][prompt_name] = { "metrics": metrics, "qa_records": model_qa_records } - + final_results = { "evaluation_summary": evaluation_summary, "all_qa_records_with_evaluations": all_qa_records } - + with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=2) - + print(f"\n💾 Detailed results saved: {report_file}") print("\n" + "=" * 80) if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser( description="Re-evaluate QA records from data directory using multiple models and prompts in parallel. " "Request rate limiting is automatically handled by base_llm.py's request_interval mechanism." @@ -539,9 +539,9 @@ if __name__ == "__main__": action="store_true", help="Use sequential execution instead of parallel (default: parallel)" ) - + args = parser.parse_args() - + asyncio.run(main( data_dir=args.data_dir, models=args.models, diff --git a/docs/todo.md b/docs/todo.md new file mode 100644 index 00000000..5d73e04d --- /dev/null +++ b/docs/todo.md @@ -0,0 +1,3 @@ +1. 如何更好的注册class +2. op的返回，使用return 还是 self.output +3. 如何把agent的东西放出来 \ No newline at end of file diff --git a/reme_ai/core/__init__.py b/reme_ai/core/__init__.py index 8eab5792..e69de29b 100644 --- a/reme_ai/core/__init__.py +++ b/reme_ai/core/__init__.py @@ -1,17 +0,0 @@ -"""Core module for ReMe AI framework.""" - -# pylint: disable=wrong-import-position -# flake8: noqa: F401 - -from . import config -from . import context -from . import embedding -from . import enumeration -from . import flow -from . import llm -from . import op -from . import schema -from . import service -from . import token_counter -from . import utils -from . import vector_store diff --git a/reme_ai/core/context/prompt_handler.py b/reme_ai/core/context/prompt_handler.py index b428b163..9639cbf8 100644 --- a/reme_ai/core/context/prompt_handler.py +++ b/reme_ai/core/context/prompt_handler.py @@ -1,24 +1,100 @@ -"""Module for managing and formatting prompt templates from files or dictionaries.""" +"""Module for managing and formatting prompt templates from files or dictionaries. +This module provides a PromptHandler class that: +- Loads prompts from YAML/JSON files or dictionaries +- Supports multi-language prompts with automatic suffix handling +- Provides conditional line filtering using boolean flags +- Formats prompts with template variable substitution +- Validates format strings and provides helpful error messages +""" + +import json from pathlib import Path +from string import Formatter +from typing import Any, Dict, Optional, Union import yaml from loguru import logger from .base_context import BaseContext -from .service_context import C + + +class PromptNotFoundError(KeyError): + """Exception raised when a requested prompt template is not found.""" + + def __init__(self, prompt_name: str, available_prompts: list[str]): + self.prompt_name = prompt_name + self.available_prompts = available_prompts + super().__init__( + f"Prompt '{prompt_name}' not found. " + f"Available prompts: {', '.join(available_prompts[:10])}" + f"{'...' if len(available_prompts) > 10 else ''}" + ) + + +class PromptFormattingError(ValueError): + """Exception raised when prompt formatting fails.""" + pass class PromptHandler(BaseContext): - """A context-aware handler for loading, retrieving, and formatting prompt templates.""" + """A context-aware handler for loading, retrieving, and formatting prompt templates. + + This handler supports: + - Loading prompts from YAML/JSON files or dictionaries + - Multi-language prompt support with automatic language suffix + - Conditional line filtering using boolean flags (e.g., [debug], [verbose]) + - Template variable substitution with validation + - Method chaining for fluent API + + Examples: + >>> handler = PromptHandler(language="en") + >>> handler.load_prompt_dict({ + ... "greeting_en": "Hello, {name}!", + ... "farewell_en": "[debug]Debug mode\\nGoodbye, {name}!" + ... }) + >>> handler.prompt_format("greeting", name="Alice") + 'Hello, Alice!' + >>> handler.prompt_format("farewell", name="Bob", debug=False) + 'Goodbye, Bob!' + """ def __init__(self, language: str = "", **kwargs): - """Initialize the handler with a specific language and optional context data.""" + """Initialize the PromptHandler with optional language configuration. + + Args: + language: Language code to append as suffix (e.g., "en", "zh", "ja"). + If provided, get_prompt will automatically try to find + prompts with this suffix (e.g., "greeting" -> "greeting_en"). + **kwargs: Additional key-value pairs to initialize the context. + """ super().__init__(**kwargs) - self.language: str = language or C.language + self.language: str = language.strip() - def load_prompt_by_file(self, prompt_file_path: Path | str = None): - """Load prompt configurations from a YAML file into the context.""" + def load_prompt_by_file( + self, + prompt_file_path: Optional[Union[Path, str]] = None, + overwrite: bool = True + ) -> "PromptHandler": + """Load prompt configurations from a YAML or JSON file into the context. + + Supports both YAML (.yaml, .yml) and JSON (.json) file formats. + Non-existent files are silently skipped. + + Args: + prompt_file_path: Path to the prompt configuration file. + If None, returns self without changes. + overwrite: If True, allows overwriting existing prompts with warnings. + If False, skips existing prompts without overwriting. + + Returns: + Self for method chaining. + + Raises: + ValueError: If file format is not supported. + yaml.YAMLError: If YAML parsing fails. + json.JSONDecodeError: If JSON parsing fails. + """ if prompt_file_path is None: return self @@ -26,70 +102,272 @@ class PromptHandler(BaseContext): prompt_file_path = Path(prompt_file_path) if not prompt_file_path.exists(): + logger.warning(f"Prompt file not found: {prompt_file_path}") return self - with prompt_file_path.open(encoding="utf-8") as f: - # Load YAML content using the full loader - prompt_dict = yaml.load(f, yaml.FullLoader) - self.load_prompt_dict(prompt_dict) + suffix = prompt_file_path.suffix.lower() + + try: + with prompt_file_path.open(encoding="utf-8") as f: + if suffix in [".yaml", ".yml"]: + prompt_dict = yaml.safe_load(f) + elif suffix == ".json": + prompt_dict = json.load(f) + else: + raise ValueError( + f"Unsupported file format: {suffix}. " + f"Supported formats: .yaml, .yml, .json" + ) + + logger.info(f"Loaded {len(prompt_dict or {})} prompts from {prompt_file_path}") + self.load_prompt_dict(prompt_dict, overwrite=overwrite) + + except (yaml.YAMLError, json.JSONDecodeError) as e: + logger.error(f"Failed to parse prompt file {prompt_file_path}: {e}") + raise + return self - def load_prompt_dict(self, prompt_dict: dict = None): - """Merge a dictionary of prompt strings into the current context.""" + def load_prompt_dict( + self, + prompt_dict: Optional[Dict[str, Any]] = None, + overwrite: bool = True + ) -> "PromptHandler": + """Merge a dictionary of prompt strings into the current context. + + Only string values are stored as prompts. Non-string values are skipped. + + Args: + prompt_dict: Dictionary mapping prompt names to prompt template strings. + overwrite: If True, allows overwriting existing prompts with warnings. + If False, skips existing prompts without overwriting. + + Returns: + Self for method chaining. + """ if not prompt_dict: return self for key, value in prompt_dict.items(): - if isinstance(value, str): - if key in self: - logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}") + if not isinstance(value, str): + logger.debug(f"Skipping non-string prompt: key={key}, type={type(value)}") + continue + + if key in self: + if overwrite: + logger.warning( + f"Overwriting prompt '{key}': " + f"old length={len(self[key])}, new length={len(value)}" + ) + self[key] = value else: - logger.debug(f"Adding new prompt key={key}, value={value}") + logger.debug(f"Skipping existing prompt: key={key}") + else: + logger.debug(f"Adding new prompt: key={key}, length={len(value)}") self[key] = value + return self - def get_prompt(self, prompt_name: str): - """Retrieve a prompt by name, automatically appending the language suffix if needed.""" - key: str = prompt_name - if self.language and not key.endswith(self.language.strip()): - key += "_" + self.language.strip() + def get_prompt(self, prompt_name: str, fallback_to_base: bool = True) -> str: + """Retrieve a prompt by name with automatic language suffix handling. + + If a language is configured, this method will: + 1. First try to find the prompt with language suffix (e.g., "greeting_en") + 2. If not found and fallback_to_base is True, try the base name (e.g., "greeting") + 3. Otherwise, raise PromptNotFoundError + + Args: + prompt_name: Name of the prompt to retrieve. + fallback_to_base: If True and language-specific prompt not found, + fallback to prompt without language suffix. + + Returns: + The prompt template string, stripped of leading/trailing whitespace. + + Raises: + PromptNotFoundError: If the prompt is not found. + """ + # Try with language suffix first + if self.language and not prompt_name.endswith(f"_{self.language}"): + key_with_lang = f"{prompt_name}_{self.language}" + if key_with_lang in self: + return self[key_with_lang].strip() - assert key in self, f"prompt_name={key} not found." - return self[key].strip() + # Try base name + if prompt_name in self: + return self[prompt_name].strip() - def prompt_format(self, prompt_name: str, **kwargs) -> str: - """Format a prompt by filtering flagged lines and filling template variables.""" - prompt = self.get_prompt(prompt_name) + # Try fallback if enabled + if fallback_to_base and self.language: + # Check if prompt_name already has language suffix, try without it + if prompt_name.endswith(f"_{self.language}"): + base_name = prompt_name[: -(len(self.language) + 1)] + if base_name in self: + return self[base_name].strip() - # Separate boolean flags from string formatting arguments - flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} - other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + # Not found, raise error with helpful message + available = list(self.keys()) + raise PromptNotFoundError(prompt_name, available) - if flag_kwargs: - split_prompt = [] - for line in prompt.strip().split("\n"): - hit = False - hit_flag = True - for key, flag in flag_kwargs.items(): - if not line.startswith(f"[{key}]"): - continue + def has_prompt(self, prompt_name: str) -> bool: + """Check if a prompt exists (with or without language suffix). + + Args: + prompt_name: Name of the prompt to check. + + Returns: + True if the prompt exists, False otherwise. + """ + try: + self.get_prompt(prompt_name) + return True + except PromptNotFoundError: + return False - hit = True - hit_flag = flag - # Remove the flag prefix from the line - line = line.strip(f"[{key}]") + def list_prompts(self, language_filter: Optional[str] = None) -> list[str]: + """List all available prompt names. + + Args: + language_filter: If provided, only return prompts for this language. + If None, return all prompts. + + Returns: + List of prompt names. + """ + if language_filter is None: + return list(self.keys()) + + suffix = f"_{language_filter.strip()}" + return [key for key in self.keys() if key.endswith(suffix)] + + @staticmethod + def _extract_format_fields(template: str) -> set[str]: + """Extract all format field names from a template string. + + Args: + template: Template string with {variable} placeholders. + + Returns: + Set of field names used in the template. + """ + return { + field_name + for _, field_name, _, _ in Formatter().parse(template) + if field_name is not None + } + + @staticmethod + def _filter_conditional_lines(prompt: str, flags: Dict[str, bool]) -> str: + """Filter lines based on boolean flags. + + Lines starting with [flag_name] are conditionally included based on + the value of flags[flag_name]. If True, the line is included (without + the flag marker). If False, the line is excluded. + + Args: + prompt: The prompt text with conditional markers. + flags: Dictionary of flag names to boolean values. + + Returns: + Filtered prompt text. + """ + filtered_lines = [] + + for line in prompt.split("\n"): + # Check each flag + matched_flag = None + for flag_name in flags: + marker = f"[{flag_name}]" + if line.startswith(marker): + matched_flag = flag_name break - # Include line if no flag is present or if the flag evaluates to True - if not hit: - split_prompt.append(line) - elif hit_flag: - split_prompt.append(line) + if matched_flag is None: + # No flag marker, always include + filtered_lines.append(line) + elif flags[matched_flag]: + # Flag is True, include without marker + marker = f"[{matched_flag}]" + filtered_lines.append(line[len(marker):]) + # else: Flag is False, skip this line - prompt = "\n".join(split_prompt) + return "\n".join(filtered_lines) - if other_kwargs: - # Apply standard Python string formatting - prompt = prompt.format(**other_kwargs) + def prompt_format( + self, + prompt_name: str, + validate: bool = True, + **kwargs + ) -> str: + """Format a prompt with conditional line filtering and variable substitution. + + This method performs two-stage formatting: + 1. Conditional line filtering: Lines marked with [flag] are included only + if the corresponding boolean kwarg is True. + 2. Variable substitution: Template variables {var} are replaced with + provided values. + + Args: + prompt_name: Name of the prompt to format. + validate: If True, check that all required template variables are provided. + **kwargs: Keyword arguments for formatting. Boolean values are treated as + conditional flags, other values are used for template substitution. + + Returns: + Formatted prompt string. + + Raises: + PromptNotFoundError: If the prompt is not found. + PromptFormattingError: If validation fails or formatting errors occur. + + Examples: + >>> handler = PromptHandler() + >>> handler["test"] = "[debug]Debug: {info}\\nResult: {value}" + >>> handler.prompt_format("test", debug=False, info="test", value=42) + 'Result: 42' + >>> handler.prompt_format("test", debug=True, info="test", value=42) + 'Debug: test\\nResult: 42' + """ + # Get the prompt template + prompt = self.get_prompt(prompt_name) - return prompt + # Separate boolean flags from format variables + flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + format_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + + # Step 1: Filter conditional lines + if flag_kwargs: + prompt = self._filter_conditional_lines(prompt, flag_kwargs) + + # Step 2: Validate required fields if requested + if validate: + required_fields = self._extract_format_fields(prompt) + missing_fields = required_fields - set(format_kwargs.keys()) + + if missing_fields: + raise PromptFormattingError( + f"Missing required format variables for prompt '{prompt_name}': " + f"{', '.join(sorted(missing_fields))}" + ) + + # Step 3: Format with variables + try: + if format_kwargs: + prompt = prompt.format(**format_kwargs) + except KeyError as e: + raise PromptFormattingError( + f"Format error in prompt '{prompt_name}': missing variable {e}" + ) from e + except (ValueError, IndexError) as e: + raise PromptFormattingError( + f"Format error in prompt '{prompt_name}': {e}" + ) from e + + return prompt.strip() + + def __repr__(self) -> str: + """Return a string representation of the PromptHandler.""" + return ( + f"PromptHandler(language='{self.language}', " + f"num_prompts={len(self)})" + ) diff --git a/reme_ai/core/context/registry.py b/reme_ai/core/context/registry.py index f403037d..9fa35d30 100644 --- a/reme_ai/core/context/registry.py +++ b/reme_ai/core/context/registry.py @@ -1,19 +1,143 @@ """Module providing a registry class for managing class-to-name mappings via decorators.""" +import inspect +from typing import Callable, TypeVar + from .base_context import BaseContext +from ..enumeration import RegistryEnum +from ...core_old.utils import singleton + +T = TypeVar("T") +@singleton class Registry(BaseContext): - """A registry container that uses decorators to map and store class references.""" + """A singleton registry manager that maintains separate registries for different component types. - def register(self, name: str = "", add_cls: bool = True): - """Return a decorator that registers a class under a specific name in the registry.""" + This class serves as the central registry hub for the entire ReMe application, providing: + - Component registration for different types (LLMs, embeddings, vector stores, etc.) + - Convenient access methods for retrieving registered classes + - Decorator-based registration API - def decorator(cls): - if add_cls: - # Use provided name or default to the class name as the key - key = name or cls.__name__ - self[key] = cls - return cls + The singleton pattern ensures only one instance exists throughout the application lifecycle, + accessible via the global `R` variable exported at the bottom of this module. + """ - return decorator + def __init__(self, **kwargs): + """Initialize the registry manager with separate registries for each component type.""" + super().__init__(**kwargs) + + # Registry system: stores class definitions for different component types + self.registry_dict: dict[RegistryEnum, dict] = { + v: {} for v in RegistryEnum.__members__.values() + } + + def register(self, name: str | type = "", register_type: RegistryEnum = None) -> Callable[[type[T]], type[T]] | type[T]: + """Return a decorator to register a component within a specific registry category. + + Can be used in multiple ways: + - @R.register_op() # with empty parentheses, uses class name + - @R.register_op # without parentheses, uses class name + - @R.register_op("custom_name") # with custom name + + Args: + name: Either a string name for the class, or the class itself when used without parentheses + register_type: The type of registry (LLM, EMBEDDING_MODEL, VECTOR_STORE, etc.) + + Returns: + Either a decorator function or the registered class itself + + Example: + @R.register("my_llm", RegistryEnum.LLM) + class MyLLM(BaseLLM): + pass + """ + if inspect.isclass(name): + # Used without parentheses: @R.register_op + self.registry_dict[register_type][name.__name__] = name + return name + else: + # Used with parentheses: @R.register_op() or @R.register_op("name") + def decorator(cls): + key = name if isinstance(name, str) and name else cls.__name__ + self.registry_dict[register_type][key] = cls + return cls + + return decorator + + def register_llm(self, name: str = ""): + """Register a Large Language Model class.""" + return self.register(name=name, register_type=RegistryEnum.LLM) + + def register_embedding_model(self, name: str = ""): + """Register an embedding model class.""" + return self.register(name=name, register_type=RegistryEnum.EMBEDDING_MODEL) + + def register_vector_store(self, name: str = ""): + """Register a vector store implementation class.""" + return self.register(name=name, register_type=RegistryEnum.VECTOR_STORE) + + def register_op(self, name: str = ""): + """Register an operation (Op) class.""" + return self.register(name=name, register_type=RegistryEnum.OP) + + def register_flow(self, name: str = ""): + """Register a workflow or logic flow class.""" + return self.register(name=name, register_type=RegistryEnum.FLOW) + + def register_service(self, name: str = ""): + """Register a backend service class.""" + return self.register(name=name, register_type=RegistryEnum.SERVICE) + + def register_token_counter(self, name: str = ""): + """Register a token counting utility class.""" + return self.register(name=name, register_type=RegistryEnum.TOKEN_COUNTER) + + def get_model_class(self, name: str, register_type: RegistryEnum): + """Retrieve a registered class by name from a specific registry category. + + Args: + name: The registration name of the class + register_type: The type of registry to search in + + Returns: + The registered class (not an instance, but the class itself) + + Raises: + AssertionError: If the class is not found in the registry + """ + assert name in self.registry_dict[register_type], f"{name} not in registry_dict[{register_type}]" + return self.registry_dict[register_type][name] + + def get_llm_class(self, name: str): + """Get the LLM class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.LLM) + + def get_embedding_model_class(self, name: str): + """Get the embedding model class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL) + + def get_vector_store_class(self, name: str): + """Get the vector store class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.VECTOR_STORE) + + def get_op_class(self, name: str): + """Get the operation class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.OP) + + def get_flow_class(self, name: str): + """Get the flow class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.FLOW) + + def get_service_class(self, name: str): + """Get the service class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.SERVICE) + + def get_token_counter_class(self, name: str): + """Get the token counter class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER) + + +# Export a global singleton instance for easy access across the application +# This is the primary way to access the registry throughout the codebase +R = Registry() diff --git a/reme_ai/core/enumeration/json_schema_enum.py b/reme_ai/core/enumeration/json_schema_enum.py index 507645f4..d66882e2 100644 --- a/reme_ai/core/enumeration/json_schema_enum.py +++ b/reme_ai/core/enumeration/json_schema_enum.py @@ -1,18 +1,38 @@ -"""Defines the standard data types supported by JSON Schema.""" +"""Defines the standard data types supported by JSON Schema. + +This enum maps common JSON Schema primitive types to their corresponding +Python runtime types, and provides a convenient string representation +compatible with JSON Schema (`"string"`, `"number"`, etc.). +""" from enum import Enum class JsonSchemaEnum(Enum): - """Enumeration of valid JSON Schema data types.""" + """Enumeration of valid JSON Schema data types. + The enum value is the corresponding Python type, while the string + representation (`str(...)`) is the canonical JSON Schema type name. + """ + + # Textual data STRING = str + + # Numeric values, including integers and floats NUMBER = float + + # Integer-only numeric values INTEGER = int + + # JSON objects (key-value mappings) OBJECT = dict + + # Ordered JSON lists/arrays ARRAY = list + + # Boolean values: true / false BOOLEAN = bool def __str__(self) -> str: - """Returns the string representation of the enum value.""" + """Return the lowercase JSON Schema type name for this enum member.""" return self.name.lower() diff --git a/reme_ai/core/enumeration/memory_type.py b/reme_ai/core/enumeration/memory_type.py index 22d35481..b9f5ed29 100644 --- a/reme_ai/core/enumeration/memory_type.py +++ b/reme_ai/core/enumeration/memory_type.py @@ -1,25 +1,33 @@ -"""Memory type enumeration for the three-layer memory architecture.""" +"""Defines the high-level categories of memory managed by ReMe. + +This enumeration is used across the system to tag, route, and store different +kinds of memories (identity, personal context, procedures, tools, etc.). +""" from enum import Enum class MemoryType(str, Enum): - """ - Three-layer memory architecture for agent memory management. + """Enumeration of memory categories used by the memory subsystem. - Layer 1 - High-level Abstraction Memory: - - IDENTITY: Self-cognition (identity, personality, current state) - - PERSONAL: Person-specific memory (preferences and context about specific individuals) - - PROCEDURAL: Procedural memory (how-to knowledge, e.g., 4 steps to write financial reports) - - TOOL: Tool memory (tool usage patterns, success rates, token consumption, latency) - - Layer 2 - Summary Memory (Compressed): Summarized digest of raw message history - Layer 3 - History Memory (Raw): Raw message history + These types describe *what* a piece of memory is about, which guides + storage, retrieval, and summarization strategies. """ + # Long‑term, relatively stable attributes about the user (name, roles, etc.) IDENTITY = "identity" + + # User-specific preferences, habits, and evolving personal context PERSONAL = "personal" + + # How‑to knowledge, workflows, and step‑by‑step instructions PROCEDURAL = "procedural" + + # Information learned about tools, APIs, and their usage patterns TOOL = "tool" + + # Condensed representation of larger memory collections SUMMARY = "summary" + + # Raw chronological interaction history, typically before summarization HISTORY = "history" diff --git a/reme_ai/core_old/__init__.py b/reme_ai/core_old/__init__.py new file mode 100644 index 00000000..8eab5792 --- /dev/null +++ b/reme_ai/core_old/__init__.py @@ -0,0 +1,17 @@ +"""Core module for ReMe AI framework.""" + +# pylint: disable=wrong-import-position +# flake8: noqa: F401 + +from . import config +from . import context +from . import embedding +from . import enumeration +from . import flow +from . import llm +from . import op +from . import schema +from . import service +from . import token_counter +from . import utils +from . import vector_store diff --git a/reme_ai/core/application.py b/reme_ai/core_old/application.py similarity index 100% rename from reme_ai/core/application.py rename to reme_ai/core_old/application.py diff --git a/reme_ai/core/config/__init__.py b/reme_ai/core_old/config/__init__.py similarity index 100% rename from reme_ai/core/config/__init__.py rename to reme_ai/core_old/config/__init__.py diff --git a/reme_ai/core/config/default.yaml b/reme_ai/core_old/config/default.yaml similarity index 100% rename from reme_ai/core/config/default.yaml rename to reme_ai/core_old/config/default.yaml diff --git a/reme_ai/core/config/reme_config_parser.py b/reme_ai/core_old/config/reme_config_parser.py similarity index 100% rename from reme_ai/core/config/reme_config_parser.py rename to reme_ai/core_old/config/reme_config_parser.py diff --git a/reme_ai/core_old/context/__init__.py b/reme_ai/core_old/context/__init__.py new file mode 100644 index 00000000..7f26d600 --- /dev/null +++ b/reme_ai/core_old/context/__init__.py @@ -0,0 +1,16 @@ +"""context""" + +from .base_context import BaseContext +from .prompt_handler import PromptHandler +from .registry import Registry +from .runtime_context import RuntimeContext +from .service_context import ServiceContext, C + +__all__ = [ + "BaseContext", + "PromptHandler", + "Registry", + "RuntimeContext", + "ServiceContext", + "C", +] diff --git a/reme_ai/core_old/context/base_context.py b/reme_ai/core_old/context/base_context.py new file mode 100644 index 00000000..dabd8cdb --- /dev/null +++ b/reme_ai/core_old/context/base_context.py @@ -0,0 +1,41 @@ +"""Module providing a dictionary subclass with attribute-style access and pickling support.""" + +from typing import Generic, TypeVar + +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") + + +class BaseContext(dict, Generic[_KT, _VT]): + """A dictionary subclass that enables accessing and modifying keys as attributes.""" + + def __getattr__(self, name: str) -> _VT: + """Retrieve a dictionary item as an attribute.""" + try: + return self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __setattr__(self, name: str, value: _VT) -> None: + """Assign a value to a dictionary item using attribute syntax.""" + self[name] = value + + def __delattr__(self, name: str) -> None: + """Remove a dictionary item using attribute syntax.""" + try: + # Delete item from dict via key + del self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __getstate__(self) -> dict: + """Return the dictionary representation for pickling.""" + return dict(self) + + def __setstate__(self, state: dict) -> None: + """Restore the dictionary state from a pickled object.""" + self.update(state) + + def __reduce__(self): + """Define the reconstruction logic for pickling processes.""" + return self.__class__, (), self.__getstate__() diff --git a/reme_ai/core_old/context/prompt_handler.py b/reme_ai/core_old/context/prompt_handler.py new file mode 100644 index 00000000..b428b163 --- /dev/null +++ b/reme_ai/core_old/context/prompt_handler.py @@ -0,0 +1,95 @@ +"""Module for managing and formatting prompt templates from files or dictionaries.""" + +from pathlib import Path + +import yaml +from loguru import logger + +from .base_context import BaseContext +from .service_context import C + + +class PromptHandler(BaseContext): + """A context-aware handler for loading, retrieving, and formatting prompt templates.""" + + def __init__(self, language: str = "", **kwargs): + """Initialize the handler with a specific language and optional context data.""" + super().__init__(**kwargs) + self.language: str = language or C.language + + def load_prompt_by_file(self, prompt_file_path: Path | str = None): + """Load prompt configurations from a YAML file into the context.""" + if prompt_file_path is None: + return self + + if isinstance(prompt_file_path, str): + prompt_file_path = Path(prompt_file_path) + + if not prompt_file_path.exists(): + return self + + with prompt_file_path.open(encoding="utf-8") as f: + # Load YAML content using the full loader + prompt_dict = yaml.load(f, yaml.FullLoader) + self.load_prompt_dict(prompt_dict) + return self + + def load_prompt_dict(self, prompt_dict: dict = None): + """Merge a dictionary of prompt strings into the current context.""" + if not prompt_dict: + return self + + for key, value in prompt_dict.items(): + if isinstance(value, str): + if key in self: + logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}") + else: + logger.debug(f"Adding new prompt key={key}, value={value}") + self[key] = value + return self + + def get_prompt(self, prompt_name: str): + """Retrieve a prompt by name, automatically appending the language suffix if needed.""" + key: str = prompt_name + if self.language and not key.endswith(self.language.strip()): + key += "_" + self.language.strip() + + assert key in self, f"prompt_name={key} not found." + return self[key].strip() + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Format a prompt by filtering flagged lines and filling template variables.""" + prompt = self.get_prompt(prompt_name) + + # Separate boolean flags from string formatting arguments + flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + + if flag_kwargs: + split_prompt = [] + for line in prompt.strip().split("\n"): + hit = False + hit_flag = True + for key, flag in flag_kwargs.items(): + if not line.startswith(f"[{key}]"): + continue + + hit = True + hit_flag = flag + # Remove the flag prefix from the line + line = line.strip(f"[{key}]") + break + + # Include line if no flag is present or if the flag evaluates to True + if not hit: + split_prompt.append(line) + elif hit_flag: + split_prompt.append(line) + + prompt = "\n".join(split_prompt) + + if other_kwargs: + # Apply standard Python string formatting + prompt = prompt.format(**other_kwargs) + + return prompt diff --git a/reme_ai/core_old/context/registry.py b/reme_ai/core_old/context/registry.py new file mode 100644 index 00000000..08ea2271 --- /dev/null +++ b/reme_ai/core_old/context/registry.py @@ -0,0 +1,46 @@ +"""Module providing a registry class for managing class-to-name mappings via decorators.""" + +import inspect +from typing import Callable, TypeVar + +from .base_context import BaseContext + +T = TypeVar('T') + + +class Registry(BaseContext): + """A registry container that uses decorators to map and store class references.""" + + def register(self, name: str | type = "", add_cls: bool = True) -> Callable[[type[T]], type[T]] | type[T]: + """Return a decorator that registers a class under a specific name in the registry. + + Can be used in three ways: + - @C.register_op() # with empty parentheses, uses class name + - @C.register_op # without parentheses, uses class name + - @C.register_op("custom_name") # with custom name + + Args: + name: Either a string name for the class, or the class itself when used without parentheses + add_cls: Whether to actually add the class to the registry + + Returns: + Either a decorator function or the registered class itself + """ + + def decorator(cls): + if add_cls: + # Use provided name or default to the class name as the key + key = name if isinstance(name, str) and name else cls.__name__ + self[key] = cls + return cls + + # If used without parentheses: @C.register_op + if inspect.isclass(name): + cls = name + # Register with class name as key + if add_cls: + self[cls.__name__] = cls + return cls + + # If used with parentheses: @C.register_op() or @C.register_op("name") + return decorator diff --git a/reme_ai/core_old/context/runtime_context.py b/reme_ai/core_old/context/runtime_context.py new file mode 100644 index 00000000..d7112e1c --- /dev/null +++ b/reme_ai/core_old/context/runtime_context.py @@ -0,0 +1,79 @@ +"""Runtime context for managing response states and asynchronous data streaming.""" + +import asyncio + +from .base_context import BaseContext +from ..enumeration import ChunkEnum +from ..schema import Response, StreamChunk + + +class RuntimeContext(BaseContext): + """Context for execution state, response metadata, and stream queues.""" + + def __init__( + self, + response: Response | None = None, + stream_queue: asyncio.Queue | None = None, + **kwargs, + ): + """Initialize the context with optional response and queue.""" + super().__init__(**kwargs) + self.response = response or Response() + self.stream_queue = stream_queue + + @classmethod + def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext": + """Create a new context from an existing instance or keywords.""" + if context is None: + return cls(**kwargs) + + context.update(kwargs) + return context + + async def _enqueue(self, chunk: StreamChunk) -> None: + """Internal helper to put a chunk into the queue if it exists.""" + if self.stream_queue: + await self.stream_queue.put(chunk) + + async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext": + """Enqueue a stream chunk from a raw string and type.""" + await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk)) + return self + + async def add_stream_chunk(self, stream_chunk: StreamChunk) -> "RuntimeContext": + """Enqueue an existing stream chunk.""" + await self._enqueue(stream_chunk) + return self + + async def add_stream_done(self) -> "RuntimeContext": + """Enqueue a termination chunk to signal the end of the stream.""" + await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True)) + return self + + def add_response_error(self, e: Exception) -> "RuntimeContext": + """Record an exception into the response object.""" + self.response.success = False + self.response.answer = str(e) + return self + + def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext": + """Copy internal values based on a source-to-target key map.""" + if not mapping: + return self + + for source, target in mapping.items(): + if source in self: + self[target] = self[source] + return self + + def validate_required_keys(self, required_keys: dict[str, bool], context_name: str = "context") -> "RuntimeContext": + """Ensure all required keys are present in the context. + + Args: + required_keys: Dictionary mapping key names to boolean indicating if required + context_name: Name of the context for error messages (e.g., operator name) + """ + for key, is_required in required_keys.items(): + if is_required and key not in self: + raise ValueError(f"{context_name}: missing required input '{key}'") + return self diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core_old/context/service_context.py similarity index 100% rename from reme_ai/core/context/service_context.py rename to reme_ai/core_old/context/service_context.py diff --git a/reme_ai/core/embedding/__init__.py b/reme_ai/core_old/embedding/__init__.py similarity index 100% rename from reme_ai/core/embedding/__init__.py rename to reme_ai/core_old/embedding/__init__.py diff --git a/reme_ai/core/embedding/base_embedding_model.py b/reme_ai/core_old/embedding/base_embedding_model.py similarity index 100% rename from reme_ai/core/embedding/base_embedding_model.py rename to reme_ai/core_old/embedding/base_embedding_model.py diff --git a/reme_ai/core/embedding/openai_embedding_model.py b/reme_ai/core_old/embedding/openai_embedding_model.py similarity index 100% rename from reme_ai/core/embedding/openai_embedding_model.py rename to reme_ai/core_old/embedding/openai_embedding_model.py diff --git a/reme_ai/core/embedding/openai_embedding_model_sync.py b/reme_ai/core_old/embedding/openai_embedding_model_sync.py similarity index 100% rename from reme_ai/core/embedding/openai_embedding_model_sync.py rename to reme_ai/core_old/embedding/openai_embedding_model_sync.py diff --git a/reme_ai/core_old/enumeration/__init__.py b/reme_ai/core_old/enumeration/__init__.py new file mode 100644 index 00000000..7202a949 --- /dev/null +++ b/reme_ai/core_old/enumeration/__init__.py @@ -0,0 +1,17 @@ +"""enumeration""" + +from .chunk_enum import ChunkEnum +from .http_enum import HttpEnum +from .json_schema_enum import JsonSchemaEnum +from .memory_type import MemoryType +from .registry_enum import RegistryEnum +from .role import Role + +__all__ = [ + "ChunkEnum", + "HttpEnum", + "JsonSchemaEnum", + "MemoryType", + "RegistryEnum", + "Role", +] diff --git a/reme_ai/core_old/enumeration/chunk_enum.py b/reme_ai/core_old/enumeration/chunk_enum.py new file mode 100644 index 00000000..dbe37106 --- /dev/null +++ b/reme_ai/core_old/enumeration/chunk_enum.py @@ -0,0 +1,25 @@ +"""Defines the types of data chunks used in streaming responses.""" + +from enum import Enum + + +class ChunkEnum(str, Enum): + """Enumeration of possible chunk categories for stream processing.""" + + # Internal reasoning or chain-of-thought process + THINK = "think" + + # The final generated response content + ANSWER = "answer" + + # Metadata or calls related to external tools + TOOL = "tool" + + # Resource consumption and token usage statistics + USAGE = "usage" + + # Error messages or exception details + ERROR = "error" + + # Final signal indicating the completion of the stream + DONE = "done" diff --git a/reme_ai/core_old/enumeration/http_enum.py b/reme_ai/core_old/enumeration/http_enum.py new file mode 100644 index 00000000..19622242 --- /dev/null +++ b/reme_ai/core_old/enumeration/http_enum.py @@ -0,0 +1,22 @@ +"""Provides a collection of standard HTTP request methods.""" + +from enum import Enum + + +class HttpEnum(str, Enum): + """Enumeration of supported HTTP methods for network requests.""" + + # Retrieves data from a specified resource + GET = "get" + + # Submits data to be processed to a specified resource + POST = "post" + + # Identical to GET but only retrieves the response headers + HEAD = "head" + + # Uploads or replaces the representation of a target resource + PUT = "put" + + # Deletes the specified resource from the server + DELETE = "delete" diff --git a/reme_ai/core_old/enumeration/json_schema_enum.py b/reme_ai/core_old/enumeration/json_schema_enum.py new file mode 100644 index 00000000..507645f4 --- /dev/null +++ b/reme_ai/core_old/enumeration/json_schema_enum.py @@ -0,0 +1,18 @@ +"""Defines the standard data types supported by JSON Schema.""" + +from enum import Enum + + +class JsonSchemaEnum(Enum): + """Enumeration of valid JSON Schema data types.""" + + STRING = str + NUMBER = float + INTEGER = int + OBJECT = dict + ARRAY = list + BOOLEAN = bool + + def __str__(self) -> str: + """Returns the string representation of the enum value.""" + return self.name.lower() diff --git a/reme_ai/core_old/enumeration/memory_type.py b/reme_ai/core_old/enumeration/memory_type.py new file mode 100644 index 00000000..22d35481 --- /dev/null +++ b/reme_ai/core_old/enumeration/memory_type.py @@ -0,0 +1,25 @@ +"""Memory type enumeration for the three-layer memory architecture.""" + +from enum import Enum + + +class MemoryType(str, Enum): + """ + Three-layer memory architecture for agent memory management. + + Layer 1 - High-level Abstraction Memory: + - IDENTITY: Self-cognition (identity, personality, current state) + - PERSONAL: Person-specific memory (preferences and context about specific individuals) + - PROCEDURAL: Procedural memory (how-to knowledge, e.g., 4 steps to write financial reports) + - TOOL: Tool memory (tool usage patterns, success rates, token consumption, latency) + + Layer 2 - Summary Memory (Compressed): Summarized digest of raw message history + Layer 3 - History Memory (Raw): Raw message history + """ + + IDENTITY = "identity" + PERSONAL = "personal" + PROCEDURAL = "procedural" + TOOL = "tool" + SUMMARY = "summary" + HISTORY = "history" diff --git a/reme_ai/core_old/enumeration/registry_enum.py b/reme_ai/core_old/enumeration/registry_enum.py new file mode 100644 index 00000000..876c06b8 --- /dev/null +++ b/reme_ai/core_old/enumeration/registry_enum.py @@ -0,0 +1,28 @@ +"""Defines the registry categories for core components of the system.""" + +from enum import Enum + + +class RegistryEnum(str, Enum): + """Enumeration of component types registered within the application lifecycle.""" + + # Large Language Model interfaces + LLM = "llm" + + # Models used for generating vector embeddings + EMBEDDING_MODEL = "embedding_model" + + # Databases or storage systems for vector search + VECTOR_STORE = "vector_store" + + # Atomic operations or functional units + OP = "op" + + # Orchestrated sequences of operations or workflows + FLOW = "flow" + + # External APIs or shared internal services + SERVICE = "service" + + # Utilities for tracking and limiting token consumption + TOKEN_COUNTER = "token_counter" diff --git a/reme_ai/core_old/enumeration/role.py b/reme_ai/core_old/enumeration/role.py new file mode 100644 index 00000000..4acad7e5 --- /dev/null +++ b/reme_ai/core_old/enumeration/role.py @@ -0,0 +1,19 @@ +"""Defines the participant roles in a chat completion sequence.""" + +from enum import Enum + + +class Role(str, Enum): + """Enumeration of standard personas involved in a conversation flow.""" + + # High-level instructions to guide the model's behavior + SYSTEM = "system" + + # Input or queries provided by the human user + USER = "user" + + # Responses or messages generated by the AI model + ASSISTANT = "assistant" + + # Output or results returned from external tool executions + TOOL = "tool" diff --git a/reme_ai/core/flow/__init__.py b/reme_ai/core_old/flow/__init__.py similarity index 100% rename from reme_ai/core/flow/__init__.py rename to reme_ai/core_old/flow/__init__.py diff --git a/reme_ai/core/flow/base_flow.py b/reme_ai/core_old/flow/base_flow.py similarity index 100% rename from reme_ai/core/flow/base_flow.py rename to reme_ai/core_old/flow/base_flow.py diff --git a/reme_ai/core/flow/cmd_flow.py b/reme_ai/core_old/flow/cmd_flow.py similarity index 100% rename from reme_ai/core/flow/cmd_flow.py rename to reme_ai/core_old/flow/cmd_flow.py diff --git a/reme_ai/core/flow/expression_flow.py b/reme_ai/core_old/flow/expression_flow.py similarity index 100% rename from reme_ai/core/flow/expression_flow.py rename to reme_ai/core_old/flow/expression_flow.py diff --git a/reme_ai/core/flow/simple_flow.py b/reme_ai/core_old/flow/simple_flow.py similarity index 100% rename from reme_ai/core/flow/simple_flow.py rename to reme_ai/core_old/flow/simple_flow.py diff --git a/reme_ai/core/llm/__init__.py b/reme_ai/core_old/llm/__init__.py similarity index 100% rename from reme_ai/core/llm/__init__.py rename to reme_ai/core_old/llm/__init__.py diff --git a/reme_ai/core/llm/base_llm.py b/reme_ai/core_old/llm/base_llm.py similarity index 98% rename from reme_ai/core/llm/base_llm.py rename to reme_ai/core_old/llm/base_llm.py index 8f743ce6..fcd543b6 100644 --- a/reme_ai/core/llm/base_llm.py +++ b/reme_ai/core_old/llm/base_llm.py @@ -19,7 +19,7 @@ class BaseLLM(ABC): def __init__(self, model_name: str, max_retries: int = 10, raise_exception: bool = False, request_interval: float = 0.0, **kwargs): """Initialize the LLM client with model configurations and retry policies. - + Args: model_name: The name of the model to use max_retries: Maximum number of retry attempts on failure @@ -32,7 +32,7 @@ class BaseLLM(ABC): self.raise_exception: bool = raise_exception self.request_interval: float = request_interval self.kwargs: dict = kwargs - + # Request rate control for async operations self._last_request_time: float = 0.0 self._request_lock: asyncio.Lock = asyncio.Lock() @@ -87,7 +87,7 @@ class BaseLLM(ABC): **kwargs, ) -> dict: """Construct provider-specific parameters for streaming API requests. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -122,7 +122,7 @@ class BaseLLM(ABC): **kwargs, ) -> AsyncGenerator[StreamChunk, None]: """Public async interface for streaming chat completions with retries. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -138,10 +138,10 @@ class BaseLLM(ABC): sleep_time = self.request_interval - elapsed await asyncio.sleep(sleep_time) self._last_request_time = time.time() - + async for chunk in self._stream_chat_impl(messages, tools, model_name, **kwargs): yield chunk - + async def _stream_chat_impl( self, messages: list[Message], @@ -178,7 +178,7 @@ class BaseLLM(ABC): **kwargs, ) -> Generator[StreamChunk, None, None]: """Public synchronous interface for streaming chat completions with retries. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -213,7 +213,7 @@ class BaseLLM(ABC): **kwargs, ) -> Message: """Internal async method to aggregate a full response by consuming the stream. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -281,7 +281,7 @@ class BaseLLM(ABC): **kwargs, ) -> Message: """Internal synchronous method to aggregate a full response by consuming the stream. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -351,7 +351,7 @@ class BaseLLM(ABC): **kwargs, ) -> Message | Any: """Perform an async chat completion with integrated retries and error handling. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -370,9 +370,9 @@ class BaseLLM(ABC): sleep_time = self.request_interval - elapsed await asyncio.sleep(sleep_time) self._last_request_time = time.time() - + return await self._chat_impl(messages, tools, enable_stream_print, callback_fn, default_value, model_name, **kwargs) - + async def _chat_impl( self, messages: list[Message], @@ -386,7 +386,7 @@ class BaseLLM(ABC): """Internal implementation of chat with retry and error handling logic.""" # Use the provided model_name or fall back to self.model_name effective_model = model_name if model_name is not None else self.model_name - + for i in range(self.max_retries): try: result = await self._chat( @@ -407,7 +407,7 @@ class BaseLLM(ABC): "exceeded your current quota" in error_message.lower() or "insufficient_quota" in error_message.lower() ) - + if is_inappropriate_content: logger.error(f"chat with model={effective_model} detected inappropriate content error") logger.error("=" * 80) @@ -424,12 +424,12 @@ class BaseLLM(ABC): logger.error("=" * 80) # Return empty Message immediately without retrying return Message(role=Role.ASSISTANT, content="") - + if is_rate_limit_error: logger.warning(f"chat with model={effective_model} hit rate limit, sleeping for 60s before retry (attempt {i + 1}/{self.max_retries})") await asyncio.sleep(60) continue - + logger.exception(f"chat with model={effective_model} encounter error with e={e.args}") if i == self.max_retries - 1: @@ -451,7 +451,7 @@ class BaseLLM(ABC): **kwargs, ) -> Message | Any: """Perform a synchronous chat completion with integrated retries and error handling. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -463,7 +463,7 @@ class BaseLLM(ABC): """ # Use the provided model_name or fall back to self.model_name effective_model = model_name if model_name is not None else self.model_name - + for i in range(self.max_retries): try: result = self._chat_sync( @@ -484,7 +484,7 @@ class BaseLLM(ABC): "exceeded your current quota" in error_message.lower() or "insufficient_quota" in error_message.lower() ) - + if is_inappropriate_content: logger.error(f"chat sync with model={effective_model} detected inappropriate content error") logger.error("=" * 80) @@ -501,12 +501,12 @@ class BaseLLM(ABC): logger.error("=" * 80) # Return empty Message immediately without retrying return Message(role=Role.ASSISTANT, content="") - + if is_rate_limit_error: logger.warning(f"chat sync with model={effective_model} hit rate limit, sleeping for 60s before retry (attempt {i + 1}/{self.max_retries})") time.sleep(60) continue - + logger.exception(f"chat sync with model={effective_model} encounter error with e={e.args}") if i == self.max_retries - 1: diff --git a/reme_ai/core/llm/lite_llm.py b/reme_ai/core_old/llm/lite_llm.py similarity index 99% rename from reme_ai/core/llm/lite_llm.py rename to reme_ai/core_old/llm/lite_llm.py index 934fa52d..458fe1c0 100644 --- a/reme_ai/core/llm/lite_llm.py +++ b/reme_ai/core_old/llm/lite_llm.py @@ -40,7 +40,7 @@ class LiteLLM(BaseLLM): **kwargs, ) -> dict: """Construct and log the parameters dictionary for LiteLLM API calls. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -50,7 +50,7 @@ class LiteLLM(BaseLLM): """ # Use the provided model_name or fall back to self.model_name effective_model = model_name if model_name is not None else self.model_name - + # Construct the API parameters by merging multiple sources llm_kwargs = { "model": effective_model, diff --git a/reme_ai/core/llm/lite_llm_sync.py b/reme_ai/core_old/llm/lite_llm_sync.py similarity index 100% rename from reme_ai/core/llm/lite_llm_sync.py rename to reme_ai/core_old/llm/lite_llm_sync.py diff --git a/reme_ai/core/llm/openai_llm.py b/reme_ai/core_old/llm/openai_llm.py similarity index 99% rename from reme_ai/core/llm/openai_llm.py rename to reme_ai/core_old/llm/openai_llm.py index 7b645036..f15dcb8e 100644 --- a/reme_ai/core/llm/openai_llm.py +++ b/reme_ai/core_old/llm/openai_llm.py @@ -45,7 +45,7 @@ class OpenAILLM(BaseLLM): **kwargs, ) -> dict: """Construct the parameter dictionary for the OpenAI Chat Completions API call. - + Args: messages: List of conversation messages tools: Optional list of tool calls @@ -55,7 +55,7 @@ class OpenAILLM(BaseLLM): """ # Use the provided model_name or fall back to self.model_name effective_model = model_name if model_name is not None else self.model_name - + # Construct the API parameters by merging multiple sources llm_kwargs = { "model": effective_model, diff --git a/reme_ai/core/llm/openai_llm_sync.py b/reme_ai/core_old/llm/openai_llm_sync.py similarity index 100% rename from reme_ai/core/llm/openai_llm_sync.py rename to reme_ai/core_old/llm/openai_llm_sync.py diff --git a/reme_ai/core/main.py b/reme_ai/core_old/main.py similarity index 100% rename from reme_ai/core/main.py rename to reme_ai/core_old/main.py diff --git a/reme_ai/core/op/__init__.py b/reme_ai/core_old/op/__init__.py similarity index 100% rename from reme_ai/core/op/__init__.py rename to reme_ai/core_old/op/__init__.py diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core_old/op/base_op.py similarity index 100% rename from reme_ai/core/op/base_op.py rename to reme_ai/core_old/op/base_op.py diff --git a/reme_ai/core/op/base_ray_op.py b/reme_ai/core_old/op/base_ray_op.py similarity index 100% rename from reme_ai/core/op/base_ray_op.py rename to reme_ai/core_old/op/base_ray_op.py diff --git a/reme_ai/core/op/mcp_tool.py b/reme_ai/core_old/op/mcp_tool.py similarity index 100% rename from reme_ai/core/op/mcp_tool.py rename to reme_ai/core_old/op/mcp_tool.py diff --git a/reme_ai/core/op/parallel_op.py b/reme_ai/core_old/op/parallel_op.py similarity index 100% rename from reme_ai/core/op/parallel_op.py rename to reme_ai/core_old/op/parallel_op.py diff --git a/reme_ai/core/op/sequential_op.py b/reme_ai/core_old/op/sequential_op.py similarity index 100% rename from reme_ai/core/op/sequential_op.py rename to reme_ai/core_old/op/sequential_op.py diff --git a/reme_ai/reme.py b/reme_ai/core_old/reme.py similarity index 90% rename from reme_ai/reme.py rename to reme_ai/core_old/reme.py index d0a71e6f..b5581159 100644 --- a/reme_ai/reme.py +++ b/reme_ai/core_old/reme.py @@ -1,14 +1,14 @@ """ReMe classes for simplified configuration and execution.""" -from .core.application import Application -from .core.config import ReMeConfigParser -from .core.context import C -from .core.embedding import BaseEmbeddingModel -from .core.enumeration import Role -from .core.llm import BaseLLM -from .core.schema import Message -from .core.utils import singleton -from .core.vector_store import BaseVectorStore +from .core_old.application import Application +from .core_old.config import ReMeConfigParser +from .core_old.context import C +from .core_old.embedding import BaseEmbeddingModel +from .core_old.enumeration import Role +from .core_old.llm import BaseLLM +from .core_old.schema import Message +from .core_old.utils import singleton +from .core_old.vector_store import BaseVectorStore from .mem_agent.retriever import ReMeRetriever from .mem_agent.retriever_v2 import ReMeRetrieverV2 from .mem_agent.summarizer import ReMeSummarizer, PersonalSummarizer @@ -177,7 +177,6 @@ class ReMe(Application): except Exception as e: print(f"Warning: reme_summarizer.call failed: {e}") return [] - else: raise NotImplementedError @@ -228,18 +227,17 @@ class ReMe(Application): except Exception as e: print(f"Warning: reme_retriever.call failed: {e}") return "error, not retrieved" - else: raise NotImplementedError async def summary_v2( - self, - messages: list[dict], - description: str = "", - user_id: str = "", - assistant_id: str = "", - **kwargs, + self, + messages: list[dict], + description: str = "", + user_id: str = "", + assistant_id: str = "", + **kwargs, ): """Summarizes messages using V2 workflow with simplified tools.""" @@ -293,14 +291,14 @@ class ReMe(Application): raise NotImplementedError async def retrieve_v2( - self, - query: str = "", - messages: list[dict] | None = None, - description: str = "", - user_id: str = "", - assistant_id: str = "", - top_k: int = 20, - **kwargs, + self, + query: str = "", + messages: list[dict] | None = None, + description: str = "", + user_id: str = "", + assistant_id: str = "", + top_k: int = 20, + **kwargs, ): """Retrieves relevant memories using V2 workflow with autonomous retrieval.""" @@ -342,12 +340,12 @@ class ReMe(Application): raise NotImplementedError async def summary_v3( - self, - messages: list[dict], - description: str = "", - user_id: str = "", - assistant_id: str = "", - **kwargs, + self, + messages: list[dict], + description: str = "", + user_id: str = "", + assistant_id: str = "", + **kwargs, ): """Summarizes messages using V3 workflow with user profile management.""" @@ -384,14 +382,14 @@ class ReMe(Application): raise NotImplementedError async def retrieve_v3( - self, - query: str = "", - messages: list[dict] | None = None, - description: str = "", - user_id: str = "", - assistant_id: str = "", - top_k: int = 20, - **kwargs, + self, + query: str = "", + messages: list[dict] | None = None, + description: str = "", + user_id: str = "", + assistant_id: str = "", + top_k: int = 20, + **kwargs, ): """Retrieves relevant memories using V3 workflow with user profile support.""" @@ -425,13 +423,13 @@ class ReMe(Application): raise NotImplementedError async def summary_v4( - self, - messages: list[dict], - description: str = "", - user_id: str = "", - assistant_id: str = "", - enable_thinking_params: bool = False, - **kwargs, + 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.""" @@ -464,15 +462,15 @@ class ReMe(Application): 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, + 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.""" diff --git a/reme_ai/core/schema/__init__.py b/reme_ai/core_old/schema/__init__.py similarity index 100% rename from reme_ai/core/schema/__init__.py rename to reme_ai/core_old/schema/__init__.py diff --git a/reme_ai/core/schema/memory_node.py b/reme_ai/core_old/schema/memory_node.py similarity index 100% rename from reme_ai/core/schema/memory_node.py rename to reme_ai/core_old/schema/memory_node.py diff --git a/reme_ai/core/schema/message.py b/reme_ai/core_old/schema/message.py similarity index 100% rename from reme_ai/core/schema/message.py rename to reme_ai/core_old/schema/message.py diff --git a/reme_ai/core/schema/request.py b/reme_ai/core_old/schema/request.py similarity index 100% rename from reme_ai/core/schema/request.py rename to reme_ai/core_old/schema/request.py diff --git a/reme_ai/core/schema/response.py b/reme_ai/core_old/schema/response.py similarity index 100% rename from reme_ai/core/schema/response.py rename to reme_ai/core_old/schema/response.py diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core_old/schema/service_config.py similarity index 100% rename from reme_ai/core/schema/service_config.py rename to reme_ai/core_old/schema/service_config.py diff --git a/reme_ai/core/schema/stream_chunk.py b/reme_ai/core_old/schema/stream_chunk.py similarity index 100% rename from reme_ai/core/schema/stream_chunk.py rename to reme_ai/core_old/schema/stream_chunk.py diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core_old/schema/tool_call.py similarity index 99% rename from reme_ai/core/schema/tool_call.py rename to reme_ai/core_old/schema/tool_call.py index 21dde495..3c01cdcd 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core_old/schema/tool_call.py @@ -177,7 +177,7 @@ class ToolCall(BaseModel): return True except Exception: return False - + def sanitize_and_check_argument(self) -> bool: """ Attempt to sanitize and validate arguments JSON. @@ -188,17 +188,17 @@ class ToolCall(BaseModel): """ 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: @@ -212,7 +212,7 @@ class ToolCall(BaseModel): sanitized = sanitized[:-1].rstrip() else: break - + return False def simple_output_dump(self) -> dict: diff --git a/reme_ai/core/schema/vector_node.py b/reme_ai/core_old/schema/vector_node.py similarity index 100% rename from reme_ai/core/schema/vector_node.py rename to reme_ai/core_old/schema/vector_node.py diff --git a/reme_ai/core/service/__init__.py b/reme_ai/core_old/service/__init__.py similarity index 100% rename from reme_ai/core/service/__init__.py rename to reme_ai/core_old/service/__init__.py diff --git a/reme_ai/core/service/base_service.py b/reme_ai/core_old/service/base_service.py similarity index 100% rename from reme_ai/core/service/base_service.py rename to reme_ai/core_old/service/base_service.py diff --git a/reme_ai/core/service/cmd_service.py b/reme_ai/core_old/service/cmd_service.py similarity index 100% rename from reme_ai/core/service/cmd_service.py rename to reme_ai/core_old/service/cmd_service.py diff --git a/reme_ai/core/service/http_service.py b/reme_ai/core_old/service/http_service.py similarity index 100% rename from reme_ai/core/service/http_service.py rename to reme_ai/core_old/service/http_service.py diff --git a/reme_ai/core/service/mcp_service.py b/reme_ai/core_old/service/mcp_service.py similarity index 100% rename from reme_ai/core/service/mcp_service.py rename to reme_ai/core_old/service/mcp_service.py diff --git a/reme_ai/core/token_counter/__init__.py b/reme_ai/core_old/token_counter/__init__.py similarity index 100% rename from reme_ai/core/token_counter/__init__.py rename to reme_ai/core_old/token_counter/__init__.py diff --git a/reme_ai/core/token_counter/base_token_counter.py b/reme_ai/core_old/token_counter/base_token_counter.py similarity index 100% rename from reme_ai/core/token_counter/base_token_counter.py rename to reme_ai/core_old/token_counter/base_token_counter.py diff --git a/reme_ai/core/token_counter/hf_token_counter.py b/reme_ai/core_old/token_counter/hf_token_counter.py similarity index 100% rename from reme_ai/core/token_counter/hf_token_counter.py rename to reme_ai/core_old/token_counter/hf_token_counter.py diff --git a/reme_ai/core/token_counter/openai_token_counter.py b/reme_ai/core_old/token_counter/openai_token_counter.py similarity index 100% rename from reme_ai/core/token_counter/openai_token_counter.py rename to reme_ai/core_old/token_counter/openai_token_counter.py diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core_old/utils/__init__.py similarity index 100% rename from reme_ai/core/utils/__init__.py rename to reme_ai/core_old/utils/__init__.py diff --git a/reme_ai/core/utils/cache_handler.py b/reme_ai/core_old/utils/cache_handler.py similarity index 100% rename from reme_ai/core/utils/cache_handler.py rename to reme_ai/core_old/utils/cache_handler.py diff --git a/reme_ai/core/utils/case_converter.py b/reme_ai/core_old/utils/case_converter.py similarity index 100% rename from reme_ai/core/utils/case_converter.py rename to reme_ai/core_old/utils/case_converter.py diff --git a/reme_ai/core/utils/common_utils.py b/reme_ai/core_old/utils/common_utils.py similarity index 100% rename from reme_ai/core/utils/common_utils.py rename to reme_ai/core_old/utils/common_utils.py diff --git a/reme_ai/core/utils/env_utils.py b/reme_ai/core_old/utils/env_utils.py similarity index 100% rename from reme_ai/core/utils/env_utils.py rename to reme_ai/core_old/utils/env_utils.py diff --git a/reme_ai/core/utils/execute_tuils.py b/reme_ai/core_old/utils/execute_tuils.py similarity index 100% rename from reme_ai/core/utils/execute_tuils.py rename to reme_ai/core_old/utils/execute_tuils.py diff --git a/reme_ai/core/utils/http_client.py b/reme_ai/core_old/utils/http_client.py similarity index 100% rename from reme_ai/core/utils/http_client.py rename to reme_ai/core_old/utils/http_client.py diff --git a/reme_ai/core/utils/llm_utils.py b/reme_ai/core_old/utils/llm_utils.py similarity index 100% rename from reme_ai/core/utils/llm_utils.py rename to reme_ai/core_old/utils/llm_utils.py diff --git a/reme_ai/core/utils/logger_utils.py b/reme_ai/core_old/utils/logger_utils.py similarity index 100% rename from reme_ai/core/utils/logger_utils.py rename to reme_ai/core_old/utils/logger_utils.py diff --git a/reme_ai/core/utils/logo_utils.py b/reme_ai/core_old/utils/logo_utils.py similarity index 100% rename from reme_ai/core/utils/logo_utils.py rename to reme_ai/core_old/utils/logo_utils.py diff --git a/reme_ai/core/utils/mcp_client.py b/reme_ai/core_old/utils/mcp_client.py similarity index 100% rename from reme_ai/core/utils/mcp_client.py rename to reme_ai/core_old/utils/mcp_client.py diff --git a/reme_ai/core/utils/pydantic_config_parser.py b/reme_ai/core_old/utils/pydantic_config_parser.py similarity index 100% rename from reme_ai/core/utils/pydantic_config_parser.py rename to reme_ai/core_old/utils/pydantic_config_parser.py diff --git a/reme_ai/core/utils/pydantic_utils.py b/reme_ai/core_old/utils/pydantic_utils.py similarity index 100% rename from reme_ai/core/utils/pydantic_utils.py rename to reme_ai/core_old/utils/pydantic_utils.py diff --git a/reme_ai/core/utils/singleton.py b/reme_ai/core_old/utils/singleton.py similarity index 100% rename from reme_ai/core/utils/singleton.py rename to reme_ai/core_old/utils/singleton.py diff --git a/reme_ai/core/utils/time.py b/reme_ai/core_old/utils/time.py similarity index 100% rename from reme_ai/core/utils/time.py rename to reme_ai/core_old/utils/time.py diff --git a/reme_ai/core/vector_store/__init__.py b/reme_ai/core_old/vector_store/__init__.py similarity index 100% rename from reme_ai/core/vector_store/__init__.py rename to reme_ai/core_old/vector_store/__init__.py diff --git a/reme_ai/core/vector_store/base_vector_store.py b/reme_ai/core_old/vector_store/base_vector_store.py similarity index 97% rename from reme_ai/core/vector_store/base_vector_store.py rename to reme_ai/core_old/vector_store/base_vector_store.py index a4a8ca8e..158af0d9 100644 --- a/reme_ai/core/vector_store/base_vector_store.py +++ b/reme_ai/core_old/vector_store/base_vector_store.py @@ -5,9 +5,9 @@ from abc import ABC, abstractmethod from collections.abc import Callable from functools import partial -from reme_ai.core.context import C -from reme_ai.core.embedding import BaseEmbeddingModel -from reme_ai.core.schema import VectorNode +from reme_ai.core_old.context import C +from reme_ai.core_old.embedding import BaseEmbeddingModel +from reme_ai.core_old.schema import VectorNode class BaseVectorStore(ABC): diff --git a/reme_ai/core/vector_store/chroma_vector_store.py b/reme_ai/core_old/vector_store/chroma_vector_store.py similarity index 99% rename from reme_ai/core/vector_store/chroma_vector_store.py rename to reme_ai/core_old/vector_store/chroma_vector_store.py index 3a483c23..567ac3b1 100644 --- a/reme_ai/core/vector_store/chroma_vector_store.py +++ b/reme_ai/core_old/vector_store/chroma_vector_store.py @@ -118,7 +118,7 @@ class ChromaVectorStore(BaseVectorStore): @staticmethod def _generate_where_clause(filters: dict | None) -> dict | None: """Convert the universal filter format to a ChromaDB-compatible where clause. - + Supports two filter formats: 1. Range query: {"field": [start_value, end_value]} - filters for field >= start_value AND field <= end_value 2. Exact match: {"field": value} - filters for field == value @@ -128,7 +128,7 @@ class ChromaVectorStore(BaseVectorStore): def convert_condition(k: str, v: Any) -> dict | list | None: """Convert a single filter condition to ChromaDB operator format. - + Returns: - dict for simple conditions - list of dicts for range queries (which need to be wrapped in $and) diff --git a/reme_ai/core/vector_store/es_vector_store.py b/reme_ai/core_old/vector_store/es_vector_store.py similarity index 100% rename from reme_ai/core/vector_store/es_vector_store.py rename to reme_ai/core_old/vector_store/es_vector_store.py diff --git a/reme_ai/core/vector_store/local_vector_store.py b/reme_ai/core_old/vector_store/local_vector_store.py similarity index 99% rename from reme_ai/core/vector_store/local_vector_store.py rename to reme_ai/core_old/vector_store/local_vector_store.py index cfeec74e..7533fa89 100644 --- a/reme_ai/core/vector_store/local_vector_store.py +++ b/reme_ai/core_old/vector_store/local_vector_store.py @@ -92,7 +92,7 @@ class LocalVectorStore(BaseVectorStore): @staticmethod def _match_filters(node: VectorNode, filters: dict | None) -> bool: """Check if a vector node matches the provided metadata filters. - + Supports two filter formats: 1. Range query: {"field": [start_value, end_value]} - filters for field >= start_value AND field <= end_value 2. Exact match: {"field": value} - filters for field == value diff --git a/reme_ai/core/vector_store/pgvector_store.py b/reme_ai/core_old/vector_store/pgvector_store.py similarity index 99% rename from reme_ai/core/vector_store/pgvector_store.py rename to reme_ai/core_old/vector_store/pgvector_store.py index 576d8e6c..882ced3e 100644 --- a/reme_ai/core/vector_store/pgvector_store.py +++ b/reme_ai/core_old/vector_store/pgvector_store.py @@ -29,7 +29,7 @@ class PGVectorStore(BaseVectorStore): @staticmethod def _validate_table_name(name: str) -> None: """Validate table name to prevent SQL injection. - + PostgreSQL table names must: - Contain only alphanumeric characters and underscores - Not start with a digit @@ -279,11 +279,11 @@ class PGVectorStore(BaseVectorStore): @staticmethod def _build_filter_clause(filters: dict | None) -> tuple[str, list]: """Generate an SQL WHERE clause and parameter list from a filter dictionary. - + Supports two filter formats: 1. Range query: {"field": [start_value, end_value]} - filters for field >= start_value AND field <= end_value 2. Exact match: {"field": value} - filters for field == value - + Range queries support both numeric and string (e.g., timestamp strings) comparisons. """ if not filters: @@ -297,7 +297,7 @@ class PGVectorStore(BaseVectorStore): # Sanitize key to prevent SQL injection (only allow alphanumeric and underscore) if not key.replace('_', '').replace('.', '').isalnum(): raise ValueError(f"Invalid metadata key: {key}. Only alphanumeric characters, underscore and dot are allowed.") - + # New syntax: [start, end] represents a range query if isinstance(value, list) and len(value) == 2: # Range query: field >= value[0] AND field <= value[1] diff --git a/reme_ai/core/vector_store/qdrant_vector_store.py b/reme_ai/core_old/vector_store/qdrant_vector_store.py similarity index 99% rename from reme_ai/core/vector_store/qdrant_vector_store.py rename to reme_ai/core_old/vector_store/qdrant_vector_store.py index 5d9fa4f5..d0b4a8aa 100644 --- a/reme_ai/core/vector_store/qdrant_vector_store.py +++ b/reme_ai/core_old/vector_store/qdrant_vector_store.py @@ -247,7 +247,7 @@ class QdrantVectorStore(BaseVectorStore): @staticmethod def _create_filter(filters: dict) -> Filter | None: """Convert a dictionary of filter conditions into a Qdrant Filter object. - + Supports two filter formats: 1. Range query: {"field": [start_value, end_value]} - filters for field >= start_value AND field <= end_value 2. Exact match: {"field": value} - filters for field == value @@ -295,7 +295,7 @@ class QdrantVectorStore(BaseVectorStore): f"Qdrant range filter for key '{key}' requires numeric lte value, got {type(value['lte']).__name__}. Skipping." ) continue - + if range_params: # Only add condition if we have valid numeric parameters conditions.append( FieldCondition( diff --git a/reme_ai/mem_agent/base_memory_agent.py b/reme_ai/mem_agent/base_memory_agent.py index d53cc8d9..5ba0cad3 100644 --- a/reme_ai/mem_agent/base_memory_agent.py +++ b/reme_ai/mem_agent/base_memory_agent.py @@ -6,9 +6,9 @@ from abc import ABCMeta from loguru import logger -from ..core.enumeration import Role, MemoryType -from ..core.op import BaseOp -from ..core.schema import Message, ToolCall, MemoryNode +from ..core_old.enumeration import Role, MemoryType +from ..core_old.op import BaseOp +from ..core_old.schema import Message, ToolCall, MemoryNode from ..mem_tool import BaseMemoryTool, ThinkTool @@ -155,13 +155,13 @@ class BaseMemoryAgent(BaseOp, metaclass=ABCMeta): tool_call_id=op.tool_call.id, ) tool_result_messages.append(tool_message) - + # # Collect tool call information to meta_info # tool_info = f"\n## Tool Call {step + 1}.{j + 1}: {op.tool_call.name}\n" # tool_info += f"Arguments: {json.dumps(assistant_message.tool_calls[j].argument_dict, ensure_ascii=False)}\n" # tool_info += f"Result: {tool_result}\n" self.meta_info += tool_result + "\n" - + logger.info(f"[{self.__class__.__name__}{stage_prefix}] step{step + 1}.{j} join tool_result={tool_result[:2000]}...\n\n") return tool_result_messages diff --git a/reme_ai/mem_agent/chat/remy_agent.py b/reme_ai/mem_agent/chat/remy_agent.py index 1eaa9ac7..c7806617 100644 --- a/reme_ai/mem_agent/chat/remy_agent.py +++ b/reme_ai/mem_agent/chat/remy_agent.py @@ -3,10 +3,10 @@ from typing import List from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role -from ...core.schema import Message -from ...core.utils import get_now_time +from ...core_old.context import C +from ...core_old.enumeration import Role +from ...core_old.schema import Message +from ...core_old.utils import get_now_time @C.register_op() diff --git a/reme_ai/mem_agent/chat/simple_chat.py b/reme_ai/mem_agent/chat/simple_chat.py index 8a71c7a8..9e09ffec 100644 --- a/reme_ai/mem_agent/chat/simple_chat.py +++ b/reme_ai/mem_agent/chat/simple_chat.py @@ -2,10 +2,10 @@ from loguru import logger -from ...core.context import C -from ...core.enumeration import Role -from ...core.op import BaseOp -from ...core.schema import Message, ToolCall +from ...core_old.context import C +from ...core_old.enumeration import Role +from ...core_old.op import BaseOp +from ...core_old.schema import Message, ToolCall @C.register_op() diff --git a/reme_ai/mem_agent/chat/stream_chat.py b/reme_ai/mem_agent/chat/stream_chat.py index 470e4647..2b121446 100644 --- a/reme_ai/mem_agent/chat/stream_chat.py +++ b/reme_ai/mem_agent/chat/stream_chat.py @@ -2,10 +2,10 @@ from loguru import logger -from ...core.context import C -from ...core.enumeration import Role, ChunkEnum -from ...core.op import BaseOp -from ...core.schema import Message, ToolCall +from ...core_old.context import C +from ...core_old.enumeration import Role, ChunkEnum +from ...core_old.op import BaseOp +from ...core_old.schema import Message, ToolCall @C.register_op() diff --git a/reme_ai/mem_agent/retriever/reme_retriever.py b/reme_ai/mem_agent/retriever/reme_retriever.py index f3700b1e..400e0d01 100644 --- a/reme_ai/mem_agent/retriever/reme_retriever.py +++ b/reme_ai/mem_agent/retriever/reme_retriever.py @@ -3,10 +3,10 @@ from typing import List from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role -from ...core.schema import Message -from ...core.utils import get_now_time, format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role +from ...core_old.schema import Message +from ...core_old.utils import get_now_time, format_messages @C.register_op() diff --git a/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.py b/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.py index ee5c0aae..fedfe188 100644 --- a/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.py +++ b/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.py @@ -3,16 +3,16 @@ from typing import List from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role -from ...core.schema import Message -from ...core.utils import format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role +from ...core_old.schema import Message +from ...core_old.utils import format_messages @C.register_op() class ReMeRetrieverV2(BaseMemoryAgent): """Memory agent that autonomously retrieves memories from multiple angles. - + This retriever: - Directly queries memories based on user questions without time constraints - Tries multiple retrieval strategies: direct vector search, metadata filtering, partial filtering @@ -24,13 +24,13 @@ class ReMeRetrieverV2(BaseMemoryAgent): # Check if ReadHistory tool is available in the tools list tools = kwargs.get('tools', []) has_read_history = any(tool.__class__.__name__ == 'ReadHistory' for tool in tools) - + # Use simple prompt if ReadHistory is not available if not has_read_history: super().__init__(prompt_name="reme_retriever_v2_simple", **kwargs) else: super().__init__(**kwargs) - + self.meta_memories: list[dict] = meta_memories or [] async def _read_meta_memories(self) -> str: diff --git a/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.yaml b/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.yaml index 281796d6..755fd6b7 100644 --- a/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.yaml +++ b/reme_ai/mem_agent/retriever_v2/reme_retriever_v2.yaml @@ -24,16 +24,16 @@ system_prompt: | 1. **Multi-Angle Vector Retrieval** (REQUIRED - at least 3 attempts): You must try AT LEAST 3 different retrieval approaches using `retrieve_memories`: - + a) **Direct Vector Search**: Use the user's question directly or with minimal reformulation - Query the most relevant memory_type and memory_target - Use straightforward query phrasing - + b) **Alternative Phrasing**: Reformulate the query from a different angle - Use synonyms or different expressions - Break down complex questions into simpler components - Try more specific or more general queries - + c) **Metadata-Filtered Search**: Add metadata filters to narrow down results - **Time-based filtering**: Use year/month/day metadata fields to filter by time periods * Example: {{"year": 2024}} for memories from 2024 @@ -41,33 +41,33 @@ system_prompt: | * Example: {{"year": 2024, "month": 5, "day": 15}} for memories from a specific date - Combine vector search with metadata constraints - Try partial metadata filtering if full filtering yields nothing (e.g., only year, or year+month) - + d) **Cross-Memory-Type Search**: If applicable, search across different memory types - Try different memory_type and memory_target combinations - Some information might be stored in unexpected memory categories - + e) **Keyword Extraction**: Extract key entities/concepts and search for them - Identify important names, places, concepts - Search for each key element separately - + 2. **Evaluate Retrieval Results** (After each attempt): - Review what memories were returned - Assess if they contain sufficient information to answer the question - If insufficient, identify what's missing and adjust your next query accordingly - Track which retrieval strategies you've already tried - + 3. **Persist Through Failures**: - DO NOT give up after 1-2 failed attempts - If a retrieval returns no results or irrelevant results, try a different approach - Consider that the information might be phrased differently than expected - Be creative with query reformulation - + 4. **Fallback to History Reading** (Only after 3+ vector retrieval attempts): - If after at least 3 different vector retrieval attempts you still lack sufficient information: * If any retrieved memories contain `ref_memory_id`, use `read_history` to read the original conversation * Use `read_history` with the `ref_memory_id` to get complete context * This can reveal details that weren't captured in the memory summaries - + 5. **Answer the Question**: - Once you have sufficient information, provide a direct answer based ONLY on retrieved memories - DO NOT fabricate, guess, or infer information not present in the memories @@ -95,30 +95,30 @@ system_prompt: | **Example 1: Simple Query** Attempt 1: Direct query "user's favorite food" → Result: No relevant memories found - + Attempt 2: Reformulated query "what does user like to eat" → Result: Some memories about meals, but not specific preferences - + Attempt 3: Keyword search "food preferences" with metadata filter → Result: Found relevant memory with ref_memory_id - + Attempt 4: Use read_history with ref_memory_id to get full context → Result: Found detailed conversation about favorite foods - + Answer: [Provide answer based on retrieved information] **Example 2: Time-based Query** Question: "What did the user do last summer?" - + Attempt 1: Direct query "user activities summer" with metadata {{"year": 2025, "month": [6, 7, 8]}} → Result: Found some vacation memories - + Attempt 2: Broader query "user summer vacation travel" with metadata {{"year": 2025}} → Result: Found additional travel-related memories - + Attempt 3: Use read_history for memories with ref_memory_id to get detailed context → Result: Complete picture of summer activities - + Answer: [Provide answer based on retrieved information] user_message: | diff --git a/reme_ai/mem_agent/retriever_v2/reme_retriever_v2_simple.yaml b/reme_ai/mem_agent/retriever_v2/reme_retriever_v2_simple.yaml index f8a4b7f5..cb9cc578 100644 --- a/reme_ai/mem_agent/retriever_v2/reme_retriever_v2_simple.yaml +++ b/reme_ai/mem_agent/retriever_v2/reme_retriever_v2_simple.yaml @@ -23,16 +23,16 @@ system_prompt: | 1. **Multi-Angle Vector Retrieval** (REQUIRED - at least 3 attempts): You must try AT LEAST 3 different retrieval approaches using `retrieve_memories`: - + a) **Direct Vector Search**: Use the user's question directly or with minimal reformulation - Query the most relevant memory_type and memory_target - Use straightforward query phrasing - + b) **Alternative Phrasing**: Reformulate the query from a different angle - Use synonyms or different expressions - Break down complex questions into simpler components - Try more specific or more general queries - + c) **Metadata-Filtered Search**: Add metadata filters to narrow down results - **Time-based filtering**: Use year/month/day metadata fields to filter by time periods * Example: {{"year": 2024}} for memories from 2024 @@ -40,27 +40,27 @@ system_prompt: | * Example: {{"year": 2024, "month": 5, "day": 15}} for memories from a specific date - Combine vector search with metadata constraints - Try partial metadata filtering if full filtering yields nothing (e.g., only year, or year+month) - + d) **Cross-Memory-Type Search**: If applicable, search across different memory types - Try different memory_type and memory_target combinations - Some information might be stored in unexpected memory categories - + e) **Keyword Extraction**: Extract key entities/concepts and search for them - Identify important names, places, concepts - Search for each key element separately - + 2. **Evaluate Retrieval Results** (After each attempt): - Review what memories were returned - Assess if they contain sufficient information to answer the question - If insufficient, identify what's missing and adjust your next query accordingly - Track which retrieval strategies you've already tried - + 3. **Persist Through Failures**: - DO NOT give up after 1-2 failed attempts - If a retrieval returns no results or irrelevant results, try a different approach - Consider that the information might be phrased differently than expected - Be creative with query reformulation - + 4. **Answer the Question**: - Once you have sufficient information, provide a direct answer based ONLY on retrieved memories - DO NOT fabricate, guess, or infer information not present in the memories @@ -88,27 +88,27 @@ system_prompt: | **Example 1: Simple Query** Attempt 1: Direct query "user's favorite food" → Result: No relevant memories found - + Attempt 2: Reformulated query "what does user like to eat" → Result: Some memories about meals, but not specific preferences - + Attempt 3: Keyword search "food preferences" with metadata filter → Result: Found relevant memory - + Answer: [Provide answer based on retrieved information] **Example 2: Time-based Query** Question: "What did the user do last summer?" - + Attempt 1: Direct query "user activities summer" with metadata {{"year": 2025, "month": [6, 7, 8]}} → Result: Found some vacation memories - + Attempt 2: Broader query "user summer vacation travel" with metadata {{"year": 2025}} → Result: Found additional travel-related memories - + Attempt 3: More specific queries about specific activities → Result: Complete picture of summer activities - + Answer: [Provide answer based on retrieved information] user_message: | diff --git a/reme_ai/mem_agent/summarizer/identity_summarizer.py b/reme_ai/mem_agent/summarizer/identity_summarizer.py index be571cdc..0a9e1410 100644 --- a/reme_ai/mem_agent/summarizer/identity_summarizer.py +++ b/reme_ai/mem_agent/summarizer/identity_summarizer.py @@ -1,10 +1,10 @@ """Specialized agent for extracting and updating agent self-cognition memories.""" from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message -from ...core.utils import get_now_time, format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message +from ...core_old.utils import get_now_time, format_messages @C.register_op() diff --git a/reme_ai/mem_agent/summarizer/personal_summarizer.py b/reme_ai/mem_agent/summarizer/personal_summarizer.py index 352fe1ee..1c32f499 100644 --- a/reme_ai/mem_agent/summarizer/personal_summarizer.py +++ b/reme_ai/mem_agent/summarizer/personal_summarizer.py @@ -1,10 +1,10 @@ """Specialized agent for extracting and managing personal memories about specific individuals.""" from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, ToolCall -from ...core.utils import get_now_time, format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, ToolCall +from ...core_old.utils import get_now_time, format_messages @C.register_op() diff --git a/reme_ai/mem_agent/summarizer/procedural_summarizer.py b/reme_ai/mem_agent/summarizer/procedural_summarizer.py index e31422e4..24a75339 100644 --- a/reme_ai/mem_agent/summarizer/procedural_summarizer.py +++ b/reme_ai/mem_agent/summarizer/procedural_summarizer.py @@ -1,10 +1,10 @@ """Specialized agent for extracting and managing procedural knowledge and workflows.""" from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message -from ...core.utils import get_now_time, format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message +from ...core_old.utils import get_now_time, format_messages @C.register_op() diff --git a/reme_ai/mem_agent/summarizer/reme_summarizer.py b/reme_ai/mem_agent/summarizer/reme_summarizer.py index 9eb06d84..5521762c 100644 --- a/reme_ai/mem_agent/summarizer/reme_summarizer.py +++ b/reme_ai/mem_agent/summarizer/reme_summarizer.py @@ -3,10 +3,10 @@ from loguru import logger from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, MemoryNode, ToolCall -from ...core.utils import get_now_time, format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, MemoryNode, ToolCall +from ...core_old.utils import get_now_time, format_messages @C.register_op() @@ -18,19 +18,19 @@ class ReMeSummarizer(BaseMemoryAgent): super().__init__(**kwargs) self.enable_identity_memory = enable_identity_memory self.meta_memories: list[dict] = meta_memories or [] - + # Check if AddMetaMemory is in tools self.enable_add_meta_memory = self._check_add_meta_memory_in_tools() def _check_add_meta_memory_in_tools(self) -> bool: """Check if AddMetaMemory tool is present in the tools list.""" from ...mem_tool import AddMetaMemory - + for tool in self.tools: if isinstance(tool, AddMetaMemory): return True return False - + def _build_tool_call(self) -> ToolCall: return ToolCall( **{ diff --git a/reme_ai/mem_agent/summarizer/tool_summarizer.py b/reme_ai/mem_agent/summarizer/tool_summarizer.py index 50399bba..1e9e33c0 100644 --- a/reme_ai/mem_agent/summarizer/tool_summarizer.py +++ b/reme_ai/mem_agent/summarizer/tool_summarizer.py @@ -1,10 +1,10 @@ """Specialized agent for extracting and managing tool usage guidelines and best practices.""" from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message -from ...core.utils import get_now_time, format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message +from ...core_old.utils import get_now_time, format_messages @C.register_op() diff --git a/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2.py b/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2.py index 13395212..cc2794b3 100644 --- a/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2.py +++ b/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2.py @@ -1,10 +1,10 @@ """Simplified personal memory summarizer using v2 memory tools.""" from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, ToolCall -from ...core.utils import format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, ToolCall +from ...core_old.utils import format_messages @C.register_op() @@ -12,7 +12,7 @@ class PersonalSummarizerV2(BaseMemoryAgent): memory_type: MemoryType = MemoryType.PERSONAL """Simplified personal memory summarizer that uses v2 memory tools. - + This summarizer follows a three-step workflow: 1. AddMemoryDrafts: Generate initial memory drafts from context 2. RetrieveRecentAndSimilarMemories: Retrieve similar and recent memories diff --git a/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2_simple.yaml b/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2_simple.yaml index acc5bd35..0a5347bb 100644 --- a/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2_simple.yaml +++ b/reme_ai/mem_agent/summarizer_v2/personal_summarizer_v2_simple.yaml @@ -10,7 +10,7 @@ system_prompt: | ## Context: {context} - + **Context Format Explanation**: The context contains formatted conversation messages in the following structure: - Each message is formatted as: `round{index} [{timestamp}] {role/name}: {content}` diff --git a/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.py b/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.py index 1aae4ad4..aa680da9 100644 --- a/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.py +++ b/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.py @@ -3,10 +3,10 @@ from loguru import logger from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, MemoryNode, ToolCall -from ...core.utils import format_messages +from ...core_old.context import C +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, MemoryNode, ToolCall +from ...core_old.utils import format_messages @C.register_op() diff --git a/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.yaml b/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.yaml index 30792a08..58080cab 100644 --- a/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.yaml +++ b/reme_ai/mem_agent/summarizer_v2/reme_summarizer_v2.yaml @@ -18,7 +18,7 @@ system_prompt: | 2. Identify which memory dimensions need updates and specify them in `memory_tasks` (each with `memory_type` and `memory_target`). - The `memory_type` and `memory_target` must exactly match existing entries in the "Main Agent's Meta Memory" listed above. - Multiple tasks can be specified to enable parallel processing by specialized agents. - + Note: If the context contains no memorable information (e.g., simple greetings), output ``. user_message: | diff --git a/reme_ai/mem_agent/v3/personal_summarizer_v3.py b/reme_ai/mem_agent/v3/personal_summarizer_v3.py index 0093884d..3f2c9ca9 100644 --- a/reme_ai/mem_agent/v3/personal_summarizer_v3.py +++ b/reme_ai/mem_agent/v3/personal_summarizer_v3.py @@ -1,7 +1,7 @@ from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, ToolCall -from ...core.utils import format_messages +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, ToolCall +from ...core_old.utils import format_messages class PersonalSummarizerV3(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/v3/reme_retriever_v3.py b/reme_ai/mem_agent/v3/reme_retriever_v3.py index 8f5c62dc..020b5ad2 100644 --- a/reme_ai/mem_agent/v3/reme_retriever_v3.py +++ b/reme_ai/mem_agent/v3/reme_retriever_v3.py @@ -3,9 +3,9 @@ from typing import List from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role -from ...core.schema import Message -from ...core.utils import format_messages +from ...core_old.enumeration import Role +from ...core_old.schema import Message +from ...core_old.utils import format_messages class ReMeRetrieverV3(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/v3/reme_retriever_v3.yaml b/reme_ai/mem_agent/v3/reme_retriever_v3.yaml index 7a3e575f..925bee68 100644 --- a/reme_ai/mem_agent/v3/reme_retriever_v3.yaml +++ b/reme_ai/mem_agent/v3/reme_retriever_v3.yaml @@ -26,13 +26,13 @@ system_prompt: | * Direct query with user's question * Reformulated queries with different phrasing/keywords * Queries focused on specific entities or concepts - + - **Time Range Filtering** (when applicable): * Format: [start_date, end_date] in YYYYMMDD format * Example: [20200101, 20200102] means 20200101 < time < 20200102 * Single-sided: [0, 20200102] for before, [20200101, 99999999] for after * If no results, try broader time ranges or remove time constraints - + - If no results after multiple attempts, try different memory_type/memory_target combinations **STEP 3: Read Original Conversations (If Step 2 insufficient)** diff --git a/reme_ai/mem_agent/v3/reme_summarizer_v3.py b/reme_ai/mem_agent/v3/reme_summarizer_v3.py index a0f466b9..3e1f17f7 100644 --- a/reme_ai/mem_agent/v3/reme_summarizer_v3.py +++ b/reme_ai/mem_agent/v3/reme_summarizer_v3.py @@ -1,9 +1,9 @@ from loguru import logger from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, MemoryNode, ToolCall -from ...core.utils import format_messages +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, MemoryNode, ToolCall +from ...core_old.utils import format_messages class ReMeSummarizerV3(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/v3/reme_summarizer_v3.yaml b/reme_ai/mem_agent/v3/reme_summarizer_v3.yaml index 30792a08..58080cab 100644 --- a/reme_ai/mem_agent/v3/reme_summarizer_v3.yaml +++ b/reme_ai/mem_agent/v3/reme_summarizer_v3.yaml @@ -18,7 +18,7 @@ system_prompt: | 2. Identify which memory dimensions need updates and specify them in `memory_tasks` (each with `memory_type` and `memory_target`). - The `memory_type` and `memory_target` must exactly match existing entries in the "Main Agent's Meta Memory" listed above. - Multiple tasks can be specified to enable parallel processing by specialized agents. - + Note: If the context contains no memorable information (e.g., simple greetings), output ``. user_message: | diff --git a/reme_ai/mem_agent/v4/personal_retriever_v4.py b/reme_ai/mem_agent/v4/personal_retriever_v4.py index 5135acd9..2ba0dba5 100644 --- a/reme_ai/mem_agent/v4/personal_retriever_v4.py +++ b/reme_ai/mem_agent/v4/personal_retriever_v4.py @@ -1,7 +1,7 @@ from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message -from ...core.utils import format_messages +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message +from ...core_old.utils import format_messages from ...mem_tool.v4 import ReadUserProfile @@ -41,7 +41,7 @@ class PersonalRetrieverV4(BaseMemoryAgent): async def execute(self): """Execute the retriever and determine success based on output markers.""" await super().execute() - + # Check for memory found/not found markers in the output if self.output: if "" in self.output: diff --git a/reme_ai/mem_agent/v4/personal_summarizer_v4.py b/reme_ai/mem_agent/v4/personal_summarizer_v4.py index c0e1d4f1..8cda9b24 100644 --- a/reme_ai/mem_agent/v4/personal_summarizer_v4.py +++ b/reme_ai/mem_agent/v4/personal_summarizer_v4.py @@ -1,8 +1,8 @@ from loguru import logger from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, MemoryNode +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, MemoryNode class PersonalSummarizerV4(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/v4/personal_summarizer_v4.yaml b/reme_ai/mem_agent/v4/personal_summarizer_v4.yaml index 5e7d12b7..3663c992 100644 --- a/reme_ai/mem_agent/v4/personal_summarizer_v4.yaml +++ b/reme_ai/mem_agent/v4/personal_summarizer_v4.yaml @@ -13,7 +13,7 @@ user_message_phase1: | **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) @@ -34,13 +34,13 @@ user_message_phase2: | {user_profile} ## Task: Update Profile with `UpdateUserProfile` - + Synchronize profile with new information from the conversation: - `profile_ids_to_delete`: Remove conflicting, or redundant entries (array of profile IDs). - `profiles_to_add`: - `conversation_time`: Time of conversation (format: `YYYY-MM-DD HH:MM:SS`, e.g., `2024-01-15 14:30:00`) - `profile_content`: Complete, self-contained profile description with full context - + **Profile Requirements**: - One user profile entry records one dimension of the user portrait, and MUST be complete and self-contained with all necessary context (preconditions, causes, and consequences) - All profiles MUST be mutually exclusive (non-overlapping) and non-conflicting diff --git a/reme_ai/mem_agent/v4/reme_retriever_v4.py b/reme_ai/mem_agent/v4/reme_retriever_v4.py index 7dbf2b8a..db5f92ae 100644 --- a/reme_ai/mem_agent/v4/reme_retriever_v4.py +++ b/reme_ai/mem_agent/v4/reme_retriever_v4.py @@ -1,9 +1,9 @@ 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 +from ...core_old.enumeration import Role +from ...core_old.schema import Message +from ...core_old.utils import format_messages class ReMeRetrieverV4(BaseMemoryAgent): @@ -47,7 +47,7 @@ class ReMeRetrieverV4(BaseMemoryAgent): async def _acting_step(self, assistant_message: Message, step: int, **kwargs) -> list[Message]: import asyncio from ...mem_tool.v4 import HandsOff - + if not assistant_message.tool_calls: return [] @@ -98,16 +98,16 @@ class ReMeRetrieverV4(BaseMemoryAgent): tool_call_id=op.tool_call.id, ) tool_result_messages.append(tool_message) - + self.meta_info += tool_result + "\n" - + logger.info(f"[{self.__class__.__name__}{stage_prefix}] step{step + 1}.{j} join tool_result={tool_result[:2000]}...\n\n") - + return tool_result_messages async def execute(self): await super().execute() - + # Assemble meta_info_dict into output if self.meta_info_dict: output_parts = [] diff --git a/reme_ai/mem_agent/v4/reme_summarizer_v4.py b/reme_ai/mem_agent/v4/reme_summarizer_v4.py index a4069c85..7787d035 100644 --- a/reme_ai/mem_agent/v4/reme_summarizer_v4.py +++ b/reme_ai/mem_agent/v4/reme_summarizer_v4.py @@ -1,9 +1,9 @@ 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 +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, MemoryNode +from ...core_old.utils import format_messages class ReMeSummarizerV4(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/v4/reme_summarizer_v4.yaml b/reme_ai/mem_agent/v4/reme_summarizer_v4.yaml index a6d322ef..4cb6f54f 100644 --- a/reme_ai/mem_agent/v4/reme_summarizer_v4.yaml +++ b/reme_ai/mem_agent/v4/reme_summarizer_v4.yaml @@ -19,7 +19,7 @@ system_prompt: | - 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 ``. user_message: | diff --git a/reme_ai/mem_agent/wk/personal_summarizer_wk.py b/reme_ai/mem_agent/wk/personal_summarizer_wk.py index e974f99a..c95ac36b 100644 --- a/reme_ai/mem_agent/wk/personal_summarizer_wk.py +++ b/reme_ai/mem_agent/wk/personal_summarizer_wk.py @@ -1,7 +1,7 @@ from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, ToolCall -from ...core.utils import format_messages +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, ToolCall +from ...core_old.utils import format_messages class PersonalSummarizerWk(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/wk/reme_retriever_wk.py b/reme_ai/mem_agent/wk/reme_retriever_wk.py index c98de937..21a5403a 100644 --- a/reme_ai/mem_agent/wk/reme_retriever_wk.py +++ b/reme_ai/mem_agent/wk/reme_retriever_wk.py @@ -3,9 +3,9 @@ from typing import List from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role -from ...core.schema import Message -from ...core.utils import format_messages +from ...core_old.enumeration import Role +from ...core_old.schema import Message +from ...core_old.utils import format_messages class ReMeRetrieverV2(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/wk/reme_retriever_wk.yaml b/reme_ai/mem_agent/wk/reme_retriever_wk.yaml index 281796d6..755fd6b7 100644 --- a/reme_ai/mem_agent/wk/reme_retriever_wk.yaml +++ b/reme_ai/mem_agent/wk/reme_retriever_wk.yaml @@ -24,16 +24,16 @@ system_prompt: | 1. **Multi-Angle Vector Retrieval** (REQUIRED - at least 3 attempts): You must try AT LEAST 3 different retrieval approaches using `retrieve_memories`: - + a) **Direct Vector Search**: Use the user's question directly or with minimal reformulation - Query the most relevant memory_type and memory_target - Use straightforward query phrasing - + b) **Alternative Phrasing**: Reformulate the query from a different angle - Use synonyms or different expressions - Break down complex questions into simpler components - Try more specific or more general queries - + c) **Metadata-Filtered Search**: Add metadata filters to narrow down results - **Time-based filtering**: Use year/month/day metadata fields to filter by time periods * Example: {{"year": 2024}} for memories from 2024 @@ -41,33 +41,33 @@ system_prompt: | * Example: {{"year": 2024, "month": 5, "day": 15}} for memories from a specific date - Combine vector search with metadata constraints - Try partial metadata filtering if full filtering yields nothing (e.g., only year, or year+month) - + d) **Cross-Memory-Type Search**: If applicable, search across different memory types - Try different memory_type and memory_target combinations - Some information might be stored in unexpected memory categories - + e) **Keyword Extraction**: Extract key entities/concepts and search for them - Identify important names, places, concepts - Search for each key element separately - + 2. **Evaluate Retrieval Results** (After each attempt): - Review what memories were returned - Assess if they contain sufficient information to answer the question - If insufficient, identify what's missing and adjust your next query accordingly - Track which retrieval strategies you've already tried - + 3. **Persist Through Failures**: - DO NOT give up after 1-2 failed attempts - If a retrieval returns no results or irrelevant results, try a different approach - Consider that the information might be phrased differently than expected - Be creative with query reformulation - + 4. **Fallback to History Reading** (Only after 3+ vector retrieval attempts): - If after at least 3 different vector retrieval attempts you still lack sufficient information: * If any retrieved memories contain `ref_memory_id`, use `read_history` to read the original conversation * Use `read_history` with the `ref_memory_id` to get complete context * This can reveal details that weren't captured in the memory summaries - + 5. **Answer the Question**: - Once you have sufficient information, provide a direct answer based ONLY on retrieved memories - DO NOT fabricate, guess, or infer information not present in the memories @@ -95,30 +95,30 @@ system_prompt: | **Example 1: Simple Query** Attempt 1: Direct query "user's favorite food" → Result: No relevant memories found - + Attempt 2: Reformulated query "what does user like to eat" → Result: Some memories about meals, but not specific preferences - + Attempt 3: Keyword search "food preferences" with metadata filter → Result: Found relevant memory with ref_memory_id - + Attempt 4: Use read_history with ref_memory_id to get full context → Result: Found detailed conversation about favorite foods - + Answer: [Provide answer based on retrieved information] **Example 2: Time-based Query** Question: "What did the user do last summer?" - + Attempt 1: Direct query "user activities summer" with metadata {{"year": 2025, "month": [6, 7, 8]}} → Result: Found some vacation memories - + Attempt 2: Broader query "user summer vacation travel" with metadata {{"year": 2025}} → Result: Found additional travel-related memories - + Attempt 3: Use read_history for memories with ref_memory_id to get detailed context → Result: Complete picture of summer activities - + Answer: [Provide answer based on retrieved information] user_message: | diff --git a/reme_ai/mem_agent/wk/reme_summarizer_wk.py b/reme_ai/mem_agent/wk/reme_summarizer_wk.py index 02a7dbf3..a04d230d 100644 --- a/reme_ai/mem_agent/wk/reme_summarizer_wk.py +++ b/reme_ai/mem_agent/wk/reme_summarizer_wk.py @@ -1,9 +1,9 @@ from loguru import logger from ..base_memory_agent import BaseMemoryAgent -from ...core.enumeration import Role, MemoryType -from ...core.schema import Message, MemoryNode, ToolCall -from ...core.utils import format_messages +from ...core_old.enumeration import Role, MemoryType +from ...core_old.schema import Message, MemoryNode, ToolCall +from ...core_old.utils import format_messages class ReMeSummarizerWk(BaseMemoryAgent): diff --git a/reme_ai/mem_agent/wk/reme_summarizer_wk.yaml b/reme_ai/mem_agent/wk/reme_summarizer_wk.yaml index 30792a08..58080cab 100644 --- a/reme_ai/mem_agent/wk/reme_summarizer_wk.yaml +++ b/reme_ai/mem_agent/wk/reme_summarizer_wk.yaml @@ -18,7 +18,7 @@ system_prompt: | 2. Identify which memory dimensions need updates and specify them in `memory_tasks` (each with `memory_type` and `memory_target`). - The `memory_type` and `memory_target` must exactly match existing entries in the "Main Agent's Meta Memory" listed above. - Multiple tasks can be specified to enable parallel processing by specialized agents. - + Note: If the context contains no memorable information (e.g., simple greetings), output ``. user_message: | diff --git a/reme_ai/mem_tool/base_memory_tool.py b/reme_ai/mem_tool/base_memory_tool.py index 8b124496..827066b9 100644 --- a/reme_ai/mem_tool/base_memory_tool.py +++ b/reme_ai/mem_tool/base_memory_tool.py @@ -3,10 +3,10 @@ from abc import ABCMeta from pathlib import Path -from ..core.enumeration import MemoryType -from ..core.op import BaseOp -from ..core.schema import ToolCall, MemoryNode -from ..core.utils import CacheHandler +from ..core_old.enumeration import MemoryType +from ..core_old.op import BaseOp +from ..core_old.schema import ToolCall, MemoryNode +from ..core_old.utils import CacheHandler class BaseMemoryTool(BaseOp, metaclass=ABCMeta): diff --git a/reme_ai/mem_tool/hands_off_tool.py b/reme_ai/mem_tool/hands_off_tool.py index 2cb28d60..7ab65cd7 100644 --- a/reme_ai/mem_tool/hands_off_tool.py +++ b/reme_ai/mem_tool/hands_off_tool.py @@ -6,8 +6,8 @@ from typing import TYPE_CHECKING from loguru import logger from .base_memory_tool import BaseMemoryTool -from ..core.context import C -from ..core.enumeration import MemoryType +from ..core_old.context import C +from ..core_old.enumeration import MemoryType if TYPE_CHECKING: from ..mem_agent import BaseMemoryAgent diff --git a/reme_ai/mem_tool/history/add_history_memory.py b/reme_ai/mem_tool/history/add_history_memory.py index a92deca5..85e0181a 100644 --- a/reme_ai/mem_tool/history/add_history_memory.py +++ b/reme_ai/mem_tool/history/add_history_memory.py @@ -3,10 +3,10 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.enumeration import MemoryType -from ...core.schema import ToolCall, Message -from ...core.utils import format_messages +from ...core_old.context import C +from ...core_old.enumeration import MemoryType +from ...core_old.schema import ToolCall, Message +from ...core_old.utils import format_messages @C.register_op() diff --git a/reme_ai/mem_tool/history/read_history_memory.py b/reme_ai/mem_tool/history/read_history_memory.py index def2ff24..24cc0366 100644 --- a/reme_ai/mem_tool/history/read_history_memory.py +++ b/reme_ai/mem_tool/history/read_history_memory.py @@ -3,8 +3,8 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode +from ...core_old.context import C +from ...core_old.schema import MemoryNode @C.register_op() diff --git a/reme_ai/mem_tool/identity/read_identity_memory.py b/reme_ai/mem_tool/identity/read_identity_memory.py index bd9f8031..dfede68f 100644 --- a/reme_ai/mem_tool/identity/read_identity_memory.py +++ b/reme_ai/mem_tool/identity/read_identity_memory.py @@ -3,7 +3,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C +from ...core_old.context import C @C.register_op() diff --git a/reme_ai/mem_tool/identity/update_identity_memory.py b/reme_ai/mem_tool/identity/update_identity_memory.py index b0211242..0883b1a9 100644 --- a/reme_ai/mem_tool/identity/update_identity_memory.py +++ b/reme_ai/mem_tool/identity/update_identity_memory.py @@ -3,7 +3,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C +from ...core_old.context import C @C.register_op() diff --git a/reme_ai/mem_tool/meta/add_meta_memory.py b/reme_ai/mem_tool/meta/add_meta_memory.py index d7b41254..99635698 100644 --- a/reme_ai/mem_tool/meta/add_meta_memory.py +++ b/reme_ai/mem_tool/meta/add_meta_memory.py @@ -5,8 +5,8 @@ import json from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.enumeration import MemoryType +from ...core_old.context import C +from ...core_old.enumeration import MemoryType @C.register_op() diff --git a/reme_ai/mem_tool/meta/read_meta_memory.py b/reme_ai/mem_tool/meta/read_meta_memory.py index 07ad1ecf..5ee59d54 100644 --- a/reme_ai/mem_tool/meta/read_meta_memory.py +++ b/reme_ai/mem_tool/meta/read_meta_memory.py @@ -3,8 +3,8 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.enumeration import MemoryType +from ...core_old.context import C +from ...core_old.enumeration import MemoryType @C.register_op() diff --git a/reme_ai/mem_tool/read_local_memories.py b/reme_ai/mem_tool/read_local_memories.py index 98d96187..f239d5d1 100644 --- a/reme_ai/mem_tool/read_local_memories.py +++ b/reme_ai/mem_tool/read_local_memories.py @@ -38,7 +38,7 @@ class ReadLocalMemories(BaseMemoryTool): cache_key = f"{memory_type}_{memory_target}" cached_data = self.meta_memory.load(cache_key, auto_clean=False) - + if not cached_data: self.output = f"Local memory not found: {memory_type}_{memory_target}" logger.info(self.output) diff --git a/reme_ai/mem_tool/think_tool.py b/reme_ai/mem_tool/think_tool.py index 1d26446a..c54c7676 100644 --- a/reme_ai/mem_tool/think_tool.py +++ b/reme_ai/mem_tool/think_tool.py @@ -5,8 +5,8 @@ before taking actions, helping agents reason about their next steps. """ from .base_memory_tool import BaseMemoryTool -from ..core.context import C -from ..core.schema import ToolCall +from ..core_old.context import C +from ..core_old.schema import ToolCall @C.register_op() diff --git a/reme_ai/mem_tool/v2/add_memory_drafts.py b/reme_ai/mem_tool/v2/add_memory_drafts.py index 93caad0c..b815be6e 100644 --- a/reme_ai/mem_tool/v2/add_memory_drafts.py +++ b/reme_ai/mem_tool/v2/add_memory_drafts.py @@ -3,7 +3,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C +from ...core_old.context import C @C.register_op() diff --git a/reme_ai/mem_tool/v2/read_history.py b/reme_ai/mem_tool/v2/read_history.py index 15aaada3..141989c0 100644 --- a/reme_ai/mem_tool/v2/read_history.py +++ b/reme_ai/mem_tool/v2/read_history.py @@ -3,20 +3,20 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode +from ...core_old.context import C +from ...core_old.schema import MemoryNode @C.register_op() class ReadHistory(BaseMemoryTool): """Read original history dialogue by reference memory ID. - + Only supports single memory read (enable_multiple=False). """ def __init__(self, **kwargs): """Initialize ReadHistory. - + Args: **kwargs: Additional args for BaseMemoryTool. """ diff --git a/reme_ai/mem_tool/v2/retrieve_memories.py b/reme_ai/mem_tool/v2/retrieve_memories.py index abdfc377..96d4ca99 100644 --- a/reme_ai/mem_tool/v2/retrieve_memories.py +++ b/reme_ai/mem_tool/v2/retrieve_memories.py @@ -3,9 +3,9 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode, VectorNode -from ...core.utils import deduplicate_memories +from ...core_old.context import C +from ...core_old.schema import MemoryNode, VectorNode +from ...core_old.utils import deduplicate_memories @C.register_op() diff --git a/reme_ai/mem_tool/v2/retrieve_memories.yaml b/reme_ai/mem_tool/v2/retrieve_memories.yaml index f83e11cb..0e65d246 100644 --- a/reme_ai/mem_tool/v2/retrieve_memories.yaml +++ b/reme_ai/mem_tool/v2/retrieve_memories.yaml @@ -9,7 +9,7 @@ tool_multiple: | This prevents redundant information in subsequent retrievals. memory_type: | - The type of memory to search for. + The type of memory to search for. You MUST select one of the memory_type values that are explicitly provided in the Available Meta-Memories. memory_target: | diff --git a/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.py b/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.py index 38107ab0..cfed10c8 100644 --- a/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.py +++ b/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.py @@ -3,9 +3,9 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode, VectorNode -from ...core.utils import deduplicate_memories +from ...core_old.context import C +from ...core_old.schema import MemoryNode, VectorNode +from ...core_old.utils import deduplicate_memories @C.register_op() diff --git a/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.yaml b/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.yaml index ed91357e..d5791a5d 100644 --- a/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.yaml +++ b/reme_ai/mem_tool/v2/retrieve_recent_and_similar_memories.yaml @@ -1,16 +1,16 @@ tool_multiple: | Retrieve memories using both time-based and multiple vector similarity searches. - + This tool combines two retrieval strategies: 1. First retrieves the most recent memories based on modification time (recent top {recent_top_k}) 2. Then retrieves semantically similar memories for each of your queries (similar top {similar_top_k} per query) - + This is useful when you need to search for different types of information in a single operation, while also considering recent context. - + The results are automatically deduplicated, so you get a combined set of both recent and relevant memories without duplicates. - + Note: Within the same session, this tool automatically deduplicates results across multiple calls. If you call this tool multiple times, only new memories (not previously retrieved) will be returned. This prevents redundant information in subsequent retrievals. diff --git a/reme_ai/mem_tool/v2/summary_and_hands_off.py b/reme_ai/mem_tool/v2/summary_and_hands_off.py index 131b1101..631c7d6b 100644 --- a/reme_ai/mem_tool/v2/summary_and_hands_off.py +++ b/reme_ai/mem_tool/v2/summary_and_hands_off.py @@ -6,9 +6,9 @@ from typing import TYPE_CHECKING from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode, Message +from ...core_old.context import C +from ...core_old.enumeration import MemoryType +from ...core_old.schema import MemoryNode, Message if TYPE_CHECKING: from ...mem_agent import BaseMemoryAgent diff --git a/reme_ai/mem_tool/v2/update_memories.py b/reme_ai/mem_tool/v2/update_memories.py index 03a9f494..cffe74ac 100644 --- a/reme_ai/mem_tool/v2/update_memories.py +++ b/reme_ai/mem_tool/v2/update_memories.py @@ -3,8 +3,8 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode +from ...core_old.context import C +from ...core_old.schema import MemoryNode @C.register_op() diff --git a/reme_ai/mem_tool/v3/add_memory.py b/reme_ai/mem_tool/v3/add_memory.py index ee488639..4893c6a2 100644 --- a/reme_ai/mem_tool/v3/add_memory.py +++ b/reme_ai/mem_tool/v3/add_memory.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class AddMemory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v3/read_history.py b/reme_ai/mem_tool/v3/read_history.py index e9ab2a15..9506d88a 100644 --- a/reme_ai/mem_tool/v3/read_history.py +++ b/reme_ai/mem_tool/v3/read_history.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class ReadHistory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v3/read_user_profile.py b/reme_ai/mem_tool/v3/read_user_profile.py index 3dba2bcf..bab1d413 100644 --- a/reme_ai/mem_tool/v3/read_user_profile.py +++ b/reme_ai/mem_tool/v3/read_user_profile.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema.memory_node import MemoryNode +from ...core_old.schema.memory_node import MemoryNode class ReadUserProfile(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v3/retrieve_memory.py b/reme_ai/mem_tool/v3/retrieve_memory.py index 32e526d7..d5b9a7bc 100644 --- a/reme_ai/mem_tool/v3/retrieve_memory.py +++ b/reme_ai/mem_tool/v3/retrieve_memory.py @@ -3,8 +3,8 @@ import json from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode -from ...core.utils import deduplicate_memories +from ...core_old.schema import MemoryNode +from ...core_old.utils import deduplicate_memories class RetrieveMemory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v3/summary_and_hands_off.py b/reme_ai/mem_tool/v3/summary_and_hands_off.py index 19d88744..e0b2756c 100644 --- a/reme_ai/mem_tool/v3/summary_and_hands_off.py +++ b/reme_ai/mem_tool/v3/summary_and_hands_off.py @@ -4,8 +4,8 @@ from typing import TYPE_CHECKING from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode, Message +from ...core_old.enumeration import MemoryType +from ...core_old.schema import MemoryNode, Message if TYPE_CHECKING: from ...mem_agent import BaseMemoryAgent diff --git a/reme_ai/mem_tool/v3/update_user_profile.py b/reme_ai/mem_tool/v3/update_user_profile.py index 46879e2b..ef47f46b 100644 --- a/reme_ai/mem_tool/v3/update_user_profile.py +++ b/reme_ai/mem_tool/v3/update_user_profile.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema.memory_node import MemoryNode +from ...core_old.schema.memory_node import MemoryNode class UpdateUserProfile(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v4/add_summary_memory.py b/reme_ai/mem_tool/v4/add_summary_memory.py index cc4be602..6fe30cda 100644 --- a/reme_ai/mem_tool/v4/add_summary_memory.py +++ b/reme_ai/mem_tool/v4/add_summary_memory.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class AddSummaryMemory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v4/hands_off.py b/reme_ai/mem_tool/v4/hands_off.py index 17a33429..2dab4531 100644 --- a/reme_ai/mem_tool/v4/hands_off.py +++ b/reme_ai/mem_tool/v4/hands_off.py @@ -3,8 +3,8 @@ 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 +from ...core_old.enumeration import MemoryType +from ...core_old.schema import Message if TYPE_CHECKING: from ...mem_agent import BaseMemoryAgent @@ -62,14 +62,14 @@ class HandsOff(BaseMemoryTool): 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, diff --git a/reme_ai/mem_tool/v4/read_history.py b/reme_ai/mem_tool/v4/read_history.py index 78a90eb4..00097234 100644 --- a/reme_ai/mem_tool/v4/read_history.py +++ b/reme_ai/mem_tool/v4/read_history.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class ReadHistory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/v4/read_user_profile.py b/reme_ai/mem_tool/v4/read_user_profile.py index aa8a57f8..f45c0c34 100644 --- a/reme_ai/mem_tool/v4/read_user_profile.py +++ b/reme_ai/mem_tool/v4/read_user_profile.py @@ -2,7 +2,7 @@ from typing import Literal from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema.memory_node import MemoryNode +from ...core_old.schema.memory_node import MemoryNode class ReadUserProfile(BaseMemoryTool): @@ -39,11 +39,11 @@ class ReadUserProfile(BaseMemoryTool): "required": [], } - async def execute(self): + async def execute(self): # Determine which IDs to show show_profile_id = self.show_ids in ("both", "profile") show_history_id = self.show_ids in ("both", "history") - + cache_key = f"{self.memory_type}_{self.memory_target}".replace(" ", "_").lower() cached_data = self.meta_memory.load(cache_key, auto_clean=False) @@ -58,22 +58,22 @@ class ReadUserProfile(BaseMemoryTool): memory_formated = [] for node in memory_nodes: node_formated_parts = [] - + # Add profile_id if enabled if show_profile_id: node_formated_parts.append(f"profile_id={node.memory_id}") - + # Always add profile_content node_formated_parts.append(f"profile_content={node.content}") - + # Add conversation_time if available if "conversation_time" in node.metadata and node.metadata["conversation_time"]: node_formated_parts.append(f"conversation_time={node.metadata['conversation_time']}") - + # Add history_id if enabled and available if show_history_id and node.ref_memory_id: node_formated_parts.append(f"history_id={node.ref_memory_id}") - + node_formated = " ".join(node_formated_parts) memory_formated.append(node_formated.strip()) diff --git a/reme_ai/mem_tool/v4/retrieve_memory.py b/reme_ai/mem_tool/v4/retrieve_memory.py index a8138bdc..a6717bfa 100644 --- a/reme_ai/mem_tool/v4/retrieve_memory.py +++ b/reme_ai/mem_tool/v4/retrieve_memory.py @@ -3,8 +3,8 @@ import json from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode -from ...core.utils import deduplicate_memories +from ...core_old.schema import MemoryNode +from ...core_old.utils import deduplicate_memories class RetrieveMemory(BaseMemoryTool): @@ -62,7 +62,7 @@ class RetrieveMemory(BaseMemoryTool): 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: diff --git a/reme_ai/mem_tool/v4/update_user_profile.py b/reme_ai/mem_tool/v4/update_user_profile.py index a8fa04f5..1ab3276c 100644 --- a/reme_ai/mem_tool/v4/update_user_profile.py +++ b/reme_ai/mem_tool/v4/update_user_profile.py @@ -1,8 +1,8 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema.memory_node import MemoryNode -from ...core.utils import deduplicate_memories +from ...core_old.schema.memory_node import MemoryNode +from ...core_old.utils import deduplicate_memories class UpdateUserProfile(BaseMemoryTool): diff --git a/reme_ai/mem_tool/vector_store/add_memory.py b/reme_ai/mem_tool/vector_store/add_memory.py index 1f87df4b..0937fe2a 100644 --- a/reme_ai/mem_tool/vector_store/add_memory.py +++ b/reme_ai/mem_tool/vector_store/add_memory.py @@ -3,8 +3,8 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode +from ...core_old.context import C +from ...core_old.schema import MemoryNode @C.register_op() diff --git a/reme_ai/mem_tool/vector_store/add_summary_memory.py b/reme_ai/mem_tool/vector_store/add_summary_memory.py index ce1127ed..54abbd57 100644 --- a/reme_ai/mem_tool/vector_store/add_summary_memory.py +++ b/reme_ai/mem_tool/vector_store/add_summary_memory.py @@ -3,9 +3,9 @@ from loguru import logger from .add_memory import AddMemory -from ...core.context import C -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode +from ...core_old.context import C +from ...core_old.enumeration import MemoryType +from ...core_old.schema import MemoryNode @C.register_op() diff --git a/reme_ai/mem_tool/vector_store/delete_memory.py b/reme_ai/mem_tool/vector_store/delete_memory.py index 45b28632..95f4f210 100644 --- a/reme_ai/mem_tool/vector_store/delete_memory.py +++ b/reme_ai/mem_tool/vector_store/delete_memory.py @@ -3,7 +3,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C +from ...core_old.context import C @C.register_op() diff --git a/reme_ai/mem_tool/vector_store/retrieve_recent_memory.py b/reme_ai/mem_tool/vector_store/retrieve_recent_memory.py index 0896d75a..dfaa9240 100644 --- a/reme_ai/mem_tool/vector_store/retrieve_recent_memory.py +++ b/reme_ai/mem_tool/vector_store/retrieve_recent_memory.py @@ -3,9 +3,9 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode, VectorNode -from ...core.utils import deduplicate_memories +from ...core_old.context import C +from ...core_old.schema import MemoryNode, VectorNode +from ...core_old.utils import deduplicate_memories @C.register_op() diff --git a/reme_ai/mem_tool/vector_store/update_memory.py b/reme_ai/mem_tool/vector_store/update_memory.py index 4873ce24..fe08fb90 100644 --- a/reme_ai/mem_tool/vector_store/update_memory.py +++ b/reme_ai/mem_tool/vector_store/update_memory.py @@ -3,8 +3,8 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema import MemoryNode +from ...core_old.context import C +from ...core_old.schema import MemoryNode @C.register_op() diff --git a/reme_ai/mem_tool/vector_store/vector_retrieve_memory.py b/reme_ai/mem_tool/vector_store/vector_retrieve_memory.py index 655de36e..24b05e8b 100644 --- a/reme_ai/mem_tool/vector_store/vector_retrieve_memory.py +++ b/reme_ai/mem_tool/vector_store/vector_retrieve_memory.py @@ -3,10 +3,10 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode, VectorNode -from ...core.utils import deduplicate_memories +from ...core_old.context import C +from ...core_old.enumeration import MemoryType +from ...core_old.schema import MemoryNode, VectorNode +from ...core_old.utils import deduplicate_memories @C.register_op() diff --git a/reme_ai/mem_tool/wk/add_memory.py b/reme_ai/mem_tool/wk/add_memory.py index 61cb154c..2af7ba71 100644 --- a/reme_ai/mem_tool/wk/add_memory.py +++ b/reme_ai/mem_tool/wk/add_memory.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class AddMemory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/wk/read_history.py b/reme_ai/mem_tool/wk/read_history.py index 945c02e2..65e8a0cf 100644 --- a/reme_ai/mem_tool/wk/read_history.py +++ b/reme_ai/mem_tool/wk/read_history.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class ReadHistory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/wk/summary_and_hands_off.py b/reme_ai/mem_tool/wk/summary_and_hands_off.py index 9a6a76cd..d384f16a 100644 --- a/reme_ai/mem_tool/wk/summary_and_hands_off.py +++ b/reme_ai/mem_tool/wk/summary_and_hands_off.py @@ -4,8 +4,8 @@ from typing import TYPE_CHECKING from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode, Message +from ...core_old.enumeration import MemoryType +from ...core_old.schema import MemoryNode, Message if TYPE_CHECKING: from ...mem_agent import BaseMemoryAgent diff --git a/reme_ai/mem_tool/wk/update_memory.py b/reme_ai/mem_tool/wk/update_memory.py index c419261f..151eba53 100644 --- a/reme_ai/mem_tool/wk/update_memory.py +++ b/reme_ai/mem_tool/wk/update_memory.py @@ -1,7 +1,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode +from ...core_old.schema import MemoryNode class UpdateMemory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/wk/vector_retrieve_memory.py b/reme_ai/mem_tool/wk/vector_retrieve_memory.py index de7498b4..2698cd6f 100644 --- a/reme_ai/mem_tool/wk/vector_retrieve_memory.py +++ b/reme_ai/mem_tool/wk/vector_retrieve_memory.py @@ -1,9 +1,9 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode, VectorNode -from ...core.utils import deduplicate_memories +from ...core_old.enumeration import MemoryType +from ...core_old.schema import MemoryNode, VectorNode +from ...core_old.utils import deduplicate_memories class VectorRetrieveMemory(BaseMemoryTool): diff --git a/reme_ai/mem_tool/write_local_memories.py b/reme_ai/mem_tool/write_local_memories.py index ac1c5aae..f6336395 100644 --- a/reme_ai/mem_tool/write_local_memories.py +++ b/reme_ai/mem_tool/write_local_memories.py @@ -30,13 +30,13 @@ class WriteLocalMemories(BaseMemoryTool): async def execute(self): memory_nodes = self.context.get("memory_nodes", []) - + if not memory_nodes: self.output = "No memory nodes provided." return memory_nodes = [MemoryNode(**node) if isinstance(node, dict) else node for node in memory_nodes] - + grouped = {} for node in memory_nodes: key = (node.memory_type.value, node.memory_target) @@ -45,11 +45,11 @@ class WriteLocalMemories(BaseMemoryTool): grouped[key].append(node) written_keys = [] - + for (memory_type, memory_target), nodes in grouped.items(): cache_key = f"{memory_type}_{memory_target}" nodes_data = [node.model_dump() for node in nodes] - + self.meta_memory.save(cache_key, nodes_data) written_keys.append(f"{memory_type}_{memory_target}") logger.info(f"Saved {len(nodes)} nodes to cache key: {cache_key}") diff --git a/reme_ai/tool/execute/execute_code.py b/reme_ai/tool/execute/execute_code.py index ea259487..f13aab3d 100644 --- a/reme_ai/tool/execute/execute_code.py +++ b/reme_ai/tool/execute/execute_code.py @@ -4,11 +4,11 @@ This module provides an operation that can execute Python code strings and return the output or error messages. """ -from ...core.context import C -from ...core.op import BaseOp -from ...core.schema import ToolCall +from ...core_old.context import C +from ...core_old.op import BaseOp +from ...core_old.schema import ToolCall -from ...core.utils import exec_code +from ...core_old.utils import exec_code @C.register_op() diff --git a/reme_ai/tool/execute/execute_shell.py b/reme_ai/tool/execute/execute_shell.py index 6e244ddb..1b235921 100644 --- a/reme_ai/tool/execute/execute_shell.py +++ b/reme_ai/tool/execute/execute_shell.py @@ -4,11 +4,11 @@ This module provides an operation that can execute shell commands asynchronously and return the output, error, and exit code. """ -from ...core.context import C -from ...core.op import BaseOp -from ...core.schema import ToolCall +from ...core_old.context import C +from ...core_old.op import BaseOp +from ...core_old.schema import ToolCall -from ...core.utils import run_shell_command +from ...core_old.utils import run_shell_command @C.register_op() diff --git a/reme_ai/tool/search/dashscope_search.py b/reme_ai/tool/search/dashscope_search.py index 19bd8104..47c399ef 100644 --- a/reme_ai/tool/search/dashscope_search.py +++ b/reme_ai/tool/search/dashscope_search.py @@ -9,9 +9,9 @@ from typing import Literal from loguru import logger -from ...core.context import C -from ...core.op import BaseOp -from ...core.schema import ToolCall +from ...core_old.context import C +from ...core_old.op import BaseOp +from ...core_old.schema import ToolCall @C.register_op() diff --git a/reme_ai/tool/search/mock_search.py b/reme_ai/tool/search/mock_search.py index 187463dc..69ef3995 100644 --- a/reme_ai/tool/search/mock_search.py +++ b/reme_ai/tool/search/mock_search.py @@ -9,11 +9,11 @@ import random from loguru import logger -from ...core.context import C -from ...core.enumeration import Role -from ...core.op import BaseOp -from ...core.schema import ToolCall, Message -from ...core.utils import extract_content +from ...core_old.context import C +from ...core_old.enumeration import Role +from ...core_old.op import BaseOp +from ...core_old.schema import ToolCall, Message +from ...core_old.utils import extract_content @C.register_op() diff --git a/reme_ai/tool/search/tavily_search.py b/reme_ai/tool/search/tavily_search.py index 5c194bdc..bb000f16 100644 --- a/reme_ai/tool/search/tavily_search.py +++ b/reme_ai/tool/search/tavily_search.py @@ -9,9 +9,9 @@ import os from loguru import logger -from ...core.context import C -from ...core.op import BaseOp -from ...core.schema import ToolCall +from ...core_old.context import C +from ...core_old.op import BaseOp +from ...core_old.schema import ToolCall @C.register_op() diff --git a/test/http_client_test.py b/test/http_client_test.py deleted file mode 100644 index acd1254c..00000000 --- a/test/http_client_test.py +++ /dev/null @@ -1,165 +0,0 @@ -import asyncio -import json - -import aiohttp - -base_url = "http://0.0.0.0:8002" - - -async def run1(session): - workspace_id = "default1" - - async with session.post( - f"{base_url}/vector_store", - json={ - "action": "delete", - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - trajectories = [ - { - "task_id": "t1", - "messages": [ - {"role": "user", "content": "搜索可以使用websearch工具"}, - ], - "score": 1, - }, - { - "task_id": "t1", - "messages": [ - {"role": "user", "content": "搜索可以使用code工具"}, - ], - "score": 0, - }, - ] - - async with session.post( - # f"{base_url}/summary_task_memory", - f"{base_url}/summary_task_memory_simple", - json={ - "trajectories": trajectories, - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - await asyncio.sleep(2) - - async with session.post( - # f"{base_url}/retrieve_task_memory", - f"{base_url}/retrieve_task_memory_simple", - json={ - "query": "茅台怎么样？", - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - -async def run2(session): - workspace_id = "default2" - - async with session.post( - f"{base_url}/vector_store", - json={ - "action": "delete", - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - messages = [ - {"role": "user", "content": "我喜欢吃西瓜🍉"}, - {"role": "user", "content": "昨天吃了苹果，很好吃"}, - {"role": "user", "content": "我不太喜欢吃西瓜"}, - {"role": "user", "content": "上周我去了日本，得了肠胃炎"}, - {"role": "user", "content": "这周只能在家里，喝粥"}, - ] - - async with session.post( - f"{base_url}/summary_personal_memory", - json={ - "messages": messages, - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - await asyncio.sleep(2) - - async with session.post( - f"{base_url}/retrieve_personal_memory", - json={ - "query": "你知道我喜欢吃什么？", - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - -async def run3(session): - workspace_id = "default2" - - async with session.post( - f"{base_url}/add_tool_call_result", - json={ - "tool_call_results": [ - {"a": 1}, - {"a": 2}, - ], - "workspace_id": workspace_id, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - -async def run4(session): - workspace_id = "default4" - - async with session.post( - f"{base_url}/agentic_retrieve", - json={ - "messages": [ - {"role": "user", "content": "hello" * 10000}, - ], - "workspace_id": workspace_id, - "context_manage_mode": "auto", - "keep_recent_count": 0, - "max_total_tokens": 10000, - }, - headers={"Content-Type": "application/json"}, - ) as response: - result = await response.json() - print(json.dumps(result, ensure_ascii=False)) - - -async def main(): - - async with aiohttp.ClientSession() as session: - # 获取工具列表 - print("获取工具列表...") - - # await run1(session) - # await run2(session) - # await run3(session) - await run4(session) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/test/mcp_client_test.py b/test/mcp_client_test.py deleted file mode 100644 index 9bb08b8b..00000000 --- a/test/mcp_client_test.py +++ /dev/null @@ -1,45 +0,0 @@ -from fastmcp import Client -from mcp.types import CallToolResult - - -async def main(): - async with Client("http://0.0.0.0:8002/sse/") as client: - tools = await client.list_tools() - for tool in tools: - print(tool.model_dump_json()) - - workspace_id = "default" - - result: CallToolResult = await client.call_tool( - "retrieve_task_memory_simple", - arguments={ - "query": "茅台怎么样？", - "workspace_id": workspace_id, - }, - ) - print(result.content) - - trajectories = [ - { - "task_id": "t1", - "messages": [ - {"role": "user", "content": "今天天气不错"}, - ], - "score": 0.9, - }, - ] - - result: CallToolResult = await client.call_tool( - "summary_task_memory_simple", - arguments={ - "trajectories": trajectories, - "workspace_id": workspace_id, - }, - ) - print(result.content) - - -if __name__ == "__main__": - import asyncio - - asyncio.run(main()) diff --git a/tests/mcp_servers_demo.json b/test/mcp_servers_demo.json similarity index 100% rename from tests/mcp_servers_demo.json rename to test/mcp_servers_demo.json diff --git a/test/record_audio.py b/test/record_audio.py deleted file mode 100644 index d9446813..00000000 --- a/test/record_audio.py +++ /dev/null @@ -1,153 +0,0 @@ -#!/usr/bin/env python3 -""" -macOS 麦克风录音脚本 -需要安装: pip install pyaudio wave -""" - -import pyaudio -import wave -import sys -import os -from datetime import datetime - - -class AudioRecorder: - """macOS 音频录制器""" - - def __init__(self, output_dir="recordings"): - """ - 初始化录音器 - - Args: - output_dir: 录音文件保存目录 - """ - self.output_dir = output_dir - self.chunk = 1024 # 每次读取的音频块大小 - self.format = pyaudio.paInt16 # 16位深度 - self.channels = 1 # 单声道 - self.rate = 44100 # 采样率 44.1kHz - - # 创建输出目录 - if not os.path.exists(output_dir): - os.makedirs(output_dir) - - def record(self, duration=5, filename=None): - """ - 录制音频 - - Args: - duration: 录制时长（秒） - filename: 输出文件名，如果为None则自动生成 - - Returns: - str: 保存的文件路径 - """ - # 生成文件名 - if filename is None: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"recording_{timestamp}.wav" - - filepath = os.path.join(self.output_dir, filename) - - # 初始化PyAudio - audio = pyaudio.PyAudio() - - try: - # 打开音频流（这会触发macOS的麦克风权限请求） - print("正在请求麦克风权限...") - stream = audio.open( - format=self.format, - channels=self.channels, - rate=self.rate, - input=True, - frames_per_buffer=self.chunk - ) - - print(f"开始录音，时长: {duration} 秒") - print("录音中...") - - frames = [] - - # 录制音频 - for i in range(0, int(self.rate / self.chunk * duration)): - data = stream.read(self.chunk) - frames.append(data) - - # 显示进度 - progress = (i + 1) / (self.rate / self.chunk * duration) * 100 - sys.stdout.write(f"\r进度: {progress:.1f}%") - sys.stdout.flush() - - print("\n录音完成！") - - # 停止并关闭流 - stream.stop_stream() - stream.close() - - # 保存为WAV文件 - print(f"正在保存到: {filepath}") - wf = wave.open(filepath, 'wb') - wf.setnchannels(self.channels) - wf.setsampwidth(audio.get_sample_size(self.format)) - wf.setframerate(self.rate) - wf.writeframes(b''.join(frames)) - wf.close() - - print(f"✓ 文件已保存: {filepath}") - return filepath - - except Exception as e: - print(f"\n错误: {e}") - print("\n提示:") - print("1. 请确保已安装 pyaudio: pip install pyaudio") - print("2. 在macOS上，首次运行会弹出权限请求对话框") - print("3. 如果权限被拒绝，请前往 系统偏好设置 > 安全性与隐私 > 隐私 > 麦克风") - return None - - finally: - audio.terminate() - - def record_interactive(self): - """交互式录音""" - print("=" * 50) - print("macOS 麦克风录音工具") - print("=" * 50) - - try: - duration = input("\n请输入录音时长（秒，默认5秒）: ").strip() - duration = int(duration) if duration else 5 - - filename = input("请输入文件名（留空自动生成）: ").strip() - filename = filename if filename else None - if filename and not filename.endswith('.wav'): - filename += '.wav' - - print() - self.record(duration=duration, filename=filename) - - except KeyboardInterrupt: - print("\n\n录音已取消") - except ValueError: - print("输入无效，请输入数字") - - -def main(): - """主函数""" - recorder = AudioRecorder() - - if len(sys.argv) > 1: - # 命令行模式 - try: - duration = int(sys.argv[1]) - filename = sys.argv[2] if len(sys.argv) > 2 else None - recorder.record(duration=duration, filename=filename) - except ValueError: - print("用法: python record_audio.py [时长(秒)] [文件名(可选)]") - print("示例: python record_audio.py 10 my_recording.wav") - else: - # 交互式模式 - recorder.record_interactive() - - -if __name__ == "__main__": - main() diff --git a/test/test1.py b/test/test1.py deleted file mode 100644 index bb73331a..00000000 --- a/test/test1.py +++ /dev/null @@ -1,12 +0,0 @@ -# 2025年半年报点评：Q2业绩同比增长，CPU、DCU业务进展顺利 -# https://data.eastmoney.com/report/info/AP202508061722561937.html -# -# https://pdf.dfcfw.com/pdf/H3_AP202508061722561937_1.pdf - -import requests - -headers = { - "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)", -} -url = requests.get("https://data.eastmoney.com/report/stock.jshtml", headers=headers) -print(url.text) diff --git a/test/test2.py b/test/test2.py deleted file mode 100644 index acd329c3..00000000 --- a/test/test2.py +++ /dev/null @@ -1,393 +0,0 @@ -import json -import os -import random -import re -from datetime import datetime, timedelta -from io import BytesIO -from time import sleep -from urllib.parse import urljoin - -import pycurl -import requests -from PyPDF2 import PdfReader - -# 全局配置 -BASE_URL = "https://reportapi.eastmoney.com/report/list" -DETAIL_BASE_URL = "https://data.eastmoney.com/report/info/" - -# 读取config.json获取stock_code -with open("config.json", "r", encoding="utf-8") as f: - config = json.load(f) -STOCK_CODE = config.get("stock_code", "600519") -MIN_PAGES = config.get("min_pages", 20) -DOWNLOAD_DIR = config.get("download_dir", "reports_pdf") -YEARS_AGO = config.get("years_ago", 2) -os.makedirs(DOWNLOAD_DIR, exist_ok=True) - -# 随机User-Agent列表 -USER_AGENTS = [ - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0", -] - - -def get_random_user_agent(): - """获取随机User-Agent""" - import random - - return random.choice(USER_AGENTS) - - -def fetch_jsonp_data(page_no=1): - """ - 获取研究报告列表数据 - :param page_no: 页码 - :return: 解析后的数据字典 - """ - # 计算日期 - today = datetime.today() - end_time = today.strftime("%Y-%m-%d") - begin_time = (today - timedelta(days=365 * YEARS_AGO)).strftime("%Y-%m-%d") - - # 检查是否存在已保存的原始数据 - raw_data_dir = "raw_data" - raw_data_file = os.path.join(raw_data_dir, f"page_{page_no}_{STOCK_CODE}_{begin_time}_{end_time}.json") - - if os.path.exists(raw_data_file): - print(f"使用已保存的原始数据: {raw_data_file}") - try: - with open(raw_data_file, "r", encoding="utf-8") as f: - return json.load(f) - except Exception as e: - print(f"读取已保存数据失败: {e}") - - params = { - "cb": "datatable6333112", - "pageNo": page_no, - "pageSize": 50, - "code": STOCK_CODE, - "industryCode": "*", - "industry": "*", - "rating": "*", - "ratingchange": "*", - "beginTime": begin_time, - "endTime": end_time, - "fields": "", - "qType": 0, - "p": page_no, - "pageNum": page_no, - "pageNumber": page_no, - "_": int(time.time() * 1000), # 使用当前时间戳 - } - headers = { - "User-Agent": get_random_user_agent(), - "Referer": "https://data.eastmoney.com/", - } - try: - response = requests.get(BASE_URL, params=params, headers=headers) - response.raise_for_status() - # 提取JSON部分 - json_str = re.search(r"\((.*)\)", response.text).group(1) - data = json.loads(json_str) - - # 保存原始数据到本地 - if not os.path.exists(raw_data_dir): - os.makedirs(raw_data_dir, exist_ok=True) - - with open(raw_data_file, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=2) - - print(f"原始数据已保存: {raw_data_file}") - return data - except Exception as e: - print(f"获取第{page_no}页数据失败: {e}") - return None - - -def get_report_detail(info_code): - """ - 获取研究报告详情页内容 - :param info_code: 报告ID - :return: 详情页HTML内容 - """ - # 检查是否存在已保存的详情页HTML - detail_data_dir = "detail_data" - detail_html_file = os.path.join(detail_data_dir, f"detail_{info_code}.html") - - if os.path.exists(detail_html_file): - print(f"使用已保存的详情页HTML: {detail_html_file}") - try: - with open(detail_html_file, "r", encoding="utf-8") as f: - return f.read() - except Exception as e: - print(f"读取已保存详情页失败: {e}") - - url = urljoin(DETAIL_BASE_URL, f"{info_code}.html") - headers = { - "User-Agent": get_random_user_agent(), - "Referer": "https://data.eastmoney.com/", - } - - try: - response = requests.get(url, headers=headers) - response.raise_for_status() - - # 保存详情页HTML原始数据 - if not os.path.exists(detail_data_dir): - os.makedirs(detail_data_dir, exist_ok=True) - - with open(detail_html_file, "w", encoding="utf-8") as f: - f.write(response.text) - - print(f"详情页HTML已保存: {detail_html_file}") - return response.text - except Exception as e: - print(f"获取报告详情{info_code}失败: {e}") - return None - - -def parse_detail_page(html, info_code): - """ - 解析详情页获取PDF下载链接及相关信息 - :param html: 详情页HTML - :param info_code: 报告ID - :return: dict，包含PDF下载URL及命名所需字段 - """ - try: - # 使用正则提取zwinfo变量 - match = re.search(r"var zwinfo\s*=\s*({.*?});", html, re.DOTALL) - if not match: - return None - zwinfo = json.loads(match.group(1)) - - # 保存解析后的zwinfo数据 - detail_data_dir = "detail_data" - zwinfo_file = os.path.join(detail_data_dir, f"zwinfo_{info_code}.json") - with open(zwinfo_file, "w", encoding="utf-8") as f: - json.dump(zwinfo, f, ensure_ascii=False, indent=2) - - print(f"zwinfo数据已保存: {zwinfo_file}") - - # 提取所需字段 - return { - "attach_url": zwinfo.get("attach_url"), - "notice_title": zwinfo.get("notice_title", ""), - "short_name": zwinfo.get("short_name", ""), - "notice_date": zwinfo.get("notice_date", ""), - "source_sample_name": zwinfo.get("source_sample_name", ""), - "attach_pages": zwinfo.get("attach_pages", ""), - } - except Exception as e: - print(f"解析详情页失败: {e}") - return None - - -def is_pdf_complete(pdf_path, expected_pages): - """ - 检查PDF页数是否与预期一致 - :param pdf_path: PDF文件路径 - :param expected_pages: 预期页数（int） - :return: bool - """ - try: - with open(pdf_path, "rb") as f: - reader = PdfReader(f) - actual_pages = len(reader.pages) - return actual_pages == expected_pages, actual_pages - except Exception as e: - print(f"读取PDF页数失败: {e}") - return False, 0 - - -def download_pdf(pdf_url, filename): - """ - 使用pycurl下载PDF文件（模拟curl请求） - - 参数: - pdf_url (str): PDF文件的URL - filename (str): 保存文件名（不含路径） - - 返回: - bool: 是否下载成功 - """ - save_path = os.path.join(DOWNLOAD_DIR, filename) - buffer = BytesIO() - c = pycurl.Curl() - - try: - # 设置curl选项 - c.setopt(pycurl.URL, pdf_url) - c.setopt(pycurl.WRITEDATA, buffer) - c.setopt(pycurl.FOLLOWLOCATION, True) - c.setopt(pycurl.MAXREDIRS, 5) - c.setopt(pycurl.CONNECTTIMEOUT, 30) - c.setopt(pycurl.TIMEOUT, 300) - - # 设置防爬虫headers - headers = [ - f"User-Agent: {get_random_user_agent()}", - "Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", - "Referer: https://data.eastmoney.com/", - "Accept-Language: zh-CN,zh;q=0.9", - ] - c.setopt(pycurl.HTTPHEADER, headers) - - # 执行下载 - c.perform() - - # 验证响应 - if c.getinfo(pycurl.HTTP_CODE) != 200: - print(f"下载失败 HTTP {c.getinfo(pycurl.HTTP_CODE)}") - return False - - # 保存文件 - with open(save_path, "wb") as f: - f.write(buffer.getvalue()) - - print(f"✓ 成功下载 {filename}") - return True - - except pycurl.error as e: - errno, errstr = e.args - print(f"pycurl错误({errno}): {errstr}") - return False - except Exception as e: - print(f"下载异常: {str(e)}") - return False - finally: - c.close() - buffer.close() - - -def process_all_reports(): - """处理所有研究报告""" - # 获取第一页数据 - first_page_data = fetch_jsonp_data(1) - if not first_page_data: - return - - total_page = first_page_data.get("TotalPage", 1) - total_reports = first_page_data.get("hits", 0) - print(f"共发现{total_reports}篇研究报告，{total_page}页") - - # 处理所有页面 - for page in range(1, total_page + 1): - print(f"\n正在处理第{page}/{total_page}页...") - # 获取当前页数据 - if page == 1: - page_data = first_page_data - else: - page_data = fetch_jsonp_data(page) - if not page_data: - continue - # 处理每篇报告 - report_list = page_data.get("data", []) - random.shuffle(report_list) - for report in report_list: - info_code = report.get("infoCode") - if not info_code: - continue - - # 检查页数，只有大于20页的才下载 - attach_pages = report.get("attachPages", 0) - try: - attach_pages = int(attach_pages) - except (ValueError, TypeError): - attach_pages = 0 - - if attach_pages < MIN_PAGES: - print(f"跳过页数不足的报告: {report.get('title')} (页数: {attach_pages})") - continue - - print(f"\n处理报告: {report.get('title')} [{info_code}] (页数: {attach_pages})") - # 获取详情页 - detail_html = get_report_detail(info_code) - if not detail_html: - continue - # 解析PDF链接及命名信息 - detail_info = parse_detail_page(detail_html, info_code) - if not detail_info or not detail_info.get("attach_url"): - print("未找到PDF链接") - continue - # 组装文件名，避免重复拼接 - notice_title = detail_info.get("notice_title", "").strip().replace("/", "_") - short_name = detail_info.get("short_name", "").strip().replace("/", "_") - notice_date = detail_info.get("notice_date", "").replace("-", "")[:8] # 只取年月日 - source_sample_name = detail_info.get("source_sample_name", "").strip().replace("/", "_") - - filename_parts = [] - filename_parts.append(notice_date) - # 判断source_sample_name是否已在notice_title中 - if source_sample_name and source_sample_name not in notice_title: - filename_parts.append(source_sample_name) - # 判断short_name是否已在notice_title中 - if short_name and short_name not in notice_title: - filename_parts.append(short_name) - filename_parts.append(notice_title) - # 分离文件名和目录 - pdf_filename = f"{'_'.join(filename_parts)}.pdf" - pdf_subdir = f"{short_name}" - - # 判断是否为深度报告（页数大于20页） - if attach_pages >= 20: - pdf_subdir = f"{short_name}/深度报告" - - pdf_full_path = os.path.join(DOWNLOAD_DIR, pdf_subdir, pdf_filename) - - # 检查并创建目录 - pdf_dir = os.path.join(DOWNLOAD_DIR, pdf_subdir) - if not os.path.exists(pdf_dir): - os.makedirs(pdf_dir, exist_ok=True) - print(f"创建目录: {pdf_dir}") - - # 检查文件是否已存在 - if os.path.exists(pdf_full_path): - print(f"文件已存在，跳过下载: {pdf_full_path}") - continue - - # 下载PDF并校验页数，最多重试3次 - max_retries = 5 - for attempt in range(1, max_retries + 1): - download_pdf(detail_info["attach_url"], os.path.join(pdf_subdir, pdf_filename)) - # 校验PDF页数 - try: - expected_pages = int(detail_info.get("attach_pages", 0)) - except Exception: - expected_pages = 0 - is_complete = True - actual_pages = 0 - if expected_pages > 0: - is_complete, actual_pages = is_pdf_complete(pdf_full_path, expected_pages) - if is_complete: - print(f"✓ PDF页数校验通过：{actual_pages}页") - break - else: - print( - f"✗ PDF页数不符：实际{actual_pages}页，预期{expected_pages}页，正在重试({attempt}/{max_retries})...", - ) - # 删除不完整文件 - try: - os.remove(pdf_full_path) - except Exception: - pass - sleep(1) - else: - break - sleep(60 * attempt) - - else: - print(f"!!! PDF多次下载后仍不完整：{pdf_full_path}") - # 礼貌性延迟 - sleep(30) - - -if __name__ == "__main__": - import time - - start_time = time.time() - - process_all_reports() - - end_time = time.time() - print(f"\n全部完成，耗时: {end_time - start_time:.2f}秒") diff --git a/test/test3.py b/test/test3.py deleted file mode 100644 index 2b0021a8..00000000 --- a/test/test3.py +++ /dev/null @@ -1,92 +0,0 @@ -import os -from io import BytesIO - -import pycurl -from PyPDF2 import PdfReader - -DOWNLOAD_DIR = "./" - -# 随机User-Agent列表 -USER_AGENTS = [ - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0", -] - - -def get_random_user_agent(): - """获取随机User-Agent""" - import random - - return random.choice(USER_AGENTS) - - -def download_pdf(pdf_url, filename): - """ - 使用pycurl下载PDF文件（模拟curl请求） - - 参数: - pdf_url (str): PDF文件的URL - filename (str): 保存文件名（不含路径） - - 返回: - bool: 是否下载成功 - """ - save_path = os.path.join(DOWNLOAD_DIR, filename) - buffer = BytesIO() - c = pycurl.Curl() - - try: - # 设置curl选项 - c.setopt(pycurl.URL, pdf_url) - c.setopt(pycurl.WRITEDATA, buffer) - c.setopt(pycurl.FOLLOWLOCATION, True) - c.setopt(pycurl.MAXREDIRS, 5) - c.setopt(pycurl.CONNECTTIMEOUT, 30) - c.setopt(pycurl.TIMEOUT, 300) - - # 设置防爬虫headers - headers = [ - f"User-Agent: {get_random_user_agent()}", - "Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", - "Referer: https://data.eastmoney.com/", - "Accept-Language: zh-CN,zh;q=0.9", - ] - c.setopt(pycurl.HTTPHEADER, headers) - - # 执行下载 - c.perform() - - # 验证响应 - if c.getinfo(pycurl.HTTP_CODE) != 200: - print(f"下载失败 HTTP {c.getinfo(pycurl.HTTP_CODE)}") - return False - - # 保存文件 - with open(save_path, "wb") as f: - f.write(buffer.getvalue()) - - print(f"✓ 成功下载 {filename}") - return True - - except pycurl.error as e: - errno, errstr = e.args - print(f"pycurl错误({errno}): {errstr}") - return False - except Exception as e: - print(f"下载异常: {str(e)}") - return False - finally: - c.close() - buffer.close() - - -if __name__ == "__main__": - url_list = [ - "https://pdf.dfcfw.com/pdf/H3_AP202508061722531920_1.pdf?1754495126000.pdf", - ] - - url_list = [x.split("?")[0] for x in url_list] - for url in url_list: - name = url.split("_")[1] - download_pdf(url, f"{name}.pdf") diff --git a/test/test4.py b/test/test4.py deleted file mode 100644 index bc1b703e..00000000 --- a/test/test4.py +++ /dev/null @@ -1,67 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - - -def analyze_corrupted_text(text): - """分析乱码文本的字节构成""" - print(f"分析文本: {text}") - print(f"文本长度: {len(text)}") - - # 显示每个字符的Unicode码点 - print("字符分析:") - for i, char in enumerate(text[:20]): # 只显示前20个字符 - print(f" {i}: '{char}' -> U+{ord(char):04X}") - - # 尝试不同的编码方式 - print("\n编码尝试:") - - try: - # 方法1: Latin1 -> UTF-8 - bytes_latin1 = text.encode("latin1") - result_utf8 = bytes_latin1.decode("utf-8") - print(f"Latin1->UTF-8: {result_utf8}") - except Exception as e: - print(f"Latin1->UTF-8 失败: {e}") - - try: - # 方法2: Latin1 -> GBK - bytes_latin1 = text.encode("latin1") - result_gbk = bytes_latin1.decode("gbk") - print(f"Latin1->GBK: {result_gbk}") - except Exception as e: - print(f"Latin1->GBK 失败: {e}") - - try: - # 方法3: CP1252 -> UTF-8 - bytes_cp1252 = text.encode("cp1252") - result_utf8 = bytes_cp1252.decode("utf-8") - print(f"CP1252->UTF-8: {result_utf8}") - except Exception as e: - print(f"CP1252->UTF-8 失败: {e}") - - # 显示原始字节 - try: - raw_bytes = text.encode("latin1") - print(f"\n原始字节 (Latin1): {raw_bytes}") - print(f"字节十六进制: {raw_bytes.hex()}") - except Exception as e: - print(f"获取原始字节失败: {e}") - - -def main(): - """调试主函数""" - test_texts = [ - "为ä»ä¹è¯´æçå»è¯è¿å¥ä¸­æå¸å±æç¹ï¼", - "åçäºâäºä¸âæé´ä¸­å½ç»æµå¤è¯éªçä¹è§å¤æ­", - "æçç§ææ°ï¼HSTECH.HIï¼åº¦æ¼æ¶4.45%", - ] - - for i, text in enumerate(test_texts, 1): - print(f"\n{'=' * 60}") - print(f"测试 {i}") - print("=" * 60) - analyze_corrupted_text(text) - - -if __name__ == "__main__": - main() diff --git a/test/test5.py b/test/test5.py deleted file mode 100644 index 6220633e..00000000 --- a/test/test5.py +++ /dev/null @@ -1,7 +0,0 @@ -import tiktoken - -enc = tiktoken.get_encoding("o200k_base") - -# r = enc.encode("我爱吃西瓜，你说啥") -r = enc.encode("hello world aaaaaaaaaaaa") -print(len(r)) diff --git a/test/test6.py b/test/test6.py deleted file mode 100644 index 7e555898..00000000 --- a/test/test6.py +++ /dev/null @@ -1,15 +0,0 @@ -import tiktoken - - -def count_tokens(text: str) -> int: - """计算给定文本在指定模型下的 token 数量""" - encoding = tiktoken.get_encoding("o200k_base") - tokens = encoding.encode(text) - return len(tokens) - - -# 示例使用 -text = "你好，世界！Hello, world!" -token_count = count_tokens(text) -print(f"Token 数量: {token_count}") -print(len(text) / 4) diff --git a/tests/test_base_context.py b/test/test_base_context.py similarity index 97% rename from tests/test_base_context.py rename to test/test_base_context.py index 316a6796..2b3ebc99 100644 --- a/tests/test_base_context.py +++ b/test/test_base_context.py @@ -4,7 +4,7 @@ Ensures attribute-style and dict-style access work interchangeably. """ import pickle -from reme_ai.core.context import BaseContext +from reme_ai.core_old.context import BaseContext def test_attribute_access(): diff --git a/tests/test_cache_handler.py b/test/test_cache_handler.py similarity index 97% rename from tests/test_cache_handler.py rename to test/test_cache_handler.py index ddcac86f..b741769a 100644 --- a/tests/test_cache_handler.py +++ b/test/test_cache_handler.py @@ -9,7 +9,7 @@ from pathlib import Path import pandas as pd from loguru import logger -from reme_ai.core.utils.cache_handler import CacheHandler +from reme_ai.core_old.utils.cache_handler import CacheHandler def run_tests(): diff --git a/tests/test_embedding.py b/test/test_embedding.py similarity index 98% rename from tests/test_embedding.py rename to test/test_embedding.py index d1c404f5..b769d05f 100644 --- a/tests/test_embedding.py +++ b/test/test_embedding.py @@ -18,12 +18,12 @@ import asyncio import argparse from typing import Type, List -from reme_ai.core.utils import load_env +from reme_ai.core_old.utils import load_env load_env() -from reme_ai.core.embedding import OpenAIEmbeddingModel, BaseEmbeddingModel -from reme_ai.core.schema import VectorNode +from reme_ai.core_old.embedding import OpenAIEmbeddingModel, BaseEmbeddingModel +from reme_ai.core_old.schema import VectorNode def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel: diff --git a/tests/test_embedding_sync.py b/test/test_embedding_sync.py similarity index 98% rename from tests/test_embedding_sync.py rename to test/test_embedding_sync.py index 361a42b3..f97e28ec 100644 --- a/tests/test_embedding_sync.py +++ b/test/test_embedding_sync.py @@ -17,12 +17,12 @@ Usage: import argparse from typing import Type, List -from reme_ai.core.utils import load_env +from reme_ai.core_old.utils import load_env load_env() -from reme_ai.core.embedding import OpenAIEmbeddingModelSync, BaseEmbeddingModel -from reme_ai.core.schema import VectorNode +from reme_ai.core_old.embedding import OpenAIEmbeddingModelSync, BaseEmbeddingModel +from reme_ai.core_old.schema import VectorNode def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel: diff --git a/tests/test_llm.py b/test/test_llm.py similarity index 98% rename from tests/test_llm.py rename to test/test_llm.py index 12c6eca3..819c2b2b 100644 --- a/tests/test_llm.py +++ b/test/test_llm.py @@ -18,13 +18,13 @@ import asyncio import argparse from typing import Type -from reme_ai.core.utils import load_env +from reme_ai.core_old.utils import load_env load_env() -from reme_ai.core.llm import OpenAILLM, LiteLLM, BaseLLM -from reme_ai.core.schema import Message, ToolCall -from reme_ai.core.enumeration import Role, ChunkEnum +from reme_ai.core_old.llm import OpenAILLM, LiteLLM, BaseLLM +from reme_ai.core_old.schema import Message, ToolCall +from reme_ai.core_old.enumeration import Role, ChunkEnum def get_llm(llm_class: Type[BaseLLM]) -> BaseLLM: diff --git a/tests/test_llm_sync.py b/test/test_llm_sync.py similarity index 98% rename from tests/test_llm_sync.py rename to test/test_llm_sync.py index 98751f87..07d80ed1 100644 --- a/tests/test_llm_sync.py +++ b/test/test_llm_sync.py @@ -17,13 +17,13 @@ Usage: import argparse from typing import Type -from reme_ai.core.utils import load_env +from reme_ai.core_old.utils import load_env load_env() -from reme_ai.core.llm import OpenAILLMSync, LiteLLMSync, BaseLLM -from reme_ai.core.schema import Message, ToolCall -from reme_ai.core.enumeration import Role, ChunkEnum +from reme_ai.core_old.llm import OpenAILLMSync, LiteLLMSync, BaseLLM +from reme_ai.core_old.schema import Message, ToolCall +from reme_ai.core_old.enumeration import Role, ChunkEnum def get_llm(llm_class: Type[BaseLLM]) -> BaseLLM: diff --git a/tests/test_logo.py b/test/test_logo.py similarity index 59% rename from tests/test_logo.py rename to test/test_logo.py index eeede81e..9b4cf089 100644 --- a/tests/test_logo.py +++ b/test/test_logo.py @@ -1,9 +1,9 @@ """test logo""" -from reme_ai.core.schema import ServiceConfig, MCPConfig +from reme_ai.core_old.schema import ServiceConfig, MCPConfig if __name__ == "__main__": - from reme_ai.core.utils import print_logo + from reme_ai.core_old.utils import print_logo c = ServiceConfig(app_name="reme", backend="mcp", mcp=MCPConfig(transport="sse")) print_logo(service_config=c) diff --git a/tests/test_mcp_client.py b/test/test_mcp_client.py similarity index 99% rename from tests/test_mcp_client.py rename to test/test_mcp_client.py index d2fef40e..0ae6fc54 100644 --- a/tests/test_mcp_client.py +++ b/test/test_mcp_client.py @@ -5,7 +5,7 @@ import asyncio import json -from reme_ai.core.utils import MCPClient +from reme_ai.core_old.utils import MCPClient async def main(): diff --git a/tests/test_mcp_server.py b/test/test_mcp_server.py similarity index 97% rename from tests/test_mcp_server.py rename to test/test_mcp_server.py index 67f9c542..4257b977 100644 --- a/tests/test_mcp_server.py +++ b/test/test_mcp_server.py @@ -5,8 +5,8 @@ from typing import Any from fastmcp import FastMCP from fastmcp.tools import FunctionTool -from reme_ai.core.schema import ToolCall -from reme_ai.core.utils import create_pydantic_model +from reme_ai.core_old.schema import ToolCall +from reme_ai.core_old.utils import create_pydantic_model mcp = FastMCP("DynamicSchemaServer", port=8010) diff --git a/tests/test_memory_vector_conversion.py b/test/test_memory_vector_conversion.py similarity index 100% rename from tests/test_memory_vector_conversion.py rename to test/test_memory_vector_conversion.py diff --git a/tests/test_message.py b/test/test_message.py similarity index 98% rename from tests/test_message.py rename to test/test_message.py index 141174c5..e77e673b 100644 --- a/tests/test_message.py +++ b/test/test_message.py @@ -4,8 +4,8 @@ import unittest from mcp.types import Tool -from reme_ai.core.enumeration import Role -from reme_ai.core.schema import ToolAttr, ToolCall, ContentBlock, Message +from reme_ai.core_old.enumeration import Role +from reme_ai.core_old.schema import ToolAttr, ToolCall, ContentBlock, Message class TestModelDefinitions(unittest.TestCase): diff --git a/tests/test_op_composition.py b/test/test_op_composition.py similarity index 99% rename from tests/test_op_composition.py rename to test/test_op_composition.py index 8d32b51a..319981a3 100644 --- a/tests/test_op_composition.py +++ b/test/test_op_composition.py @@ -5,8 +5,8 @@ Tests asynchronous execution mode. import asyncio -from reme_ai.core.op import BaseOp -from reme_ai.core.schema import ToolCall, ToolAttr +from reme_ai.core_old.op import BaseOp +from reme_ai.core_old.schema import ToolCall, ToolAttr class AddOp(BaseOp): diff --git a/tests/test_reme.py b/test/test_reme.py similarity index 98% rename from tests/test_reme.py rename to test/test_reme.py index c9e6843e..b23aa63e 100644 --- a/tests/test_reme.py +++ b/test/test_reme.py @@ -2,7 +2,7 @@ import asyncio -from reme_ai.core.schema import VectorNode, MemoryNode +from reme_ai.core_old.schema import VectorNode, MemoryNode from reme_ai.reme import ReMe reme = ReMe( diff --git a/tests/test_timer.py b/test/test_timer.py similarity index 97% rename from tests/test_timer.py rename to test/test_timer.py index c9714e38..1c9de938 100644 --- a/tests/test_timer.py +++ b/test/test_timer.py @@ -7,7 +7,7 @@ import time from loguru import logger -from reme_ai.core.utils import timer +from reme_ai.core_old.utils import timer @timer diff --git a/tests/test_token_counter.py b/test/test_token_counter.py similarity index 99% rename from tests/test_token_counter.py rename to test/test_token_counter.py index 3c44a298..e67fb5c2 100644 --- a/tests/test_token_counter.py +++ b/test/test_token_counter.py @@ -14,9 +14,9 @@ Usage: import argparse from typing import Type, List -from reme_ai.core.enumeration import Role -from reme_ai.core.schema import Message, ToolCall -from reme_ai.core.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter +from reme_ai.core_old.enumeration import Role +from reme_ai.core_old.schema import Message, ToolCall +from reme_ai.core_old.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter def get_token_counter(counter_class: Type[BaseTokenCounter], **kwargs) -> BaseTokenCounter: diff --git a/tests/test_tool.py b/test/test_tool.py similarity index 98% rename from tests/test_tool.py rename to test/test_tool.py index 9db5a3ed..765051b3 100644 --- a/tests/test_tool.py +++ b/test/test_tool.py @@ -169,8 +169,8 @@ async def test_stream_chat(): process and stream responses in real-time using async operations. """ from reme_ai.mem_agent.chat import StreamChat - from reme_ai.core.utils import execute_stream_task - from reme_ai.core.context import RuntimeContext + from reme_ai.core_old.utils import execute_stream_task + from reme_ai.core_old.context import RuntimeContext from asyncio import Queue op = StreamChat() diff --git a/tests/test_tool_call.py b/test/test_tool_call.py similarity index 99% rename from tests/test_tool_call.py rename to test/test_tool_call.py index 30c0a37e..2ae7a619 100644 --- a/tests/test_tool_call.py +++ b/test/test_tool_call.py @@ -2,7 +2,7 @@ import json -from reme_ai.core.schema.tool_call import ToolCall +from reme_ai.core_old.schema.tool_call import ToolCall def test_simple_schema(): diff --git a/test/test_update_insight_op.py b/test/test_update_insight_op.py deleted file mode 100644 index 6505b6a4..00000000 --- a/test/test_update_insight_op.py +++ /dev/null @@ -1,128 +0,0 @@ -#!/usr/bin/env python3 -""" -Simple test script to verify the UpdateInsightOp implementation. -This is a basic validation test to ensure the class structure is correct. -""" - -import sys - -sys.path.append("/Users/yuli/workspace/MemoryScope") - - -def test_update_insight_op_import(): - """Test that we can import the UpdateInsightOp class""" - try: - from reme_ai.summary.personal.update_insight_op import UpdateInsightOp - - print("✓ Successfully imported UpdateInsightOp") - return True - except ImportError as e: - print(f"✗ Failed to import UpdateInsightOp: {e}") - return False - - -def test_personal_memory_import(): - """Test that we can import PersonalMemory""" - try: - from reme_ai.schema.memory import PersonalMemory - - print("✓ Successfully imported PersonalMemory") - return True - except ImportError as e: - print(f"✗ Failed to import PersonalMemory: {e}") - return False - - -def test_op_utils_import(): - """Test that we can import the utility functions""" - try: - from reme_ai.utils.op_utils import parse_update_insight_response - - print("✓ Successfully imported parse_update_insight_response") - return True - except ImportError as e: - print(f"✗ Failed to import parse_update_insight_response: {e}") - return False - - -def test_personal_memory_creation(): - """Test PersonalMemory creation with reflection_subject""" - try: - from reme_ai.schema.memory import PersonalMemory - - memory = PersonalMemory( - workspace_id="test_workspace", - content="User likes playing basketball", - target="test_user", - reflection_subject="hobbies", - author="test_system", - ) - - print(f"✓ Created PersonalMemory: {memory.content}") - print(f" - Memory ID: {memory.memory_id}") - print(f" - Target: {memory.target}") - print(f" - Reflection Subject: {memory.reflection_subject}") - return True - except Exception as e: - print(f"✗ Failed to create PersonalMemory: {e}") - return False - - -def test_parse_update_insight_response(): - """Test the parse_update_insight_response function""" - try: - from reme_ai.utils.op_utils import parse_update_insight_response - - # Test Chinese format - chinese_response = "思考：用户喜欢篮球和足球\ntest_user的资料：<喜欢篮球和足球>" - result_zh = parse_update_insight_response(chinese_response, "zh") - print(f"✓ Parsed Chinese response: '{result_zh}'") - - # Test English format - english_response = ( - "Thoughts: User likes basketball and football\ntest_user's profile: " - ) - result_en = parse_update_insight_response(english_response, "en") - print(f"✓ Parsed English response: '{result_en}'") - - return True - except Exception as e: - print(f"✗ Failed to test parse_update_insight_response: {e}") - return False - - -def main(): - """Run all tests""" - print("Running UpdateInsightOp validation tests...\n") - - tests = [ - test_personal_memory_import, - test_op_utils_import, - test_update_insight_op_import, - test_personal_memory_creation, - test_parse_update_insight_response, - ] - - passed = 0 - total = len(tests) - - for test in tests: - print(f"\nRunning {test.__name__}:") - if test(): - passed += 1 - print() - - print("=" * 50) - print(f"Test Results: {passed}/{total} passed") - - if passed == total: - print("🎉 All tests passed! The UpdateInsightOp implementation looks good.") - else: - print("⚠️ Some tests failed. Please check the implementation.") - - return passed == total - - -if __name__ == "__main__": - success = main() - sys.exit(0 if success else 1) diff --git a/tests/test_vector_store.py b/test/test_vector_store.py similarity index 99% rename from tests/test_vector_store.py rename to test/test_vector_store.py index 6ef9ab0e..51edbd55 100644 --- a/tests/test_vector_store.py +++ b/test/test_vector_store.py @@ -23,9 +23,9 @@ from typing import List from loguru import logger -from reme_ai.core.embedding import OpenAIEmbeddingModel -from reme_ai.core.schema import VectorNode -from reme_ai.core.vector_store import ( +from reme_ai.core_old.embedding import OpenAIEmbeddingModel +from reme_ai.core_old.schema import VectorNode +from reme_ai.core_old.vector_store import ( BaseVectorStore, ChromaVectorStore, LocalVectorStore, @@ -388,7 +388,8 @@ async def test_search_with_multiple_filters(store: BaseVectorStore, _store_name: ) logger.info( - f"Multi-filter search (node_type=tech AND source=research AND priority=high) " f"returned {len(results)} results", + f"Multi-filter search (node_type=tech AND source=research AND priority=high) " + f"returned {len(results)} results", ) for i, r in enumerate(results, 1): node_type = r.metadata.get("node_type") @@ -1452,11 +1453,11 @@ async def test_sql_injection_protection(store: BaseVectorStore, store_name: str) # Test 1: Invalid collection name (SQL injection attempt) try: - from reme_ai.core.vector_store import PGVectorStore - from reme_ai.core.embedding import OpenAIEmbeddingModel - + from reme_ai.core_old.vector_store import PGVectorStore + from reme_ai.core_old.embedding import OpenAIEmbeddingModel + embedding_model = OpenAIEmbeddingModel() - + # This should raise ValueError due to invalid table name try: invalid_store = PGVectorStore( @@ -1467,7 +1468,7 @@ async def test_sql_injection_protection(store: BaseVectorStore, store_name: str) assert False, "Should have raised ValueError for invalid collection name" except ValueError as e: logger.info(f"✓ Invalid collection name rejected: {e}") - + # Test 2: Invalid metadata key in filters try: results = await store.search( @@ -1481,9 +1482,9 @@ async def test_sql_injection_protection(store: BaseVectorStore, store_name: str) assert False, "Should have raised ValueError for invalid metadata key" except ValueError as e: logger.info(f"✓ Invalid metadata key rejected: {e}") - + logger.info("✓ SQL injection protection validated") - + except Exception as e: logger.error(f"SQL injection protection test failed: {e}") raise