From de28c8e243a2b623eb1c5e2b7e12c87cea9b0d38 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 7 Feb 2026 03:37:12 +0800 Subject: [PATCH 1/2] refactor(core): rename memory_storage to memory_store and add base fs tool --- reme/__init__.py | 2 + reme/agent/fs/fs_compactor.py | 198 +++-- reme/agent/fs/fs_compactor.yaml | 16 + reme/agent/fs/fs_summarizer.py | 35 +- reme/agent/fs/fs_summarizer.yaml | 14 +- reme/core/__init__.py | 4 + reme/core/application.py | 2 +- reme/core/context/service_context.py | 116 +-- reme/core/file_watcher/base_file_watcher.py | 2 +- reme/core/llm/lite_llm.py | 3 +- reme/core/llm/lite_llm_sync.py | 4 +- .../__init__.py | 2 +- .../base_memory_store.py | 0 .../sqlite_memory_store.py | 26 +- reme/core/op/base_op.py | 2 +- reme/core/schema/__init__.py | 2 + reme/core/schema/compaction_result.py | 15 + reme/core/schema/memory_search_result.py | 4 +- reme/reme.py | 2 +- reme/reme_fs.py | 106 ++- reme/tool/fs/__init__.py | 6 + reme/tool/fs/base_fs_tool.py | 41 + reme/tool/fs/bash_tool.py | 62 +- reme/tool/fs/edit_tool.py | 18 +- reme/tool/fs/find_tool.py | 41 +- reme/tool/fs/fs_memory_get.py | 42 +- reme/tool/fs/fs_memory_search.py | 12 +- reme/tool/fs/grep_tool.py | 19 +- reme/tool/fs/ls_tool.py | 9 +- reme/tool/fs/read_tool.py | 42 +- reme/tool/fs/write_tool.py | 4 +- tests/demo_memory_search.py | 2 +- tests/test_fs_agent.py | 323 -------- tests/test_fs_compact.py | 301 ++++++++ tests/test_fs_memory_get.py | 369 +++++++++ tests/test_fs_memory_search.py | 704 ++++++++++++++++++ tests/test_fs_summary.py | 234 ++++++ ...st_file_system_tool.py => test_fs_tool.py} | 0 tests/test_memory_store.py | 4 +- 39 files changed, 2096 insertions(+), 692 deletions(-) rename reme/core/{memory_storage => memory_store}/__init__.py (86%) rename reme/core/{memory_storage => memory_store}/base_memory_store.py (100%) rename reme/core/{memory_storage => memory_store}/sqlite_memory_store.py (96%) create mode 100644 reme/core/schema/compaction_result.py create mode 100644 reme/tool/fs/base_fs_tool.py delete mode 100644 tests/test_fs_agent.py create mode 100644 tests/test_fs_compact.py create mode 100644 tests/test_fs_memory_get.py create mode 100644 tests/test_fs_memory_search.py create mode 100644 tests/test_fs_summary.py rename tests/{test_file_system_tool.py => test_fs_tool.py} (100%) diff --git a/reme/__init__.py b/reme/__init__.py index cd5d4e6c..81fe52cf 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -6,6 +6,7 @@ from . import core from . import tool from . import workflow from .reme import ReMe +from .reme_fs import ReMeFs __all__ = [ "agent", @@ -14,6 +15,7 @@ __all__ = [ "tool", "workflow", "ReMe", + "ReMeFs", ] __version__ = "0.3.0.0a1" diff --git a/reme/agent/fs/fs_compactor.py b/reme/agent/fs/fs_compactor.py index 79d70c6f..07c2dd18 100644 --- a/reme/agent/fs/fs_compactor.py +++ b/reme/agent/fs/fs_compactor.py @@ -4,7 +4,7 @@ from loguru import logger from ...core.enumeration import Role, MemoryType from ...core.op import BaseReact -from ...core.schema import Message +from ...core.schema import CutPointResult, Message class FsCompactor(BaseReact): @@ -24,44 +24,39 @@ class FsCompactor(BaseReact): self.reserve_tokens: int = reserve_tokens self.keep_recent_tokens: int = keep_recent_tokens + @staticmethod + def _normalize_messages(messages: list[Message | dict]) -> list[Message]: + """Convert dict messages to Message objects.""" + return [Message(**m) if isinstance(m, dict) else m for m in messages] + @staticmethod def _is_user_message(message: Message) -> bool: - """Check if a message is a user-initiated message (user or tool result).""" + """Check if message is user role.""" return message.role is Role.USER def _find_turn_start_index(self, messages: list[Message], entry_index: int) -> int: - """ - Find the user message that starts the turn containing the given entry index. - Returns -1 if no turn start found before the index. - """ + """Find user message that starts the turn. Returns -1 if not found.""" + if not messages or entry_index < 0 or entry_index >= len(messages): + return -1 + for i in range(entry_index, -1, -1): if self._is_user_message(messages[i]): return i return -1 - def _find_cut_point(self, messages: list[Message]) -> dict: + def _find_cut_point(self, messages: list[Message]) -> CutPointResult: """ Find cut point with split turn detection. - A "split turn" occurs when the cut point falls in the middle of a conversation turn - rather than at a clean user message boundary. For example: - User → Assistant → [CUT HERE] → Assistant continues → User - - In this case, we need to: - 1. Summarize complete history (before turn start) - 2. Separately summarize the turn prefix (turn start to cut point) - 3. Keep the turn suffix (cut point onwards) in full - - Returns dict with: - - messages_to_summarize: Complete turns before the current turn - - turn_prefix_messages: If split turn, messages from turn start to cut point - - is_split_turn: Whether this is a split turn - - cut_index: The actual cut point index + Split turn: User → Assistant → [CUT] → Assistant → User + Clean cut: User → [CUT] → Assistant → User """ + if not messages: + return CutPointResult() + accumulated_tokens = 0 cut_index = 0 - # Walk backwards from the newest messages, accumulating tokens until we hit the keep threshold for i in range(len(messages) - 1, -1, -1): msg = messages[i] msg_tokens = self.token_counter.count_token([msg]) @@ -69,51 +64,51 @@ class FsCompactor(BaseReact): if accumulated_tokens >= self.keep_recent_tokens: cut_index = i + logger.debug(f"Cut point at index {cut_index}, {accumulated_tokens} tokens") break if cut_index == 0: - return { - "messages_to_summarize": [], - "turn_prefix_messages": [], - "is_split_turn": False, - "cut_index": 0, - } + return CutPointResult(left_messages=messages) - # Check if cut point is a user message (clean turn boundary) or assistant/other (mid-turn) cut_message = messages[cut_index] is_user_cut = self._is_user_message(cut_message) if is_user_cut: - # Clean cut: cut point is at a turn boundary, summarize everything before - return { - "messages_to_summarize": messages[:cut_index], - "turn_prefix_messages": [], - "is_split_turn": False, - "cut_index": cut_index, - } + return CutPointResult( + messages_to_summarize=messages[:cut_index], + left_messages=messages[cut_index:], + cut_index=cut_index, + ) - # Split turn detected: find where the current turn started turn_start_index = self._find_turn_start_index(messages, cut_index) if turn_start_index == -1: - # No turn start found (shouldn't happen), treat as clean cut - return { - "messages_to_summarize": messages[:cut_index], - "turn_prefix_messages": [], - "is_split_turn": False, - "cut_index": cut_index, - } + logger.warning("Split turn detected but no turn start found, treating as clean cut") + return CutPointResult( + messages_to_summarize=messages[:cut_index], + left_messages=messages[cut_index:], + cut_index=cut_index, + ) - # Split turn: separate complete history from turn prefix - # History: [0, turn_start_index) - complete turns to summarize - # Turn prefix: [turn_start_index, cut_index) - needs special context summary - # Turn suffix: [cut_index, end) - kept in full (recent work) - return { - "messages_to_summarize": messages[:turn_start_index], - "turn_prefix_messages": messages[turn_start_index:cut_index], - "is_split_turn": True, - "cut_index": cut_index, - } + return CutPointResult( + messages_to_summarize=messages[:turn_start_index], + turn_prefix_messages=messages[turn_start_index:cut_index], + left_messages=messages[cut_index:], + is_split_turn=True, + cut_index=cut_index, + ) + + async def _generate_summary(self, prompt_messages: list[Message]) -> str: + """Generate summary via LLM. Returns empty string if no messages.""" + if not prompt_messages: + return "" + + try: + assistant_message = await self.llm.chat(prompt_messages) + return assistant_message.content if assistant_message.content else "" + except Exception as e: + logger.error(f"Failed to generate summary: {e}") + raise RuntimeError(f"Summarization failed: {e}") from e @staticmethod def _serialize_conversation(messages: list[Message]) -> str: @@ -135,20 +130,15 @@ class FsCompactor(BaseReact): return "\n".join(lines) def build_messages_s1(self) -> list[Message]: - """ - Build messages for compaction summarization. - - This creates the prompt for the main history summary. If split turn is detected, - a separate turn prefix summary will be generated later in execute(). - """ - messages: list[Message] = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] - + """Build prompt for main history summary.""" + messages = self._normalize_messages(self.context.messages) cut_result = self._find_cut_point(messages) - messages_to_summarize = cut_result["messages_to_summarize"] - self.context.is_split_turn = cut_result["is_split_turn"] - self.context.turn_prefix_messages = cut_result["turn_prefix_messages"] - if not messages_to_summarize: + self.context.is_split_turn = cut_result.is_split_turn + self.context.turn_prefix_messages = cut_result.turn_prefix_messages + self.context.left_messages = cut_result.left_messages + + if not cut_result.messages_to_summarize: logger.info("No messages to summarize") return [] @@ -157,7 +147,7 @@ class FsCompactor(BaseReact): user_prompt = self.prompt_format("update_user_message", previous_summary=self.context.previous_summary) else: user_prompt = self.get_prompt("initial_user_message") - conversation_text = self._serialize_conversation(messages_to_summarize) + conversation_text = self._serialize_conversation(cut_result.messages_to_summarize) return [ Message(role=Role.SYSTEM, content=system_prompt), @@ -165,16 +155,7 @@ class FsCompactor(BaseReact): ] def build_messages_s2(self) -> list[Message]: - """ - Generate summary for turn prefix when splitting a turn. - - This provides context for the retained turn suffix. The summary focuses on: - - What the user originally asked for in this turn - - Key decisions and early progress made in the prefix - - Information needed to understand the kept suffix - - This is shorter and more focused than the full history summary. - """ + """Build prompt for turn prefix summary (split turn only).""" if not self.context.turn_prefix_messages: return [] @@ -191,55 +172,54 @@ class FsCompactor(BaseReact): """ Execute compaction if needed. - Compaction process: - 1. Check if token count exceeds threshold - 2. Find cut point and detect if it's a split turn - 3. Generate history summary (complete turns before cut point) - 4. If split turn: generate turn prefix summary (partial turn before cut point) - 5. Merge summaries and update context - - Final context structure after compaction: - - Summary (history + optional turn prefix context) - - Recent messages kept in full (from cut point onwards) + Returns: [summary_message, ...left_messages] if compacted, else original messages. """ - messages: list[Message] = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] - token_count: int = self.token_counter.count_token(messages) + original_messages = self._normalize_messages(self.context.messages) + token_count: int = self.token_counter.count_token(original_messages) threshold = self.context_window_tokens - self.reserve_tokens + if token_count < threshold: - logger.info(f"Token count {token_count} below threshold, skipping compaction") + logger.info(f"Token count {token_count} below threshold ({threshold}), skipping compaction") return { - "answer": "", - "success": True, - "messages": [], - "tools": [], - "skipped": True, + "compacted": False, + "tokens_before": token_count, + "is_split_turn": False, + "messages": original_messages, } - logger.info(f"Starting compaction, token count: {token_count}") - messages: list[Message] = self.build_messages_s1() - if messages: - assistant_message = await self.llm.chat(messages) - history_summary = assistant_message.content - else: - history_summary = "" + logger.info(f"Starting compaction, token count: {token_count}, threshold: {threshold}") + + history_prompt_messages = self.build_messages_s1() + + if not history_prompt_messages and not self.context.get("is_split_turn"): + logger.warning("No messages to summarize and not a split turn, returning original messages") + return { + "compacted": False, + "tokens_before": token_count, + "is_split_turn": False, + "messages": original_messages, + } + + history_summary = await self._generate_summary(history_prompt_messages) if history_prompt_messages else "" if self.context.is_split_turn and self.context.turn_prefix_messages: logger.info("Split turn detected, generating turn prefix summary") - messages: list[Message] = self.build_messages_s2() - if messages: - assistant_message = await self.llm.chat(messages) - turn_prefix_summary = assistant_message.content - else: - turn_prefix_summary = "" + turn_prefix_prompt_messages = self.build_messages_s2() + turn_prefix_summary = await self._generate_summary(turn_prefix_prompt_messages) summary = f"{history_summary}\n\n---\n\n**Turn Context (split turn):**\n\n{turn_prefix_summary}" else: summary = history_summary logger.info(f"Compaction complete, summary length: {len(summary)}, split_turn: {self.context.is_split_turn}") + summary_content = self.prompt_format("compaction_summary_format", summary=summary) + summary_message = Message(role=Role.USER, content=summary_content) + left_messages = self.context.get("left_messages", []) + final_messages = [summary_message] + left_messages + return { "compacted": True, "tokens_before": token_count, - "summary": summary, "is_split_turn": self.context.is_split_turn, + "messages": final_messages, } diff --git a/reme/agent/fs/fs_compactor.yaml b/reme/agent/fs/fs_compactor.yaml index 40fb6ec4..dbad84f2 100644 --- a/reme/agent/fs/fs_compactor.yaml +++ b/reme/agent/fs/fs_compactor.yaml @@ -210,3 +210,19 @@ turn_prefix_summarization_zh: | - [理解保留的最近工作所需的信息] 保持简洁。专注于理解保留后缀所需的内容。 + +# Format for wrapping the compaction summary when presenting to LLM +# This matches the TypeScript format: COMPACTION_SUMMARY_PREFIX + summary + COMPACTION_SUMMARY_SUFFIX +compaction_summary_format: | + The conversation history before this point was compacted into the following summary: + + + {summary} + + +compaction_summary_format_zh: | + 此点之前的对话历史已被压缩为以下摘要： + + + {summary} + diff --git a/reme/agent/fs/fs_summarizer.py b/reme/agent/fs/fs_summarizer.py index 25176dc1..0b88ad5c 100644 --- a/reme/agent/fs/fs_summarizer.py +++ b/reme/agent/fs/fs_summarizer.py @@ -1,5 +1,7 @@ """Personal memory retriever agent for retrieving personal memories through vector search.""" +import datetime + from loguru import logger from ...core.enumeration import Role, MemoryType @@ -14,28 +16,25 @@ class FsSummarizer(BaseReact): def __init__( self, - memory_dir: str, + memory_dir: str = "memory", version: str = "default", - context_window_tokens: int = 128000, - reserve_tokens: int = 32000, - soft_threshold_tokens: int = 4000, **kwargs, ): super().__init__(**kwargs) self.memory_dir: str = memory_dir self.version: str = version - self.context_window_tokens: int = context_window_tokens - self.reserve_tokens: int = reserve_tokens - self.soft_threshold_tokens: int = soft_threshold_tokens async def build_messages(self) -> list[Message]: messages: list[Message] = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] + date_str: str = self.context.get("date", datetime.datetime.now().strftime("%Y-%m-%d")) + if self.version == "default": messages.append( Message( role=Role.USER, content=self.prompt_format( "user_message_v2", + date=date_str, memory_dir=self.memory_dir, ), ), @@ -47,6 +46,7 @@ class FsSummarizer(BaseReact): role=Role.USER, content=self.prompt_format( "user_message", + date=date_str, memory_dir=self.memory_dir, ), ), @@ -54,27 +54,6 @@ class FsSummarizer(BaseReact): return messages async def execute(self): - context_window = max(1, int(self.context_window_tokens)) - reserve_tokens = max(0, int(self.reserve_tokens)) - soft_threshold = max(0, int(self.soft_threshold_tokens)) - threshold = max(0, context_window - reserve_tokens - soft_threshold) - messages: list[Message] = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] - token_count: int = self.token_counter.count_token(messages) - - if token_count >= threshold: - logger.info(f"[{self.__class__.__name__}] Skipping summary execution based on threshold check") - return { - "answer": "", - "success": True, - "messages": [], - "tools": [], - "skipped": True, - } - - # Mark that we're executing a summary in this cycle - summary_count = self.context.get("summary_count", 0) - self.context["last_summary_at"] = summary_count - result = await super().execute() answer = result["answer"] logger.info(f"[{self.__class__.__name__}] answer={answer}") diff --git a/reme/agent/fs/fs_summarizer.yaml b/reme/agent/fs/fs_summarizer.yaml index f4673750..ef3de8ae 100644 --- a/reme/agent/fs/fs_summarizer.yaml +++ b/reme/agent/fs/fs_summarizer.yaml @@ -5,11 +5,19 @@ system_prompt: | user_message: | Pre-compaction memory flush. + Current Date: {date} Store durable memories now (use {memory_dir}/YYYY-MM-DD.md; create {memory_dir}/ if needed). If nothing to store, reply with [SILENT]. user_message_v2: | - Pre-compaction memory flush. + Current Date: {date} The session is near auto-compaction; capture durable memories to disk. - Store durable memories now (use {memory_dir}/YYYY-MM-DD.md; create {memory_dir}/ if needed). - If nothing to store, reply with [SILENT]. \ No newline at end of file + + Memory storage workflow: + 1. Check if {memory_dir}/ exists; if not, create it via bash + 2. Check if {memory_dir}/YYYY-MM-DD.md exists (use actual date) + 3. If file is NEW: Write memories directly (be concise) + 4. If file EXISTS: Read it first, then UPDATE with new memories (keep concise, merge/deduplicate) + 5. If NO valuable information to store: Reply with reason and [SILENT] + + Store durable memories. Keep entries concise and well-organized. \ No newline at end of file diff --git a/reme/core/__init__.py b/reme/core/__init__.py index 2b56fb08..4128cbd9 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -3,8 +3,10 @@ from . import context from . import embedding from . import enumeration +from . import file_watcher from . import flow from . import llm +from . import memory_store from . import op from . import schema from . import service @@ -18,8 +20,10 @@ __all__ = [ "context", "embedding", "enumeration", + "file_watcher", "flow", "llm", + "memory_store", "op", "schema", "service", diff --git a/reme/core/application.py b/reme/core/application.py index ff093692..fde84afe 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -7,7 +7,7 @@ from .embedding import BaseEmbeddingModel from .file_watcher import BaseFileWatcher from .flow import BaseFlow from .llm import BaseLLM -from .memory_storage import BaseMemoryStore +from .memory_store import BaseMemoryStore from .schema import Response from .token_counter import BaseTokenCounter from .utils import execute_stream_task, PydanticConfigParser diff --git a/reme/core/context/service_context.py b/reme/core/context/service_context.py index a74eac40..f6f07efa 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/context/service_context.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from ..llm import BaseLLM from ..embedding import BaseEmbeddingModel from ..vector_store import BaseVectorStore - from ..memory_storage import BaseMemoryStore + from ..memory_store import BaseMemoryStore from ..token_counter import BaseTokenCounter from ..flow import BaseFlow from ..service import BaseService @@ -45,24 +45,39 @@ class ServiceContext(BaseContext): **kwargs, ): super().__init__() - self.service_config: ServiceConfig = self._build_service_config( - *args, - llm_api_key=llm_api_key, - llm_api_base=llm_api_base, - embedding_api_key=embedding_api_key, - embedding_api_base=embedding_api_base, - service_config=service_config, - parser=parser, - config_path=config_path, - enable_logo=enable_logo, - llm=llm, - embedding_model=embedding_model, - vector_store=vector_store, - memory_store=memory_store, - token_counter=token_counter, - file_watcher=file_watcher, - **kwargs, - ) + + load_env() + self._update_env("REME_LLM_API_KEY", llm_api_key) + self._update_env("REME_LLM_BASE_URL", llm_api_base) + self._update_env("REME_EMBEDDING_API_KEY", embedding_api_key) + self._update_env("REME_EMBEDDING_BASE_URL", embedding_api_base) + + if service_config is None: + parser_class = parser if parser is not None else PydanticConfigParser + parser = parser_class(ServiceConfig) + input_args = [] + if config_path: + input_args.append(f"config={config_path}") + if args: + input_args.extend(args) + + if llm: + self._update_section_config(kwargs, "llm", **llm) + if embedding_model: + self._update_section_config(kwargs, "embedding_model", **embedding_model) + if token_counter: + self._update_section_config(kwargs, "token_counter", **token_counter) + if vector_store: + self._update_section_config(kwargs, "vector_store", **vector_store) + if memory_store: + self._update_section_config(kwargs, "memory_store", **memory_store) + if file_watcher: + self._update_section_config(kwargs, "file_watcher", **file_watcher) + kwargs["enable_logo"] = enable_logo + logger.info(f"update with args: {input_args} kwargs: {kwargs}") + service_config = parser.parse_args(*input_args, **kwargs) + + self.service_config: ServiceConfig = service_config if self.service_config.init_logger: init_logger() @@ -92,57 +107,6 @@ class ServiceContext(BaseContext): self._build_flows() - def _build_service_config( - self, - *args, - llm_api_key: str | None = None, - llm_api_base: str | None = None, - embedding_api_key: str | None = None, - embedding_api_base: str | None = None, - service_config: ServiceConfig | None = None, - parser: type[PydanticConfigParser] | None = None, - config_path: str | None = None, - enable_logo: bool = True, - llm: dict | None = None, - embedding_model: dict | None = None, - vector_store: dict | None = None, - memory_store: dict | None = None, - token_counter: dict | None = None, - file_watcher: dict | None = None, - **kwargs, - ) -> ServiceConfig: - - load_env() - self._update_env("REME_LLM_API_KEY", llm_api_key) - self._update_env("REME_LLM_BASE_URL", llm_api_base) - self._update_env("REME_EMBEDDING_API_KEY", embedding_api_key) - self._update_env("REME_EMBEDDING_BASE_URL", embedding_api_base) - - if service_config is None: - parser_class = parser if parser is not None else PydanticConfigParser - parser = parser_class(ServiceConfig) - input_args = [] - if config_path: - input_args.append(f"config={config_path}") - if args: - input_args.extend(args) - service_config = parser.parse_args(*input_args, **kwargs) - - service_config.enable_logo = enable_logo - if llm: - self._update_section_config(service_config, "llm", **llm) - if embedding_model: - self._update_section_config(service_config, "embedding_model", **embedding_model) - if token_counter: - self._update_section_config(service_config, "token_counter", **token_counter) - if vector_store: - self._update_section_config(service_config, "vector_store", **vector_store) - if memory_store: - self._update_section_config(service_config, "memory_store", **memory_store) - if file_watcher: - self._update_section_config(service_config, "file_watcher", **file_watcher) - return service_config - @staticmethod def _update_env(key: str, value: str | None): """Update environment variable if value is provided.""" @@ -150,13 +114,13 @@ class ServiceContext(BaseContext): os.environ[key] = value @staticmethod - def _update_section_config(service_config: ServiceConfig, section_name: str, **kwargs): + def _update_section_config(config: dict, section_name: str, **kwargs): """Update a specific section of the service config with new values.""" - section_dict: dict = getattr(service_config, section_name) - if "default" not in section_dict: - raise KeyError(f"Default `{section_name}` config not found") - current_config = section_dict["default"] - section_dict["default"] = current_config.model_copy(update=kwargs, deep=True) + if section_name not in config: + config[section_name] = {} + if "default" not in config[section_name]: + config[section_name]["default"] = {} + config[section_name]["default"].update(kwargs) def _build_flows(self): expression_flow_cls = None diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 620d2e17..a4e8af36 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -11,7 +11,7 @@ from typing import Any, Callable from loguru import logger from watchfiles import awatch, Change -from ..memory_storage import BaseMemoryStore +from ..memory_store import BaseMemoryStore class BaseFileWatcher: diff --git a/reme/core/llm/lite_llm.py b/reme/core/llm/lite_llm.py index 5663702b..7950a8c2 100644 --- a/reme/core/llm/lite_llm.py +++ b/reme/core/llm/lite_llm.py @@ -3,7 +3,6 @@ import os from typing import AsyncGenerator -import litellm from loguru import logger from .base_llm import BaseLLM @@ -88,6 +87,8 @@ class LiteLLM(BaseLLM): stream_kwargs: dict | None = None, ) -> AsyncGenerator[StreamChunk, None]: """Execute async streaming chat requests and yield processed response chunks.""" + import litellm + stream_kwargs = stream_kwargs or {} completion = await litellm.acompletion(**stream_kwargs) ret_tool_calls: list[ToolCall] = [] diff --git a/reme/core/llm/lite_llm_sync.py b/reme/core/llm/lite_llm_sync.py index 778eaed0..9951d223 100644 --- a/reme/core/llm/lite_llm_sync.py +++ b/reme/core/llm/lite_llm_sync.py @@ -2,8 +2,6 @@ from typing import Generator -import litellm - from .lite_llm import LiteLLM from ..enumeration import ChunkEnum from ..schema import Message @@ -21,6 +19,8 @@ class LiteLLMSync(LiteLLM): stream_kwargs: dict | None = None, ) -> Generator[StreamChunk, None, None]: """Internal synchronous generator for processing streaming chat completion chunks.""" + import litellm + stream_kwargs = stream_kwargs or {} completion = litellm.completion(**stream_kwargs) ret_tool_calls: list[ToolCall] = [] diff --git a/reme/core/memory_storage/__init__.py b/reme/core/memory_store/__init__.py similarity index 86% rename from reme/core/memory_storage/__init__.py rename to reme/core/memory_store/__init__.py index 371d4e0f..6d43e423 100644 --- a/reme/core/memory_storage/__init__.py +++ b/reme/core/memory_store/__init__.py @@ -1,4 +1,4 @@ -"""Memory storage module for persistent memory management. +"""Memory store module for persistent memory management. This module provides storage backends for memory chunks and file metadata, including SQLite-based implementations with vector and full-text search. diff --git a/reme/core/memory_storage/base_memory_store.py b/reme/core/memory_store/base_memory_store.py similarity index 100% rename from reme/core/memory_storage/base_memory_store.py rename to reme/core/memory_store/base_memory_store.py diff --git a/reme/core/memory_storage/sqlite_memory_store.py b/reme/core/memory_store/sqlite_memory_store.py similarity index 96% rename from reme/core/memory_storage/sqlite_memory_store.py rename to reme/core/memory_store/sqlite_memory_store.py index 756c2b3b..7e1519b8 100644 --- a/reme/core/memory_storage/sqlite_memory_store.py +++ b/reme/core/memory_store/sqlite_memory_store.py @@ -208,14 +208,17 @@ class SqliteMemoryStore(BaseMemoryStore): ), ) - # Insert vector + # Insert vector (vec0 doesn't support OR REPLACE, use DELETE + INSERT) if self.vector_available: assert chunk.embedding, "Embedding is required for vector insert" + # Delete existing vector first cursor.execute( - f""" - INSERT OR REPLACE INTO {self.vector_table_name} (id, embedding) - VALUES (?, ?) - """, + f"DELETE FROM {self.vector_table_name} WHERE id = ?", + (chunk.id,), + ) + # Then insert new vector + cursor.execute( + f"INSERT INTO {self.vector_table_name} (id, embedding) VALUES (?, ?)", (chunk.id, self.vector_to_blob(chunk.embedding)), ) @@ -372,14 +375,17 @@ class SqliteMemoryStore(BaseMemoryStore): ), ) - # Insert/update vector + # Insert/update vector (vec0 doesn't support OR REPLACE, use DELETE + INSERT) if self.vector_available: assert chunk.embedding, "Embedding is required for vector insert" + # Delete existing vector first cursor.execute( - f""" - INSERT OR REPLACE INTO {self.vector_table_name} (id, embedding) - VALUES (?, ?) - """, + f"DELETE FROM {self.vector_table_name} WHERE id = ?", + (chunk.id,), + ) + # Then insert new vector + cursor.execute( + f"INSERT INTO {self.vector_table_name} (id, embedding) VALUES (?, ?)", (chunk.id, self.vector_to_blob(chunk.embedding)), ) diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index 015d0fbb..8911cead 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -13,7 +13,7 @@ from tqdm import tqdm from ..context import RuntimeContext, PromptHandler, ServiceContext from ..embedding import BaseEmbeddingModel from ..llm import BaseLLM -from ..memory_storage import BaseMemoryStore +from ..memory_store import BaseMemoryStore from ..schema import Response from ..token_counter import BaseTokenCounter from ..utils import camel_to_snake, CacheHandler, timer diff --git a/reme/core/schema/__init__.py b/reme/core/schema/__init__.py index d4c5c051..ae6b8392 100644 --- a/reme/core/schema/__init__.py +++ b/reme/core/schema/__init__.py @@ -1,5 +1,6 @@ """schema""" +from .compaction_result import CutPointResult from .file_metadata import FileMetadata from .memory_chunk import MemoryChunk from .memory_node import MemoryNode @@ -26,6 +27,7 @@ from .vector_node import VectorNode __all__ = [ "CmdConfig", "ContentBlock", + "CutPointResult", "EmbeddingModelConfig", "FileMetadata", "FlowConfig", diff --git a/reme/core/schema/compaction_result.py b/reme/core/schema/compaction_result.py new file mode 100644 index 00000000..0cc1d86e --- /dev/null +++ b/reme/core/schema/compaction_result.py @@ -0,0 +1,15 @@ +"""Compaction result schemas for context window management.""" + +from pydantic import BaseModel, Field + +from .message import Message + + +class CutPointResult(BaseModel): + """Cut point detection result for conversation compaction.""" + + messages_to_summarize: list[Message] = Field(default_factory=list, description="Complete turns before cut point") + turn_prefix_messages: list[Message] = Field(default_factory=list, description="Turn prefix if split turn") + left_messages: list[Message] = Field(default_factory=list, description="Messages to keep from cut point onwards") + is_split_turn: bool = Field(default=False, description="Whether cut point is mid-turn") + cut_index: int = Field(default=0, description="Index of cut point in original message list") diff --git a/reme/core/schema/memory_search_result.py b/reme/core/schema/memory_search_result.py index d08ca12b..784fe7c7 100644 --- a/reme/core/schema/memory_search_result.py +++ b/reme/core/schema/memory_search_result.py @@ -1,7 +1,5 @@ """Memory search result schema.""" -from typing import Any, Dict - from pydantic import BaseModel, Field from ..enumeration import MemorySource @@ -16,7 +14,7 @@ class MemorySearchResult(BaseModel): score: float = Field(..., description="Relevance score of the search result") snippet: str = Field(..., description="Text snippet from the matched content") source: MemorySource = Field(..., description="Source of the memory data") - metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata") + metadata: dict = Field(default_factory=dict, description="Additional metadata") @property def merge_key(self) -> str: diff --git a/reme/reme.py b/reme/reme.py index f08858ec..4f197dd6 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -251,7 +251,7 @@ class ReMe(Application): enable_when_to_use=False, enable_multiple=True, ), - UpdateMemoryV1( + AddMemory( enable_thinking_params=enable_thinking_params, enable_memory_target=False, enable_when_to_use=False, diff --git a/reme/reme_fs.py b/reme/reme_fs.py index f5ef42d2..51875ffc 100644 --- a/reme/reme_fs.py +++ b/reme/reme_fs.py @@ -1,7 +1,24 @@ """ReMe File System""" +from pathlib import Path + +from .agent.fs import FsCompactor, FsSummarizer from .config import ReMeConfigParser from .core import Application +from .core.enumeration import MemorySource +from .core.op import BaseTool +from .core.schema import Message +from .tool.fs import ( + BashTool, + EditTool, + FindTool, + FsMemoryGet, + FsMemorySearch, + GrepTool, + LsTool, + ReadTool, + WriteTool, +) class ReMeFs(Application): @@ -17,9 +34,10 @@ class ReMeFs(Application): enable_logo: bool = True, llm: dict | None = None, embedding_model: dict | None = None, - vector_store: dict | None = None, + memory_store: dict | None = None, token_counter: dict | None = None, - working_dir: str = "./agent", + file_watcher: dict | None = None, + working_dir: str = ".reme", **kwargs, ): """Initialize ReMe with config.""" @@ -33,9 +51,91 @@ class ReMeFs(Application): parser=ReMeConfigParser, llm=llm, embedding_model=embedding_model, - vector_store=vector_store, + memory_store=memory_store, token_counter=token_counter, + file_watcher=file_watcher, **kwargs, ) self.working_dir: str = working_dir + self.fs_tools: list[BaseTool] = [ + BashTool(cwd=self.working_dir), + EditTool(cwd=self.working_dir), + FindTool(cwd=self.working_dir), + GrepTool(cwd=self.working_dir), + LsTool(cwd=self.working_dir), + ReadTool(cwd=self.working_dir), + WriteTool(cwd=self.working_dir), + ] + self.working_path: Path = Path(self.working_dir) + self.working_path.mkdir(parents=True, exist_ok=True) + + async def compact( + self, + messages: list[Message | dict], + context_window_tokens: int = 128000, + reserve_tokens: int = 36000, + keep_recent_tokens: int = 20000, + ): + """Compact messages.""" + messages = [Message(**message) if isinstance(message, dict) else message for message in messages] + compactor = FsCompactor( + context_window_tokens=context_window_tokens, + reserve_tokens=reserve_tokens, + keep_recent_tokens=keep_recent_tokens, + ) + + return await compactor.call(messages=messages, service_context=self.service_context) + + async def summary( + self, + messages: list[Message | dict], + date: str, + version: str = "default", + context_window_tokens: int = 128000, + reserve_tokens: int = 32000, + soft_threshold_tokens: int = 4000, + ): + """Summarize messages.""" + messages = [Message(**message) if isinstance(message, dict) else message for message in messages] + summarizer = FsSummarizer( + tools=self.fs_tools, + version=version, + context_window_tokens=context_window_tokens, + reserve_tokens=reserve_tokens, + soft_threshold_tokens=soft_threshold_tokens, + ) + + return await summarizer.call(messages=messages, date=date, service_context=self.service_context) + + async def memory_search( + self, + query: str, + max_results: int = 20, + min_score: float = 0.1, + sources: list[MemorySource] | None = None, + hybrid_enabled: bool = True, + hybrid_vector_weight: float = 0.7, + hybrid_text_weight: float = 0.3, + hybrid_candidate_multiplier: float = 3.0, + ) -> str: + """Semantically search memory files.""" + search_tool = FsMemorySearch( + sources=sources, + hybrid_enabled=hybrid_enabled, + hybrid_vector_weight=hybrid_vector_weight, + hybrid_text_weight=hybrid_text_weight, + hybrid_candidate_multiplier=hybrid_candidate_multiplier, + ) + + return await search_tool.call( + query=query, + max_results=max_results, + min_score=min_score, + service_context=self.service_context, + ) + + async def memory_get(self, path: str, offset: int | None = None, limit: int | None = None) -> str: + """Read specific snippets from memory files.""" + get_tool = FsMemoryGet(workspace_dir=self.working_dir, memory_store=self.memory_store) + return await get_tool.call(path=path, offset=offset, limit=limit, service_context=self.service_context) diff --git a/reme/tool/fs/__init__.py b/reme/tool/fs/__init__.py index 090e1a1f..3a85838d 100644 --- a/reme/tool/fs/__init__.py +++ b/reme/tool/fs/__init__.py @@ -1,8 +1,11 @@ """File system tools.""" +from .base_fs_tool import BaseFsTool from .bash_tool import BashTool from .edit_tool import EditTool from .find_tool import FindTool +from .fs_memory_get import FsMemoryGet +from .fs_memory_search import FsMemorySearch from .grep_tool import GrepTool from .ls_tool import LsTool from .read_tool import ReadTool @@ -10,9 +13,12 @@ from .write_tool import WriteTool from ...core import R __all__ = [ + "BaseFsTool", "BashTool", "EditTool", "FindTool", + "FsMemoryGet", + "FsMemorySearch", "GrepTool", "LsTool", "ReadTool", diff --git a/reme/tool/fs/base_fs_tool.py b/reme/tool/fs/base_fs_tool.py new file mode 100644 index 00000000..92c9092a --- /dev/null +++ b/reme/tool/fs/base_fs_tool.py @@ -0,0 +1,41 @@ +"""Base class for file system tools with unified error handling.""" + +from loguru import logger + +from ...core.context import RuntimeContext +from ...core.op import BaseTool + + +class BaseFsTool(BaseTool): + """Base class for file system tools. + + Features: + - No retry logic (max_retries=1) + - Catches all exceptions and returns error messages to LLM + - Simplifies error handling in subclasses + """ + + def __init__(self, **kwargs): + """Initialize fs tool with no retry.""" + kwargs.setdefault("max_retries", 1) + kwargs.setdefault("raise_exception", False) + super().__init__(**kwargs) + + async def call(self, context: RuntimeContext = None, **kwargs): + """Execute the tool with unified error handling. + + This method catches all exceptions and returns error messages + to the LLM instead of raising them. + """ + self.context = RuntimeContext.from_context(context, **kwargs) + + try: + await self.before_execute() + response = await self.execute() + response = await self.after_execute(response) + return response + except Exception as e: + # Return error message to LLM instead of raising + error_msg = f"{self.__class__.__name__} failed: {str(e)}" + logger.error(error_msg) + return await self.after_execute(error_msg) diff --git a/reme/tool/fs/bash_tool.py b/reme/tool/fs/bash_tool.py index d8aea81c..1b084323 100644 --- a/reme/tool/fs/bash_tool.py +++ b/reme/tool/fs/bash_tool.py @@ -11,8 +11,8 @@ import platform import signal from pathlib import Path +from .base_fs_tool import BaseFsTool from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncate_tail -from ...core.op import BaseTool from ...core.schema import ToolCall, TruncationResult @@ -53,7 +53,7 @@ def kill_process_tree(pid: int) -> None: pass # Best effort -class BashTool(BaseTool): +class BashTool(BaseFsTool): """Production-grade tool for executing bash commands. Features: @@ -118,41 +118,35 @@ class BashTool(BaseTool): shell, shell_args = get_shell_config() # Start process - try: - process = await asyncio.create_subprocess_exec( - shell, - *shell_args, - command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=self.cwd, - # Create process group for clean termination - preexec_fn=os.setpgrp if platform.system() != "Windows" else None, - ) - except Exception as e: - raise RuntimeError(f"Failed to start process: {e}") from e + process = await asyncio.create_subprocess_exec( + shell, + *shell_args, + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self.cwd, + # Create process group for clean termination + preexec_fn=os.setpgrp if platform.system() != "Windows" else None, + ) # Execute command with optional timeout - try: - if timeout and timeout > 0: + if timeout and timeout > 0: + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), + timeout=timeout, + ) + except asyncio.TimeoutError as e: + # Kill process tree on timeout + if process.pid: + kill_process_tree(process.pid) try: - stdout, stderr = await asyncio.wait_for( - process.communicate(), - timeout=timeout, - ) - except asyncio.TimeoutError as e: - # Kill process tree on timeout - if process.pid: - kill_process_tree(process.pid) - try: - await asyncio.wait_for(process.wait(), timeout=1.0) - except asyncio.TimeoutError: - process.kill() - raise TimeoutError(f"Command timed out after {timeout} seconds") from e - else: - stdout, stderr = await process.communicate() - except TimeoutError as e: - raise RuntimeError(str(e)) from e + await asyncio.wait_for(process.wait(), timeout=1.0) + except asyncio.TimeoutError: + process.kill() + raise TimeoutError(f"Command timed out after {timeout} seconds") from e + else: + stdout, stderr = await process.communicate() # Decode output full_output = stdout.decode("utf-8", errors="ignore") diff --git a/reme/tool/fs/edit_tool.py b/reme/tool/fs/edit_tool.py index fe73bd1f..8733c2af 100644 --- a/reme/tool/fs/edit_tool.py +++ b/reme/tool/fs/edit_tool.py @@ -3,6 +3,7 @@ import os from pathlib import Path +from .base_fs_tool import BaseFsTool from .edit_diff import ( detect_line_ending, fuzzy_find_text, @@ -12,11 +13,10 @@ from .edit_diff import ( restore_line_endings, strip_bom, ) -from ...core.op import BaseTool from ...core.schema import ToolCall -class EditTool(BaseTool): +class EditTool(BaseFsTool): """Edit a file by replacing exact text.""" def __init__(self, cwd: str | None = None): @@ -77,11 +77,8 @@ class EditTool(BaseTool): raise PermissionError(f"File not readable/writable: {path}") # Read file - try: - with open(absolute_path, "r", encoding="utf-8") as f: - raw_content = f.read() - except Exception as e: - raise IOError(f"Failed to read file {path}: {e}") from e + with open(absolute_path, "r", encoding="utf-8") as f: + raw_content = f.read() # Strip BOM (LLM won't include invisible BOM in oldText) bom, content = strip_bom(raw_content) @@ -129,11 +126,8 @@ class EditTool(BaseTool): # Write file final_content = bom + restore_line_endings(new_content, original_ending) - try: - with open(absolute_path, "w", encoding="utf-8") as f: - f.write(final_content) - except Exception as e: - raise IOError(f"Failed to write file {path}: {e}") from e + with open(absolute_path, "w", encoding="utf-8") as f: + f.write(final_content) # Generate diff diff_result = generate_diff_string(base_content, new_content) diff --git a/reme/tool/fs/find_tool.py b/reme/tool/fs/find_tool.py index 7b855586..731722cd 100644 --- a/reme/tool/fs/find_tool.py +++ b/reme/tool/fs/find_tool.py @@ -3,12 +3,12 @@ import os from pathlib import Path +from .base_fs_tool import BaseFsTool from .truncate import FIND_MAX_BYTES, FIND_MAX_LINES, format_size, truncate_head -from ...core.op import BaseTool from ...core.schema import ToolCall -class FindTool(BaseTool): +class FindTool(BaseFsTool): """Search for files by glob pattern, respecting .gitignore.""" def __init__(self, cwd: str | None = None): @@ -131,28 +131,25 @@ class FindTool(BaseTool): # Search for files results = [] - try: - for file_path in search_path.glob(pattern): - if len(results) >= limit: - break + for file_path in search_path.glob(pattern): + if len(results) >= limit: + break - # Skip if matches ignore patterns - if self._should_ignore(file_path, ignore_patterns): - continue + # Skip if matches ignore patterns + if self._should_ignore(file_path, ignore_patterns): + continue - # Get relative path - try: - rel_path = file_path.relative_to(search_path) - # Add trailing slash for directories - if file_path.is_dir(): - results.append(f"{rel_path}/") - else: - results.append(str(rel_path)) - except ValueError: - # If relative_to fails, use the path as-is - results.append(str(file_path)) - except Exception as e: - raise RuntimeError(f"Error searching for files: {e}") from e + # Get relative path + try: + rel_path = file_path.relative_to(search_path) + # Add trailing slash for directories + if file_path.is_dir(): + results.append(f"{rel_path}/") + else: + results.append(str(rel_path)) + except ValueError: + # If relative_to fails, use the path as-is + results.append(str(file_path)) # Handle no results if not results: diff --git a/reme/tool/fs/fs_memory_get.py b/reme/tool/fs/fs_memory_get.py index c872aabf..c4100507 100644 --- a/reme/tool/fs/fs_memory_get.py +++ b/reme/tool/fs/fs_memory_get.py @@ -3,11 +3,11 @@ import os from pathlib import Path -from reme.core.op import BaseTool from reme.core.schema import ToolCall +from .base_fs_tool import BaseFsTool -class FsMemoryGet(BaseTool): +class FsMemoryGet(BaseFsTool): """Read specific snippets from memory files.""" def __init__(self, workspace_dir: str | None = None, **kwargs): @@ -20,7 +20,7 @@ class FsMemoryGet(BaseTool): return ToolCall( **{ "description": ( - "Safe snippet read from MEMORY.md, memory/*.md with optional from/lines; " + "Safe snippet read from MEMORY.md, memory/*.md with optional offset/limit; " "use after memory_search to pull only the needed lines and keep context small." ), "parameters": { @@ -30,11 +30,11 @@ class FsMemoryGet(BaseTool): "type": "string", "description": "Path to the memory file to read (relative or absolute)", }, - "from": { + "offset": { "type": "integer", "description": "Starting line number (1-indexed, optional)", }, - "lines": { + "limit": { "type": "integer", "description": "Number of lines to read from the starting line (optional)", }, @@ -47,8 +47,8 @@ class FsMemoryGet(BaseTool): async def execute(self) -> str: """Execute the memory get operation.""" raw_path: str = self.context.path.strip() - from_param: int | None = self.context.get("from", None) - lines_param: int | None = self.context.get("lines", None) + offset: int | None = self.context.get("offset", None) + limit: int | None = self.context.get("limit", None) if os.path.isabs(raw_path): abs_path = os.path.abspath(raw_path) @@ -65,15 +65,25 @@ class FsMemoryGet(BaseTool): with open(abs_path, "r", encoding="utf-8") as f: content = f.read() - if from_param is None and lines_param is None: + if offset is None and limit is None: return content - else: - lines = content.split("\n") - start = max(1, from_param if from_param is not None else 1) - count = max(1, lines_param if lines_param is not None else len(lines)) + lines = content.split("\n") + total_lines = len(lines) - # Extract slice (1-indexed to 0-indexed conversion) - selected = lines[start - 1 : start - 1 + count] - text = "\n".join(selected) - return text + # Validate and normalize offset (1-indexed) + start = offset if offset is not None else 1 + assert start >= 1, f"offset must be >= 1, got {start}" + assert start <= total_lines, f"offset {start} exceeds total lines {total_lines}" + + # Validate and calculate count + if limit is not None: + assert limit > 0, f"limit must be positive, got {limit}" + count = limit + else: + # Read from start to end of file + count = total_lines - start + 1 + + # Extract slice (1-indexed to 0-indexed conversion) + selected = lines[start - 1 : start - 1 + count] + return "\n".join(selected) diff --git a/reme/tool/fs/fs_memory_search.py b/reme/tool/fs/fs_memory_search.py index af245a99..924be3d2 100644 --- a/reme/tool/fs/fs_memory_search.py +++ b/reme/tool/fs/fs_memory_search.py @@ -3,11 +3,11 @@ import json from reme.core.enumeration import MemorySource -from reme.core.op import BaseTool from reme.core.schema import MemorySearchResult, ToolCall +from .base_fs_tool import BaseFsTool -class FsMemorySearch(BaseTool): +class FsMemorySearch(BaseFsTool): """Semantically search MEMORY.md and memory files.""" def __init__( @@ -48,11 +48,11 @@ class FsMemorySearch(BaseTool): "type": "string", "description": "The semantic search query to find relevant memory snippets", }, - "maxResults": { + "max_results": { "type": "integer", "description": "Maximum number of search results to return (optional)", }, - "minScore": { + "min_score": { "type": "number", "description": "Minimum similarity score threshold for results (optional)", }, @@ -65,8 +65,8 @@ class FsMemorySearch(BaseTool): async def execute(self) -> str: """Execute the memory search operation.""" query: str = self.context.query.strip() - min_score = self.context.get("minScore", self.min_score) - max_results = self.context.get("maxResults", self.max_results) + min_score = self.context.get("min_score", self.min_score) + max_results = self.context.get("max_results", self.max_results) candidates = min(200, max(1, int(max_results * self.hybrid_candidate_multiplier))) # Perform hybrid search (vector + keyword) diff --git a/reme/tool/fs/grep_tool.py b/reme/tool/fs/grep_tool.py index 9df9f8f2..0546a495 100644 --- a/reme/tool/fs/grep_tool.py +++ b/reme/tool/fs/grep_tool.py @@ -13,6 +13,7 @@ import os import shutil from pathlib import Path +from .base_fs_tool import BaseFsTool from .truncate import ( DEFAULT_MAX_BYTES, GREP_MAX_LINE_LENGTH, @@ -20,14 +21,13 @@ from .truncate import ( truncate_head, truncate_line, ) -from ...core.op import BaseTool from ...core.schema import ToolCall # Default limits DEFAULT_LIMIT = 100 # Maximum number of matches -class GrepTool(BaseTool): +class GrepTool(BaseFsTool): """Tool for searching file contents using ripgrep. Features: @@ -145,15 +145,12 @@ class GrepTool(BaseTool): args.extend([pattern, search_path]) # Execute ripgrep - try: - process = await asyncio.create_subprocess_exec( - *args, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=self.cwd, - ) - except Exception as e: - raise RuntimeError(f"Failed to run ripgrep: {e}") from e + process = await asyncio.create_subprocess_exec( + *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self.cwd, + ) stdout, stderr = await process.communicate() diff --git a/reme/tool/fs/ls_tool.py b/reme/tool/fs/ls_tool.py index 1c054c62..f231119c 100644 --- a/reme/tool/fs/ls_tool.py +++ b/reme/tool/fs/ls_tool.py @@ -3,14 +3,14 @@ import os from pathlib import Path +from .base_fs_tool import BaseFsTool from .truncate import DEFAULT_MAX_BYTES, truncate_head -from ...core.op import BaseTool from ...core.schema import ToolCall DEFAULT_LIMIT = 500 -class LsTool(BaseTool): +class LsTool(BaseFsTool): """List directory contents with smart truncation. Features: @@ -75,10 +75,7 @@ class LsTool(BaseTool): raise NotADirectoryError(f"Not a directory: {dir_path}") # Read directory entries - try: - entries = list(dir_path.iterdir()) - except Exception as e: - raise PermissionError(f"Cannot read directory: {e}") from e + entries = list(dir_path.iterdir()) # Sort alphabetically (case-insensitive) entries.sort(key=lambda e: e.name.lower()) diff --git a/reme/tool/fs/read_tool.py b/reme/tool/fs/read_tool.py index b236d4a4..ace2339a 100644 --- a/reme/tool/fs/read_tool.py +++ b/reme/tool/fs/read_tool.py @@ -9,8 +9,8 @@ Features: import os from pathlib import Path +from .base_fs_tool import BaseFsTool from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, format_size, truncate_head -from ...core.op import BaseTool from ...core.schema import ToolCall # Supported image extensions @@ -29,7 +29,7 @@ def is_image_file(path: str) -> bool: return Path(path).suffix.lower() in IMAGE_EXTENSIONS -class ReadTool(BaseTool): +class ReadTool(BaseFsTool): """Read file contents with smart truncation. Features: @@ -163,24 +163,32 @@ class ReadTool(BaseTool): all_lines = content.split("\n") total_file_lines = len(all_lines) - # Apply offset if specified (convert 1-indexed to 0-indexed) - start_line = max(0, (offset - 1)) if offset else 0 + # Validate and apply offset (1-indexed to 0-indexed) + if offset is not None: + if offset < 1: + raise ValueError(f"offset must be >= 1, got {offset}") + start_line = offset - 1 + else: + start_line = 0 + start_line_display = start_line + 1 # Check offset bounds - if start_line >= len(all_lines): + if start_line >= total_file_lines: raise IndexError( - f"Offset {offset} is beyond end of file ({len(all_lines)} lines total)", + f"Offset {offset} is beyond end of file ({total_file_lines} lines total)", ) - # Apply user limit if specified + # Validate and apply limit if limit is not None: - end_line = min(start_line + limit, len(all_lines)) - selected_content = "\n".join(all_lines[start_line:end_line]) - user_limited_lines = end_line - start_line + if limit <= 0: + raise ValueError(f"limit must be positive, got {limit}") + end_line = min(start_line + limit, total_file_lines) else: - selected_content = "\n".join(all_lines[start_line:]) - user_limited_lines = None + end_line = total_file_lines + + # Extract selected lines + selected_content = "\n".join(all_lines[start_line:end_line]) # Apply truncation truncation = truncate_head(selected_content) @@ -205,15 +213,15 @@ class ReadTool(BaseTool): f"of {total_file_lines} ({max_kb}KB limit). " f"Use offset={next_offset} to continue.]" ) - elif user_limited_lines is not None and start_line + user_limited_lines < len(all_lines): - # User limit exceeded, but no truncation - remaining = len(all_lines) - (start_line + user_limited_lines) - next_offset = start_line + user_limited_lines + 1 + elif end_line < total_file_lines: + # User limit reached but no truncation + remaining = total_file_lines - end_line + next_offset = end_line + 1 output_text = truncation.content output_text += f"\n\n[{remaining} more lines in file. " f"Use offset={next_offset} to continue.]" else: - # No truncation or user limit exceeded + # No truncation or limit output_text = truncation.content return output_text diff --git a/reme/tool/fs/write_tool.py b/reme/tool/fs/write_tool.py index 55a4a0c1..a291e0ac 100644 --- a/reme/tool/fs/write_tool.py +++ b/reme/tool/fs/write_tool.py @@ -8,11 +8,11 @@ This module provides a tool for writing content to files with: import os -from ...core.op import BaseTool +from .base_fs_tool import BaseFsTool from ...core.schema import ToolCall -class WriteTool(BaseTool): +class WriteTool(BaseFsTool): """Tool for writing content to files. Features: diff --git a/tests/demo_memory_search.py b/tests/demo_memory_search.py index 104d333a..a1d37091 100644 --- a/tests/demo_memory_search.py +++ b/tests/demo_memory_search.py @@ -24,7 +24,7 @@ from reme.core.embedding import OpenAIEmbeddingModel from reme.core.enumeration import MemorySource from reme.core.file_watcher.delta_file_watcher import DeltaFileWatcher from reme.core.file_watcher.full_file_watcher import FullFileWatcher -from reme.core.memory_storage import SqliteMemoryStore +from reme.core.memory_store import SqliteMemoryStore from reme.core.utils import load_env load_env() diff --git a/tests/test_fs_agent.py b/tests/test_fs_agent.py deleted file mode 100644 index 47995785..00000000 --- a/tests/test_fs_agent.py +++ /dev/null @@ -1,323 +0,0 @@ -"""Tests for fs (full-session) agents including compactor and summarizer. - -This module contains test functions for FsCompactor and FsSummarizer operations. -""" - -import asyncio -import os -import tempfile -from pathlib import Path - -from reme import ReMe -from reme.agent.fs.fs_compactor import FsCompactor -from reme.agent.fs.fs_summarizer import FsSummarizer -from reme.core.enumeration import Role -from reme.core.schema import Message -from reme.tool.fs import ReadTool, WriteTool, EditTool - - -def create_test_messages(num_messages: int = 10) -> list[Message]: - """Create a list of test messages for testing. - - Args: - num_messages: Number of messages to create - - Returns: - List of Message objects alternating between user and assistant - """ - messages = [] - for i in range(num_messages): - if i % 2 == 0: - # User messages - messages.append( - Message( - role=Role.USER, - content=f"User message {i}: Can you help me with task {i}?", - ), - ) - else: - # Assistant messages - messages.append( - Message( - role=Role.ASSISTANT, - content=f"Assistant message {i}: Sure, I'd be happy to help you with task {i - 1}. " - f"Let me explain the solution in detail. " * 10, # Make it longer - ), - ) - return messages - - -def create_long_conversation() -> list[Message]: - """Create a long conversation that exceeds token thresholds.""" - messages = [ - Message( - role=Role.USER, - content="I need help building a complete web application with authentication, database, and API endpoints.", - ), - Message( - role=Role.ASSISTANT, - content="""I'll help you build a complete web application. Here's what we'll do: - -1. Set up the project structure -2. Implement authentication system -3. Design and create database schema -4. Build API endpoints -5. Add frontend components -6. Test and deploy - -Let me start with the project structure...""", - ), - ] - - # Initial user request - - # Assistant response with detailed steps - - # Continue with multiple turns - for i in range(15): - messages.append( - Message( - role=Role.USER, - content=f"What about step {i + 1}? Can you provide more details?", - ), - ) - messages.append( - Message( - role=Role.ASSISTANT, - content=f"""For step {i + 1}, here's a detailed explanation: - -First, we need to consider the architecture. """ - + "This is important context. " * 50 - + """ - -Then we implement the following: -- Component A -- Component B -- Component C - -Let me show you the code for this part...""" - + "\n\ncode_example = 'example'" * 20, - ), - ) - - return messages - - -async def test_compactor_basic(reme: ReMe): - """Test basic FsCompactor functionality without triggering compaction. - - Tests that the compactor correctly skips compaction when token count - is below the threshold. - """ - print("\n" + "=" * 60) - print("Testing FsCompactor - Basic (Below Threshold)") - print("=" * 60) - - # Create a small conversation that won't trigger compaction - messages = create_test_messages(num_messages=6) - - # Create compactor with high threshold so it won't trigger - compactor = FsCompactor( - context_window_tokens=128000, - reserve_tokens=10000, - keep_recent_tokens=5000, - ) - - print(f"Number of messages: {len(messages)}") - output = await compactor.call(messages=messages, service_context=reme.service_context) - print(f"test_compactor_basic output: {output}") - - -async def test_compactor_with_compaction(reme: ReMe): - """Test FsCompactor with a long conversation that triggers compaction. - - Tests that the compactor correctly summarizes old messages when - the conversation exceeds the token threshold. - """ - print("\n" + "=" * 60) - print("Testing FsCompactor - With Compaction") - print("=" * 60) - - # Create a long conversation - messages = create_long_conversation() - - # Create compactor with low threshold to trigger compaction - compactor = FsCompactor( - context_window_tokens=10000, # Low threshold - reserve_tokens=2000, - keep_recent_tokens=2000, - ) - - print(f"Number of messages: {len(messages)}") - output = await compactor.call(messages=messages, service_context=reme.service_context) - print(f"test_compactor_with_compaction output: {output}") - - -async def test_compactor_split_turn(reme: ReMe): - """Test FsCompactor with a split turn scenario. - - Tests the scenario where the cut point falls in the middle of a turn, - requiring special handling to maintain context. - """ - print("\n" + "=" * 60) - print("Testing FsCompactor - Split Turn Detection") - print("=" * 60) - - messages = [] - - # Add some initial conversation - for i in range(5): - messages.append(Message(role=Role.USER, content=f"Question {i}")) - messages.append(Message(role=Role.ASSISTANT, content=f"Answer {i}. " * 30)) - - # Add a very long assistant response that will be split - messages.append(Message(role=Role.USER, content="Please explain this in great detail.")) - messages.append( - Message( - role=Role.ASSISTANT, - content="This is the first part of a very long response. " * 100, - ), - ) - messages.append( - Message( - role=Role.ASSISTANT, - content="This is the continuation of the response. " * 100, - ), - ) - messages.append( - Message( - role=Role.ASSISTANT, - content="And here's the final part with the conclusion. " * 50, - ), - ) - - compactor = FsCompactor( - context_window_tokens=8000, - reserve_tokens=1000, - keep_recent_tokens=2000, - ) - - print(f"Number of messages: {len(messages)}") - output = await compactor.call(messages=messages, service_context=reme.service_context) - print(f"test_compactor_split_turn output: {output}") - - -async def test_summarizer_basic(reme: ReMe): - """Test basic FsSummarizer functionality. - - Tests that the summarizer correctly skips when below threshold - and executes when above threshold. - """ - print("\n" + "=" * 60) - print("Testing FsSummarizer - Basic") - print("=" * 60) - - # Create a temporary directory for memory storage - with tempfile.TemporaryDirectory() as temp_dir: - memory_dir = os.path.join(temp_dir, "memories") - Path(memory_dir).mkdir(parents=True, exist_ok=True) - - # Create a small conversation (below threshold) - messages = create_test_messages(num_messages=4) - - summarizer = FsSummarizer( - tools=[ReadTool(), WriteTool(), EditTool()], - memory_dir=memory_dir, - context_window_tokens=128000, - reserve_tokens=32000, - soft_threshold_tokens=4000, - ) - - print(f"Memory directory: {memory_dir}") - print(f"Number of messages: {len(messages)}") - output = await summarizer.call(messages=messages, service_context=reme.service_context) - print(f"test_summarizer_basic output: {output}") - - -async def test_summarizer_with_execution(reme: ReMe): - """Test FsSummarizer with execution triggered. - - Tests that the summarizer executes when token count is within - the soft threshold range before compaction. - """ - print("\n" + "=" * 60) - print("Testing FsSummarizer - With Execution") - print("=" * 60) - - with tempfile.TemporaryDirectory() as temp_dir: - memory_dir = os.path.join(temp_dir, "memories") - Path(memory_dir).mkdir(parents=True, exist_ok=True) - - # Create messages that will trigger summarizer but not compactor - messages = create_test_messages(num_messages=10) - - # Set low thresholds to trigger execution - summarizer = FsSummarizer( - tools=[ReadTool(), WriteTool(), EditTool()], - memory_dir=memory_dir, - context_window_tokens=5000, - reserve_tokens=1000, - soft_threshold_tokens=500, - ) - - print(f"Memory directory: {memory_dir}") - print(f"Number of messages: {len(messages)}") - output = await summarizer.call(messages=messages, service_context=reme.service_context) - print(f"test_summarizer_with_execution output: {output}") - - -def test_compactor_serialization(): - """Test message serialization in FsCompactor. - - Tests that messages are correctly serialized to text format - for summarization. - """ - print("\n" + "=" * 60) - print("Testing FsCompactor - Message Serialization") - print("=" * 60) - - messages = [ - Message(role=Role.USER, content="Hello, how are you?", name="Alice"), - Message(role=Role.ASSISTANT, content="I'm doing great, thanks!"), - Message(role=Role.USER, content="Can you help me?"), - ] - - # Access static method for testing serialization - serialized = FsCompactor._serialize_conversation(messages) # pylint: disable=protected-access - - print("Serialized conversation:") - print(serialized) - print("\n✓ Serialization completed") - - # Check that it contains expected markers - assert "[Alice]" in serialized - assert "[assistant]" in serialized - assert "Hello, how are you?" in serialized - print("✓ Serialization format is correct") - - -async def main(): - """Run all tests.""" - # Run basic tests first - reme = ReMe() - await reme.start() - test_compactor_serialization() - await test_compactor_basic(reme) - await test_summarizer_basic(reme) - - # Run tests that require LLM calls (commented out by default) - # Uncomment these if you want to test with actual LLM calls - # await test_compactor_with_compaction(reme) - # await test_compactor_split_turn(reme) - # await test_summarizer_with_execution(reme) - - print("\n" + "=" * 60) - print("All basic tests completed!") - print("=" * 60) - print("\nNote: Tests requiring LLM calls are commented out.") - print("Uncomment them in the main() function to run with actual LLM.") - await reme.close() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests/test_fs_compact.py b/tests/test_fs_compact.py new file mode 100644 index 00000000..8824ea4c --- /dev/null +++ b/tests/test_fs_compact.py @@ -0,0 +1,301 @@ +"""Tests for ReMeFs compact interface. + +This module tests the compact() method of ReMeFs class which provides +a high-level interface for conversation compaction. +""" + +import asyncio + +from reme import ReMeFs +from reme.core.enumeration import Role +from reme.core.schema import Message + + +def print_messages(messages: list[Message], title: str = "Messages", max_content_len: int = 150): + """Print messages with their role and content. + + Args: + messages: List of messages to print + title: Title for the message list + max_content_len: Maximum content length to display (truncate if longer) + """ + print(f"\n{title}: (count: {len(messages)})") + print("-" * 80) + for i, msg in enumerate(messages): + content = str(msg.content) + if len(content) > max_content_len: + content = content[:max_content_len] + "..." + print(f" [{i}] {msg.role.value:10s}: {content}") + print("-" * 80) + + +def create_test_messages(num_messages: int = 10) -> list[Message]: + """Create a list of test messages. + + Args: + num_messages: Number of messages to create + + Returns: + List of Message objects alternating between user and assistant + """ + messages = [] + for i in range(num_messages): + if i % 2 == 0: + messages.append( + Message( + role=Role.USER, + content=f"User message {i}: Can you help me with task {i}?", + ), + ) + else: + messages.append( + Message( + role=Role.ASSISTANT, + content=f"Assistant message {i}: Sure, I'd be happy to help you with task {i - 1}. " + f"Let me explain the solution in detail. " * 10, + ), + ) + return messages + + +def create_long_conversation() -> list[Message]: + """Create a long conversation that exceeds token thresholds.""" + messages = [ + Message( + role=Role.USER, + content="I need help building a complete web application with authentication, database, and API endpoints.", + ), + Message( + role=Role.ASSISTANT, + content="""I'll help you build a complete web application. Here's what we'll do: + +1. Set up the project structure +2. Implement authentication system +3. Design and create database schema +4. Build API endpoints +5. Add frontend components +6. Test and deploy + +Let me start with the project structure...""", + ), + ] + + for i in range(15): + messages.append( + Message( + role=Role.USER, + content=f"What about step {i + 1}? Can you provide more details?", + ), + ) + messages.append( + Message( + role=Role.ASSISTANT, + content=f"""For step {i + 1}, here's a detailed explanation: + +First, we need to consider the architecture. """ + + "This is important context. " * 50 + + """ + +Then we implement the following: +- Component A +- Component B +- Component C + +Let me show you the code for this part...""" + + "\n\ncode_example = 'example'" * 20, + ), + ) + + return messages + + +async def test_compact_below_threshold(): + """Test compact() when messages are below threshold. + + Expects: compacted=False, returns original messages + """ + print("\n" + "=" * 80) + print("TEST 1: Compact - Below Threshold (No Compaction)") + print("=" * 80) + + reme_fs = ReMeFs(enable_logo=False, vector_store=None) + await reme_fs.start() + + messages = create_test_messages(num_messages=4) + print_messages(messages, "INPUT MESSAGES", max_content_len=80) + + print("\nParameters:") + print(" context_window_tokens: 5000") + print(" reserve_tokens: 2000 (threshold = 3000)") + print(" keep_recent_tokens: 1000") + + result = await reme_fs.compact( + messages=messages, + context_window_tokens=5000, + reserve_tokens=2000, + keep_recent_tokens=1000, + ) + + print(f"\n{'='*80}") + print("RESULT:") + print(f" compacted: {result.get('compacted')}") + print(f" tokens_before: {result.get('tokens_before')}") + print(f" is_split_turn: {result.get('is_split_turn')}") + + result_messages = result.get("messages", []) + print_messages(result_messages, "OUTPUT MESSAGES", max_content_len=80) + + assert result.get("compacted") is False, "Should not compact below threshold" + assert len(result_messages) == len(messages), "Should return all original messages" + print("\n✓ TEST PASSED: No compaction below threshold\n") + + await reme_fs.close() + + +async def test_compact_above_threshold(): + """Test compact() when messages exceed threshold. + + Expects: compacted=True, returns summary + left_messages + """ + print("\n" + "=" * 80) + print("TEST 2: Compact - Above Threshold (With Compaction & LLM Summary)") + print("=" * 80) + + reme_fs = ReMeFs(enable_logo=False, vector_store=None) + await reme_fs.start() + + messages = create_test_messages(num_messages=12) + print_messages(messages, "INPUT MESSAGES", max_content_len=60) + + print("\nParameters:") + print(" context_window_tokens: 3000") + print(" reserve_tokens: 1500 (threshold = 1500)") + print(" keep_recent_tokens: 500 (keep only recent messages)") + + result = await reme_fs.compact( + messages=messages, + context_window_tokens=3000, + reserve_tokens=1500, + keep_recent_tokens=500, + ) + + print(f"\n{'='*80}") + print("RESULT:") + print(f" compacted: {result.get('compacted')}") + print(f" tokens_before: {result.get('tokens_before')}") + print(f" is_split_turn: {result.get('is_split_turn')}") + + result_messages = result.get("messages", []) + if result.get("compacted") and result_messages: + has_summary = "" in str(result_messages[0].content) + print(f"\n *** First message contains summary: {has_summary}") + + print_messages(result_messages, "OUTPUT MESSAGES (Summary + Recent)", max_content_len=1500) + + assert result.get("compacted") is True, "Should compact above threshold" + assert len(result_messages) < len(messages), "Should reduce message count" + print("\n✓ TEST PASSED: Compaction triggered and summary generated\n") + + await reme_fs.close() + + +async def test_compact_split_turn_scenario(): + """Test compact() with split turn scenario. + + Expects: is_split_turn=True when cut point is mid-turn + """ + print("\n" + "=" * 80) + print("TEST 3: Compact - Split Turn Scenario (Cut in Middle of Assistant Response)") + print("=" * 80) + + reme_fs = ReMeFs(enable_logo=False, vector_store=None) + await reme_fs.start() + + messages = [] + + # Add initial conversation + for i in range(3): + messages.append(Message(role=Role.USER, content=f"Question {i}")) + messages.append(Message(role=Role.ASSISTANT, content=f"Answer {i}. " * 30)) + + # Add a very long multi-part assistant response + messages.append(Message(role=Role.USER, content="Please explain this in great detail.")) + messages.append( + Message( + role=Role.ASSISTANT, + content="This is the first part of a very long response. " * 50, + ), + ) + messages.append( + Message( + role=Role.ASSISTANT, + content="This is the continuation of the response. " * 50, + ), + ) + messages.append( + Message( + role=Role.ASSISTANT, + content="And here's the final part with the conclusion. " * 30, + ), + ) + + print_messages(messages, "INPUT MESSAGES", max_content_len=80) + + print("\nParameters:") + print(" context_window_tokens: 3000") + print(" reserve_tokens: 1000 (threshold = 2000)") + print(" keep_recent_tokens: 800 (should cut in middle of assistant responses)") + + result = await reme_fs.compact( + messages=messages, + context_window_tokens=3000, + reserve_tokens=1000, + keep_recent_tokens=800, + ) + + print(f"\n{'='*80}") + print("RESULT:") + print(f" compacted: {result.get('compacted')}") + print(f" tokens_before: {result.get('tokens_before')}") + print(f" is_split_turn: {result.get('is_split_turn')} *** (should be True)") + + result_messages = result.get("messages", []) + print_messages(result_messages, "OUTPUT MESSAGES (Summary with Turn Context + Recent)", max_content_len=150) + + if result.get("is_split_turn"): + print("\n✓ TEST PASSED: Split turn correctly detected and handled\n") + else: + print("\n⚠ WARNING: Split turn not detected (parameters may need adjustment)\n") + + await reme_fs.close() + + +async def main(): + """Run core compact interface tests.""" + print("\n" + "=" * 80) + print("ReMeFs Compact Interface - Core Test Suite") + print("=" * 80) + print("\nThis test suite demonstrates the three key scenarios of conversation compaction:") + print(" 1. Below threshold - no compaction needed") + print(" 2. Above threshold - full compaction with LLM summary") + print(" 3. Split turn - cut point falls in middle of assistant response") + print("=" * 80) + + # Test 1: No compaction (below threshold) + await test_compact_below_threshold() + + # Test 2: Full compaction (requires LLM) + await test_compact_above_threshold() + + # Test 3: Split turn compaction (requires LLM) + await test_compact_split_turn_scenario() + + print("\n" + "=" * 80) + print("All basic tests completed!") + print("=" * 80) + print("\nNote: Tests requiring LLM calls are commented out.") + print("Uncomment them in the main() function to run with actual LLM.") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_fs_memory_get.py b/tests/test_fs_memory_get.py new file mode 100644 index 00000000..246a695e --- /dev/null +++ b/tests/test_fs_memory_get.py @@ -0,0 +1,369 @@ +"""Tests for ReMeFs memory_get interface. + +This module tests the memory_get() method of ReMeFs class which provides +a high-level interface for reading specific snippets from memory files. + +The memory_get function should enable the LLM to: +1. Read entire memory files (MEMORY.md, memory/*.md) +2. Read specific line ranges using offset and limit parameters +3. Extract only the needed content to keep context small +""" + +import asyncio +import os +from pathlib import Path + +from reme import ReMeFs + + +def print_result(content: str, title: str = "RESULT", max_len: int = 300): + """Print the result of memory_get() call. + + Args: + content: Content returned from memory_get() + title: Title for the result section + max_len: Maximum content length to display (truncate if longer) + """ + print(f"\n{'=' * 80}") + print(f"{title}:") + + lines = content.split("\n") + print(f" total_lines: {len(lines)}") + print(f" total_chars: {len(content)}") + + if len(content) > max_len: + preview = content[:max_len] + "..." + else: + preview = content + + print("\n Content Preview:") + print("-" * 80) + print(preview) + print("-" * 80) + print(f"{'=' * 80}") + + +def create_test_memory_file(workspace_dir: str) -> str: + """Create a test memory file with numbered lines. + + Args: + workspace_dir: Directory to create the test file in + + Returns: + Path to the created test file (relative to workspace_dir) + """ + workspace_path = Path(workspace_dir) + workspace_path.mkdir(parents=True, exist_ok=True) + + memory_dir = workspace_path / "memory" + memory_dir.mkdir(parents=True, exist_ok=True) + + test_file = memory_dir / "test_profile.md" + + content = """# User Profile + +## Personal Information +Name: Alice Johnson +Age: 28 +Location: San Francisco, CA + +## Professional Background +Occupation: Software Engineer +Company: Tech Innovations Inc. +Years of Experience: 5 + +## Skills +- Python Programming +- Machine Learning +- Natural Language Processing +- Docker & Kubernetes + +## Interests +- Reading sci-fi novels +- Hiking in national parks +- Photography +- Cooking international cuisines + +## Preferences +- Prefers detailed technical explanations +- Likes to see code examples +- Values efficiency and clean code +- Appreciates constructive feedback +""" + + test_file.write_text(content, encoding="utf-8") + return "memory/test_profile.md" + + +async def test_memory_get_full_file(): + """Test memory_get() reads entire file without offset/limit. + + Expects: Returns complete file content + """ + print("\n" + "=" * 80) + print("TEST 1: Memory Get - Read Entire File") + print("=" * 80) + + workspace_dir = ".reme_test_get" + reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + await reme_fs.start() + + # Create test file + test_file_path = create_test_memory_file(workspace_dir) + print("\nTest Setup:") + print(f" test_file: {test_file_path}") + print(f" workspace: {workspace_dir}") + + print("\nParameters:") + print(f" path: {test_file_path}") + print(" offset: None") + print(" limit: None") + print(" Expected: Read entire file content") + + # Call memory_get + result = await reme_fs.memory_get(path=test_file_path) + + print_result(result, "MEMORY_GET RESULT", max_len=500) + + # Verify + lines = result.split("\n") + print("\nVerification:") + print(f" ✓ Got {len(lines)} lines") + print(f" ✓ Content starts with: {result[:50].strip()}") + + await reme_fs.close() + + +async def test_memory_get_with_offset(): + """Test memory_get() reads from specific line to end. + + Expects: Returns content from line 10 to end of file + """ + print("\n" + "=" * 80) + print("TEST 2: Memory Get - Read with Offset") + print("=" * 80) + + workspace_dir = ".reme_test_get" + reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + await reme_fs.start() + + test_file_path = "memory/test_profile.md" + offset = 10 + + print("\nParameters:") + print(f" path: {test_file_path}") + print(f" offset: {offset}") + print(" limit: None") + print(f" Expected: Read from line {offset} to end of file") + + # Call memory_get + result = await reme_fs.memory_get(path=test_file_path, offset=offset) + + print_result(result, "MEMORY_GET RESULT", max_len=400) + + # Verify + lines = result.split("\n") + print("\nVerification:") + print(f" ✓ Got {len(lines)} lines starting from line {offset}") + print(f" ✓ First line of result: {lines[0]}") + + await reme_fs.close() + + +async def test_memory_get_with_offset_and_limit(): + """Test memory_get() reads specific line range. + + Expects: Returns exactly 5 lines starting from line 5 + """ + print("\n" + "=" * 80) + print("TEST 3: Memory Get - Read with Offset and Limit") + print("=" * 80) + + workspace_dir = ".reme_test_get" + reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + await reme_fs.start() + + test_file_path = "memory/test_profile.md" + offset = 5 + limit = 5 + + print("\nParameters:") + print(f" path: {test_file_path}") + print(f" offset: {offset}") + print(f" limit: {limit}") + print(f" Expected: Read exactly {limit} lines starting from line {offset}") + + # Call memory_get + result = await reme_fs.memory_get(path=test_file_path, offset=offset, limit=limit) + + print_result(result, "MEMORY_GET RESULT", max_len=400) + + # Verify + lines = result.split("\n") + print("\nVerification:") + print(f" ✓ Got {len(lines)} lines (expected {limit})") + print(" ✓ Lines content:") + for i, line in enumerate(lines, start=offset): + print(f" Line {i}: {line}") + + await reme_fs.close() + + +async def test_memory_get_beginning_lines(): + """Test memory_get() reads first few lines. + + Expects: Returns first 3 lines of the file + """ + print("\n" + "=" * 80) + print("TEST 4: Memory Get - Read Beginning Lines") + print("=" * 80) + + workspace_dir = ".reme_test_get" + reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + await reme_fs.start() + + test_file_path = "memory/test_profile.md" + offset = 1 + limit = 3 + + print("\nParameters:") + print(f" path: {test_file_path}") + print(f" offset: {offset}") + print(f" limit: {limit}") + print(f" Expected: Read first {limit} lines") + + # Call memory_get + result = await reme_fs.memory_get(path=test_file_path, offset=offset, limit=limit) + + print_result(result, "MEMORY_GET RESULT", max_len=400) + + # Verify + lines = result.split("\n") + print("\nVerification:") + print(f" ✓ Got {len(lines)} lines (expected {limit})") + print(" ✓ Should contain '# User Profile' header") + assert "# User Profile" in result, "Expected header not found" + + await reme_fs.close() + + +async def test_memory_get_single_line(): + """Test memory_get() reads a single specific line. + + Expects: Returns exactly 1 line + """ + print("\n" + "=" * 80) + print("TEST 5: Memory Get - Read Single Line") + print("=" * 80) + + workspace_dir = ".reme_test_get" + reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + await reme_fs.start() + + test_file_path = "memory/test_profile.md" + offset = 3 + limit = 1 + + print("\nParameters:") + print(f" path: {test_file_path}") + print(f" offset: {offset}") + print(f" limit: {limit}") + print(f" Expected: Read exactly 1 line at position {offset}") + + # Call memory_get + result = await reme_fs.memory_get(path=test_file_path, offset=offset, limit=limit) + + print_result(result, "MEMORY_GET RESULT", max_len=400) + + # Verify + lines = result.split("\n") + print("\nVerification:") + print(f" ✓ Got {len(lines)} line (expected {limit})") + print(f" ✓ Line {offset}: {result}") + + await reme_fs.close() + + +async def test_memory_get_with_absolute_path(): + """Test memory_get() with absolute path. + + Expects: Works with both relative and absolute paths + """ + print("\n" + "=" * 80) + print("TEST 6: Memory Get - Read with Absolute Path") + print("=" * 80) + + workspace_dir = ".reme_test_get" + reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + await reme_fs.start() + + # Get absolute path + abs_path = os.path.abspath(os.path.join(workspace_dir, "memory/test_profile.md")) + limit = 5 + + print("\nParameters:") + print(f" path: {abs_path}") + print(" offset: None") + print(f" limit: {limit}") + print(f" Expected: Read first {limit} lines using absolute path") + + # Call memory_get with absolute path + result = await reme_fs.memory_get(path=abs_path, limit=limit) + + print_result(result, "MEMORY_GET RESULT", max_len=400) + + # Verify + lines = result.split("\n") + print("\nVerification:") + print(f" ✓ Got {len(lines)} lines using absolute path") + print(" ✓ Absolute path handling works correctly") + + await reme_fs.close() + + +async def main(): + """Run core memory_get interface tests.""" + print("\n" + "=" * 80) + print("ReMeFs Memory Get Interface - Tests") + print("=" * 80) + print("\nThis test suite validates that the memory_get() function:") + print(" 1. Reads entire memory files without parameters") + print(" 2. Reads from specific line (offset) to end of file") + print(" 3. Reads specific line ranges (offset + limit)") + print(" 4. Handles both relative and absolute paths") + print("\nTest Scenarios:") + print(" 1. Full file read (no offset/limit)") + print(" 2. Read with offset (from line N to end)") + print(" 3. Read with offset and limit (specific range)") + print(" 4. Read beginning lines (first N lines)") + print(" 5. Read single line") + print(" 6. Read with absolute path") + print("=" * 80) + + # Test 1: Read entire file + await test_memory_get_full_file() + + # Test 2: Read with offset + await test_memory_get_with_offset() + + # Test 3: Read with offset and limit + await test_memory_get_with_offset_and_limit() + + # Test 4: Read beginning lines + await test_memory_get_beginning_lines() + + # Test 5: Read single line + await test_memory_get_single_line() + + # Test 6: Read with absolute path + await test_memory_get_with_absolute_path() + + print("\n" + "=" * 80) + print("All memory_get tests completed!") + print("=" * 80) + print("\nNote: Test files are created in .reme_test_get/ directory") + print("You can manually inspect them or delete the directory after testing.") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_fs_memory_search.py b/tests/test_fs_memory_search.py new file mode 100644 index 00000000..88e6da8b --- /dev/null +++ b/tests/test_fs_memory_search.py @@ -0,0 +1,704 @@ +"""Tests for ReMeFs memory_search interface. + +This module tests the memory_search() method of ReMeFs class which provides +a high-level interface for searching personal information stored in memory files. + +The memory_search function should enable: +1. Vector similarity search across memory chunks +2. Keyword/FTS (full-text search) if enabled +3. Hybrid search combining vector and keyword results +4. Source filtering (MEMORY, SESSIONS, etc.) +5. Score-based filtering and result limiting +""" + +import asyncio +import hashlib +import shutil +from pathlib import Path + +from reme import ReMeFs +from reme.core.enumeration import MemorySource +from reme.core.schema import FileMetadata, MemoryChunk + + +# ==================== Test Configuration ==================== + + +class TestConfig: + """Test configuration settings.""" + + WORKING_DIR = ".reme_test_search" + + +# ==================== Sample Data Generator ==================== + + +class SampleDataGenerator: + """Generator for sample test data.""" + + @staticmethod + def create_personal_info_chunks(test_name: str = "") -> list[MemoryChunk]: + """Create sample chunks with personal information.""" + base_path = "memory/personal_info.md" + prefix = f"{test_name}_" if test_name else "" + return [ + MemoryChunk( + id=f"{prefix}personal_1", + path=base_path, + source=MemorySource.MEMORY, + start_line=1, + end_line=3, + text="My name is Alice Chen. I am a software engineer working at TechCorp.", + hash=hashlib.md5(b"personal_1").hexdigest(), + embedding=None, + metadata={"category": "personal", "type": "basic_info"}, + ), + MemoryChunk( + id=f"{prefix}personal_2", + path=base_path, + source=MemorySource.MEMORY, + start_line=4, + end_line=6, + text=( + "I love Python programming and machine learning. " + "My favorite frameworks are PyTorch and scikit-learn." + ), + hash=hashlib.md5(b"personal_2").hexdigest(), + embedding=None, + metadata={"category": "personal", "type": "interests"}, + ), + MemoryChunk( + id=f"{prefix}personal_3", + path=base_path, + source=MemorySource.MEMORY, + start_line=7, + end_line=9, + text="In my free time, I enjoy reading science fiction novels and hiking in the mountains.", + hash=hashlib.md5(b"personal_3").hexdigest(), + embedding=None, + metadata={"category": "personal", "type": "hobbies"}, + ), + ] + + @staticmethod + def create_technical_chunks(test_name: str = "") -> list[MemoryChunk]: + """Create sample chunks with technical information.""" + base_path = "memory/technical_notes.md" + prefix = f"{test_name}_" if test_name else "" + return [ + MemoryChunk( + id=f"{prefix}tech_1", + path=base_path, + source=MemorySource.MEMORY, + start_line=1, + end_line=3, + text=( + "Artificial intelligence is transforming software development " + "with automated code generation and testing." + ), + hash=hashlib.md5(b"tech_1").hexdigest(), + embedding=None, + metadata={"category": "tech", "topic": "AI"}, + ), + MemoryChunk( + id=f"{prefix}tech_2", + path=base_path, + source=MemorySource.MEMORY, + start_line=4, + end_line=6, + text=( + "Machine learning models require careful tuning of hyperparameters " + "to achieve optimal performance." + ), + hash=hashlib.md5(b"tech_2").hexdigest(), + embedding=None, + metadata={"category": "tech", "topic": "ML"}, + ), + MemoryChunk( + id=f"{prefix}tech_3", + path=base_path, + source=MemorySource.MEMORY, + start_line=7, + end_line=9, + text="Deep learning neural networks excel at image recognition and natural language processing tasks.", + hash=hashlib.md5(b"tech_3").hexdigest(), + embedding=None, + metadata={"category": "tech", "topic": "DL"}, + ), + ] + + @staticmethod + def create_session_chunks(test_name: str = "") -> list[MemoryChunk]: + """Create sample session chunks.""" + base_path = "sessions/2024-01-15.jsonl" + prefix = f"{test_name}_" if test_name else "" + return [ + MemoryChunk( + id=f"{prefix}session_1", + path=base_path, + source=MemorySource.SESSIONS, + start_line=1, + end_line=2, + text="User asked about Python best practices for async programming.", + hash=hashlib.md5(b"session_1").hexdigest(), + embedding=None, + metadata={"session_id": "sess_001", "date": "2024-01-15"}, + ), + MemoryChunk( + id=f"{prefix}session_2", + path=base_path, + source=MemorySource.SESSIONS, + start_line=3, + end_line=4, + text="Discussed asyncio event loop and common pitfalls in concurrent Python code.", + hash=hashlib.md5(b"session_2").hexdigest(), + embedding=None, + metadata={"session_id": "sess_001", "date": "2024-01-15"}, + ), + ] + + @staticmethod + def create_file_metadata(path: str, chunk_count: int) -> FileMetadata: + """Create file metadata.""" + return FileMetadata( + path=path, + hash=hashlib.md5(path.encode()).hexdigest(), + mtime_ms=1704067200000, # 2024-01-01 00:00:00 + size=1000, + chunk_count=chunk_count, + ) + + +# ==================== Helper Functions ==================== + + +def print_search_results(results: list[dict], query: str, title: str = "SEARCH RESULTS"): + """Pretty print search results.""" + print(f"\n{'=' * 80}") + print(f"{title}") + print(f"Query: '{query}'") + print(f"Found {len(results)} results") + print(f"{'=' * 80}") + + for i, result in enumerate(results, 1): + print(f"\n[{i}] Source: {result.get('source', 'N/A')}") + print(f" Path: {result.get('path', 'N/A')}") + print(f" Lines: {result.get('start_line', 'N/A')}-{result.get('end_line', 'N/A')}") + print(f" Score: {result.get('score', 0):.6f}") + snippet = result.get("snippet", result.get("text", "")) + if len(snippet) > 100: + snippet = snippet[:100] + "..." + print(f" Snippet: {snippet}") + if result.get("metadata"): + print(f" Metadata: {result.get('metadata')}") + + print(f"{'=' * 80}\n") + + +# ==================== Test Functions ==================== + + +async def test_memory_search_basic(): + """Test basic memory search functionality. + + Insert sample data and perform a simple search query. + """ + print("\n" + "=" * 80) + print("TEST 1: Basic Memory Search") + print("=" * 80) + + # Initialize ReMeFs with unique store name + reme_fs = ReMeFs( + enable_logo=False, + working_dir=TestConfig.WORKING_DIR, + memory_store={ + "backend": "sqlite", + "store_name": "test_basic", + "embedding_model": "default", + "fts_enabled": True, + "snippet_max_chars": 700, + }, + ) + await reme_fs.start() + + # Insert personal info chunks + personal_chunks = SampleDataGenerator.create_personal_info_chunks("test_basic") + personal_chunks = await reme_fs.memory_store.get_chunk_embeddings(personal_chunks) + + file_meta = SampleDataGenerator.create_file_metadata( + "memory/personal_info.md", + len(personal_chunks), + ) + await reme_fs.memory_store.upsert_file( + file_meta, + MemorySource.MEMORY, + personal_chunks, + ) + print(f"✓ Inserted {len(personal_chunks)} personal info chunks") + + # Perform search + query = "What programming languages does the user like?" + print(f"\nSearching for: '{query}'") + + result_json = await reme_fs.memory_search( + query=query, + max_results=5, + min_score=0.0, + ) + + # Parse results + import json + + results = json.loads(result_json) + print_search_results(results, query, "BASIC SEARCH RESULTS") + + # Verify results + assert len(results) > 0, "Should find at least one result" + assert results[0]["score"] > 0, "Top result should have positive score" + print("✓ Basic memory search test passed") + + await reme_fs.close() + + +async def test_memory_search_technical_content(): + """Test memory search with technical content. + + Insert technical chunks and search for ML/AI related queries. + """ + print("\n" + "=" * 80) + print("TEST 2: Technical Content Search") + print("=" * 80) + + # Initialize ReMeFs with unique store name + reme_fs = ReMeFs( + enable_logo=False, + working_dir=TestConfig.WORKING_DIR, + memory_store={ + "backend": "sqlite", + "store_name": "test_technical", + "embedding_model": "default", + "fts_enabled": True, + "snippet_max_chars": 700, + }, + ) + await reme_fs.start() + + # Insert technical chunks + tech_chunks = SampleDataGenerator.create_technical_chunks("test_technical") + tech_chunks = await reme_fs.memory_store.get_chunk_embeddings(tech_chunks) + + file_meta = SampleDataGenerator.create_file_metadata( + "memory/technical_notes.md", + len(tech_chunks), + ) + await reme_fs.memory_store.upsert_file( + file_meta, + MemorySource.MEMORY, + tech_chunks, + ) + print(f"✓ Inserted {len(tech_chunks)} technical chunks") + + # Test multiple queries + queries = [ + "artificial intelligence and machine learning", + "neural networks for image processing", + "hyperparameter tuning in ML models", + ] + + for query in queries: + print(f"\n--- Searching for: '{query}' ---") + result_json = await reme_fs.memory_search( + query=query, + max_results=3, + min_score=0.0, + ) + + import json + + results = json.loads(result_json) + print(f"Found {len(results)} results") + + for i, result in enumerate(results, 1): + print(f" [{i}] Score: {result['score']:.6f} | {result['path']}") + + assert len(results) > 0, f"Should find results for query: {query}" + + print("\n✓ Technical content search test passed") + await reme_fs.close() + + +async def test_memory_search_with_source_filter(): + """Test memory search with source filtering. + + Insert data from different sources and test source-specific searches. + """ + print("\n" + "=" * 80) + print("TEST 3: Memory Search with Source Filter") + print("=" * 80) + + # Initialize ReMeFs with unique store name + reme_fs = ReMeFs( + enable_logo=False, + working_dir=TestConfig.WORKING_DIR, + memory_store={ + "backend": "sqlite", + "store_name": "test_source_filter", + "embedding_model": "default", + "fts_enabled": True, + "snippet_max_chars": 700, + }, + ) + await reme_fs.start() + + # Insert MEMORY source data + personal_chunks = SampleDataGenerator.create_personal_info_chunks("test_source") + personal_chunks = await reme_fs.memory_store.get_chunk_embeddings(personal_chunks) + personal_meta = SampleDataGenerator.create_file_metadata( + "memory/personal_info.md", + len(personal_chunks), + ) + await reme_fs.memory_store.upsert_file( + personal_meta, + MemorySource.MEMORY, + personal_chunks, + ) + print(f"✓ Inserted {len(personal_chunks)} MEMORY chunks") + + # Insert SESSIONS source data + session_chunks = SampleDataGenerator.create_session_chunks("test_source") + session_chunks = await reme_fs.memory_store.get_chunk_embeddings(session_chunks) + session_meta = SampleDataGenerator.create_file_metadata( + "sessions/2024-01-15.jsonl", + len(session_chunks), + ) + await reme_fs.memory_store.upsert_file( + session_meta, + MemorySource.SESSIONS, + session_chunks, + ) + print(f"✓ Inserted {len(session_chunks)} SESSIONS chunks") + + query = "Python programming and async" + + # Search only MEMORY source + print(f"\n--- Searching MEMORY source for: '{query}' ---") + result_json = await reme_fs.memory_search( + query=query, + max_results=5, + sources=[MemorySource.MEMORY], + ) + import json + + memory_results = json.loads(result_json) + print(f"Found {len(memory_results)} results in MEMORY source") + for result in memory_results: + assert result["source"] == MemorySource.MEMORY.value, "Should only return MEMORY source results" + + # Search only SESSIONS source + print(f"\n--- Searching SESSIONS source for: '{query}' ---") + result_json = await reme_fs.memory_search( + query=query, + max_results=5, + sources=[MemorySource.SESSIONS], + ) + session_results = json.loads(result_json) + print(f"Found {len(session_results)} results in SESSIONS source") + for result in session_results: + assert result["source"] == MemorySource.SESSIONS.value, "Should only return SESSIONS source results" + + # Search all sources + print(f"\n--- Searching ALL sources for: '{query}' ---") + result_json = await reme_fs.memory_search( + query=query, + max_results=10, + ) + all_results = json.loads(result_json) + print(f"Found {len(all_results)} results across all sources") + + sources_found = {result["source"] for result in all_results} + print(f"Sources found: {sources_found}") + + print("\n✓ Source filter search test passed") + await reme_fs.close() + + +async def test_memory_search_score_filtering(): + """Test memory search with score threshold. + + Test min_score parameter to filter low-relevance results. + """ + print("\n" + "=" * 80) + print("TEST 4: Memory Search with Score Filtering") + print("=" * 80) + + # Initialize ReMeFs with unique store name + reme_fs = ReMeFs( + enable_logo=False, + working_dir=TestConfig.WORKING_DIR, + memory_store={ + "backend": "sqlite", + "store_name": "test_score_filter", + "embedding_model": "default", + "fts_enabled": True, + "snippet_max_chars": 700, + }, + ) + await reme_fs.start() + + # Insert test data + chunks = SampleDataGenerator.create_technical_chunks("test_score") + chunks = await reme_fs.memory_store.get_chunk_embeddings(chunks) + file_meta = SampleDataGenerator.create_file_metadata( + "memory/technical_notes.md", + len(chunks), + ) + await reme_fs.memory_store.upsert_file( + file_meta, + MemorySource.MEMORY, + chunks, + ) + print(f"✓ Inserted {len(chunks)} test chunks") + + query = "machine learning algorithms" + + # Search with different min_score thresholds + thresholds = [0.0, 0.1, 0.3, 0.5] + + for min_score in thresholds: + print(f"\n--- Searching with min_score={min_score} ---") + result_json = await reme_fs.memory_search( + query=query, + max_results=10, + min_score=min_score, + ) + + import json + + results = json.loads(result_json) + print(f"Found {len(results)} results with min_score >= {min_score}") + + # Verify all results meet threshold + for result in results: + assert result["score"] >= min_score, f"Result score {result['score']:.6f} should be >= {min_score}" + + if results: + print(f" Top score: {results[0]['score']:.6f}") + print(f" Lowest score: {results[-1]['score']:.6f}") + + print("\n✓ Score filtering search test passed") + await reme_fs.close() + + +async def test_memory_search_max_results(): + """Test memory search with result limiting. + + Test max_results parameter to limit number of returned results. + """ + print("\n" + "=" * 80) + print("TEST 5: Memory Search with Result Limiting") + print("=" * 80) + + # Initialize ReMeFs with unique store name + reme_fs = ReMeFs( + enable_logo=False, + working_dir=TestConfig.WORKING_DIR, + memory_store={ + "backend": "sqlite", + "store_name": "test_max_results", + "embedding_model": "default", + "fts_enabled": True, + "snippet_max_chars": 700, + }, + ) + await reme_fs.start() + + # Insert multiple chunks + personal_chunks = SampleDataGenerator.create_personal_info_chunks("test_max") + tech_chunks = SampleDataGenerator.create_technical_chunks("test_max") + all_chunks = personal_chunks + tech_chunks + + all_chunks = await reme_fs.memory_store.get_chunk_embeddings(all_chunks) + + # Insert as one file for simplicity + combined_meta = SampleDataGenerator.create_file_metadata( + "memory/combined.md", + len(all_chunks), + ) + await reme_fs.memory_store.upsert_file( + combined_meta, + MemorySource.MEMORY, + all_chunks, + ) + + print(f"✓ Inserted {len(all_chunks)} total chunks") + + query = "programming and technology" + + # Test different max_results values + result_limits = [1, 2, 3, 5, 20] + + for max_results in result_limits: + print(f"\n--- Searching with max_results={max_results} ---") + result_json = await reme_fs.memory_search( + query=query, + max_results=max_results, + min_score=0.0, + ) + + import json + + results = json.loads(result_json) + print(f"Requested {max_results}, got {len(results)} results") + + assert len(results) <= max_results, f"Should return at most {max_results} results, got {len(results)}" + + print("\n✓ Result limiting search test passed") + await reme_fs.close() + + +async def test_memory_search_hybrid_mode(): + """Test memory search with hybrid mode (vector + keyword). + + Test different hybrid configurations and weights. + """ + print("\n" + "=" * 80) + print("TEST 6: Memory Search with Hybrid Mode") + print("=" * 80) + + # Initialize ReMeFs with unique store name + reme_fs = ReMeFs( + enable_logo=False, + working_dir=TestConfig.WORKING_DIR, + memory_store={ + "backend": "sqlite", + "store_name": "test_hybrid", + "embedding_model": "default", + "fts_enabled": True, + "snippet_max_chars": 700, + }, + ) + await reme_fs.start() + + # Insert test data + chunks = SampleDataGenerator.create_technical_chunks("test_hybrid") + chunks = await reme_fs.memory_store.get_chunk_embeddings(chunks) + file_meta = SampleDataGenerator.create_file_metadata( + "memory/technical_notes.md", + len(chunks), + ) + await reme_fs.memory_store.upsert_file( + file_meta, + MemorySource.MEMORY, + chunks, + ) + print(f"✓ Inserted {len(chunks)} test chunks") + + query = "neural networks" + + # Test with hybrid enabled + print(f"\n--- Hybrid search (enabled) for: '{query}' ---") + result_json_hybrid = await reme_fs.memory_search( + query=query, + max_results=5, + hybrid_enabled=True, + hybrid_vector_weight=0.7, + hybrid_text_weight=0.3, + ) + + import json + + hybrid_results = json.loads(result_json_hybrid) + print(f"Hybrid search found {len(hybrid_results)} results") + print_search_results(hybrid_results, query, "HYBRID SEARCH RESULTS") + + # Test with hybrid disabled (vector only) + print(f"\n--- Vector-only search for: '{query}' ---") + result_json_vector = await reme_fs.memory_search( + query=query, + max_results=5, + hybrid_enabled=False, + ) + + vector_results = json.loads(result_json_vector) + print(f"Vector search found {len(vector_results)} results") + print_search_results(vector_results, query, "VECTOR-ONLY SEARCH RESULTS") + + # Test different weight configurations + print("\n--- Testing different hybrid weights ---") + weight_configs = [ + (0.9, 0.1), # Mostly vector + (0.5, 0.5), # Balanced + (0.3, 0.7), # Mostly text + ] + + for vec_weight, text_weight in weight_configs: + result_json = await reme_fs.memory_search( + query=query, + max_results=5, + hybrid_enabled=True, + hybrid_vector_weight=vec_weight, + hybrid_text_weight=text_weight, + ) + results = json.loads(result_json) + print(f" Vector:{vec_weight}/Text:{text_weight} -> {len(results)} results") + + print("\n✓ Hybrid mode search test passed") + await reme_fs.close() + + +async def cleanup_test_data(): + """Clean up test data directory.""" + print("\n" + "=" * 80) + print("CLEANUP: Removing test data") + print("=" * 80) + + test_dir = Path(TestConfig.WORKING_DIR) + if test_dir.exists(): + shutil.rmtree(test_dir) + print(f"✓ Removed test directory: {test_dir}") + else: + print(f"⊘ Test directory does not exist: {test_dir}") + + +# ==================== Main Entry Point ==================== + + +async def main(): + """Run all memory search tests.""" + print("\n" + "=" * 80) + print("ReMeFs Memory Search Interface Tests") + print("=" * 80) + print("\nThis test suite validates the memory_search() function:") + print(" 1. Basic semantic search functionality") + print(" 2. Technical content search") + print(" 3. Source filtering (MEMORY, SESSIONS)") + print(" 4. Score threshold filtering") + print(" 5. Result limiting (max_results)") + print(" 6. Hybrid mode (vector + keyword search)") + print("=" * 80) + + try: + # Run tests + await test_memory_search_basic() + await test_memory_search_technical_content() + await test_memory_search_with_source_filter() + await test_memory_search_score_filtering() + await test_memory_search_max_results() + await test_memory_search_hybrid_mode() + + print("\n" + "=" * 80) + print("✓ All memory search tests passed!") + print("=" * 80) + + finally: + # Cleanup + await cleanup_test_data() + + print("\nNote: These tests require:") + print(" - Valid API keys for embedding model") + print(" - sqlite-vec extension for vector search") + print(" - FTS5 enabled for keyword search") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_fs_summary.py b/tests/test_fs_summary.py new file mode 100644 index 00000000..ff9ec210 --- /dev/null +++ b/tests/test_fs_summary.py @@ -0,0 +1,234 @@ +"""Tests for ReMeFs summary interface. + +This module tests the summary() method of ReMeFs class which provides +a high-level interface for storing user's personal information into memory files. + +The summary function should enable the LLM to: +1. Extract personal information from user messages (name, preferences, requirements) +2. Call file system tools (WriteTool, EditTool) to store this information +3. Maintain personalized memory for future conversations +""" + +import asyncio + +from reme import ReMeFs +from reme.core.enumeration import Role +from reme.core.schema import Message + + +def print_messages(messages: list[Message], title: str = "Messages", max_content_len: int = 150): + """Print messages with their role and content. + + Args: + messages: List of messages to print + title: Title for the message list + max_content_len: Maximum content length to display (truncate if longer) + """ + print(f"\n{title}: (count: {len(messages)})") + print("-" * 80) + for i, msg in enumerate(messages): + content = str(msg.content) + if len(content) > max_content_len: + content = content[:max_content_len] + "..." + print(f" [{i}] {msg.role.value:10s}: {content}") + print("-" * 80) + + +def print_result(result: dict, title: str = "RESULT"): + """Print the result of summary() call. + + Args: + result: Result dictionary from summary() + title: Title for the result section + """ + print(f"\n{'=' * 80}") + print(f"{title}:") + print(f" success: {result.get('success')}") + print(f" skipped: {result.get('skipped', False)}") + + tools_used = result.get("tools", []) + print(f" tools_called: {len(tools_used)}") + + if tools_used: + print("\n Tool Usage Details:") + for i, tool in enumerate(tools_used): + print(f" [{i}] Tool: {tool.name} Arguments: {tool.tool_call.arguments}") + + answer = result.get("answer", "") + if answer: + answer_preview = answer[:300] + "..." if len(answer) > 300 else answer + print(f"\n answer: {answer_preview}") + + print(f"{'=' * 80}") + + +def create_personal_info_introduction() -> list[Message]: + """Create a conversation where user introduces personal information.""" + return [ + Message( + role=Role.USER, + content="Hi! My name is Alice, and I'm a software engineer.", + ), + Message( + role=Role.ASSISTANT, + content="Nice to meet you, Alice! How can I help you today?", + ), + Message( + role=Role.USER, + content=( + "I love Python programming and working on AI projects. " + "I also enjoy reading sci-fi novels in my free time." + ), + ), + Message( + role=Role.ASSISTANT, + content="That's great! Python and AI are exciting fields. What kind of AI projects are you interested in?", + ), + ] + + +def create_detailed_profile_conversation() -> list[Message]: + """Create a comprehensive conversation with multiple personal details.""" + return [ + Message( + role=Role.USER, + content=( + "Let me tell you about myself. My name is Charlie Chen, " "I'm a data scientist based in San Francisco." + ), + ), + Message( + role=Role.ASSISTANT, + content="Hello Charlie! It's nice to meet you. Tell me more about your work.", + ), + Message( + role=Role.USER, + content=( + "I specialize in machine learning and natural language processing. " + "My favorite tools are PyTorch and Hugging Face transformers." + ), + ), + Message( + role=Role.ASSISTANT, + content="Those are excellent tools for NLP work. What kind of projects do you work on?", + ), + Message( + role=Role.USER, + content="I work on chatbots and sentiment analysis. Outside of work, I love hiking and photography.", + ), + Message( + role=Role.ASSISTANT, + content="That's a great combination of technical and creative interests!", + ), + Message( + role=Role.USER, + content=( + "Oh, and one more thing - please be more assertive when you think " + "I'm making a mistake. I want you to challenge my ideas." + ), + ), + Message( + role=Role.ASSISTANT, + content="Absolutely, I'll make sure to provide critical feedback when needed.", + ), + ] + + +async def test_summary_personal_info_storage(): + """Test summary() stores user's personal information to memory. + + User introduces name and basic info. + Expects: LLM should call WriteTool to store this information + """ + print("\n" + "=" * 80) + print("TEST 1: Summary - Personal Information Storage") + print("=" * 80) + + reme_fs = ReMeFs(enable_logo=False) + await reme_fs.start() + + messages = create_personal_info_introduction() + print_messages(messages, "INPUT MESSAGES", max_content_len=200) + + print("\nParameters:") + print(" version: default") + print(" Expected: LLM should call WriteTool to save user's name and interests") + + result = await reme_fs.summary( + messages=messages, + version="default", + date="2023-09-01", + ) + + print_result(result, "SUMMARY RESULT") + await reme_fs.close() + + +async def test_summary_detailed_profile(): + """Test summary() with comprehensive user profile information. + + User provides detailed personal, professional, and preference information. + Expects: LLM should organize and store all relevant information + """ + print("\n" + "=" * 80) + print("TEST 2: Summary - Detailed User Profile Storage") + print("=" * 80) + + reme_fs = ReMeFs(enable_logo=False, vector_store=None) + await reme_fs.start() + + messages = create_detailed_profile_conversation() + print_messages(messages, "INPUT MESSAGES", max_content_len=150) + + print("\nParameters:") + print(" version: default") + print(" Expected: LLM should extract and store:") + print(" - Name: Charlie Chen") + print(" - Profession: Data Scientist") + print(" - Location: San Francisco") + print(" - Skills: ML, NLP, PyTorch, Hugging Face") + print(" - Hobbies: Hiking, Photography") + print(" - Assistant behavior: Be assertive and critical") + + result = await reme_fs.summary( + messages=messages, + version="default", + date="2023-10-01", + ) + + print_result(result, "SUMMARY RESULT") + await reme_fs.close() + + +async def main(): + """Run core summary interface tests for personal memory storage.""" + print("\n" + "=" * 80) + print("ReMeFs Summary Interface - Personal Memory Storage Tests") + print("=" * 80) + print("\nThis test suite validates that the summary() function:") + print(" 1. Extracts personal information from user messages") + print(" 2. Calls file system tools (WriteTool/EditTool) to store the info") + print(" 3. Organizes information for future retrieval") + print("\nTest Scenarios:") + print(" 1. Personal info - name, profession, interests") + print(" 2. Detailed profile - comprehensive user information") + print("=" * 80) + + # Test 1: Basic personal information storage + await test_summary_personal_info_storage() + + # Test 2: Comprehensive user profile + await test_summary_detailed_profile() + + print("\n" + "=" * 80) + print("All summary tests completed!") + print("=" * 80) + print("\nNote: These tests require LLM calls to:") + print(" - Analyze user messages for personal information") + print(" - Decide what information to store") + print(" - Call appropriate tools (WriteTool/EditTool) to save to memory") + print("\nMake sure your API keys are properly configured before running.") + print("The LLM will autonomously decide what to store based on the conversation.") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_file_system_tool.py b/tests/test_fs_tool.py similarity index 100% rename from tests/test_file_system_tool.py rename to tests/test_fs_tool.py diff --git a/tests/test_memory_store.py b/tests/test_memory_store.py index a057c897..0052049c 100644 --- a/tests/test_memory_store.py +++ b/tests/test_memory_store.py @@ -21,8 +21,8 @@ from loguru import logger from reme.core.embedding import OpenAIEmbeddingModel from reme.core.enumeration.memory_source import MemorySource -from reme.core.memory_storage.base_memory_store import BaseMemoryStore -from reme.core.memory_storage.sqlite_memory_store import SqliteMemoryStore +from reme.core.memory_store.base_memory_store import BaseMemoryStore +from reme.core.memory_store.sqlite_memory_store import SqliteMemoryStore from reme.core.schema.file_metadata import FileMetadata from reme.core.schema.memory_chunk import MemoryChunk from reme.core.utils import load_env From 55d61f1dc5acd8059f5eb08ec38fc52f9b4f7d0c Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 7 Feb 2026 15:05:02 +0800 Subject: [PATCH 2/2] refactor(core): update registry naming and application configuration --- .pre-commit-config.yaml | 10 +- benchmark/halumem/eval_reme.py | 325 ++++++++++-------- reme/agent/chat/__init__.py | 4 +- reme/agent/memory/__init__.py | 2 +- reme/config/default.yaml | 17 +- reme/core/application.py | 72 ++-- reme/core/context/registry_factory.py | 18 +- reme/core/context/service_context.py | 66 ++-- reme/core/embedding/__init__.py | 4 +- reme/core/file_watcher/__init__.py | 4 +- reme/core/flow/__init__.py | 2 +- reme/core/flow/base_flow.py | 4 +- reme/core/llm/__init__.py | 8 +- reme/core/memory_store/__init__.py | 2 +- reme/core/op/__init__.py | 2 +- reme/core/schema/service_config.py | 14 +- reme/core/service/__init__.py | 6 +- reme/core/token_counter/__init__.py | 6 +- reme/core/vector_store/__init__.py | 10 +- reme/reme.py | 42 ++- reme/reme_fs.py | 22 +- reme/tool/fs/__init__.py | 2 +- reme/tool/gallery/__init__.py | 2 +- reme/tool/memory/__init__.py | 2 +- reme/tool/search/__init__.py | 2 +- reme/workflow/gallery/__init__.py | 2 +- .../procedural_memory/summarizer/__init__.py | 2 +- tests/test_fs_memory_search.py | 40 +-- tests/test_reme.py | 6 +- 29 files changed, 386 insertions(+), 312 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e29ebf69..d44e7f60 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/|reme_ai/|bench) + exclude: ^(test/|cookbook/|reme_ai/) - 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/|reme_ai/|bench) + exclude: ^(test/|cookbook/|reme_ai/) - repo: https://github.com/psf/black rev: 25.9.0 hooks: - id: black - exclude: ^(test/|cookbook/|reme_ai/|bench) + exclude: ^(test/|cookbook/|reme_ai/) args: [--line-length=120] - repo: https://github.com/PyCQA/flake8 rev: 7.3.0 hooks: - id: flake8 - exclude: ^(test/|cookbook/|reme_ai/|bench) + exclude: ^(test/|cookbook/|reme_ai/) args: [ "--extend-ignore=E203", "--max-line-length=120" @@ -45,7 +45,6 @@ repos: | \.md$ | \.html$ | reme_ai/ - | bench ) args: [ --disable=W0511, @@ -80,6 +79,7 @@ repos: --disable=R0912, --max-statements=120, --max-line-length=120, + --max-module-lines=1500, ] - repo: https://github.com/regebro/pyroma rev: "5.0" diff --git a/benchmark/halumem/eval_reme.py b/benchmark/halumem/eval_reme.py index ce6631d8..229cfb85 100644 --- a/benchmark/halumem/eval_reme.py +++ b/benchmark/halumem/eval_reme.py @@ -35,6 +35,7 @@ from reme.reme import ReMe @dataclass class EvalConfig: """Evaluation configuration parameters.""" + data_path: str top_k: int = 20 user_num: int = 1 @@ -49,6 +50,7 @@ class EvalConfig: # ==================== Utilities ==================== + class DataLoader: """Handles loading and parsing of HaluMem data.""" @@ -74,7 +76,8 @@ class DataLoader: "role": turn["role"], "content": turn["content"], "time_created": datetime.strptime( - turn["timestamp"], "%b %d, %Y, %H:%M:%S" + turn["timestamp"], + "%b %d, %Y, %H:%M:%S", ) .replace(tzinfo=timezone.utc) .strftime("%Y-%m-%d %H:%M:%S"), @@ -88,17 +91,20 @@ class DataLoader: """Format dialogue into string for evaluation.""" formatted_turns = [] for turn in dialogue: - timestamp = datetime.strptime( - turn["timestamp"], "%b %d, %Y, %H:%M:%S" - ).replace(tzinfo=timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + 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'] + role = user_name if turn["role"] == "user" and user_name else turn["role"] formatted_turns.append( - f"Role: {role}\n" - f"Content: {turn['content']}\n" - f"Time: {timestamp}" + f"Role: {role}\n" f"Content: {turn['content']}\n" f"Time: {timestamp}", ) return "\n\n".join(formatted_turns) @@ -139,8 +145,7 @@ class FileManager: def user_has_cache(self, user_name: str) -> bool: """Check if user has cached results.""" user_dir = self.get_user_dir(user_name) - return any(f.name.startswith("session_") and f.suffix == ".json" - for f in user_dir.iterdir()) + 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.""" @@ -149,10 +154,9 @@ class FileManager: if not user_dir.is_dir(): continue - session_files = sorted([ - f for f in user_dir.iterdir() - if f.name.startswith("session_") and f.suffix == ".json" - ]) + session_files = sorted( + [f for f in user_dir.iterdir() if f.name.startswith("session_") and f.suffix == ".json"], + ) if not session_files: continue @@ -164,7 +168,7 @@ class FileManager: user_data = { "uuid": first_session["uuid"], "user_name": first_session["user_name"], - "sessions": [] + "sessions": [], } # Load all sessions @@ -181,12 +185,13 @@ class FileManager: # ==================== Evaluation Functions ==================== + async def answer_question_with_memories( - reme: ReMe, - question: str, - memories: str, - user_id: str = None, - model_name: str = "qwen3-30b-a3b-instruct-2507" + reme: ReMe, + question: str, + memories: str, + user_id: str = None, + model_name: str = "qwen3-30b-a3b-instruct-2507", ): """ Answer a question using retrieved memories with PROMPT_MEMZERO_JSON template. @@ -206,7 +211,7 @@ async def answer_question_with_memories( context = reme.prompt_handler.prompt_format( "TEMPLATE_MEMOS", user_id=user_id, - memories=memories + memories=memories, ) else: context = f"Memories:\n{memories}" @@ -215,10 +220,10 @@ async def answer_question_with_memories( prompt = reme.prompt_handler.prompt_format( "PROMPT_MEMZERO_JSON", context=context, - question=question + question=question, ) - result = await reme.llm.simple_request_for_json( + result = await reme.get_llm("qwen3_max_instruct").simple_request_for_json( prompt=prompt, model_name=model_name, ) @@ -227,13 +232,13 @@ async def answer_question_with_memories( async def evaluation_for_question( - reme: ReMe, - question: str, - reference_answer: str, - key_memory_points: str, - response: str, - dialogue: str = None, - model_name: str = "qwen3-max" + reme: ReMe, + question: str, + reference_answer: str, + key_memory_points: str, + response: str, + dialogue: str = None, + model_name: str = "qwen3-max", ): """ Question-Answering Evaluation with optional Dialogue Context. @@ -256,12 +261,12 @@ async def evaluation_for_question( reference_answer=reference_answer, key_memory_points=key_memory_points, response=response, - dialogue=dialogue if dialogue else "" + dialogue=dialogue if dialogue else "", ) - result = await reme.llm.simple_request_for_json( + result = await reme.get_llm("qwen3_max_instruct").simple_request_for_json( prompt=prompt, - model_name=model_name + model_name=model_name, ) return result @@ -269,10 +274,18 @@ async def evaluation_for_question( # ==================== Memory Operations ==================== + class MemoryProcessor: """Handles ReMe memory operations.""" - def __init__(self, reme: ReMe, reme_model_name:str="qwen3-max",eval_model_name: str = "qwen3-max", algo_version: str = "halumem", enable_thinking_params: bool = False): + def __init__( + self, + reme: ReMe, + reme_model_name: str = "qwen3-max", + eval_model_name: str = "qwen3-max", + algo_version: str = "halumem", + enable_thinking_params: bool = False, + ): self.reme = reme self.reme_model_name = reme_model_name self.eval_model_name = eval_model_name @@ -280,10 +293,10 @@ class MemoryProcessor: self.enable_thinking_params = enable_thinking_params async def add_memories( - self, - user_id: str, - messages: list[dict], - batch_size: int = 10000 + self, + user_id: str, + messages: list[dict], + batch_size: int = 10000, ) -> tuple[list[str], list, float]: """ Add memories in batches using ReMe and return extracted memory contents. @@ -296,7 +309,7 @@ class MemoryProcessor: total_duration_ms = 0 for i in range(0, len(messages), batch_size): - batch = messages[i:i + batch_size] + batch = messages[i : i + batch_size] start = time.time() # Use new summary API @@ -306,7 +319,8 @@ class MemoryProcessor: version=self.algo_version, return_dict=True, enable_time_filter=True, - enable_thinking_params=self.enable_thinking_params + enable_thinking_params=self.enable_thinking_params, + llm_config_name="qwen-plus-t", ) duration_ms = (time.time() - start) * 1000 @@ -318,10 +332,10 @@ class MemoryProcessor: return extracted_memories, summary_messages, total_duration_ms async def search_memory( - self, - query: str, - user_id: str, - top_k: int = 20 + self, + query: str, + user_id: str, + top_k: int = 20, ) -> tuple[dict, list, float]: """ Search memory using ReMe and return structured answer with reasoning. @@ -340,7 +354,8 @@ class MemoryProcessor: version=self.algo_version, return_dict=True, enable_time_filter=True, - enable_thinking_params=self.enable_thinking_params + enable_thinking_params=self.enable_thinking_params, + llm_config_name="qwen-plus-t", ) # Extract memories from response @@ -367,6 +382,7 @@ class MemoryProcessor: # ==================== Evaluation ==================== + class QuestionAnsweringEvaluator: """Evaluates question answering performance.""" @@ -377,12 +393,12 @@ class QuestionAnsweringEvaluator: self.eval_model_name = eval_model_name async def evaluate_questions( - self, - questions: list[dict], - user_name: str, - uuid: str, - session_id: int, - formatted_dialogue: str + self, + questions: list[dict], + user_name: str, + uuid: str, + session_id: int, + formatted_dialogue: str, ) -> list[dict]: """Evaluate all questions for a session.""" results = [] @@ -391,7 +407,7 @@ class QuestionAnsweringEvaluator: answer_dict, agent_messages, duration_ms = await self.memory_processor.search_memory( query=qa["question"], user_id=user_name, - top_k=self.top_k + top_k=self.top_k, ) # Extract answer and reasoning from the structured response @@ -409,7 +425,7 @@ class QuestionAnsweringEvaluator: key_memory_points=evidence_text, response=system_answer, dialogue=formatted_dialogue, - model_name=self.eval_model_name + model_name=self.eval_model_name, ) eval_result_original_answer = await evaluation_for_question( @@ -419,7 +435,7 @@ class QuestionAnsweringEvaluator: key_memory_points=evidence_text, response=retrieved_memories, dialogue=formatted_dialogue, - model_name=self.eval_model_name + model_name=self.eval_model_name, ) # Build result record @@ -459,7 +475,7 @@ class MetricsAggregator: "hallucination_qa_ratio(valid)": 0, "omission_qa_ratio(valid)": 0, "qa_valid_num": 0, - "qa_num": 0 + "qa_num": 0, } correct = 0 @@ -484,21 +500,25 @@ class MetricsAggregator: "hallucination_qa_ratio(all)": hallucination / total, "omission_qa_ratio(all)": omission / total, "qa_valid_num": valid, - "qa_num": total + "qa_num": total, } if valid > 0: - metrics.update({ - "correct_qa_ratio(valid)": correct / valid, - "hallucination_qa_ratio(valid)": hallucination / valid, - "omission_qa_ratio(valid)": omission / valid - }) + metrics.update( + { + "correct_qa_ratio(valid)": correct / valid, + "hallucination_qa_ratio(valid)": hallucination / valid, + "omission_qa_ratio(valid)": omission / valid, + }, + ) else: - metrics.update({ - "correct_qa_ratio(valid)": 0, - "hallucination_qa_ratio(valid)": 0, - "omission_qa_ratio(valid)": 0 - }) + metrics.update( + { + "correct_qa_ratio(valid)": 0, + "hallucination_qa_ratio(valid)": 0, + "omission_qa_ratio(valid)": 0, + }, + ) return metrics @@ -507,7 +527,7 @@ class MetricsAggregator: """Compute question answering metrics for both result_type and original_result_type.""" return { "with_llm_answer": MetricsAggregator._compute_single_metric(qa_records, "result_type"), - "with_original_memories": MetricsAggregator._compute_single_metric(qa_records, "original_result_type") + "with_original_memories": MetricsAggregator._compute_single_metric(qa_records, "original_result_type"), } @staticmethod @@ -533,18 +553,32 @@ class MetricsAggregator: return { "add_dialogue_duration_time": add_duration / 1000 / 60, "search_memory_duration_time": search_duration / 1000 / 60, - "total_duration_time": (add_duration + search_duration) / 1000 / 60 + "total_duration_time": (add_duration + search_duration) / 1000 / 60, } # ==================== Main Pipeline ==================== + class HaluMemEvaluator: """HaluMem evaluator with proper resource management.""" def __init__(self, config: EvalConfig): self.config = config - self.reme = ReMe(llm={"model_name": self.config.reme_model_name}) + self.reme = ReMe( + default_llm_config={ + "model_name": self.config.reme_model_name, + }, + llms={ + "qwen-plus-t": { + "backend": "openai", + "model_name": "qwen-plus", + "extra_body": { + "enable_thinking": True, + }, + }, + }, + ) # Load evaluation prompts into ReMe's prompt handler prompts_yaml_path = Path(__file__).parent / "eval_reme.yaml" @@ -556,13 +590,13 @@ class HaluMemEvaluator: config.reme_model_name, config.eval_model_name, config.algo_version, - config.enable_thinking_params + config.enable_thinking_params, ) self.qa_evaluator = QuestionAnsweringEvaluator( self.memory_processor, self.reme, config.top_k, - config.eval_model_name + config.eval_model_name, ) self.data_loader = DataLoader() @@ -581,18 +615,18 @@ class HaluMemEvaluator: return False async def process_session( - self, - session: dict, - session_id: int, - user_name: str, - uuid: str + self, + session: dict, + session_id: int, + user_name: str, + uuid: str, ) -> dict: """Process a single session using ReMe.""" session_data = { "uuid": uuid, "user_name": user_name, "session_id": session_id, - "memory_points": session["memory_points"] + "memory_points": session["memory_points"], } # Skip generated QA sessions @@ -606,15 +640,17 @@ class HaluMemEvaluator: 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 + batch_size=self.config.batch_size, ) - session_data.update({ - "dialogue": dialogue, - "extracted_memories": extracted_memories, - "summary_messages": agent_messages, - "add_dialogue_duration_ms": duration_ms - }) + session_data.update( + { + "dialogue": dialogue, + "extracted_memories": extracted_memories, + "summary_messages": agent_messages, + "add_dialogue_duration_ms": duration_ms, + }, + ) # Evaluate questions if present if "questions" in session: @@ -624,11 +660,11 @@ class HaluMemEvaluator: user_name=user_name, uuid=uuid, session_id=session_id, - formatted_dialogue=formatted_dialogue + formatted_dialogue=formatted_dialogue, ) session_data["evaluation_results"] = { - "question_answering_records": qa_results + "question_answering_records": qa_results, } return session_data @@ -647,7 +683,7 @@ class HaluMemEvaluator: session=session, session_id=idx, user_name=user_name, - uuid=uuid + uuid=uuid, ) self.file_manager.save_session(user_name, idx, session_data) @@ -672,23 +708,20 @@ class HaluMemEvaluator: # Load user data first to get user names all_users = self.data_loader.load_jsonl(self.config.data_path) - users_to_process = all_users[:self.config.user_num] + users_to_process = all_users[: self.config.user_num] # Extract all user names and delete all profiles - all_user_names = [ - self.data_loader.extract_user_name(user_data["persona_info"]) - for user_data in all_users - ] + all_user_names = [self.data_loader.extract_user_name(user_data["persona_info"]) for user_data in all_users] if all_user_names: for user_name in all_user_names: self.reme.get_profile_handler(user_name).delete_all() logger.info(f"Deleted all profiles for {len(all_user_names)} users") # Clear existing data - await self.reme.vector_store.delete_all() + await self.reme.default_vector_store.delete_all() # Clear meta_memory directory - meta_memory_path = Path(f"meta_memory/{self.reme.vector_store.collection_name}") + meta_memory_path = Path(f"meta_memory/{self.reme.default_vector_store.collection_name}") if meta_memory_path.exists(): shutil.rmtree(meta_memory_path) logger.info(f"Cleared meta_memory directory: {meta_memory_path}") @@ -725,10 +758,7 @@ class HaluMemEvaluator: return result - tasks = [ - process_with_cache_check(idx, user) - for idx, user in enumerate(users_to_process, 1) - ] + tasks = [process_with_cache_check(idx, user) for idx, user in enumerate(users_to_process, 1)] await asyncio.gather(*tasks) elapsed = time.time() - start_time @@ -758,7 +788,7 @@ class HaluMemEvaluator: eval_results = session.get("evaluation_results", {}) qa_records.extend( - eval_results.get("question_answering_records", []) + eval_results.get("question_answering_records", []), ) except (json.JSONDecodeError, KeyError): return @@ -773,9 +803,9 @@ class HaluMemEvaluator: final_results = { "overall_score": { "question_answering": qa_metrics, - "time_consuming": time_metrics + "time_consuming": time_metrics, }, - "question_answering_records": qa_records + "question_answering_records": qa_records, } # Save statistics @@ -803,7 +833,7 @@ class HaluMemEvaluator: eval_results = session.get("evaluation_results", {}) qa_records.extend( - eval_results.get("question_answering_records", []) + eval_results.get("question_answering_records", []), ) # Compute metrics @@ -813,9 +843,9 @@ class HaluMemEvaluator: final_results = { "overall_score": { "question_answering": qa_metrics, - "time_consuming": time_metrics + "time_consuming": time_metrics, }, - "question_answering_records": qa_records + "question_answering_records": qa_records, } # Save final report @@ -856,7 +886,7 @@ class HaluMemEvaluator: print(f" Omission (valid): {orig_metrics['omission_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {orig_metrics['qa_valid_num']}/{orig_metrics['qa_num']}") - print(f"\n⏱️ Time Metrics:") + print("\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") @@ -865,16 +895,17 @@ class HaluMemEvaluator: # ==================== Entry Point ==================== + async def main_async( - data_path: str, - top_k: int, - batch_size: int, - user_num: int, - max_concurrency: int, - reme_model_name: str= "qwen-flash", - eval_model_name: str = "qwen3-max", - algo_version: str = "halumem", - enable_thinking_params: bool = False + data_path: str, + top_k: int, + batch_size: int, + user_num: int, + max_concurrency: int, + reme_model_name: str = "qwen-flash", + eval_model_name: str = "qwen3-max", + algo_version: str = "halumem", + enable_thinking_params: bool = False, ): """Main async entry point for ReMe evaluation with proper resource cleanup.""" config = EvalConfig( @@ -886,7 +917,7 @@ async def main_async( reme_model_name=reme_model_name, eval_model_name=eval_model_name, algo_version=algo_version, - enable_thinking_params=enable_thinking_params + enable_thinking_params=enable_thinking_params, ) # Use async context manager for automatic cleanup @@ -895,90 +926,92 @@ async def main_async( def main( - data_path: str, - top_k: int, - batch_size: int, - user_num: int, - max_concurrency: int, - reme_model_name: str= "qwen-flash", - eval_model_name: str = "qwen3-max", - algo_version: str = "halumem", - enable_thinking_params: bool = False + data_path: str, + top_k: int, + batch_size: int, + user_num: int, + max_concurrency: int, + reme_model_name: str = "qwen-flash", + eval_model_name: str = "qwen3-max", + algo_version: str = "halumem", + enable_thinking_params: bool = False, ): """Main entry point for ReMe evaluation.""" - asyncio.run(main_async( - data_path=data_path, - top_k=top_k, - batch_size=batch_size, - user_num=user_num, - max_concurrency=max_concurrency, - reme_model_name=reme_model_name, - eval_model_name=eval_model_name, - algo_version=algo_version, - enable_thinking_params=enable_thinking_params - )) + asyncio.run( + main_async( + data_path=data_path, + top_k=top_k, + batch_size=batch_size, + user_num=user_num, + max_concurrency=max_concurrency, + reme_model_name=reme_model_name, + eval_model_name=eval_model_name, + algo_version=algo_version, + enable_thinking_params=enable_thinking_params, + ), + ) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( - description="Evaluate ReMe on HaluMem benchmark (Question Answering)" + description="Evaluate ReMe on HaluMem benchmark (Question Answering)", ) parser.add_argument( "--data_path", type=str, # required=True, default="/Users/zhouwk/PycharmProjects/MemAgent/dataset/halumem/HaluMem-Medium.jsonl", - help="Path to HaluMem JSONL file" + help="Path to HaluMem JSONL file", ) parser.add_argument( "--top_k", type=int, default=20, - help="Number of memories to retrieve (default: 20)" + help="Number of memories to retrieve (default: 20)", ) parser.add_argument( "--user_num", type=int, default=1, - help="Number of users to evaluate (default: 1)" + help="Number of users to evaluate (default: 1)", ) parser.add_argument( "--max_concurrency", type=int, default=1, - help="Maximum concurrent user processing (default: 100)" + help="Maximum concurrent user processing (default: 100)", ) parser.add_argument( "--batch_size", type=int, default=40, - help="Batch size for memory summary processing of each conversation (default: 40)" + help="Batch size for memory summary processing of each conversation (default: 40)", ) parser.add_argument( "--reme_model_name", type=str, default="qwen-flash", - help="Model name for ReMe (default: qwen-flash)" + help="Model name for ReMe (default: qwen-flash)", ) parser.add_argument( "--eval_model_name", type=str, default="qwen3-max", - help="Model name for evaluation (default: qwen3-max)" + help="Model name for evaluation (default: qwen3-max)", ) parser.add_argument( "--algo_version", type=str, default="v1", - help="Algorithm version for summary and retrieval (default: v1)" + help="Algorithm version for summary and retrieval (default: v1)", ) parser.add_argument( "--enable_thinking_params", action="store_true", default=False, - help="Enable thinking parameters for summary and retrieval (default: False)" + help="Enable thinking parameters for summary and retrieval (default: False)", ) args = parser.parse_args() @@ -993,5 +1026,5 @@ if __name__ == "__main__": reme_model_name=args.reme_model_name, eval_model_name=args.eval_model_name, algo_version=args.algo_version, - enable_thinking_params=args.enable_thinking_params + enable_thinking_params=args.enable_thinking_params, ) diff --git a/reme/agent/chat/__init__.py b/reme/agent/chat/__init__.py index a04e8542..506bce7a 100644 --- a/reme/agent/chat/__init__.py +++ b/reme/agent/chat/__init__.py @@ -9,5 +9,5 @@ __all__ = [ "SimpleChat", ] -R.op.register(SimpleChat) -R.op.register(StreamChat) +R.ops.register(SimpleChat) +R.ops.register(StreamChat) diff --git a/reme/agent/memory/__init__.py b/reme/agent/memory/__init__.py index b1443803..ae158f8b 100644 --- a/reme/agent/memory/__init__.py +++ b/reme/agent/memory/__init__.py @@ -38,4 +38,4 @@ for name in __all__: and issubclass(agent_class, BaseMemoryAgent) and agent_class is not BaseMemoryAgent ): - R.op.register(agent_class) + R.ops.register(agent_class) diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 561f931a..24eee5cf 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -12,12 +12,12 @@ http: timeout_keep_alive: 600 limit_concurrency: 64 -flow: +flows: test: flow_content: TestOp() description: "test" -llm: +llms: default: backend: openai model_name: qwen3-30b-a3b-instruct-2507 @@ -29,20 +29,27 @@ llm: model_name: qwen3-max request_interval: 2 -embedding_model: + qwen-plus-thinking: + backend: openai + model_name: qwen-plus + request_interval: 1 + extra_body: + enable_thinking: True + +embedding_models: default: backend: openai model_name: text-embedding-v4 dimensions: 1024 -vector_store: +vector_stores: default: backend: chroma # backend: local embedding_model: default collection_name: reme -token_counter: +token_counters: default: backend: base diff --git a/reme/core/application.py b/reme/core/application.py index fde84afe..68465771 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -26,12 +26,12 @@ class Application: embedding_api_base: str | None = None, enable_logo: bool = True, parser: type[PydanticConfigParser] | None = None, - llm: dict | None = None, - embedding_model: dict | None = None, - vector_store: dict | None = None, - memory_store: dict | None = None, - token_counter: dict | None = None, - file_watcher: dict | None = None, + default_llm_config: dict | None = None, + default_embedding_model_config: dict | None = None, + default_vector_store_config: dict | None = None, + default_memory_store_config: dict | None = None, + default_token_counter_config: dict | None = None, + default_file_watcher_config: dict | None = None, **kwargs, ): self.service_context = ServiceContext( @@ -44,12 +44,12 @@ class Application: parser=parser, config_path=None, enable_logo=enable_logo, - llm=llm, - embedding_model=embedding_model, - vector_store=vector_store, - memory_store=memory_store, - token_counter=token_counter, - file_watcher=file_watcher, + default_llm_config=default_llm_config, + default_embedding_model_config=default_embedding_model_config, + default_vector_store_config=default_vector_store_config, + default_memory_store_config=default_memory_store_config, + default_token_counter_config=default_token_counter_config, + default_file_watcher_config=default_file_watcher_config, **kwargs, ) self.prompt_handler = PromptHandler(language=self.service_context.language) @@ -82,12 +82,12 @@ class Application: embedding_api_base=embedding_api_base, enable_logo=enable_logo, parser=parser, - llm=llm, - embedding_model=embedding_model, - vector_store=vector_store, - memory_store=memory_store, - token_counter=token_counter, - file_watcher=file_watcher, + default_llm_config=llm, + default_embedding_model_config=embedding_model, + default_vector_store_config=vector_store, + default_memory_store_config=memory_store, + default_token_counter_config=token_counter, + default_file_watcher_config=file_watcher, **kwargs, ) await instance.start() @@ -141,35 +141,59 @@ class Application: yield chunk @property - def llm(self) -> BaseLLM: + def default_llm(self) -> BaseLLM: """Get the default LLM instance.""" return self.service_context.llms.get("default") + def get_llm(self, name: str): + """Get an LLM instance by name.""" + return self.service_context.llms.get(name) + @property - def embedding_model(self) -> BaseEmbeddingModel: + def default_embedding_model(self) -> BaseEmbeddingModel: """Get the default embedding model instance.""" return self.service_context.embedding_models.get("default") + def get_embedding_model(self, name: str): + """Get an embedding model instance by name.""" + return self.service_context.embedding_models.get(name) + @property - def vector_store(self) -> BaseVectorStore: + def default_vector_store(self) -> BaseVectorStore: """Get the default vector store instance.""" return self.service_context.vector_stores.get("default") + def get_vector_store(self, name: str): + """Get a vector store instance by name.""" + return self.service_context.vector_stores.get(name) + @property - def memory_store(self) -> BaseMemoryStore: + def default_memory_store(self) -> BaseMemoryStore: """Get the default memory store instance.""" return self.service_context.memory_stores.get("default") + def get_memory_store(self, name: str): + """Get a memory store instance by name.""" + return self.service_context.memory_stores.get(name) + @property - def file_watcher(self) -> BaseFileWatcher: + def default_file_watcher(self) -> BaseFileWatcher: """Get the default file watcher instance.""" return self.service_context.file_watchers.get("default") + def get_file_watcher(self, name: str): + """Get a file watcher instance by name.""" + return self.service_context.file_watchers.get(name) + @property - def token_counter(self) -> BaseTokenCounter: + def default_token_counter(self) -> BaseTokenCounter: """Get the default token counter instance.""" return self.service_context.token_counters.get("default") + def get_token_counter(self, name: str): + """Get a token counter instance by name.""" + return self.service_context.token_counters.get(name) + def run_service(self): """Run the configured service (HTTP, MCP, or CMD).""" import warnings diff --git a/reme/core/context/registry_factory.py b/reme/core/context/registry_factory.py index 8cb51e61..e0633cfd 100644 --- a/reme/core/context/registry_factory.py +++ b/reme/core/context/registry_factory.py @@ -33,15 +33,15 @@ class RegistryFactory: """A factory class for creating registries.""" def __init__(self): - self.llm = Registry() - self.embedding_model = Registry() - self.vector_store = Registry() - self.memory_store = Registry() - self.op = Registry() - self.flow = Registry() - self.service = Registry() - self.token_counter = Registry() - self.file_watcher = Registry() + self.llms = Registry() + self.embedding_models = Registry() + self.vector_stores = Registry() + self.memory_stores = Registry() + self.ops = Registry() + self.flows = Registry() + self.services = Registry() + self.token_counters = Registry() + self.file_watchers = Registry() R = RegistryFactory() diff --git a/reme/core/context/service_context.py b/reme/core/context/service_context.py index f6f07efa..b17745bd 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/context/service_context.py @@ -36,12 +36,12 @@ class ServiceContext(BaseContext): parser: type[PydanticConfigParser] | None = None, config_path: str | None = None, enable_logo: bool = True, - llm: dict | None = None, - embedding_model: dict | None = None, - vector_store: dict | None = None, - memory_store: dict | None = None, - token_counter: dict | None = None, - file_watcher: dict | None = None, + default_llm_config: dict | None = None, + default_embedding_model_config: dict | None = None, + default_vector_store_config: dict | None = None, + default_memory_store_config: dict | None = None, + default_token_counter_config: dict | None = None, + default_file_watcher_config: dict | None = None, **kwargs, ): super().__init__() @@ -61,18 +61,18 @@ class ServiceContext(BaseContext): if args: input_args.extend(args) - if llm: - self._update_section_config(kwargs, "llm", **llm) - if embedding_model: - self._update_section_config(kwargs, "embedding_model", **embedding_model) - if token_counter: - self._update_section_config(kwargs, "token_counter", **token_counter) - if vector_store: - self._update_section_config(kwargs, "vector_store", **vector_store) - if memory_store: - self._update_section_config(kwargs, "memory_store", **memory_store) - if file_watcher: - self._update_section_config(kwargs, "file_watcher", **file_watcher) + if default_llm_config: + self._update_section_config(kwargs, "llms", **default_llm_config) + if default_embedding_model_config: + self._update_section_config(kwargs, "embedding_models", **default_embedding_model_config) + if default_token_counter_config: + self._update_section_config(kwargs, "token_counters", **default_token_counter_config) + if default_vector_store_config: + self._update_section_config(kwargs, "vector_stores", **default_vector_store_config) + if default_memory_store_config: + self._update_section_config(kwargs, "memory_stores", **default_memory_store_config) + if default_file_watcher_config: + self._update_section_config(kwargs, "file_watchers", **default_file_watcher_config) kwargs["enable_logo"] = enable_logo logger.info(f"update with args: {input_args} kwargs: {kwargs}") service_config = parser.parse_args(*input_args, **kwargs) @@ -103,7 +103,7 @@ class ServiceContext(BaseContext): self.flows: dict[str, "BaseFlow"] = {} self.mcp_server_mapping: dict[str, dict] = {} - self.service: "BaseService" = R.service[self.service_config.backend](service_context=self) + self.service: "BaseService" = R.services[self.service_config.backend](service_context=self) self._build_flows() @@ -124,7 +124,7 @@ class ServiceContext(BaseContext): def _build_flows(self): expression_flow_cls = None - for name, flow_cls in R.flow.items(): + for name, flow_cls in R.flows.items(): if not self._filter_flows(name): continue @@ -135,7 +135,7 @@ class ServiceContext(BaseContext): self.flows[flow.name] = flow if expression_flow_cls is not None: - for name, flow_config in self.service_config.flow.items(): + for name, flow_config in self.service_config.flows.items(): if not self._filter_flows(name): continue flow_config.name = name @@ -146,23 +146,23 @@ class ServiceContext(BaseContext): async def start(self): """Start the service context by initializing all configured components.""" - for name, config in self.service_config.llm.items(): - self.llms[name] = R.llm[config.backend](model_name=config.model_name, **config.model_extra) + for name, config in self.service_config.llms.items(): + self.llms[name] = R.llms[config.backend](model_name=config.model_name, **config.model_extra) - for name, config in self.service_config.embedding_model.items(): - self.embedding_models[name] = R.embedding_model[config.backend]( + for name, config in self.service_config.embedding_models.items(): + self.embedding_models[name] = R.embedding_models[config.backend]( model_name=config.model_name, **config.model_extra, ) - for name, config in self.service_config.token_counter.items(): - self.token_counters[name] = R.token_counter[config.backend]( + for name, config in self.service_config.token_counters.items(): + self.token_counters[name] = R.token_counters[config.backend]( model_name=config.model_name, **config.model_extra, ) - for name, config in self.service_config.vector_store.items(): - self.vector_stores[name] = R.vector_store[config.backend]( + for name, config in self.service_config.vector_stores.items(): + self.vector_stores[name] = R.vector_stores[config.backend]( collection_name=config.collection_name, embedding_model=self.embedding_models[config.embedding_model], thread_pool=self.thread_pool, @@ -170,8 +170,8 @@ class ServiceContext(BaseContext): ) await self.vector_stores[name].create_collection(config.collection_name) - for name, config in self.service_config.memory_store.items(): - self.memory_stores[name] = R.memory_store[config.backend]( + for name, config in self.service_config.memory_stores.items(): + self.memory_stores[name] = R.memory_stores[config.backend]( store_name=config.store_name, embedding_model=self.embedding_models[config.embedding_model], fts_enabled=config.fts_enabled, @@ -180,8 +180,8 @@ class ServiceContext(BaseContext): ) await self.memory_stores[name].start() - for name, config in self.service_config.file_watcher.items(): - self.file_watchers[name] = R.file_watcher[config.backend]( + for name, config in self.service_config.file_watchers.items(): + self.file_watchers[name] = R.file_watchers[config.backend]( watch_paths=config.watch_paths, suffix_filters=config.suffix_filters, recursive=config.recursive, diff --git a/reme/core/embedding/__init__.py b/reme/core/embedding/__init__.py index f694d065..71fc4f43 100644 --- a/reme/core/embedding/__init__.py +++ b/reme/core/embedding/__init__.py @@ -11,5 +11,5 @@ __all__ = [ "OpenAIEmbeddingModelSync", ] -R.embedding_model.register("openai")(OpenAIEmbeddingModel) -R.embedding_model.register("openai_sync")(OpenAIEmbeddingModelSync) +R.embedding_models.register("openai")(OpenAIEmbeddingModel) +R.embedding_models.register("openai_sync")(OpenAIEmbeddingModelSync) diff --git a/reme/core/file_watcher/__init__.py b/reme/core/file_watcher/__init__.py index 74eb0fd7..e1d72fdb 100644 --- a/reme/core/file_watcher/__init__.py +++ b/reme/core/file_watcher/__init__.py @@ -15,5 +15,5 @@ __all__ = [ "FullFileWatcher", ] -R.file_watcher.register("full")(FullFileWatcher) -R.file_watcher.register("delta")(DeltaFileWatcher) +R.file_watchers.register("full")(FullFileWatcher) +R.file_watchers.register("delta")(DeltaFileWatcher) diff --git a/reme/core/flow/__init__.py b/reme/core/flow/__init__.py index 75ecd7a9..0258d63c 100644 --- a/reme/core/flow/__init__.py +++ b/reme/core/flow/__init__.py @@ -11,4 +11,4 @@ __all__ = [ "ExpressionFlow", ] -R.flow.register(ExpressionFlow) +R.flows.register(ExpressionFlow) diff --git a/reme/core/flow/base_flow.py b/reme/core/flow/base_flow.py index af907e22..7e0e1f0c 100644 --- a/reme/core/flow/base_flow.py +++ b/reme/core/flow/base_flow.py @@ -129,9 +129,9 @@ class BaseFlow(ABC): raise ValueError("Expression is empty") if len(lines) > 1: - exec("\n".join(lines[:-1]), {"__builtins__": {}}, R.op) + exec("\n".join(lines[:-1]), {"__builtins__": {}}, R.ops) - result = eval(lines[-1], {"__builtins__": {}}, R.op) + result = eval(lines[-1], {"__builtins__": {}}, R.ops) if not isinstance(result, BaseOp): raise TypeError(f"Expression evaluated to {type(result)}, expected BaseOp") return result diff --git a/reme/core/llm/__init__.py b/reme/core/llm/__init__.py index 1b80641e..1cb7e03a 100644 --- a/reme/core/llm/__init__.py +++ b/reme/core/llm/__init__.py @@ -15,7 +15,7 @@ __all__ = [ "OpenAILLMSync", ] -R.llm.register("litellm")(LiteLLM) -R.llm.register("litellm_sync")(LiteLLMSync) -R.llm.register("openai")(OpenAILLM) -R.llm.register("openai_sync")(OpenAILLMSync) +R.llms.register("litellm")(LiteLLM) +R.llms.register("litellm_sync")(LiteLLMSync) +R.llms.register("openai")(OpenAILLM) +R.llms.register("openai_sync")(OpenAILLMSync) diff --git a/reme/core/memory_store/__init__.py b/reme/core/memory_store/__init__.py index 6d43e423..2226e1b0 100644 --- a/reme/core/memory_store/__init__.py +++ b/reme/core/memory_store/__init__.py @@ -13,4 +13,4 @@ __all__ = [ "SqliteMemoryStore", ] -R.memory_store.register("sqlite")(SqliteMemoryStore) +R.memory_stores.register("sqlite")(SqliteMemoryStore) diff --git a/reme/core/op/__init__.py b/reme/core/op/__init__.py index 6e1e97e4..f07b161c 100644 --- a/reme/core/op/__init__.py +++ b/reme/core/op/__init__.py @@ -19,4 +19,4 @@ __all__ = [ "SequentialOp", ] -R.op.register(MCPTool) +R.ops.register(MCPTool) diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py index 11c8fca6..30797d48 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -131,10 +131,10 @@ class ServiceConfig(BaseModel): mcp: MCPConfig = Field(default_factory=MCPConfig) http: HttpConfig = Field(default_factory=HttpConfig) cmd: CmdConfig = Field(default_factory=CmdConfig) - flow: dict[str, FlowConfig] = Field(default_factory=dict) - llm: dict[str, LLMConfig] = Field(default_factory=dict) - embedding_model: dict[str, EmbeddingModelConfig] = Field(default_factory=dict) - vector_store: dict[str, VectorStoreConfig] = Field(default_factory=dict) - memory_store: dict[str, MemoryStoreConfig] = Field(default_factory=dict) - token_counter: dict[str, TokenCounterConfig] = Field(default_factory=dict) - file_watcher: dict[str, FileWatcherConfig] = Field(default_factory=dict) + flows: dict[str, FlowConfig] = Field(default_factory=dict) + llms: dict[str, LLMConfig] = Field(default_factory=dict) + embedding_models: dict[str, EmbeddingModelConfig] = Field(default_factory=dict) + vector_stores: dict[str, VectorStoreConfig] = Field(default_factory=dict) + memory_stores: dict[str, MemoryStoreConfig] = Field(default_factory=dict) + token_counters: dict[str, TokenCounterConfig] = Field(default_factory=dict) + file_watchers: dict[str, FileWatcherConfig] = Field(default_factory=dict) diff --git a/reme/core/service/__init__.py b/reme/core/service/__init__.py index 1cd4ae58..4cd60264 100644 --- a/reme/core/service/__init__.py +++ b/reme/core/service/__init__.py @@ -13,6 +13,6 @@ __all__ = [ "MCPService", ] -R.service.register("cmd")(CmdService) -R.service.register("http")(HttpService) -R.service.register("mcp")(MCPService) +R.services.register("cmd")(CmdService) +R.services.register("http")(HttpService) +R.services.register("mcp")(MCPService) diff --git a/reme/core/token_counter/__init__.py b/reme/core/token_counter/__init__.py index f0cb5e28..7792046f 100644 --- a/reme/core/token_counter/__init__.py +++ b/reme/core/token_counter/__init__.py @@ -11,6 +11,6 @@ __all__ = [ "OpenAITokenCounter", ] -R.token_counter.register("base")(BaseTokenCounter) -R.token_counter.register("hf")(HFTokenCounter) -R.token_counter.register("openai")(OpenAITokenCounter) +R.token_counters.register("base")(BaseTokenCounter) +R.token_counters.register("hf")(HFTokenCounter) +R.token_counters.register("openai")(OpenAITokenCounter) diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py index bd8a1869..84b68c3c 100644 --- a/reme/core/vector_store/__init__.py +++ b/reme/core/vector_store/__init__.py @@ -17,8 +17,8 @@ __all__ = [ "QdrantVectorStore", ] -R.vector_store.register("chroma")(ChromaVectorStore) -R.vector_store.register("es")(ESVectorStore) -R.vector_store.register("local")(LocalVectorStore) -R.vector_store.register("pgvector")(PGVectorStore) -R.vector_store.register("qdrant")(QdrantVectorStore) +R.vector_stores.register("chroma")(ChromaVectorStore) +R.vector_stores.register("es")(ESVectorStore) +R.vector_stores.register("local")(LocalVectorStore) +R.vector_stores.register("pgvector")(PGVectorStore) +R.vector_stores.register("qdrant")(QdrantVectorStore) diff --git a/reme/reme.py b/reme/reme.py index 4f197dd6..9a216fcd 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -52,10 +52,10 @@ class ReMe(Application): embedding_api_key: str | None = None, embedding_api_base: str | None = None, enable_logo: bool = True, - llm: dict | None = None, - embedding_model: dict | None = None, - vector_store: dict | None = None, - token_counter: dict | None = None, + default_llm_config: dict | None = None, + default_embedding_model_config: dict | None = None, + default_vector_store_config: dict | None = None, + default_token_counter_config: dict | None = None, target_user_names: list[str] | None = None, target_task_names: list[str] | None = None, target_tool_names: list[str] | None = None, @@ -71,10 +71,10 @@ class ReMe(Application): embedding_api_key: API key for embedding provider embedding_api_base: API base for embedding provider enable_logo: Enable logo - llm: LLM configuration - embedding_model: Embedding model configuration - vector_store: Vector store configuration - token_counter: Token counter configuration + default_llm_config: LLM configuration + default_embedding_model_config: Embedding model configuration + default_vector_store_config: Vector store configuration + default_token_counter_config: Token counter configuration target_user_names: List of user names for personal memory target_task_names: List of task names for procedural memory target_tool_names: List of tool names for tool memory @@ -102,10 +102,10 @@ class ReMe(Application): embedding_api_base=embedding_api_base, enable_logo=enable_logo, parser=ReMeConfigParser, - llm=llm, - embedding_model=embedding_model, - vector_store=vector_store, - token_counter=token_counter, + default_llm_config=default_llm_config, + default_embedding_model_config=default_embedding_model_config, + default_vector_store_config=default_vector_store_config, + default_token_counter_config=default_token_counter_config, **kwargs, ) memory_target_type_mapping: dict[str, MemoryType] = {} @@ -184,6 +184,7 @@ class ReMe(Application): version: str = "default", retrieve_top_k: int = 20, return_dict: bool = False, + llm_config_name: str = "default", **kwargs, ) -> str | dict: """Summarize personal, procedural and tool memories for the given context.""" @@ -197,6 +198,7 @@ class ReMe(Application): personal_summarizer: BaseMemoryAgent if version == "default": personal_summarizer = PersonalSummarizer( + llm=llm_config_name, tools=[ AddAndRetrieveSimilarMemory( enable_thinking_params=enable_thinking_params, @@ -216,6 +218,7 @@ class ReMe(Application): elif version == "v1": personal_summarizer = PersonalV1Summarizer( + llm=llm_config_name, tools=[ AddDraftAndRetrieveSimilarMemory( enable_thinking_params=enable_thinking_params, @@ -244,6 +247,7 @@ class ReMe(Application): ) elif version == "v2": personal_summarizer = PersonalV1Summarizer( + llm=llm_config_name, tools=[ AddDraftAndRetrieveSimilarMemory( enable_thinking_params=enable_thinking_params, @@ -272,6 +276,7 @@ class ReMe(Application): ) elif version == "halumem": personal_summarizer = PersonalHalumemSummarizer( + llm=llm_config_name, tools=[ AddAndRetrieveSimilarMemory( enable_thinking_params=enable_thinking_params, @@ -382,6 +387,7 @@ class ReMe(Application): retrieve_top_k: int = 20, enable_time_filter: bool = True, return_dict: bool = False, + llm_config_name: str = "default", **kwargs, ) -> str | dict: """Retrieve relevant personal, procedural and tool memories for a query.""" @@ -389,6 +395,7 @@ class ReMe(Application): personal_retriever: BaseMemoryAgent if version == "default": personal_retriever = PersonalRetriever( + llm=llm_config_name, tools=[ ReadAllProfiles( enable_thinking_params=enable_thinking_params, @@ -405,6 +412,7 @@ class ReMe(Application): elif version == "v1": personal_retriever = PersonalV1Retriever( + llm=llm_config_name, return_memory_nodes=False, tools=[ ReadAllProfiles( @@ -426,6 +434,7 @@ class ReMe(Application): ) elif version == "v2": personal_retriever = PersonalV1Retriever( + llm=llm_config_name, return_memory_nodes=False, tools=[ ReadAllProfiles( @@ -449,6 +458,7 @@ class ReMe(Application): ) elif version == "halumem": personal_retriever = PersonalHalumemRetriever( + llm=llm_config_name, tools=[ ReadAllProfiles( enable_thinking_params=enable_thinking_params, @@ -598,7 +608,7 @@ class ReMe(Application): Returns: MemoryNode: The retrieved memory node """ - vector_node = await self.vector_store.get(memory_id) + vector_node = await self.default_vector_store.get(memory_id) return MemoryNode.from_vector_node(vector_node) async def delete_memory( @@ -610,11 +620,11 @@ class ReMe(Application): Args: memory_id: The ID of the memory to delete """ - await self.vector_store.delete(memory_id) + await self.default_vector_store.delete(memory_id) async def delete_all(self): """Delete all memory nodes in the vector store.""" - await self.vector_store.delete_all() + await self.default_vector_store.delete_all() async def update_memory( self, @@ -707,7 +717,7 @@ class ReMe(Application): @property def profile_path(self) -> Path: """Get the path to the profile directory.""" - return Path(self.profile_dir) / self.vector_store.collection_name + return Path(self.profile_dir) / self.default_vector_store.collection_name def get_profile_handler(self, user_name: str) -> ProfileHandler: """Get the profile handler for the specified user.""" diff --git a/reme/reme_fs.py b/reme/reme_fs.py index 51875ffc..402e66e0 100644 --- a/reme/reme_fs.py +++ b/reme/reme_fs.py @@ -32,11 +32,11 @@ class ReMeFs(Application): embedding_api_key: str | None = None, embedding_api_base: str | None = None, enable_logo: bool = True, - llm: dict | None = None, - embedding_model: dict | None = None, - memory_store: dict | None = None, - token_counter: dict | None = None, - file_watcher: dict | None = None, + default_llm_config: dict | None = None, + default_embedding_model_config: dict | None = None, + default_memory_store_config: dict | None = None, + default_token_counter_config: dict | None = None, + default_file_watcher_config: dict | None = None, working_dir: str = ".reme", **kwargs, ): @@ -49,11 +49,11 @@ class ReMeFs(Application): embedding_api_base=embedding_api_base, enable_logo=enable_logo, parser=ReMeConfigParser, - llm=llm, - embedding_model=embedding_model, - memory_store=memory_store, - token_counter=token_counter, - file_watcher=file_watcher, + default_llm_config=default_llm_config, + default_embedding_model_config=default_embedding_model_config, + default_memory_store_config=default_memory_store_config, + default_token_counter_config=default_token_counter_config, + default_file_watcher_config=default_file_watcher_config, **kwargs, ) @@ -137,5 +137,5 @@ class ReMeFs(Application): async def memory_get(self, path: str, offset: int | None = None, limit: int | None = None) -> str: """Read specific snippets from memory files.""" - get_tool = FsMemoryGet(workspace_dir=self.working_dir, memory_store=self.memory_store) + get_tool = FsMemoryGet(workspace_dir=self.working_dir) return await get_tool.call(path=path, offset=offset, limit=limit, service_context=self.service_context) diff --git a/reme/tool/fs/__init__.py b/reme/tool/fs/__init__.py index 3a85838d..a9f003ae 100644 --- a/reme/tool/fs/__init__.py +++ b/reme/tool/fs/__init__.py @@ -27,4 +27,4 @@ __all__ = [ for name in __all__: tool_class = globals()[name] - R.op.register(tool_class) + R.ops.register(tool_class) diff --git a/reme/tool/gallery/__init__.py b/reme/tool/gallery/__init__.py index bf3b21b3..2f3ab74f 100644 --- a/reme/tool/gallery/__init__.py +++ b/reme/tool/gallery/__init__.py @@ -13,4 +13,4 @@ __all__ = [ for name in __all__: tool_class = globals()[name] - R.op.register(tool_class) + R.ops.register(tool_class) diff --git a/reme/tool/memory/__init__.py b/reme/tool/memory/__init__.py index b6403f3f..d7ee2e93 100644 --- a/reme/tool/memory/__init__.py +++ b/reme/tool/memory/__init__.py @@ -56,4 +56,4 @@ __all__ = [ for name in __all__: tool_class = globals()[name] if isinstance(tool_class, type) and issubclass(tool_class, BaseMemoryTool) and tool_class is not BaseMemoryTool: - R.op.register(tool_class) + R.ops.register(tool_class) diff --git a/reme/tool/search/__init__.py b/reme/tool/search/__init__.py index 9c68f8cb..69995c37 100644 --- a/reme/tool/search/__init__.py +++ b/reme/tool/search/__init__.py @@ -13,4 +13,4 @@ __all__ = [ for name in __all__: tool_class = globals()[name] - R.op.register(tool_class) + R.ops.register(tool_class) diff --git a/reme/workflow/gallery/__init__.py b/reme/workflow/gallery/__init__.py index 8d18015c..6203089d 100644 --- a/reme/workflow/gallery/__init__.py +++ b/reme/workflow/gallery/__init__.py @@ -11,4 +11,4 @@ __all__ = [ for name in __all__: agent_class = globals()[name] - R.op.register(agent_class) + R.ops.register(agent_class) diff --git a/reme/workflow/procedural_memory/summarizer/__init__.py b/reme/workflow/procedural_memory/summarizer/__init__.py index eb4b194e..34644cd0 100644 --- a/reme/workflow/procedural_memory/summarizer/__init__.py +++ b/reme/workflow/procedural_memory/summarizer/__init__.py @@ -12,4 +12,4 @@ __all__ = ["TrajectoryPreprocess", "SuccessExtraction"] for name in __all__: tool_class = globals()[name] - R.op.register(tool_class) + R.ops.register(tool_class) diff --git a/tests/test_fs_memory_search.py b/tests/test_fs_memory_search.py index 88e6da8b..cb88fa76 100644 --- a/tests/test_fs_memory_search.py +++ b/tests/test_fs_memory_search.py @@ -211,7 +211,7 @@ async def test_memory_search_basic(): reme_fs = ReMeFs( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - memory_store={ + default_memory_store_config={ "backend": "sqlite", "store_name": "test_basic", "embedding_model": "default", @@ -223,13 +223,13 @@ async def test_memory_search_basic(): # Insert personal info chunks personal_chunks = SampleDataGenerator.create_personal_info_chunks("test_basic") - personal_chunks = await reme_fs.memory_store.get_chunk_embeddings(personal_chunks) + personal_chunks = await reme_fs.default_memory_store.get_chunk_embeddings(personal_chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/personal_info.md", len(personal_chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( file_meta, MemorySource.MEMORY, personal_chunks, @@ -273,7 +273,7 @@ async def test_memory_search_technical_content(): reme_fs = ReMeFs( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - memory_store={ + default_memory_store_config={ "backend": "sqlite", "store_name": "test_technical", "embedding_model": "default", @@ -285,13 +285,13 @@ async def test_memory_search_technical_content(): # Insert technical chunks tech_chunks = SampleDataGenerator.create_technical_chunks("test_technical") - tech_chunks = await reme_fs.memory_store.get_chunk_embeddings(tech_chunks) + tech_chunks = await reme_fs.default_memory_store.get_chunk_embeddings(tech_chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/technical_notes.md", len(tech_chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( file_meta, MemorySource.MEMORY, tech_chunks, @@ -340,7 +340,7 @@ async def test_memory_search_with_source_filter(): reme_fs = ReMeFs( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - memory_store={ + default_memory_store_config={ "backend": "sqlite", "store_name": "test_source_filter", "embedding_model": "default", @@ -352,12 +352,12 @@ async def test_memory_search_with_source_filter(): # Insert MEMORY source data personal_chunks = SampleDataGenerator.create_personal_info_chunks("test_source") - personal_chunks = await reme_fs.memory_store.get_chunk_embeddings(personal_chunks) + personal_chunks = await reme_fs.default_memory_store.get_chunk_embeddings(personal_chunks) personal_meta = SampleDataGenerator.create_file_metadata( "memory/personal_info.md", len(personal_chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( personal_meta, MemorySource.MEMORY, personal_chunks, @@ -366,12 +366,12 @@ async def test_memory_search_with_source_filter(): # Insert SESSIONS source data session_chunks = SampleDataGenerator.create_session_chunks("test_source") - session_chunks = await reme_fs.memory_store.get_chunk_embeddings(session_chunks) + session_chunks = await reme_fs.default_memory_store.get_chunk_embeddings(session_chunks) session_meta = SampleDataGenerator.create_file_metadata( "sessions/2024-01-15.jsonl", len(session_chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( session_meta, MemorySource.SESSIONS, session_chunks, @@ -435,7 +435,7 @@ async def test_memory_search_score_filtering(): reme_fs = ReMeFs( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - memory_store={ + default_memory_store_config={ "backend": "sqlite", "store_name": "test_score_filter", "embedding_model": "default", @@ -447,12 +447,12 @@ async def test_memory_search_score_filtering(): # Insert test data chunks = SampleDataGenerator.create_technical_chunks("test_score") - chunks = await reme_fs.memory_store.get_chunk_embeddings(chunks) + chunks = await reme_fs.default_memory_store.get_chunk_embeddings(chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/technical_notes.md", len(chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( file_meta, MemorySource.MEMORY, chunks, @@ -502,7 +502,7 @@ async def test_memory_search_max_results(): reme_fs = ReMeFs( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - memory_store={ + default_memory_store_config={ "backend": "sqlite", "store_name": "test_max_results", "embedding_model": "default", @@ -517,14 +517,14 @@ async def test_memory_search_max_results(): tech_chunks = SampleDataGenerator.create_technical_chunks("test_max") all_chunks = personal_chunks + tech_chunks - all_chunks = await reme_fs.memory_store.get_chunk_embeddings(all_chunks) + all_chunks = await reme_fs.default_memory_store.get_chunk_embeddings(all_chunks) # Insert as one file for simplicity combined_meta = SampleDataGenerator.create_file_metadata( "memory/combined.md", len(all_chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( combined_meta, MemorySource.MEMORY, all_chunks, @@ -569,7 +569,7 @@ async def test_memory_search_hybrid_mode(): reme_fs = ReMeFs( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - memory_store={ + default_memory_store_config={ "backend": "sqlite", "store_name": "test_hybrid", "embedding_model": "default", @@ -581,12 +581,12 @@ async def test_memory_search_hybrid_mode(): # Insert test data chunks = SampleDataGenerator.create_technical_chunks("test_hybrid") - chunks = await reme_fs.memory_store.get_chunk_embeddings(chunks) + chunks = await reme_fs.default_memory_store.get_chunk_embeddings(chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/technical_notes.md", len(chunks), ) - await reme_fs.memory_store.upsert_file( + await reme_fs.default_memory_store.upsert_file( file_meta, MemorySource.MEMORY, chunks, diff --git a/tests/test_reme.py b/tests/test_reme.py index e78867ca..3bdd10c2 100644 --- a/tests/test_reme.py +++ b/tests/test_reme.py @@ -9,11 +9,11 @@ from reme.core.schema import VectorNode, MemoryNode async def test_reme(): """Tests ReMe memory system with personal information storage and retrieval.""" # 构建一段包含个人信息的对话 - reme = ReMe(vector_store={"collection_name": "reme"}) + reme = ReMe(default_vector_store_config={"collection_name": "reme"}) await reme.start() # reme = await ReMe.create(vector_store={"collection_name": "reme"}) - await reme.vector_store.delete_all() + await reme.default_vector_store.delete_all() messages = [ { @@ -76,7 +76,7 @@ async def test_reme(): print("=" * 60) # 列出所有存储的记忆节点 - nodes: list[VectorNode] = await reme.vector_store.list() + nodes: list[VectorNode] = await reme.default_vector_store.list() for i, node in enumerate(nodes, 1): memory_node = MemoryNode.from_vector_node(node) print(f"{i} {memory_node.model_dump_json()}")