From 72eabfa858bd331d021310435f77bb718931ad47 Mon Sep 17 00:00:00 2001 From: Zhouwk <57825291+nitwtog@users.noreply.github.com> Date: Thu, 30 Apr 2026 10:19:36 +0800 Subject: [PATCH] fix(user profile): update locomo benchmark and update vector based profile code (#225) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(reme): 添加配置选项以启用或禁用个人资料功能 - 在 ReMe 初始化方法中添加 enable_profile 参数，默认值为 True - 根据 enable_profile 设置决定是否创建 profile 目录和设置 profile_dir - 在 PersonalSummarizer 中根据 enable_profile 条件性地添加个人资料相关工具 - 在 PersonalRetriever 中根据 enable_profile 条件性地添加 ReadAllProfiles 工具 - 修改 profile_path 属性以在禁用个人资料时返回 None - 修改 get_profile_handler 方法以在禁用个人资料时返回 None - 为 enable_profile 参数添加文档说明其用于云向量存储场景 * refactor(benchmark): 重构LongMemEval基准测试中的ReMe实例管理 - 移除未使用的shutil导入 - 将固定的ReMe实例改为每个问题创建独立实例以实现隔离 - 更新LLM配置名称从qwen3-max-think到qwen-max-t - 修改模型调用逻辑使用正确的model_name参数 - 添加qwen-flash和GPT-4o-mini等新模型配置 - 统一使用"User"作为用户名，通过集合名实现隔离 - 调整并发处理数从4降至1，批处理大小从10增至30 - 每个问题类型采样数从2增至4 - 添加异步上下文管理确保资源正确释放 * reformat 2 files * refactor(benchmark): 重构长记忆评估中的模型配置 - 将原有的 eval_model_name 替换为专门的 retrieve_model_name 用于检索操作 - 添加对 qwen-max 模型配置的支持 - 更新参数解析器以支持新的检索模型参数 - 修改最大并发数默认值从 1 提升到 4 - 调整样本数量默认值从 4 减少到 1 - 统一模型参数命名规范，区分摘要、检索和评估模型 - 优化内存处理器初始化逻辑，支持独立的检索模型配置 * fix(benchmark): 移除数据路径默认值并设为必填参数 - 将LongMemEval评估脚本中的data_path参数改为必需参数 - 将HaluMem评估脚本中的data_path参数改为必需参数 - 删除了硬编码的默认文件路径配置 - 强制用户显式指定数据集文件路径以避免路径错误 * Update __init__.py * Update __init__.py * fix(benchmark): 修复ReMe评估中的模型配置和空值处理问题 - 移除了retrieve_memory调用中不需要的llm_config_name参数 - 修复了长字符串打印的换行格式问题 - 添加了eval_result为空时的初始化处理 - 在accuracy评估中加入了eval_model_name参数传递 * style(benchmark): 格式化模型名称打印输出 - 移除了多行字符串中的换行符和多余空格 - 将模型名称信息合并为单行连续显示 - 保持了原有的打印格式和信息完整性 * docs(readme): 更新文档添加实验结果表格 - 在英文版 README 中添加 🧪 Experiments 章节 - 添加 LoCoMo 和 HaluMem 两个基准测试的结果表格 - 在中文版 README_ZH 中添加 🧪 实验 章节 - 添加 LoCoMo 和 HaluMem 测试集的实验配置说明 - 添加完整的实验数据对比表格和评估协议说明 * docs(readme): 更新文档中的内存系统链接 - 为基于文件的记忆系统添加锚点链接 - 为基于向量库的记忆系统添加锚点链接 - 修复英文文档中的链接格式 - 修复中文文档中的链接格式和空行问题 * docs(readme): update experimental results section in documentation - Remove outdated experimental data placeholder "Coming soon..." - Add complete evaluation results for LoCoMo and HaluMem benchmarks - Include detailed performance metrics tables for all memory methods - Update experimental settings description with ReMe backbone details - Align evaluation protocol information with LLM-as-a-Judge approach - Maintain consistent formatting between English and Chinese documentation * docs(benchmark): add quick start guides for halumem and longmemeval experiments - Created HaluMem experiment quick start guide with ReMe integration setup - Added detailed steps for installing ReMe environment using conda - Included repository cloning instructions for HaluMem benchmark - Provided complete command examples for running HaluMem experiments - Created LongMeMEval quick start guide with data download procedures - Added wget commands for downloading cleaned dataset files - Included evaluation script instructions for computing experiment statistics - Documented parameter configurations for different model types and batch sizes * docs(longmemeval): update quickstart guide documentation - Changed project name from Halumem to Longmemeval in title - Updated description to reference Longmemeval experiments instead of Halumem - Maintained existing ReMe integration instructions unchanged * chore(logger): add test comment to logger configuration - Added test comment in logger utility function - Removed duplicate log handling by keeping the remove() call * chore(logger): add test comment to logger configuration - Added test comment in logger utility function - Removed duplicate log handling by keeping the remove() call * feat(core): add file logging capability to application - Added log_to_file parameter to Application class constructor - Integrated log_to_file option in logger initialization - Updated ServiceContext to support file logging configuration - Modified init_logger function to conditionally enable file logging - Added log_to_file field to ServiceConfig schema - Updated ReMe class to include file logging option - Wrapped file logging setup in conditional check to prevent unnecessary operations * docs(benchmark): update HaluMem quickstart guide with dataset download instructions - Replace repository cloning with direct dataset download using curl - Add commands to download HaluMem-Medium.jsonl and HaluMem-Long.jsonl files - Include both official Hugging Face and mirror download sources - Update data path reference from nested directory to local data folder - Add dataset page link and mirror usage instructions for mainland China access * feat(memory): add profile retrieval tool and refactor profile management - Introduce RetrieveProfile tool for fetching specific user profiles - Refactor ProfileHandler to support both filesystem and vector backends - Add async methods to ProfileHandler with synchronous fallbacks - Update PersonalRetriever to support two-stage profile and memory retrieval - Enhance PersonalSummarizer with improved tool partitioning logic - Add profile_backend, profile_store_name, and profile_max_capacity configuration options - Replace direct ProfileHandler imports with get_profile_handler method - Implement profile search functionality with dedicated prompts and workflows - Add FileProfileBackend and VectorProfileBackend implementations - Update base memory tool with new profile configuration parameters * feat(profile): add custom profile collection name support - Add profile_collection_name parameter to Application constructor - Allow custom database collection name for vector profiles instead of default suffix - Update profile vector store configuration logic to use custom collection name - Modify _ensure_profile_vector_store_config to handle custom collection names - Update docstring with detailed parameter descriptions for profile configuration options * test(history): add single history id acceptance test for multiple mode - Add test case to verify multiple-mode history lookup accepts a single history_id string - Create FakeVectorStore stub with minimal implementation for ReadHistory tests - Return requested history node from vector store mock - Initialize ReadHistory tool with multiple mode enabled - Add pylint disable comment for protected access to vector store property * refactor(memory): update profile handler and vector tools with improved formatting and error handling - Add module docstring to profiles/__init__.py - Add pylint disable comments for no-name-in-module and missing-function-docstring - Format long error message in ProfileHandler.sync_run method for better readability - Reformat parameters in ProfileHandler.aadd method to separate lines - Update model_copy call in reme.py to span multiple lines for better readability - Format aadd_batch call in update_profile.py to span multiple lines * feat(profiles): add profile management system with file and vector storage backends - Add FileProfileBackend for filesystem-based profile persistence - Add VectorProfileBackend for vector store-based profile management - Create abstract BaseProfileBackend interface for profile operations - Implement ProfileVectorHandler for vector-backed profile storage - Add RetrieveProfile tool for semantic profile retrieval - Update eval_reme.py to use user_message_s2 for retriever prompt - Modify eval_reme.yaml to use {profiles} instead of {user_profile} - Implement complete CRUD operations for profile management - Add batch operations for efficient profile handling - Include search functionality with semantic matching capabilities - Add capacity limits and automatic cleanup for profile storage * docs(profiles): add comprehensive docstrings for profile backend and handler methods - Added documentation for get_all_sync, get_by_sync, delete_sync, delete_all_sync methods - Documented add_sync and add_batch_sync functionality with deduping behavior - Added docstrings for update_sync and search_sync operations - Updated ProfileHandler.format_node method with proper documentation - Refactored private _format_node to public format_node method - Added comprehensive documentation for profile vector handler operations - Documented _vector_profile_matches, _get_by_profile_id, _get_by_profile_key helper methods - Added docstrings for retrieve_profile functionality and formatting methods --- benchmark/locomo/eval_reme.py | 2 +- benchmark/locomo/eval_reme.yaml | 2 +- .../profiles/file_profile_backend.py | 234 +++++++++++++++++ .../vector_tools/profiles/profile_backend.py | 51 ++++ .../vector_tools/profiles/profile_handler.py | 7 +- .../profiles/profile_vector_handler.py | 245 ++++++++++++++++++ .../vector_tools/profiles/retrieve_profile.py | 102 ++++++++ .../profiles/vector_profile_backend.py | 55 ++++ 8 files changed, 693 insertions(+), 5 deletions(-) create mode 100644 reme/memory/vector_tools/profiles/file_profile_backend.py create mode 100644 reme/memory/vector_tools/profiles/profile_backend.py create mode 100644 reme/memory/vector_tools/profiles/profile_vector_handler.py create mode 100644 reme/memory/vector_tools/profiles/retrieve_profile.py create mode 100644 reme/memory/vector_tools/profiles/vector_profile_backend.py diff --git a/benchmark/locomo/eval_reme.py b/benchmark/locomo/eval_reme.py index 82d91031..bb424aee 100644 --- a/benchmark/locomo/eval_reme.py +++ b/benchmark/locomo/eval_reme.py @@ -643,7 +643,7 @@ class LocomoEvaluator: }, "personal_retriever": { "prompt_dict": { - "user_message": self.retriever_prompt, + "user_message_s2": self.retriever_prompt, }, "params": { "return_memory_nodes": True, diff --git a/benchmark/locomo/eval_reme.yaml b/benchmark/locomo/eval_reme.yaml index 113431b3..5228884a 100644 --- a/benchmark/locomo/eval_reme.yaml +++ b/benchmark/locomo/eval_reme.yaml @@ -132,7 +132,7 @@ user_message_retrieve: | You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}. ## User Profile - {user_profile} + {profiles} ## User Question {context} diff --git a/reme/memory/vector_tools/profiles/file_profile_backend.py b/reme/memory/vector_tools/profiles/file_profile_backend.py new file mode 100644 index 00000000..e6226459 --- /dev/null +++ b/reme/memory/vector_tools/profiles/file_profile_backend.py @@ -0,0 +1,234 @@ +"""Filesystem-backed profile storage.""" + +from pathlib import Path + +from loguru import logger + +from .profile_backend import BaseProfileBackend +from ....core.enumeration import MemoryType +from ....core.schema import MemoryNode +from ....core.utils import CacheHandler, deduplicate_memories + + +class FileProfileBackend(BaseProfileBackend): + """Persist user profiles in local JSONL cache files.""" + + def __init__(self, profile_path: str | Path, memory_target: str, max_capacity: int = 50): + super().__init__(memory_target=memory_target, max_capacity=max_capacity) + self.cache_key: str = self.memory_target.replace(" ", "_").lower() + self.cache_handler: CacheHandler = CacheHandler(profile_path) + + def _load_nodes(self) -> list[MemoryNode]: + cached_data = self.cache_handler.load(self.cache_key, auto_clean=False) + if not cached_data: + return [] + return [MemoryNode(**data) for data in cached_data] + + def _save_nodes(self, nodes: list[MemoryNode], apply_limits: bool = True): + if apply_limits: + nodes = deduplicate_memories(nodes) + + if len(nodes) > self.max_capacity: + sorted_nodes = sorted(nodes, key=lambda n: n.message_time) + removed_count = len(sorted_nodes) - self.max_capacity + nodes = sorted_nodes[removed_count:] + logger.info( + f"Capacity limit reached: removed {removed_count} oldest profiles " + f"(kept {len(nodes)}/{self.max_capacity})", + ) + + nodes_data = [node.model_dump(exclude_none=True) for node in nodes] + self.cache_handler.save(self.cache_key, nodes_data) + logger.info(f"Saved {len(nodes)} profiles to {self.cache_key}") + + def get_all_sync(self) -> list[MemoryNode]: + """Load all profile nodes from cache, ordered by ``message_time``.""" + nodes = self._load_nodes() + nodes.sort(key=lambda n: n.message_time) + return nodes + + def get_by_sync(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + """Return the first node matching ``profile_id`` or ``profile_key``.""" + if not profile_id and not profile_key: + raise ValueError("Must provide either profile_id or profile_key") + + for node in self._load_nodes(): + if profile_id and node.memory_id == profile_id: + return node + if profile_key and node.when_to_use == profile_key: + return node + return None + + def delete_sync(self, profile_id: str | list[str]) -> bool | int: + """Remove one id, many ids, or none; returns bool, count, or 0/false if nothing removed.""" + nodes = self._load_nodes() + original_count = len(nodes) + + if isinstance(profile_id, list): + profile_ids_set = set(profile_id) + nodes = [n for n in nodes if n.memory_id not in profile_ids_set] + deleted_count = original_count - len(nodes) + if deleted_count == 0: + logger.warning(f"No profiles found to delete from {len(profile_id)} IDs") + return 0 + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Batch deleted {deleted_count} profiles") + return deleted_count + + nodes = [n for n in nodes if n.memory_id != profile_id] + if len(nodes) == original_count: + logger.warning(f"Profile {profile_id} not found") + return False + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Deleted profile {profile_id}") + return True + + def delete_all_sync(self) -> int: + """Clear every cached profile for this target; returns how many were stored.""" + nodes = self._load_nodes() + count = len(nodes) + self._save_nodes([], apply_limits=False) + logger.info(f"Deleted all {count} profiles") + return count + + def add_sync(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode: + """Append a profile row, replacing any existing row with the same key.""" + nodes = self._load_nodes() + + new_node = MemoryNode( + memory_type=MemoryType.PERSONAL, + memory_target=self.memory_target, + when_to_use=profile_key, + content=profile_value, + message_time=message_time, + ref_memory_id=ref_memory_id, + ) + + original_count = len(nodes) + nodes = [n for n in nodes if n.when_to_use != profile_key] + if len(nodes) < original_count: + logger.info(f"Removed {original_count - len(nodes)} duplicate profile(s) with key: {profile_key}") + + nodes.append(new_node) + self._save_nodes(nodes) + logger.info(f"Added profile: {profile_key}={profile_value}") + return new_node + + def add_batch_sync(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + """Insert many profiles in one write, deduping by key against existing rows.""" + if not profiles: + return [] + + nodes = self._load_nodes() + new_nodes = [ + MemoryNode( + memory_type=MemoryType.PERSONAL, + memory_target=self.memory_target, + when_to_use=p.get("profile_key", ""), + content=p.get("profile_value", ""), + message_time=p.get("message_time", ""), + ref_memory_id=ref_memory_id, + ) + for p in profiles + ] + + new_keys = {n.when_to_use for n in new_nodes} + original_count = len(nodes) + nodes = [n for n in nodes if n.when_to_use not in new_keys] + if len(nodes) < original_count: + logger.info(f"Removed {original_count - len(nodes)} duplicate profile(s) with matching keys") + + nodes.extend(new_nodes) + self._save_nodes(nodes) + logger.info(f"Batch added {len(new_nodes)} profiles") + return new_nodes + + def update_sync( + self, + profile_id: str, + message_time: str, + profile_key: str, + profile_value: str, + ) -> MemoryNode | None: + """Update fields for ``profile_id``; return ``None`` if that id is missing.""" + nodes = self._load_nodes() + target_node = None + for node in nodes: + if node.memory_id == profile_id: + node.when_to_use = profile_key + node.content = profile_value + node.message_time = message_time + target_node = node + break + + if target_node is None: + logger.warning(f"Profile {profile_id} not found") + return None + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Updated profile {profile_id}: {profile_key}={profile_value}") + return target_node + + def search_sync(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]: + """Simple substring/token match over key and content, best matches first.""" + queries = [query] if isinstance(query, str) else query + query_terms = [q.strip().lower() for q in queries if q and q.strip()] + if not query_terms: + return [] + + scored_nodes = [] + for node in self.get_all_sync(): + profile_key = str(node.metadata.get("profile_key", node.when_to_use)).lower() + haystack = f"{profile_key}: {node.content}".lower() + score = 0 + for term in query_terms: + if term in haystack: + score += len(term) + 10 + else: + token_hits = sum(1 for token in term.split() if token and token in haystack) + score += token_hits + + if score > 0: + node.score = float(score) + scored_nodes.append(node) + + scored_nodes.sort(key=lambda n: (n.score, n.message_time), reverse=True) + return scored_nodes[:limit] + + async def get_all(self) -> list[MemoryNode]: + return self.get_all_sync() + + async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + return self.get_by_sync(profile_id=profile_id, profile_key=profile_key) + + async def delete(self, profile_id: str | list[str]) -> bool | int: + return self.delete_sync(profile_id) + + async def delete_all(self) -> int: + return self.delete_all_sync() + + async def add( + self, + message_time: str, + profile_key: str, + profile_value: str, + ref_memory_id: str = "", + ) -> MemoryNode: + return self.add_sync(message_time, profile_key, profile_value, ref_memory_id) + + async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + return self.add_batch_sync(profiles, ref_memory_id) + + async def update( + self, + profile_id: str, + message_time: str, + profile_key: str, + profile_value: str, + ) -> MemoryNode | None: + return self.update_sync(profile_id, message_time, profile_key, profile_value) + + async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]: + return self.search_sync(query, limit) diff --git a/reme/memory/vector_tools/profiles/profile_backend.py b/reme/memory/vector_tools/profiles/profile_backend.py new file mode 100644 index 00000000..8ffcb9bc --- /dev/null +++ b/reme/memory/vector_tools/profiles/profile_backend.py @@ -0,0 +1,51 @@ +"""Profile backend abstractions.""" + +from abc import ABC, abstractmethod + +from ....core.schema import MemoryNode + + +class BaseProfileBackend(ABC): + """Abstract interface for profile storage backends.""" + + def __init__(self, memory_target: str, max_capacity: int = 50): + self.memory_target = memory_target + self.max_capacity = max_capacity + + @abstractmethod + async def get_all(self) -> list[MemoryNode]: + """Return all profile rows for the current user.""" + + @abstractmethod + async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + """Return one profile row by id or key.""" + + @abstractmethod + async def delete(self, profile_id: str | list[str]) -> bool | int: + """Delete one or more profile rows.""" + + @abstractmethod + async def delete_all(self) -> int: + """Delete all profile rows for the current user.""" + + @abstractmethod + async def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode: + """Add a single profile row.""" + + @abstractmethod + async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + """Add multiple profile rows.""" + + @abstractmethod + async def update( + self, + profile_id: str, + message_time: str, + profile_key: str, + profile_value: str, + ) -> MemoryNode | None: + """Update one profile row.""" + + @abstractmethod + async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]: + """Search profile rows relevant to the query.""" diff --git a/reme/memory/vector_tools/profiles/profile_handler.py b/reme/memory/vector_tools/profiles/profile_handler.py index 0df3707a..145caa6e 100644 --- a/reme/memory/vector_tools/profiles/profile_handler.py +++ b/reme/memory/vector_tools/profiles/profile_handler.py @@ -115,7 +115,8 @@ class ProfileHandler: return await self.backend.search(query=query, limit=limit) @staticmethod - def _format_node(node: MemoryNode, add_profile_id: bool = False, add_history_id: bool = False) -> str: + def format_node(node: MemoryNode, add_profile_id: bool = False, add_history_id: bool = False) -> str: + """Render a profile ``MemoryNode`` as a single-line string for tools/logs.""" parts = [] profile_key = str(node.metadata.get("profile_key", node.when_to_use)) @@ -134,7 +135,7 @@ class ProfileHandler: async def aread_all(self, add_profile_id: bool = False, add_history_id: bool = False) -> str: nodes = await self.aget_all() - formatted_profiles = [self._format_node(node, add_profile_id, add_history_id) for node in nodes] + formatted_profiles = [self.format_node(node, add_profile_id, add_history_id) for node in nodes] logger.info(f"Read {len(formatted_profiles)} profiles from {self.cache_key}") return "\n".join(formatted_profiles).strip() @@ -146,7 +147,7 @@ class ProfileHandler: add_history_id: bool = False, ) -> tuple[list[MemoryNode], str]: nodes = await self.asearch(query=query, limit=limit) - formatted_profiles = [self._format_node(node, add_profile_id, add_history_id) for node in nodes] + formatted_profiles = [self.format_node(node, add_profile_id, add_history_id) for node in nodes] return nodes, "\n".join(formatted_profiles).strip() def delete(self, profile_id: str | list[str]) -> bool | int: diff --git a/reme/memory/vector_tools/profiles/profile_vector_handler.py b/reme/memory/vector_tools/profiles/profile_vector_handler.py new file mode 100644 index 00000000..f03e6a96 --- /dev/null +++ b/reme/memory/vector_tools/profiles/profile_vector_handler.py @@ -0,0 +1,245 @@ +"""Vector-backed handler for bounded user profiles.""" + +import hashlib + +from loguru import logger + +from ....core import ServiceContext +from ....core.enumeration import MemoryType +from ....core.schema import MemoryNode +from ....core.vector_store import BaseVectorStore + + +class ProfileVectorHandler: + """Manage profile rows stored in a dedicated vector collection.""" + + PROFILE_KIND = "profile" + + def __init__( + self, + memory_target: str, + service_context: ServiceContext, + vector_store_name: str = "profile", + max_capacity: int = 50, + ): + self.memory_target = memory_target + self.service_context = service_context + self.vector_store_name = vector_store_name + self.max_capacity = max_capacity + self.vector_store: BaseVectorStore = service_context.vector_stores[vector_store_name] + + @staticmethod + def build_retrieval_text(profile_key: str, profile_value: str) -> str: + """Build the text that will be embedded for semantic profile retrieval.""" + return f"{profile_key}: {profile_value}".strip(": ") + + def build_profile_id(self, profile_key: str) -> str: + """Build a stable id from user and key.""" + hash_obj = hashlib.sha256(f"{self.memory_target}\n{profile_key}".encode("utf-8")) + return hash_obj.hexdigest()[:16] + + def _base_filters(self) -> dict: + """Filters shared by all profile rows in the vector collection.""" + return { + "memory_type": MemoryType.IDENTITY.value, + "memory_target": self.memory_target, + "profile_kind": self.PROFILE_KIND, + } + + def _build_profile_node(self, profile: dict, ref_memory_id: str = "") -> MemoryNode: + """Turn a profile dict into a ``MemoryNode`` for upsert into the vector store.""" + profile_key = profile.get("profile_key", "").strip() + profile_value = profile.get("profile_value", "").strip() + message_time = profile.get("message_time", "") + ref_id = profile.get("ref_memory_id", ref_memory_id) + metadata = dict(profile.get("metadata", {})) + metadata.update( + { + "profile_key": profile_key, + "profile_kind": self.PROFILE_KIND, + "profile_backend": "vector", + }, + ) + return MemoryNode( + memory_id=self.build_profile_id(profile_key), + memory_type=MemoryType.IDENTITY, + memory_target=self.memory_target, + when_to_use=self.build_retrieval_text(profile_key, profile_value), + content=profile_value, + message_time=message_time, + ref_memory_id=ref_id, + metadata=metadata, + ) + + def _vector_profile_matches(self, memory_node: MemoryNode) -> bool: + """True if ``memory_node`` belongs to this handler's target and profile kind.""" + if memory_node.memory_target != self.memory_target: + return False + if memory_node.memory_type is not MemoryType.IDENTITY: + return False + if memory_node.metadata.get("profile_kind") != self.PROFILE_KIND: + return False + return True + + async def _get_by_profile_id(self, profile_id: str) -> MemoryNode | None: + """Load by vector id and validate filters.""" + try: + vector_node = await self.vector_store.get(profile_id) + except KeyError: + logger.warning(f"Profile {profile_id} not found in vector store") + return None + if vector_node is None: + logger.warning(f"Profile {profile_id} not found in vector store") + return None + memory_node = MemoryNode.from_vector_node(vector_node) + if not self._vector_profile_matches(memory_node): + return None + return memory_node + + async def _get_by_profile_key(self, profile_key: str) -> MemoryNode | None: + """Load the single row matching ``profile_key`` under base filters.""" + filters = {**self._base_filters(), "profile_key": profile_key} + vector_nodes = await self.vector_store.list(filters=filters, limit=1) + if not vector_nodes: + return None + return MemoryNode.from_vector_node(vector_nodes[0]) + + async def get_all(self) -> list[MemoryNode]: + """List every profile row for this memory target, sorted by store.""" + vector_nodes = await self.vector_store.list( + filters=self._base_filters(), + sort_key="message_time", + reverse=False, + ) + return [MemoryNode.from_vector_node(node) for node in vector_nodes] + + async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + """Return one profile by stable id or by logical profile key.""" + if not profile_id and not profile_key: + raise ValueError("Must provide either profile_id or profile_key") + if profile_id: + return await self._get_by_profile_id(profile_id) + return await self._get_by_profile_key(profile_key or "") + + async def delete(self, profile_id: str | list[str]) -> bool | int: + """Delete one id, many ids, or report zero/false when nothing matched.""" + if isinstance(profile_id, list): + profile_ids = list(dict.fromkeys(pid for pid in profile_id if pid)) + if not profile_ids: + return 0 + existing_nodes = [] + for pid in profile_ids: + node = await self.get_by(profile_id=pid) + if node is not None: + existing_nodes.append(node) + if not existing_nodes: + return 0 + await self.vector_store.delete([node.memory_id for node in existing_nodes]) + return len(existing_nodes) + + existing_node = await self.get_by(profile_id=profile_id) + if existing_node is None: + return False + await self.vector_store.delete(existing_node.memory_id) + return True + + async def delete_all(self) -> int: + """Remove all profile vectors for this target; returns how many were deleted.""" + nodes = await self.get_all() + if not nodes: + return 0 + await self.vector_store.delete([node.memory_id for node in nodes]) + return len(nodes) + + async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + """Upsert many profiles at once (last dict wins per key), then enforce capacity.""" + if not profiles: + return [] + + deduped_profiles: dict[str, dict] = {} + for profile in profiles: + profile_key = profile.get("profile_key", "").strip() + if not profile_key: + continue + deduped_profiles[profile_key] = profile + + new_nodes = [ + self._build_profile_node(profile, ref_memory_id=ref_memory_id) for profile in deduped_profiles.values() + ] + if not new_nodes: + return [] + + await self.vector_store.delete([node.memory_id for node in new_nodes]) + await self.vector_store.insert([node.to_vector_node() for node in new_nodes]) + await self.enforce_capacity() + return new_nodes + + async def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode: + """Insert or replace a single profile row.""" + nodes = await self.add_batch( + [ + { + "message_time": message_time, + "profile_key": profile_key, + "profile_value": profile_value, + }, + ], + ref_memory_id=ref_memory_id, + ) + return nodes[0] + + async def update( + self, + profile_id: str, + message_time: str, + profile_key: str, + profile_value: str, + ) -> MemoryNode | None: + """Replace content and key for ``profile_id``; return ``None`` if missing.""" + existing_node = await self.get_by(profile_id=profile_id) + if existing_node is None: + return None + + new_node = self._build_profile_node( + { + "message_time": message_time, + "profile_key": profile_key, + "profile_value": profile_value, + "ref_memory_id": existing_node.ref_memory_id, + "metadata": existing_node.metadata, + }, + ) + + if existing_node.memory_id != new_node.memory_id: + await self.vector_store.delete(existing_node.memory_id) + else: + await self.vector_store.delete(new_node.memory_id) + + await self.vector_store.insert(new_node.to_vector_node()) + await self.enforce_capacity() + return new_node + + async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]: + """Semantic search with de-duplication across multiple query strings.""" + queries = [query] if isinstance(query, str) else query + seen_nodes: dict[str, MemoryNode] = {} + for item in queries: + if not item or not item.strip(): + continue + vector_nodes = await self.vector_store.search(item, limit=limit, filters=self._base_filters()) + for vector_node in vector_nodes: + memory_node = MemoryNode.from_vector_node(vector_node) + seen_nodes[memory_node.memory_id] = memory_node + nodes = list(seen_nodes.values()) + nodes.sort(key=lambda node: (node.score, node.message_time), reverse=True) + return nodes[:limit] + + async def enforce_capacity(self): + """Drop oldest rows when count exceeds ``max_capacity``.""" + nodes = await self.get_all() + overflow = len(nodes) - self.max_capacity + if overflow <= 0: + return + + to_delete = [node.memory_id for node in nodes[:overflow]] + await self.vector_store.delete(to_delete) diff --git a/reme/memory/vector_tools/profiles/retrieve_profile.py b/reme/memory/vector_tools/profiles/retrieve_profile.py new file mode 100644 index 00000000..38337bd8 --- /dev/null +++ b/reme/memory/vector_tools/profiles/retrieve_profile.py @@ -0,0 +1,102 @@ +"""Retrieve relevant profile rows.""" + +from loguru import logger + +from .profile_handler import ProfileHandler +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import MemoryNode, ToolCall + + +class RetrieveProfile(BaseMemoryTool): + """Tool to retrieve relevant profiles using the configured backend.""" + + def __init__(self, top_k: int = 5, enable_memory_target: bool = False, **kwargs): + super().__init__(**kwargs) + self.top_k = top_k + self.enable_memory_target = enable_memory_target + + def _build_query_parameters(self) -> dict: + properties = { + "query": { + "type": "string", + "description": "query", + }, + } + required = ["query"] + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "memory_target", + } + required.append("memory_target") + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "Retrieve relevant user profiles using semantic matching.", + "parameters": self._build_query_parameters(), + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "Retrieve relevant user profiles using semantic matching.", + "parameters": { + "type": "object", + "properties": { + "query_items": { + "type": "array", + "description": "List of query items.", + "items": self._build_query_parameters(), + }, + }, + "required": ["query_items"], + }, + }, + ) + + async def execute(self): + if self.enable_multiple: + query_items = self.context.get("query_items", []) + else: + query_items = [self.context] + + queries_by_target: dict[str, list[str]] = {} + for item in query_items: + target = item["memory_target"] if self.enable_memory_target else self.memory_target + queries_by_target.setdefault(target, []).append(item["query"]) + + profile_nodes: list[MemoryNode] = [] + for target, queries in queries_by_target.items(): + profile_handler = self.get_profile_handler(target) + nodes, _ = await profile_handler.aretrieve( + query=queries, + limit=self.top_k, + add_profile_id=True, + add_history_id=True, + ) + profile_nodes.extend(nodes) + + seen_ids = {node.memory_id: node for node in self.retrieved_nodes if node.memory_id} + new_nodes = [] + for node in profile_nodes: + if node.memory_id not in seen_ids: + seen_ids[node.memory_id] = node + new_nodes.append(node) + self.retrieved_nodes.extend(new_nodes) + + if not new_nodes: + output = "No new profiles found." + else: + output = "\n".join( + [ProfileHandler.format_node(node, add_profile_id=True, add_history_id=True) for node in new_nodes], + ) + + logger.info(f"Retrieved {len(profile_nodes)} profiles, {len(new_nodes)} new after deduplication") + return output diff --git a/reme/memory/vector_tools/profiles/vector_profile_backend.py b/reme/memory/vector_tools/profiles/vector_profile_backend.py new file mode 100644 index 00000000..3147cc92 --- /dev/null +++ b/reme/memory/vector_tools/profiles/vector_profile_backend.py @@ -0,0 +1,55 @@ +"""Vector-backed profile storage.""" + +from .profile_backend import BaseProfileBackend +from .profile_vector_handler import ProfileVectorHandler +from ....core import ServiceContext +from ....core.schema import MemoryNode + + +class VectorProfileBackend(BaseProfileBackend): + """Persist user profiles in a dedicated vector store.""" + + def __init__( + self, + memory_target: str, + service_context: ServiceContext, + vector_store_name: str = "profile", + max_capacity: int = 50, + ): + super().__init__(memory_target=memory_target, max_capacity=max_capacity) + self.handler = ProfileVectorHandler( + memory_target=memory_target, + service_context=service_context, + vector_store_name=vector_store_name, + max_capacity=max_capacity, + ) + + async def get_all(self) -> list[MemoryNode]: + return await self.handler.get_all() + + async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + return await self.handler.get_by(profile_id=profile_id, profile_key=profile_key) + + async def delete(self, profile_id: str | list[str]) -> bool | int: + return await self.handler.delete(profile_id) + + async def delete_all(self) -> int: + return await self.handler.delete_all() + + async def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode: + return await self.handler.add(message_time, profile_key, profile_value, ref_memory_id) + + async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + return await self.handler.add_batch(profiles, ref_memory_id) + + async def update( + self, + profile_id: str, + message_time: str, + profile_key: str, + profile_value: str, + ) -> MemoryNode | None: + return await self.handler.update(profile_id, message_time, profile_key, profile_value) + + async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]: + return await self.handler.search(query, limit)