diff --git a/benchmark/halumem/eval_reme.py b/benchmark/halumem/eval_reme.py index 229cfb85..9c71478a 100644 --- a/benchmark/halumem/eval_reme.py +++ b/benchmark/halumem/eval_reme.py @@ -320,7 +320,6 @@ class MemoryProcessor: return_dict=True, enable_time_filter=True, enable_thinking_params=self.enable_thinking_params, - llm_config_name="qwen-plus-t", ) duration_ms = (time.time() - start) * 1000 @@ -355,7 +354,6 @@ class MemoryProcessor: return_dict=True, enable_time_filter=True, enable_thinking_params=self.enable_thinking_params, - llm_config_name="qwen-plus-t", ) # Extract memories from response @@ -569,15 +567,6 @@ class HaluMemEvaluator: 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 diff --git a/pyproject.toml b/pyproject.toml index 028f4fa8..19df32d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,5 +113,11 @@ remecli = "reme.reme_cli:main" [tool.pytest.ini_options] asyncio_default_fixture_loop_scope = "function" +# Script-style tests that need to be run with `python test_*.py` +testpaths = ["tests"] +python_files = ["test_*.py"] +python_functions = ["test_*"] +# Exclude script-style tests that require manual execution +addopts = "--ignore=tests/test_embedding.py --ignore=tests/test_embedding_cache.py --ignore=tests/test_embedding_sync.py --ignore=tests/test_file_store.py" # python -m build && twine upload dist/* diff --git a/reme/__init__.py b/reme/__init__.py index b02aaa94..f7f90a29 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -1,27 +1,27 @@ """ReMe""" -from . import agent from . import config from . import core -from . import tool +from . import extension +from . import memory from . import workflow from .reme import ReMe from .reme_cli import ReMeCli -from .reme_fs import ReMeFs +from .reme_fb import ReMeFb + +__version__ = "0.3.0.0b5" __all__ = [ - "agent", "config", "core", - "tool", + "extension", + "memory", "workflow", "ReMe", "ReMeCli", - "ReMeFs", + "ReMeFb", ] -__version__ = "0.3.0.0b4" - """ conda create -n fl_test2 python=3.10 conda activate fl_test2 diff --git a/reme/agent/__init__.py b/reme/agent/__init__.py deleted file mode 100644 index a2a0aa17..00000000 --- a/reme/agent/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""A simple chatbot.""" - -from . import chat -from . import memory - -__all__ = [ - "chat", - "memory", -] diff --git a/reme/agent/chat/__init__.py b/reme/agent/chat/__init__.py deleted file mode 100644 index 18661cfa..00000000 --- a/reme/agent/chat/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -"""chat agent""" - -from .fs_cli import FsCli -from .simple_chat import SimpleChat -from .stream_chat import StreamChat -from ...core import R - -__all__ = [ - "FsCli", - "StreamChat", - "SimpleChat", -] - -R.ops.register(FsCli) -R.ops.register(SimpleChat) -R.ops.register(StreamChat) diff --git a/reme/agent/fs/__init__.py b/reme/agent/fs/__init__.py deleted file mode 100644 index f649ce34..00000000 --- a/reme/agent/fs/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""File system agents for memory management.""" - -from .fs_compactor import FsCompactor -from .fs_context_checker import FsContextChecker -from .fs_summarizer import FsSummarizer - -__all__ = [ - "FsSummarizer", - "FsCompactor", - "FsContextChecker", -] diff --git a/reme/agent/memory/personal/personal_halumem_retriever.py b/reme/agent/memory/personal/personal_halumem_retriever.py deleted file mode 100644 index c2974ead..00000000 --- a/reme/agent/memory/personal/personal_halumem_retriever.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Personal memory retriever agent for retrieving personal memories through vector search.""" - -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import Role, MemoryType -from ....core.op import BaseTool -from ....core.schema import Message -from ....core.utils import format_messages - - -class PersonalHalumemRetriever(BaseMemoryAgent): - """Retrieve personal memories through vector search and history reading.""" - - memory_type: MemoryType = MemoryType.PERSONAL - - async def build_messages(self) -> list[Message]: - if self.context.get("query"): - context = self.context.query - elif self.context.get("messages"): - context = self.description + "\n" + format_messages(self.context.messages) - else: - raise ValueError("input must have either `query` or `messages`") - - read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles") - if read_all_profiles_tool is not None: - all_profiles = await read_all_profiles_tool.call( - memory_target=self.memory_target, - service_context=self.service_context, - ) - else: - all_profiles = "" - - return [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt", - memory_type=self.memory_type.value, - memory_target=self.memory_target, - user_profile=all_profiles, - context=context.strip(), - ), - ), - Message( - role=Role.USER, - content=self.prompt_format( - prompt_name="user_message", - memory_type=self.memory_type.value, - memory_target=self.memory_target, - user_profile=all_profiles, - context=context.strip(), - ), - ), - ] - - async def _acting_step( - self, - assistant_message: Message, - tools: list[BaseTool], - step: int, - stage: str = "", - **kwargs, - ) -> tuple[list[BaseTool], list[Message]]: - """Execute tool calls with memory context.""" - return await super()._acting_step( - assistant_message, - tools, - step, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - retrieved_nodes=self.retrieved_nodes, - **kwargs, - ) - - async def execute(self): - result = await super().execute() - answer = result["answer"] - if "MEMORY_NOT_FOUND" in answer: - result["answer"] = "\n".join( - [ - n.format( - include_memory_id=False, - include_when_to_use=False, - include_content=True, - include_message_time=False, - ref_memory_id_key="", - ) - for n in self.retrieved_nodes - ], - ) - - result["retrieved_nodes"] = self.retrieved_nodes - return result diff --git a/reme/agent/memory/personal/personal_halumem_retriever.yaml b/reme/agent/memory/personal/personal_halumem_retriever.yaml deleted file mode 100644 index 41395dbb..00000000 --- a/reme/agent/memory/personal/personal_halumem_retriever.yaml +++ /dev/null @@ -1,91 +0,0 @@ -system_prompt: | - # Role Definition: - You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}. - - ## Multi-Phase Retrieval Strategy - Follow these phases sequentially to gather comprehensive information: - - ### Tool Rules - **Tool**: `retrieve_memory` (without time constraints) - **Objective**: Cast a wide net to find potentially relevant memories - **Approach**: - - Execute 3-5 diverse search queries using different formulations: - * Original question verbatim - * Rephrased variations (different wording, synonyms) - * Entity-focused queries (extract and search specific names, places, events) - * Keyword-based searches (core concepts, topics) - * Related context queries (broader themes) - - Review all results before proceeding to next phase - - **Tool**: `retrieve_memory` (with time filter) - **When to use**: Only if the user question contains temporal references - **Time Filter Format**: - - Single date: `20200101` - - Date range: `20200101,20200102` (inclusive: 20200101 ≤ time ≤ 20200102) - - Before date: `0,20200102` (up to and including 20200102) - - After date: `20200101,99999999` (from 20200101 onwards) - **Approach**: - - Identify temporal constraints from the user question - - Refine Phase 1 queries with appropriate time filters - - Try multiple time ranges if initial searches yield no results - - **Tool**: `read_history` - **When to use**: After exhausting retrieval attempts OR when specific conversation context is needed - **Approach**: - - Extract `history_id` from retrieved memory references - - Prioritize histories that are most relevant or recent - - Read multiple histories if necessary for complete context - - Use this to understand the full conversation surrounding a memory - - -user_message: | - ## User Profile - {user_profile} - - ## User Question - {context} - - # Core Objective: - Before responding to the user, you must strictly follow the **[Memory Retrieve -> Original Source Tracing -> Broad Search Fallback]** retrieval strategy. It is strictly forbidden to directly opt for an indiscriminate search of massive historical original texts. - - ## Retrieval Strategy & Workflow (Strictly Enforced Chain of Thought) - ### Phase 1: Intent Decomposition and Primary Retrieval (Summary First) - 1. **Analyze Intent**: Analyze the user's current Query, decomposing it into 1-3 core search intents.2. **Summary Priority**: First, retrieve from **high-level memories**. - - **Action**: Call `vector_retrieve_memory` using at least two different `query`. - - **Filter**: (Optional) Set metadata filter {{"timestamp": "YYYY-MM-DD"}} - - **Goal**: Obtain refined conclusions such as entity attributes, task status, user preferences, or environmental information. - - ### Phase 2: Memory Evaluation and Deep Tracing (Drill Down) - Check the retrieval results of Phase 1: - - **Case A (Sufficient Information)**: If the summarized memory contains all the details needed for the answer, proceed directly to Phase 4 for the response. - - **Case B (Vague/Complex Information)**: If summarized memory exists (e.g., 'discussed project architecture') but lacks specific details (e.g., 'specific parameter configuration'), use clues from the summary to trace the original text. - - **Action**: Call `read_history` using ref_memory_id from the retrieved memory. - - **Goal**: Obtain the specific conversation context at that time. - - ### Phase 3: Fallback Retrieval and Strategy Adjustment (Fallback & Expand) - If no valid information is found in both Phase 1 and Phase 2 (result is empty or similarity is too low): Rewrite the Query based on the context (remove non-keywords, synonym substitution), and search again. - - ### Phase 4: Result Compilation and Response - - Combine the retrieved content (summary or original text) with the current conversation context. - - If all retrieved results are irrelevant, **it is strictly forbidden to fabricate memories**; directly inform the user that no relevant information was found. - - ## Output Format - Before the final reply, ensure at least 3 tool calls for retrieval, and then output your answer in ten words. - - Base your answer EXCLUSIVELY on retrieved memories, user profile, and history data - - Never infer, assume, or hallucinate information - - Always cite sources with timestamps: `[timestamp] Memory content` - - Present conflicting information transparently with respective timestamps - - Exhaust all search strategies before concluding information doesn't exist - - Before the final reply, ensure at least 3 tool calls for retrieval, and then output the most relevant JSON retrieval summary, followed by your answer: - ```json - {{ - "retrieved_memories": [ - {{"type": "profile", "timestamp":"...", "content": "..."}}, - {{"type": "personal", "timestamp":"...", "content": "..."}}, - {{"type": "history", "timestamp":"...", "content": "..."}}, - .... - ], - "summary": "Fill in your summarized answer here." - }} - ``` diff --git a/reme/agent/memory/personal/personal_halumem_summarizer.py b/reme/agent/memory/personal/personal_halumem_summarizer.py deleted file mode 100644 index 4541d9de..00000000 --- a/reme/agent/memory/personal/personal_halumem_summarizer.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Personal memory summarizer agent for two-phase personal memory processing.""" - -from loguru import logger - -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import Role, MemoryType -from ....core.op import BaseTool -from ....core.schema import Message - - -class PersonalHalumemSummarizer(BaseMemoryAgent): - """Two-phase personal memory processor: retrieve/add memories then update profile.""" - - memory_type: MemoryType = MemoryType.PERSONAL - - async def _build_s1_messages(self) -> list[Message]: - return [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt_s1", - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - ), - ), - Message( - role=Role.USER, - # content=self.get_prompt("user_message_s1"), - content=self.prompt_format( - prompt_name="user_message_s1", - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - ), - ), - ] - - async def _build_s2_messages(self, user_profile: str) -> list[Message]: - return [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt_s2", - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - user_profile=user_profile, - ), - ), - Message( - role=Role.USER, - content=self.prompt_format( - prompt_name="user_message_s2", - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - user_profile=user_profile, - ), - ), - ] - - async def _acting_step( - self, - assistant_message: Message, - tools: list[BaseTool], - step: int, - stage: str = "", - **kwargs, - ) -> tuple[list[BaseTool], list[Message]]: - """Execute tool calls with memory context.""" - return await super()._acting_step( - assistant_message, - tools, - step, - stage=stage, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - history_node=self.history_node, - author=self.author, - retrieved_nodes=self.retrieved_nodes, - **kwargs, - ) - - async def execute(self): - memory_tools = [] - profile_tools = [] - for i, tool in enumerate(self.tools): - tool_name = tool.tool_call.name - if "_memory" in tool_name: - memory_tools.append(tool) - elif "_profile" in tool_name: - profile_tools.append(tool) - else: - raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={tool_name}") - logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}") - - stage = "s1-memory" - messages_s1 = await self._build_s1_messages() - for i, message in enumerate(messages_s1): - role = message.name or message.role - logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") - tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage) - - if profile_tools: - - read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles") - if read_all_profiles_tool is not None: - all_profiles = await read_all_profiles_tool.call( - memory_target=self.memory_target, - service_context=self.service_context, - ) - else: - all_profiles = "" - - stage = "s2-profile" - messages_s2 = await self._build_s2_messages(user_profile=all_profiles) - for i, message in enumerate(messages_s2): - role = message.name or message.role - logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") - tools_s2, messages_s2, success_s2 = await self.react(messages_s2, profile_tools, stage=stage) - else: - tools_s2, messages_s2, success_s2 = [], [], True - - answer = (messages_s1[-1].content if success_s1 else "") + (messages_s2[-1].content if success_s2 else "") - success = success_s1 and success_s2 - messages = messages_s1 + messages_s2 - tools = tools_s1 + tools_s2 - memory_nodes = [] - for tool in tools: - if tool.memory_nodes: - memory_nodes.extend(tool.memory_nodes) - - return { - "answer": answer, - "success": success, - "messages": messages, - "tools": tools, - "memory_nodes": memory_nodes, - } diff --git a/reme/agent/memory/personal/personal_halumem_summarizer.yaml b/reme/agent/memory/personal/personal_halumem_summarizer.yaml deleted file mode 100644 index 6a55a241..00000000 --- a/reme/agent/memory/personal/personal_halumem_summarizer.yaml +++ /dev/null @@ -1,96 +0,0 @@ -system_prompt_s1: | - You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}. - - ## Tool Rules - 1. `add_and_retrieve_similar_memory`: Create a memory in the vector store. - - Use this tool to add memories, and it will return the relevant content related to the added memories. - - Use actual names from the conversation (e.g., "Bob likes apples") instead of generic references (e.g., "user likes apples") - - The tool will retrieve similar historical memories via vector search to help you consolidate in Step 2 - - 2. `update_memory`: Update memories in the vector store. - **What to Delete** (via `memory_ids_to_delete`): - - Duplicate memories with identical or highly similar content - - Memories that should be merged into a single consolidated entry - - **What to Add** (via `memories_to_add` with message_time and memory_content): - - For each topic with changes: add ONE consolidated memory that merges related information - - New distinct memories that don't overlap with existing ones - - Updated memories that capture the latest state while preserving temporal evolution - -user_message_s1: | - ## Latest Conversation - Format: round [] : - {context} - - ## Task - ### Step 1: Create Memory - - At this step, you can call the tool multiple times to store memories, or you can call it once to store multiple memories. - - ### Step 2: Update Memory Store - - Update the vector store using `update_memory` to keep it well-organized and consolidated. - - ## Storage Scope (Biographical & Behavioral ONLY) - - Personal Biography: Significant milestones, past experiences, and life events. - - Behavioral Patterns: How the agent reacts, specific actions taken, and recurring habits. - - **EXCLUSION**: DO NOT record objective world facts, general knowledge, or user-specific health/states. - - Extraction & Formatting Rules - - Fact Filtering: Only extract information that builds the biography of **{memory_target}**. - - Subject Splitting: If a conversation mentions multiple subject (e.g., the User's childhood and their Father's career), create separate memory entries for each subject. - - Atomic Content: Each entry should focus on one specific event or trait. Keep descriptions concise to ensure efficient retrieval. - - -system_prompt_s2: | - You are a Profile Agent responsible for managing profiles about {memory_target}. - - ## Tool Rules - 1. Update Profile with `update_profile` - Synchronize profile with new information from the conversation: - - `profile_ids_to_delete`: Remove conflicting, or redundant entries (array of profile IDs). - - `profiles_to_add`: - - `conversation_time`: Time of conversation (format: `YYYY-MM-DD HH:MM:SS`, e.g., `2024-01-15 14:30:00`) - - `profile_content`: Complete, self-contained profile description with full context - Update user profile using `update_profile` based on the conversation and current profile. - - -user_message_s2: | - You are a memory agent managing **{memory_type}** memories about **{memory_target}**. - - ## Latest Conversation: - {context} - - Message format: `round [] : ` (timestamp: YYYY-MM-DD HH:MM:SS). - - **CRITICAL**: Extract ONLY explicitly stated information. DO NOT infer, assume, or fabricate. - - ## Current User Profile: - {user_profile} - - ## Task - ### Step 1: ADD Profile - - Add new, relevant, and up-to-date information to the user profile using `update_profile` (via `profiles_to_add`). - - ### Step 2: DELETE Profile - - Delete outdated, redundant, or resolved states from the user profile using `update_profile` (via `profile_ids_to_delete`). - - ## Storage Scope (Current States ONLY) - - **EXCLUSION PRINCIPLE**: DO NOT record any user *actions*, *requests*, *queries*, or *interactions with the system* (e.g., "asked for code", "solved a puzzle", "requested translation"). These are interaction logs, not user states. - - Identity: Geography, job title, work content, income. - - Background: Education, family, relationships, hobbies, interests, and other personal preferences. - - Temporary States: Physical health (e.g., "Has a cold"), emotional mood, stress levels, and specific prohibitions (e.g., "Cannot drink alcohol due to medication"). - - ## Profile Management Rules - - Subject Splitting (CRITICAL): If the conversation mentions multiple subjects (e.g., the User's job and their Spouse's health), you MUST create separate profile entries for each unique subject. - - Conflict Resolution: Use profile_ids_to_delete to remove outdated, redundant, or resolved states (e.g., if a user is "Recovered," delete the "Illness" entry). - - Each profile entry MUST describe a **persistent or temporary state of the user themselves** (e.g., who they are, what they like, what they’re dealing with), NOT an event they participated in or a request they made. - - ## Profile Format - - **ONLY record what the user EXPLICITLY STATES about themselves as a state or preference.** - - The key in the record represents the category of memory, and the value should record the specific content. For example： - {{"message_time": "YYYY-MM-DD HH:MM:SS", "profile_key": "the category of memory", "profile_value": "content" }} - - When there is no information conflict or outdated information, you don't need to delete any of the memory. If there is no information that meets the requirements, it is also acceptable not to add it. - - ## Forbidden Case - 1.There is no need to record user behavior: {{ "profile_key": "workouts", "profile_content": "confident in new running shoes' suitability for chosen route; they have significantly improved morning jogs"}} - 2. There is no need to record the users' plans or requirements: {{ "profile_key": "plans.vacation", "profile_content": "planning to go to a nearby city for a week and ask for a job change"}} - diff --git a/reme/agent/memory/personal/personal_retriever.py b/reme/agent/memory/personal/personal_retriever.py deleted file mode 100644 index 5b2fa88c..00000000 --- a/reme/agent/memory/personal/personal_retriever.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Personal memory retriever agent for retrieving personal memories through vector search.""" - -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import Role, MemoryType -from ....core.op import BaseTool -from ....core.schema import Message -from ....core.utils import format_messages - - -class PersonalRetriever(BaseMemoryAgent): - """Retrieve personal memories through vector search and history reading.""" - - memory_type: MemoryType = MemoryType.PERSONAL - - async def build_messages(self) -> list[Message]: - if self.context.get("query"): - context = self.context.query - elif self.context.get("messages"): - context = self.description + "\n" + format_messages(self.context.messages) - else: - raise ValueError("input must have either `query` or `messages`") - - read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles") - if read_all_profiles_tool is not None: - all_profiles = await read_all_profiles_tool.call( - memory_target=self.memory_target, - service_context=self.service_context, - ) - else: - all_profiles = "" - - return [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt", - memory_type=self.memory_type.value, - memory_target=self.memory_target, - user_profile=all_profiles, - context=context.strip(), - ), - ), - Message( - role=Role.USER, - content=self.get_prompt("user_message"), - ), - ] - - async def _acting_step( - self, - assistant_message: Message, - tools: list[BaseTool], - step: int, - stage: str = "", - **kwargs, - ) -> tuple[list[BaseTool], list[Message]]: - """Execute tool calls with memory context.""" - return await super()._acting_step( - assistant_message, - tools, - step, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - retrieved_nodes=self.retrieved_nodes, - **kwargs, - ) - - async def execute(self): - result = await super().execute() - answer = result["answer"] - if "MEMORY_NOT_FOUND" in answer: - result["answer"] = "\n".join( - [ - n.format( - include_memory_id=False, - include_when_to_use=False, - include_content=True, - include_message_time=False, - ref_memory_id_key="", - ) - for n in self.retrieved_nodes - ], - ) - - result["retrieved_nodes"] = self.retrieved_nodes - return result diff --git a/reme/agent/memory/personal/personal_retriever.yaml b/reme/agent/memory/personal/personal_retriever.yaml deleted file mode 100644 index 32b85e45..00000000 --- a/reme/agent/memory/personal/personal_retriever.yaml +++ /dev/null @@ -1,66 +0,0 @@ -system_prompt: | - You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}. - - ## User Profile - {user_profile} - - ## User Question - {context} - - ## Multi-Phase Retrieval Strategy - Follow these phases sequentially to gather comprehensive information: - - ### Phase 1: Semantic Search (No Time Filter) - **Tool**: `retrieve_memory` (without time constraints) - **Objective**: Cast a wide net to find potentially relevant memories - **Approach**: - - Execute 3-5 diverse search queries using different formulations: - * Original question verbatim - * Rephrased variations (different wording, synonyms) - * Entity-focused queries (extract and search specific names, places, events) - * Keyword-based searches (core concepts, topics) - * Related context queries (broader themes) - - Review all results before proceeding to next phase - - ### Phase 2: Temporal Search (Optional) - **Tool**: `retrieve_memory` (with time filter) - **When to use**: Only if the user question contains temporal references - **Time Filter Format**: - - Single date: `20200101` - - Date range: `20200101,20200102` (inclusive: 20200101 ≤ time ≤ 20200102) - - Before date: `0,20200102` (up to and including 20200102) - - After date: `20200101,99999999` (from 20200101 onwards) - **Approach**: - - Identify temporal constraints from the user question - - Refine Phase 1 queries with appropriate time filters - - Try multiple time ranges if initial searches yield no results - - ### Phase 3: Deep Dive into History - **Tool**: `read_history` - **When to use**: After exhausting retrieval attempts OR when specific conversation context is needed - **Approach**: - - Extract `history_id` from retrieved memory references - - Prioritize histories that are most relevant or recent - - Read multiple histories if necessary for complete context - - Use this to understand the full conversation surrounding a memory - - ## Response Guidelines - **Critical Rules**: - - Base your answer EXCLUSIVELY on retrieved memories, user profile, and history data - - Never infer, assume, or hallucinate information - - Always cite sources with timestamps: `[timestamp] Memory content` - - Present conflicting information transparently with respective timestamps - - Exhaust all search strategies before concluding information doesn't exist - - **Output Format**: - - When information is found: - - [timestamp] All relevant memory/profile/history content - - - When no information is found after thorough search (5+ queries across phases): - - No relevant information found after exhaustive search using multiple query strategies and retrieval phases. - - -user_message: | - Retrieve relevant memories following the multi-phase strategy outlined above. \ No newline at end of file diff --git a/reme/agent/memory/personal/personal_summarizer.yaml b/reme/agent/memory/personal/personal_summarizer.yaml deleted file mode 100644 index f30dabf6..00000000 --- a/reme/agent/memory/personal/personal_summarizer.yaml +++ /dev/null @@ -1,65 +0,0 @@ -system_prompt_s1: | - You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}. - - ## Latest Conversation - Format: round [] : - {context} - - ## Task - ### Step 1: Create Memory Draft - Create a memory draft in `add_draft_and_retrieve_similar_memory` based on the latest conversation. - - Use actual names from the conversation (e.g., "Bob likes apples") instead of generic references (e.g., "user likes apples") - - Always record memories with real names - - The tool will retrieve similar historical memories via vector search to help you consolidate in Step 2 - - ### Step 2: Update Memory Store - Update the vector store using `update_memory` to keep it well-organized and consolidated: - - **What to Delete** (via `memory_ids_to_delete`): - - Duplicate memories with identical or highly similar content - - Memories that should be merged into a single consolidated entry - - **What to Add** (via `memories_to_add` with message_time and memory_content): - - For each topic with changes: add ONE consolidated memory that merges related information - - New distinct memories that don't overlap with existing ones - - Updated memories that capture the latest state while preserving temporal evolution - - ## Requirements - - Extract only what's explicitly stated—no inferences, assumptions, or fabrications - - Preserve temporal evolution: capture how things change over time within the same topic - - Maintain organization: group related memories by topic and eliminate all redundancy - -user_message_s1: | - Complete the task by following Step 1 and Step 2 in order - -system_prompt_s2: | - You are a Profile Agent responsible for managing profiles about {memory_target}. - - ## Latest Conversation - Format: round [] : - {context} - - ## Task - ### Step 1: Create Profile Draft - Create a profile draft in `add_draft_and_read_all_profiles` based on the latest conversation. - - The tool will return all existing profiles to help you maintain the profile store in Step 2 - - ### Step 2: Update Profile Store - Update the profile store using `update_profile` to keep it well-organized and consolidated: - - **What to Delete** (via `profile_ids_to_delete`): - - Duplicate profiles with identical keys or values - - Conflicting profiles that contradict the new information - - Profiles that should be merged into a single consolidated entry - - **What to Add** (via `profiles_to_add` with message_time, profile_key, and profile_value): - - For each profile key with changes: add ONE consolidated profile that merges related information - - New distinct profiles that don't overlap with existing ones - - Updated profiles that capture the latest state - - ## Requirements - - Extract only what's explicitly stated—no inferences, assumptions, or fabrications - - Maintain organization: group related profiles by key and eliminate all redundancy - -user_message_s2: | - Complete the task by following Step 1 and Step 2 in order diff --git a/reme/agent/memory/personal/personal_v1_summarizer.py b/reme/agent/memory/personal/personal_v1_summarizer.py deleted file mode 100644 index 72f08434..00000000 --- a/reme/agent/memory/personal/personal_v1_summarizer.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Personal memory summarizer agent for two-phase personal memory processing.""" - -from loguru import logger - -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import Role, MemoryType -from ....core.op import BaseTool -from ....core.schema import Message - - -class PersonalV1Summarizer(BaseMemoryAgent): - """Two-phase personal memory processor: retrieve/add memories then update profile.""" - - memory_type: MemoryType = MemoryType.PERSONAL - - async def _build_s1_messages(self) -> list[Message]: - return [ - Message( - role=Role.USER, - content=self.prompt_format( - prompt_name="user_message_s1", - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - ), - ), - ] - - async def _build_s2_messages(self, profiles: str) -> list[Message]: - return [ - Message( - role=Role.USER, - content=self.prompt_format( - prompt_name="user_message_s2", - profiles=profiles, - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - ), - ), - ] - - async def _acting_step( - self, - assistant_message: Message, - tools: list[BaseTool], - step: int, - stage: str = "", - **kwargs, - ) -> tuple[list[BaseTool], list[Message]]: - """Execute tool calls with memory context.""" - return await super()._acting_step( - assistant_message, - tools, - step, - stage=stage, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - history_node=self.history_node, - author=self.author, - retrieved_nodes=self.retrieved_nodes, - **kwargs, - ) - - async def execute(self): - memory_tools = [] - profile_tools = [] - read_all_profiles_tool: BaseTool | None = None - for i, tool in enumerate(self.tools): - tool_name = tool.tool_call.name - if tool_name == "read_all_profiles": - read_all_profiles_tool = tool - elif "_memory" in tool_name: - memory_tools.append(tool) - elif "_profile" in tool_name: - profile_tools.append(tool) - else: - raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={tool_name}") - logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}") - - stage = "s1-memory" - messages_s1 = await self._build_s1_messages() - for i, message in enumerate(messages_s1): - role = message.name or message.role - logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") - tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage) - - if read_all_profiles_tool is not None: - profiles = await read_all_profiles_tool.call( - memory_target=self.memory_target, - service_context=self.service_context, - ) - else: - profiles = "" - - if profile_tools: - stage = "s2-profile" - messages_s2 = await self._build_s2_messages(profiles) - for i, message in enumerate(messages_s2): - role = message.name or message.role - logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") - tools_s2, messages_s2, success_s2 = await self.react(messages_s2, profile_tools, stage=stage) - else: - tools_s2, messages_s2, success_s2 = [], [], True - - answer = (messages_s1[-1].content if success_s1 and messages_s1 else "") + ( - messages_s2[-1].content if success_s2 and messages_s2 else "" - ) - success = success_s1 and success_s2 - messages = messages_s1 + messages_s2 - tools = tools_s1 + tools_s2 - memory_nodes = [] - for tool in tools: - if tool.memory_nodes: - memory_nodes.extend(tool.memory_nodes) - - return { - "answer": answer, - "success": success, - "messages": messages, - "tools": tools, - "memory_nodes": memory_nodes, - } diff --git a/reme/config/cli.yaml b/reme/config/cli.yaml index d368506d..35c9c2f6 100644 --- a/reme/config/cli.yaml +++ b/reme/config/cli.yaml @@ -11,7 +11,7 @@ metadata: llms: default: backend: openai -# model_name: qwen3-235b-a22b-thinking-2507 + # model_name: qwen3-235b-a22b-thinking-2507 model_name: qwen3.5-plus request_interval: 1 @@ -20,12 +20,12 @@ embedding_models: backend: openai model_name: text-embedding-v4 dimensions: 1024 + enable_cache: true -memory_stores: +file_stores: default: backend: chroma # backend: local - db_name: reme.db store_name: reme embedding_model: default fts_enabled: true @@ -34,7 +34,7 @@ memory_stores: file_watchers: default: backend: full - memory_store: default + file_store: default watch_paths: [ ".reme", ".reme/memory" ] suffix_filters: [ ".md" ] recursive: false diff --git a/reme/config/fs.yaml b/reme/config/file.yaml similarity index 71% rename from reme/config/fs.yaml rename to reme/config/file.yaml index e4c191bb..163fe79d 100644 --- a/reme/config/fs.yaml +++ b/reme/config/file.yaml @@ -1,26 +1,23 @@ backend: cmd +working_dir: .reme llms: default: backend: openai - # model_name: qwen3-30b-a3b-instruct-2507 - # model_name: qwen3-30b-a3b-thinking-2507 - model_name: qwen3-235b-a22b-thinking-2507 + model_name: qwen3.5-plus request_interval: 1 -# temperature: 0.0001 embedding_models: default: backend: openai model_name: text-embedding-v4 dimensions: 1024 + enable_cache: true -memory_stores: +file_stores: default: - # backend: sqlite backend: chroma # backend: local - db_name: reme.db store_name: reme embedding_model: default fts_enabled: true @@ -29,7 +26,7 @@ memory_stores: file_watchers: default: backend: full - memory_store: default + file_store: default watch_paths: [ ".reme", ".reme/memory" ] suffix_filters: [ ".md" ] recursive: false diff --git a/reme/config/service.yaml b/reme/config/service.yaml new file mode 100644 index 00000000..7d1b531a --- /dev/null +++ b/reme/config/service.yaml @@ -0,0 +1,64 @@ +backend: http +working_dir: .reme +thread_pool_max_workers: 64 + +mcp: + transport: sse + host: "0.0.0.0" + port: 8001 + +http: + host: "0.0.0.0" + port: 8002 + timeout_keep_alive: 600 + limit_concurrency: 64 + +flows: + test: + flow_content: TestOp() + description: "test" + +# curl -X POST http://localhost:8002/simple_chat \ +# -H "Content-Type: application/json" \ +# -d '{ +# "query": "hello" +# }' + simple_chat: + flow_content: SimpleChat() + description: "test" + + stream_chat: + flow_content: StreamChat() + description: "test" + stream: true + +llms: + default: + backend: openai + model_name: qwen3.5-plus + request_interval: 1 +# temperature: 0.0001 + +embedding_models: + default: + backend: openai + model_name: text-embedding-v4 + dimensions: 1024 + enable_cache: false + +vector_stores: + default: + backend: chroma + # backend: local + collection_name: reme + embedding_model: default + +token_counters: + default: + backend: base + + hf: + backend: hf + model_name: Qwen/Qwen3-Coder-30B-A3B-Instruct + use_mirror: true + diff --git a/reme/config/default.yaml b/reme/config/vector.yaml similarity index 94% rename from reme/config/default.yaml rename to reme/config/vector.yaml index 5ce200b7..540e5df9 100644 --- a/reme/config/default.yaml +++ b/reme/config/vector.yaml @@ -1,21 +1,5 @@ -backend: http -thread_pool_max_workers: 64 - -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -http: - host: "0.0.0.0" - port: 8002 - timeout_keep_alive: 600 - limit_concurrency: 64 - -flows: - test: - flow_content: TestOp() - description: "test" +backend: cmd +working_dir: .reme retrieve_task_memory: flow_content: BuildQuery() >> MemoryRetrieval() >> RerankMemory() >> RewriteMemory() @@ -162,8 +146,7 @@ flows: llms: default: backend: openai - model_name: qwen3-30b-a3b-instruct-2507 - # model_name: qwen3-30b-a3b-thinking-2507 + model_name: qwen3.5-plus request_interval: 1 # temperature: 0.0001 @@ -184,13 +167,14 @@ embedding_models: backend: openai model_name: text-embedding-v4 dimensions: 1024 + enable_cache: false vector_stores: default: backend: chroma # backend: local - embedding_model: default collection_name: reme + embedding_model: default token_counters: default: diff --git a/reme/core/__init__.py b/reme/core/__init__.py index 4128cbd9..5872e2ad 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -1,12 +1,11 @@ """Core""" -from . import context from . import embedding from . import enumeration +from . import file_store from . import file_watcher from . import flow from . import llm -from . import memory_store from . import op from . import schema from . import service @@ -14,22 +13,33 @@ from . import token_counter from . import utils from . import vector_store from .application import Application -from .context import R +from .base_dict import BaseDict +from .prompt_handler import PromptHandler +from .registry_factory import R, Registry, RegistryFactory +from .runtime_context import RuntimeContext +from .service_context import ServiceContext __all__ = [ - "context", + # Submodules "embedding", "enumeration", "file_watcher", "flow", "llm", - "memory_store", + "file_store", "op", "schema", "service", "token_counter", "utils", "vector_store", + # Classes "Application", + "BaseDict", + "PromptHandler", "R", + "Registry", + "RegistryFactory", + "RuntimeContext", + "ServiceContext", ] diff --git a/reme/core/application.py b/reme/core/application.py index 0d85fc29..49537807 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -6,15 +6,17 @@ from pathlib import Path from loguru import logger -from .context import PromptHandler, ServiceContext, R from .embedding import BaseEmbeddingModel +from .file_store import BaseFileStore from .file_watcher import BaseFileWatcher from .flow import BaseFlow from .llm import BaseLLM -from .memory_store import BaseMemoryStore +from .prompt_handler import PromptHandler +from .registry_factory import R from .schema import Response, ServiceConfig +from .service_context import ServiceContext from .token_counter import BaseTokenCounter -from .utils import execute_stream_task, PydanticConfigParser, init_logger, print_logo, MCPClient +from .utils import execute_stream_task, PydanticConfigParser, init_logger, MCPClient, print_logo from .vector_store import BaseVectorStore @@ -36,7 +38,7 @@ class Application: 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_file_store_config: dict | None = None, default_token_counter_config: dict | None = None, default_file_watcher_config: dict | None = None, **kwargs, @@ -56,12 +58,17 @@ class Application: 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_file_store_config=default_file_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_config.language) + + # NOTE: flows are initialized here to start service! + self.init_flows() + self._started: bool = False @classmethod @@ -71,40 +78,8 @@ class Application: await instance.start() return instance - @property - def service_config(self) -> ServiceConfig: - """Get the service configuration.""" - return self.service_context.service_config - - async def start(self): - """Start the service context by initializing all configured components.""" - if self._started: - logger.warning("Application has already started.") - return self - - init_logger(log_to_console=self.service_config.log_to_console) - logger.info(f"Init ReMe with config: {self.service_config.model_dump_json()}") - - working_path = Path(self.service_config.working_dir) - working_path.mkdir(parents=True, exist_ok=True) - - if self.service_config.enable_logo: - print_logo(service_config=self.service_config) - - if self.service_config.ray_max_workers > 1: - import ray - - if not ray.is_initialized(): - ray.init(num_cpus=self.service_config.ray_max_workers) - - if ( - self.service_context.thread_pool is None - or self.service_context.thread_pool._shutdown # pylint: disable=protected-access - ): - self.service_context.thread_pool = ThreadPoolExecutor( - max_workers=self.service_config.thread_pool_max_workers, - ) - + def init_flows(self): + """Initialize flows.""" expression_flow_cls = None for name, flow_cls in R.flows.items(): if not self._filter_flows(name): @@ -129,12 +104,56 @@ class Application: else: logger.info("No expression flow found, please check your configuration.") + def _filter_flows(self, name: str) -> bool: + """Filter flows based on enabled_flows and disabled_flows configuration.""" + if self.service_config.enabled_flows: + return name in self.service_config.enabled_flows + elif self.service_config.disabled_flows: + return name not in self.service_config.disabled_flows + else: + return True + + @property + def service_config(self) -> ServiceConfig: + """Get the service configuration.""" + return self.service_context.service_config + + async def start(self): + """Start the service context by initializing all configured components.""" + if self._started: + logger.warning("Application has already started.") + return self + + init_logger(log_to_console=self.service_config.log_to_console) + logger.info(f"Init ReMe with config: {self.service_config.model_dump_json()}") + + working_path = Path(self.service_config.working_dir) + working_path.mkdir(parents=True, exist_ok=True) + + if self.service_config.ray_max_workers > 1: + import ray + + if not ray.is_initialized(): + ray.init(num_cpus=self.service_config.ray_max_workers) + + if ( + self.service_context.thread_pool is None + or self.service_context.thread_pool._shutdown # pylint: disable=protected-access + ): + self.service_context.thread_pool = ThreadPoolExecutor( + max_workers=self.service_config.thread_pool_max_workers, + ) + + if self.service_context.service_config.enable_logo: + print_logo(service_config=self.service_config) + for name, config in self.service_config.llms.items(): if config.backend not in R.llms: logger.warning(f"LLM backend {config.backend} is not supported.") else: config_dict = config.model_dump(exclude={"backend"}) self.service_context.llms[name] = R.llms[config.backend](**config_dict) + await self.service_context.llms[name].start() for name, config in self.service_config.embedding_models.items(): if config.backend not in R.embedding_models: @@ -143,6 +162,7 @@ class Application: config_dict = config.model_dump(exclude={"backend"}) config_dict["cache_dir"] = working_path / "embedding_cache" self.service_context.embedding_models[name] = R.embedding_models[config.backend](**config_dict) + await self.service_context.embedding_models[name].start() for name, config in self.service_config.token_counters.items(): if config.backend not in R.token_counters: @@ -159,33 +179,32 @@ class Application: config_dict.update( { "embedding_model": self.service_context.embedding_models[config.embedding_model], - "thread_pool": self.service_context.thread_pool, + "db_path": working_path / "vector_store", }, ) self.service_context.vector_stores[name] = R.vector_stores[config.backend](**config_dict) - await self.service_context.vector_stores[name].create_collection(config.collection_name) + await self.service_context.vector_stores[name].start() - for name, config in self.service_config.memory_stores.items(): - if config.backend not in R.memory_stores: - logger.warning(f"Memory store backend {config.backend} is not supported.") + for name, config in self.service_config.file_stores.items(): + if config.backend not in R.file_stores: + logger.warning(f"File store backend {config.backend} is not supported.") else: config_dict = config.model_dump(exclude={"backend", "embedding_model"}) config_dict.update( { "embedding_model": self.service_context.embedding_models[config.embedding_model], - "thread_pool": self.service_context.thread_pool, - "db_path": working_path / "memory_store", + "db_path": working_path / "file_store", }, ) - self.service_context.memory_stores[name] = R.memory_stores[config.backend](**config_dict) - await self.service_context.memory_stores[name].start() + self.service_context.file_stores[name] = R.file_stores[config.backend](**config_dict) + await self.service_context.file_stores[name].start() for name, config in self.service_config.file_watchers.items(): if config.backend not in R.file_watchers: logger.warning(f"File watcher backend {config.backend} is not supported.") else: - config_dict = config.model_dump(exclude={"backend", "memory_store"}) - config_dict["memory_store"] = self.service_context.memory_stores[config.memory_store] + config_dict = config.model_dump(exclude={"backend", "file_store"}) + config_dict["file_store"] = self.service_context.file_stores[config.file_store] self.service_context.file_watchers[name] = R.file_watchers[config.backend](**config_dict) await self.service_context.file_watchers[name].start() @@ -195,15 +214,6 @@ class Application: self._started = True return self - def _filter_flows(self, name: str) -> bool: - """Filter flows based on enabled_flows and disabled_flows configuration.""" - if self.service_config.enabled_flows: - return name in self.service_config.enabled_flows - elif self.service_config.disabled_flows: - return name not in self.service_config.disabled_flows - else: - return True - async def prepare_mcp_servers(self): """Prepare and initialize MCP server connections.""" mcp_client = MCPClient(config={"mcpServers": self.service_config.mcp_servers}) @@ -228,9 +238,9 @@ class Application: logger.info(f"Closing vector store: {name}") await vector_store.close() - for name, memory_store in self.service_context.memory_stores.items(): - logger.info(f"Closing memory store: {name}") - await memory_store.close() + for name, file_store in self.service_context.file_stores.items(): + logger.info(f"Closing file store: {name}") + await file_store.close() for name, file_watcher in self.service_context.file_watchers.items(): logger.info(f"Closing file watcher: {name}") @@ -327,13 +337,13 @@ class Application: return self.service_context.vector_stores.get(name) @property - def default_memory_store(self) -> BaseMemoryStore: - """Get the default memory store instance.""" - return self.service_context.memory_stores.get("default") + def default_file_store(self) -> BaseFileStore: + """Get the default file store instance.""" + return self.service_context.file_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) + def get_file_store(self, name: str): + """Get a file store instance by name.""" + return self.service_context.file_stores.get(name) @property def default_file_watcher(self) -> BaseFileWatcher: @@ -358,7 +368,7 @@ class Application: import warnings warnings.filterwarnings("ignore", category=DeprecationWarning) - service = R.services[self.service_config.backend](service_context=self.service_context) + service = R.services[self.service_config.backend](app=self) service.run() async def reset_default_collection(self, collection_name: str): diff --git a/reme/core/context/base_context.py b/reme/core/base_dict.py similarity index 97% rename from reme/core/context/base_context.py rename to reme/core/base_dict.py index dabd8cdb..b96b4c24 100644 --- a/reme/core/context/base_context.py +++ b/reme/core/base_dict.py @@ -6,7 +6,7 @@ _KT = TypeVar("_KT") _VT = TypeVar("_VT") -class BaseContext(dict, Generic[_KT, _VT]): +class BaseDict(dict, Generic[_KT, _VT]): """A dictionary subclass that enables accessing and modifying keys as attributes.""" def __getattr__(self, name: str) -> _VT: diff --git a/reme/core/context/__init__.py b/reme/core/context/__init__.py deleted file mode 100644 index 7bbd5869..00000000 --- a/reme/core/context/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -"""context""" - -from .base_context import BaseContext -from .prompt_handler import PromptHandler -from .registry_factory import R -from .runtime_context import RuntimeContext -from .service_context import ServiceContext - -__all__ = [ - "BaseContext", - "PromptHandler", - "R", - "RuntimeContext", - "ServiceContext", -] diff --git a/reme/core/embedding/__init__.py b/reme/core/embedding/__init__.py index 71fc4f43..6431703e 100644 --- a/reme/core/embedding/__init__.py +++ b/reme/core/embedding/__init__.py @@ -3,7 +3,7 @@ from .base_embedding_model import BaseEmbeddingModel from .openai_embedding_model import OpenAIEmbeddingModel from .openai_embedding_model_sync import OpenAIEmbeddingModelSync -from ..context import R +from ..registry_factory import R __all__ = [ "BaseEmbeddingModel", diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index da93eb8d..d29dc3fc 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -14,8 +14,7 @@ from pathlib import Path from loguru import logger -from ..schema import VectorNode -from ..schema.memory_chunk import MemoryChunk +from ..schema import VectorNode, MemoryChunk class BaseEmbeddingModel(ABC): @@ -76,9 +75,6 @@ class BaseEmbeddingModel(ABC): self.cache_path: Path = Path(self.cache_dir) self.cache_path.mkdir(parents=True, exist_ok=True) - # Load cache from disk if available - self._load_cache() - @property def api_key(self) -> str | None: """Get API key from environment variable.""" @@ -531,6 +527,14 @@ class BaseEmbeddingModel(ABC): logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(chunks)} chunks") return chunks + def start_sync(self): + """Synchronously initialize resources and load cache.""" + self._load_cache() + + async def start(self): + """Asynchronously initialize resources and load cache.""" + self._load_cache() + def close_sync(self): """Synchronously release resources and close connections.""" self._save_cache() diff --git a/reme/core/embedding/openai_embedding_model.py b/reme/core/embedding/openai_embedding_model.py index a068b35d..efc0133d 100644 --- a/reme/core/embedding/openai_embedding_model.py +++ b/reme/core/embedding/openai_embedding_model.py @@ -45,6 +45,10 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel): result_emb[emb.index] = emb.embedding return result_emb + async def start(self): + """Initialize the asynchronous OpenAI embedding model and load cache.""" + await super().start() + async def close(self): """Close the asynchronous OpenAI client and release network resources.""" if self._client is not None: diff --git a/reme/core/enumeration/registry_enum.py b/reme/core/enumeration/registry_enum.py index aec8ca93..68450e6a 100644 --- a/reme/core/enumeration/registry_enum.py +++ b/reme/core/enumeration/registry_enum.py @@ -15,8 +15,8 @@ class RegistryEnum(str, Enum): # Databases or storage systems for vector search VECTOR_STORE = "vector_store" - # Databases or storage systems for long-term memory storage - MEMORY_STORE = "memory_store" + # Databases or storage systems for long-term file storage + FILE_STORE = "file_store" # Atomic operations or functional units OP = "op" diff --git a/reme/core/file_store/__init__.py b/reme/core/file_store/__init__.py new file mode 100644 index 00000000..1358df52 --- /dev/null +++ b/reme/core/file_store/__init__.py @@ -0,0 +1,23 @@ +"""File store module for persistent memory management. + +This module provides storage backends for memory chunks and file metadata, +including SQLite-based, ChromaDB-based, and pure-Python local implementations +with vector and full-text search. +""" + +from .base_file_store import BaseFileStore +from .chroma_file_store import ChromaFileStore +from .local_file_store import LocalFileStore +from .sqlite_file_store import SqliteFileStore +from ..registry_factory import R + +__all__ = [ + "BaseFileStore", + "ChromaFileStore", + "LocalFileStore", + "SqliteFileStore", +] + +R.file_stores.register("sqlite")(SqliteFileStore) +R.file_stores.register("chroma")(ChromaFileStore) +R.file_stores.register("local")(LocalFileStore) diff --git a/reme/core/memory_store/base_memory_store.py b/reme/core/file_store/base_file_store.py similarity index 85% rename from reme/core/memory_store/base_memory_store.py rename to reme/core/file_store/base_file_store.py index c9083a3c..9c79f81f 100644 --- a/reme/core/memory_store/base_memory_store.py +++ b/reme/core/file_store/base_file_store.py @@ -1,27 +1,22 @@ -"""Base storage interface for memory manager.""" +"""Base storage interface for file store.""" -import asyncio import re from abc import ABC, abstractmethod -from concurrent.futures import ThreadPoolExecutor -from functools import partial from pathlib import Path -from typing import Callable from ..embedding import BaseEmbeddingModel from ..enumeration import MemorySource from ..schema import FileMetadata, MemoryChunk, MemorySearchResult -class BaseMemoryStore(ABC): - """Abstract base class for memory storage backends.""" +class BaseFileStore(ABC): + """Abstract base class for file storage backends.""" def __init__( self, store_name: str, db_path: str | Path, - thread_pool: ThreadPoolExecutor, - embedding_model: BaseEmbeddingModel, + embedding_model: BaseEmbeddingModel | None = None, vector_enabled: bool = False, fts_enabled: bool = True, **kwargs, @@ -36,11 +31,14 @@ class BaseMemoryStore(ABC): if not vector_enabled and not fts_enabled: raise ValueError("At least one of vector_enabled or fts_enabled must be True.") + # Ensure embedding_model is provided when vector search is enabled + if vector_enabled and embedding_model is None: + raise ValueError("embedding_model is required when vector_enabled is True.") + self.store_name: str = store_name self.db_path: Path = Path(db_path) self.db_path.mkdir(parents=True, exist_ok=True) - self.thread_pool: ThreadPoolExecutor = thread_pool - self.embedding_model: BaseEmbeddingModel = embedding_model + self.embedding_model: BaseEmbeddingModel | None = embedding_model self.vector_enabled: bool = vector_enabled self.fts_enabled: bool = fts_enabled self.kwargs: dict = kwargs @@ -48,6 +46,8 @@ class BaseMemoryStore(ABC): @property def embedding_dim(self) -> int: """Get the embedding model's dimensionality.""" + if self.embedding_model is None: + return 1024 return self.embedding_model.dimensions def _get_mock_embedding(self) -> list[float]: @@ -82,11 +82,6 @@ class BaseMemoryStore(ABC): return chunks return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs) - async def _run_sync_in_executor(self, sync_func: Callable, *args, **kwargs): - """Run a synchronous function in the context-defined thread pool executor.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor(self.thread_pool, partial(sync_func, *args, **kwargs)) # noqa - @abstractmethod async def start(self): """Initialize the storage backend.""" @@ -115,6 +110,18 @@ class BaseMemoryStore(ABC): async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None: """Get full file metadata with statistics.""" + @abstractmethod + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks. + + This is useful for incremental updates where only metadata needs to be updated + (e.g., after adding/removing chunks in delta file watcher). + + Args: + file_meta: Updated file metadata (hash, mtime_ms, size, chunk_count) + source: Memory source + """ + @abstractmethod async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]: """Get all chunks for a file.""" diff --git a/reme/core/memory_store/chroma_memory_store.py b/reme/core/file_store/chroma_file_store.py similarity index 95% rename from reme/core/memory_store/chroma_memory_store.py rename to reme/core/file_store/chroma_file_store.py index ee7c4911..390a2bc4 100644 --- a/reme/core/memory_store/chroma_memory_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -1,4 +1,4 @@ -"""ChromaDB storage backend for memory index.""" +"""ChromaDB storage backend for file store.""" import json import time @@ -6,7 +6,7 @@ from pathlib import Path from loguru import logger -from .base_memory_store import BaseMemoryStore +from .base_file_store import BaseFileStore from ..enumeration import MemorySource from ..schema import FileMetadata, MemoryChunk, MemorySearchResult @@ -21,10 +21,10 @@ except ImportError: Settings = None -class ChromaMemoryStore(BaseMemoryStore): - """ChromaDB memory storage with vector and full-text search. +class ChromaFileStore(BaseFileStore): + """ChromaDB file storage with vector and full-text search. - Inherits embedding methods from BaseMemoryStore: + Inherits embedding methods from BaseFileStore: - get_chunk_embedding / get_chunk_embeddings (async) - get_embedding / get_embeddings (async) @@ -40,7 +40,7 @@ class ChromaMemoryStore(BaseMemoryStore): ): if not CHROMADB_AVAILABLE: raise ImportError( - "chromadb package is required for ChromaMemoryStore. Install it with: pip install chromadb", + "chromadb package is required for ChromaFileStore. Install it with: pip install chromadb", ) super().__init__(**kwargs) @@ -287,6 +287,19 @@ class ChromaMemoryStore(BaseMemoryStore): return None return self._metadata_cache[source.value].get(path) + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks.""" + if source.value not in self._metadata_cache: + self._metadata_cache[source.value] = {} + + self._metadata_cache[source.value][file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + ) + async def get_file_chunks( self, path: str, diff --git a/reme/core/memory_store/local_memory_store.py b/reme/core/file_store/local_file_store.py similarity index 74% rename from reme/core/memory_store/local_memory_store.py rename to reme/core/file_store/local_file_store.py index bee6882b..37df17f4 100644 --- a/reme/core/memory_store/local_memory_store.py +++ b/reme/core/file_store/local_file_store.py @@ -1,40 +1,24 @@ -"""Pure-Python in-memory storage backend for memory index, with JSON file persistence.""" +"""Pure-Python in-memory storage backend for file store, with JSON file persistence.""" import json -import time -from dataclasses import dataclass from pathlib import Path +import numpy as np from loguru import logger -from .base_memory_store import BaseMemoryStore +from .base_file_store import BaseFileStore from ..enumeration import MemorySource from ..schema import FileMetadata, MemoryChunk, MemorySearchResult -from ..utils.common_utils import cosine_similarity +from ..utils.common_utils import batch_cosine_similarity -@dataclass -class _ChunkRecord: - """Internal in-memory representation of a stored chunk.""" - - id: str - path: str - source: str - start_line: int - end_line: int - text: str - hash: str - embedding: list[float] | None - updated_at: int - - -class LocalMemoryStore(BaseMemoryStore): - """Pure-Python in-memory memory storage with JSONL file persistence. +class LocalFileStore(BaseFileStore): + """Pure-Python in-memory file storage with JSONL file persistence. No external dependencies required. All data lives in Python dicts; writes are persisted to JSONL files on disk so state survives restarts. - Inherits embedding methods from BaseMemoryStore: + Inherits embedding methods from BaseFileStore: - get_chunk_embedding / get_chunk_embeddings (async) - get_embedding / get_embeddings (async) @@ -48,9 +32,9 @@ class LocalMemoryStore(BaseMemoryStore): super().__init__(**kwargs) self._started: bool = False # In-memory indexes - self._chunks: dict[str, _ChunkRecord] = {} + self._chunks: dict[str, MemoryChunk] = {} self._files: dict[str, dict[str, FileMetadata]] = {} # source -> path -> meta - # Persistence paths (mirror ChromaMemoryStore convention) + # Persistence paths (mirror ChromaFileStore convention) self._chunks_file: Path = self.db_path / f"{self.store_name}_chunks.jsonl" self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json" @@ -69,8 +53,8 @@ class LocalMemoryStore(BaseMemoryStore): if not line: continue rec = json.loads(line) - chunk_id = rec["id"] - self._chunks[chunk_id] = _ChunkRecord(**rec) + chunk = MemoryChunk.model_validate(rec) + self._chunks[chunk.id] = chunk logger.debug(f"Loaded {len(self._chunks)} chunks from {self._chunks_file}") except Exception as e: logger.warning(f"Failed to load chunks from {self._chunks_file}: {e}") @@ -79,18 +63,8 @@ class LocalMemoryStore(BaseMemoryStore): """Persist chunks to JSONL file.""" try: lines = [] - for rec in self._chunks.values(): - chunk_dict = { - "id": rec.id, - "path": rec.path, - "source": rec.source, - "start_line": rec.start_line, - "end_line": rec.end_line, - "text": rec.text, - "hash": rec.hash, - "embedding": rec.embedding, - "updated_at": rec.updated_at, - } + for chunk in self._chunks.values(): + chunk_dict = chunk.model_dump(mode="json") lines.append(json.dumps(chunk_dict, ensure_ascii=False)) data = "\n".join(lines) self._chunks_file.write_text(data, encoding="utf-8") @@ -145,7 +119,7 @@ class LocalMemoryStore(BaseMemoryStore): await self._load_metadata() await self._load_chunks() logger.info( - f"LocalMemoryStore '{self.store_name}' ready: " + f"LocalFileStore '{self.store_name}' ready: " f"{len(self._chunks)} chunks, metadata at {self._metadata_file}", ) @@ -177,19 +151,8 @@ class LocalMemoryStore(BaseMemoryStore): # Batch generate embeddings (base class returns mock embeddings when vector_enabled=False) chunks = await self.get_chunk_embeddings(chunks) - now = int(time.time() * 1000) for chunk in chunks: - self._chunks[chunk.id] = _ChunkRecord( - id=chunk.id, - path=file_meta.path, - source=source.value, - start_line=chunk.start_line, - end_line=chunk.end_line, - text=chunk.text, - hash=chunk.hash, - embedding=chunk.embedding, - updated_at=now, - ) + self._chunks[chunk.id] = chunk if source.value not in self._files: self._files[source.value] = {} @@ -203,7 +166,7 @@ class LocalMemoryStore(BaseMemoryStore): async def delete_file(self, path: str, source: MemorySource) -> None: """Delete file and all its chunks.""" - to_delete = [cid for cid, rec in self._chunks.items() if rec.path == path and rec.source == source.value] + to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path and chunk.source == source] for cid in to_delete: del self._chunks[cid] @@ -218,10 +181,12 @@ class LocalMemoryStore(BaseMemoryStore): for cid in chunk_ids: self._chunks.pop(cid, None) - # Recalculate chunk_count in file metadata - for source_meta in self._files.values(): + # Recalculate chunk_count in file metadata (per source) + for source_key, source_meta in self._files.items(): if path in source_meta: - source_meta[path].chunk_count = sum(1 for rec in self._chunks.values() if rec.path == path) + source_meta[path].chunk_count = sum( + 1 for chunk in self._chunks.values() if chunk.path == path and chunk.source.value == source_key + ) async def upsert_chunks( self, @@ -234,19 +199,8 @@ class LocalMemoryStore(BaseMemoryStore): chunks = await self.get_chunk_embeddings(chunks) - now = int(time.time() * 1000) for chunk in chunks: - self._chunks[chunk.id] = _ChunkRecord( - id=chunk.id, - path=chunk.path, - source=source.value, - start_line=chunk.start_line, - end_line=chunk.end_line, - text=chunk.text, - hash=chunk.hash, - embedding=chunk.embedding, - updated_at=now, - ) + self._chunks[chunk.id] = chunk # ------------------------------------------------------------------ # Read operations @@ -264,27 +218,28 @@ class LocalMemoryStore(BaseMemoryStore): """Get file metadata.""" return self._files.get(source.value, {}).get(path) + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks.""" + if source.value not in self._files: + self._files[source.value] = {} + + self._files[source.value][file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + ) + async def get_file_chunks( self, path: str, source: MemorySource, ) -> list[MemoryChunk]: """Get all chunks for a file, sorted by start_line.""" - records = [rec for rec in self._chunks.values() if rec.path == path and rec.source == source.value] - records.sort(key=lambda r: r.start_line) - return [ - MemoryChunk( - id=rec.id, - path=rec.path, - source=MemorySource(rec.source), - start_line=rec.start_line, - end_line=rec.end_line, - text=rec.text, - hash=rec.hash, - embedding=rec.embedding, - ) - for rec in records - ] + chunks = [chunk for chunk in self._chunks.values() if chunk.path == path and chunk.source == source] + chunks.sort(key=lambda c: c.start_line) + return chunks # ------------------------------------------------------------------ # Search @@ -304,26 +259,32 @@ class LocalMemoryStore(BaseMemoryStore): if not query_embedding: return [] - source_values = {s.value for s in sources} if sources else None - results = [] - for rec in self._chunks.values(): - if source_values and rec.source not in source_values: - continue - if not rec.embedding: - continue + # Collect candidate chunks with embeddings + candidates = [ + chunk for chunk in self._chunks.values() if (not sources or chunk.source in sources) and chunk.embedding + ] - similarity = cosine_similarity(query_embedding, rec.embedding) - results.append( - MemorySearchResult( - path=rec.path, - start_line=rec.start_line, - end_line=rec.end_line, - score=similarity, - snippet=rec.text, - source=MemorySource(rec.source), - raw_metric=1.0 - similarity, # distance equivalent - ), + if not candidates: + return [] + + # Build embedding matrix and compute similarities in batch + query_array = np.array([query_embedding]) # Shape: (1, emb_size) + chunk_embeddings = np.array([chunk.embedding for chunk in candidates]) # Shape: (n, emb_size) + similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] # Shape: (n,) + + # Build results + results = [ + MemorySearchResult( + path=chunk.path, + start_line=chunk.start_line, + end_line=chunk.end_line, + score=float(similarity), + snippet=chunk.text, + source=chunk.source, + raw_metric=1.0 - float(similarity), ) + for chunk, similarity in zip(candidates, similarities) + ] results.sort(key=lambda r: r.score, reverse=True) return results[:limit] @@ -346,13 +307,12 @@ class LocalMemoryStore(BaseMemoryStore): words_lower = [w.lower() for w in words] n_words = len(words) - source_values = {s.value for s in sources} if sources else None results = [] - for rec in self._chunks.values(): - if source_values and rec.source not in source_values: + for chunk in self._chunks.values(): + if sources and chunk.source not in sources: continue - text_lower = rec.text.lower() + text_lower = chunk.text.lower() match_count = sum(1 for w in words_lower if w in text_lower) if match_count == 0: continue @@ -364,12 +324,12 @@ class LocalMemoryStore(BaseMemoryStore): results.append( MemorySearchResult( - path=rec.path, - start_line=rec.start_line, - end_line=rec.end_line, + path=chunk.path, + start_line=chunk.start_line, + end_line=chunk.end_line, score=score, - snippet=rec.text, - source=MemorySource(rec.source), + snippet=chunk.text, + source=chunk.source, ), ) @@ -452,18 +412,21 @@ class LocalMemoryStore(BaseMemoryStore): merged: dict[str, MemorySearchResult] = {} for result in vector: - result.score = result.score * vector_weight + result.metadata["_weighted_score"] = result.score * vector_weight merged[result.merge_key] = result for result in keyword: key = result.merge_key if key in merged: - merged[key].score += result.score * text_weight + merged[key].metadata["_weighted_score"] += result.score * text_weight else: - result.score = result.score * text_weight + result.metadata["_weighted_score"] = result.score * text_weight merged[key] = result results = list(merged.values()) + for r in results: + r.score = r.metadata.pop("_weighted_score") + results.sort(key=lambda r: r.score, reverse=True) return results @@ -473,4 +436,4 @@ class LocalMemoryStore(BaseMemoryStore): self._files.clear() await self._save_chunks() await self._save_metadata() - logger.info(f"Cleared all data from LocalMemoryStore '{self.store_name}'") + logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'") diff --git a/reme/core/memory_store/sqlite_memory_store.py b/reme/core/file_store/sqlite_file_store.py similarity index 97% rename from reme/core/memory_store/sqlite_memory_store.py rename to reme/core/file_store/sqlite_file_store.py index e05c6adc..7fd49762 100644 --- a/reme/core/memory_store/sqlite_memory_store.py +++ b/reme/core/file_store/sqlite_file_store.py @@ -1,4 +1,4 @@ -"""SQLite storage backend for memory index.""" +"""SQLite storage backend for file store.""" import json import sqlite3 @@ -7,15 +7,15 @@ import time from loguru import logger -from .base_memory_store import BaseMemoryStore +from .base_file_store import BaseFileStore from ..enumeration import MemorySource from ..schema import FileMetadata, MemoryChunk, MemorySearchResult -class SqliteMemoryStore(BaseMemoryStore): - """SQLite memory storage with vector and full-text search. +class SqliteFileStore(BaseFileStore): + """SQLite file storage with vector and full-text search. - Inherits embedding methods from BaseMemoryStore: + Inherits embedding methods from BaseFileStore: - get_chunk_embedding / get_chunk_embeddings (async) - get_chunk_embedding_sync / get_chunk_embeddings_sync (sync) - get_embedding / get_embeddings (async) @@ -467,6 +467,24 @@ class SqliteMemoryStore(BaseMemoryStore): finally: cursor.close() + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks.""" + cursor = self.conn.cursor() + try: + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.files_table_name} (path, source, hash, mtime, size) + VALUES (?, ?, ?, ?, ?) + """, + (file_meta.path, source.value, file_meta.hash, file_meta.mtime_ms, file_meta.size), + ) + self.conn.commit() + except Exception as e: + logger.error(f"Failed to update file metadata for {file_meta.path}: {e}") + raise + finally: + cursor.close() + async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]: """Get all chunks for a file.""" cursor = self.conn.cursor() diff --git a/reme/core/file_watcher/__init__.py b/reme/core/file_watcher/__init__.py index e1d72fdb..020d0cfa 100644 --- a/reme/core/file_watcher/__init__.py +++ b/reme/core/file_watcher/__init__.py @@ -7,7 +7,7 @@ and updating memory stores accordingly. from .base_file_watcher import BaseFileWatcher from .delta_file_watcher import DeltaFileWatcher from .full_file_watcher import FullFileWatcher -from ..context import R +from ..registry_factory import R __all__ = [ "BaseFileWatcher", diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 7e8c0bd9..5030a9cb 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -13,7 +13,7 @@ from loguru import logger from watchfiles import awatch, Change from ..enumeration import MemorySource -from ..memory_store import BaseMemoryStore +from ..file_store import BaseFileStore class BaseFileWatcher: @@ -32,7 +32,7 @@ class BaseFileWatcher: debounce: int = 500, # Millisecond debounce chunk_tokens: int = 400, chunk_overlap: int = 80, - memory_store: BaseMemoryStore | None = None, + file_store: BaseFileStore | None = None, callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None, scan_on_start: bool = False, **kwargs, @@ -47,7 +47,7 @@ class BaseFileWatcher: debounce: Debounce time in milliseconds chunk_tokens: Token size for chunking chunk_overlap: Overlap size for chunks - memory_store: Memory store instance + file_store: File store instance callback: Callback function for changes scan_on_start: If True, scan existing files on start and trigger on_changes with Change.added **kwargs: Additional keyword arguments @@ -58,7 +58,7 @@ class BaseFileWatcher: self.debounce: int = debounce self.chunk_tokens: int = chunk_tokens self.chunk_overlap: int = chunk_overlap - self.memory_store: BaseMemoryStore = memory_store + self.file_store: BaseFileStore = file_store self.callback = callback self.scan_on_start: bool = scan_on_start self.kwargs: dict = kwargs @@ -140,9 +140,9 @@ class BaseFileWatcher: else: logger.info("[SCAN_ON_START] No existing files found matching watch criteria") - files: list[str] = await self.memory_store.list_files(MemorySource.MEMORY) + files: list[str] = await self.file_store.list_files(MemorySource.MEMORY) for file_path in files: - chunks = await self.memory_store.get_file_chunks(file_path, MemorySource.MEMORY) + chunks = await self.file_store.get_file_chunks(file_path, MemorySource.MEMORY) logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks") async def _watch_loop(self): diff --git a/reme/core/file_watcher/delta_file_watcher.py b/reme/core/file_watcher/delta_file_watcher.py index 0c9cae7c..6148bd07 100644 --- a/reme/core/file_watcher/delta_file_watcher.py +++ b/reme/core/file_watcher/delta_file_watcher.py @@ -158,17 +158,17 @@ class DeltaFileWatcher(BaseFileWatcher): ) if chunks: - chunks = await self.memory_store.get_chunk_embeddings(chunks) + chunks = await self.file_store.get_chunk_embeddings(chunks) file_meta.chunk_count = len(chunks) - await self.memory_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) + await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) logger.info(f"File added: {path} ({len(chunks)} chunks)") else: logger.warning(f"No chunks generated for new file {path}") elif change_type == Change.modified: # Get existing data - old_chunks = await self.memory_store.get_file_chunks(path, MemorySource.MEMORY) - old_file_meta = await self.memory_store.get_file_metadata(path, MemorySource.MEMORY) + old_chunks = await self.file_store.get_file_chunks(path, MemorySource.MEMORY) + old_file_meta = await self.file_store.get_file_metadata(path, MemorySource.MEMORY) # Read new file file_meta = await self._build_file_metadata(path) @@ -187,10 +187,10 @@ class DeltaFileWatcher(BaseFileWatcher): or [] ) if chunks: - chunks = await self.memory_store.get_chunk_embeddings(chunks) + chunks = await self.file_store.get_chunk_embeddings(chunks) file_meta.chunk_count = len(chunks) - await self.memory_store.delete_file(path, MemorySource.MEMORY) - await self.memory_store.upsert_file( + await self.file_store.delete_file(path, MemorySource.MEMORY) + await self.file_store.upsert_file( file_meta, MemorySource.MEMORY, chunks, @@ -216,10 +216,10 @@ class DeltaFileWatcher(BaseFileWatcher): or [] ) if chunks: - chunks = await self.memory_store.get_chunk_embeddings(chunks) + chunks = await self.file_store.get_chunk_embeddings(chunks) file_meta.chunk_count = len(chunks) - await self.memory_store.delete_file(path, MemorySource.MEMORY) - await self.memory_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) + await self.file_store.delete_file(path, MemorySource.MEMORY) + await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) logger.info(f"File modified (full): {path} ({len(chunks)} chunks)") else: # Append-only: incremental update @@ -247,16 +247,22 @@ class DeltaFileWatcher(BaseFileWatcher): f"{chunk.source}:{chunk.path}:{chunk.start_line}:" f"{chunk.end_line}:{chunk.hash}:{idx}", ) - new_chunks = await self.memory_store.get_chunk_embeddings(new_chunks) + new_chunks = await self.file_store.get_chunk_embeddings(new_chunks) chunks_to_delete = [c.id for c in old_chunks_sorted if c.start_line >= cutoff_line] # Apply incremental updates if chunks_to_delete: - await self.memory_store.delete_file_chunks(path, chunks_to_delete) + await self.file_store.delete_file_chunks(path, chunks_to_delete) if new_chunks: - await self.memory_store.upsert_chunks(new_chunks, MemorySource.MEMORY) + await self.file_store.upsert_chunks(new_chunks, MemorySource.MEMORY) + + # Update file metadata to reflect the changes + # Calculate new chunk count: old chunks - deleted + new chunks + new_chunk_count = len(old_chunks) - len(chunks_to_delete) + len(new_chunks) + file_meta.chunk_count = new_chunk_count + await self.file_store.update_file_metadata(file_meta, MemorySource.MEMORY) logger.info( f"File modified (incremental): {path} " @@ -265,7 +271,7 @@ class DeltaFileWatcher(BaseFileWatcher): ) elif change_type == Change.deleted: - await self.memory_store.delete_file(path, MemorySource.MEMORY) + await self.file_store.delete_file(path, MemorySource.MEMORY) logger.info(f"File deleted: {path}") else: diff --git a/reme/core/file_watcher/full_file_watcher.py b/reme/core/file_watcher/full_file_watcher.py index 1e26e936..2d365f23 100644 --- a/reme/core/file_watcher/full_file_watcher.py +++ b/reme/core/file_watcher/full_file_watcher.py @@ -58,17 +58,17 @@ class FullFileWatcher(BaseFileWatcher): or [] ) if chunks: - chunks = await self.memory_store.get_chunk_embeddings(chunks) + chunks = await self.file_store.get_chunk_embeddings(chunks) file_meta.chunk_count = len(chunks) - await self.memory_store.delete_file(file_meta.path, MemorySource.MEMORY) + await self.file_store.delete_file(file_meta.path, MemorySource.MEMORY) logger.info(f"delete_file {file_meta.path}") - await self.memory_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) + await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) logger.info(f"Upserted {file_meta.chunk_count} chunks for {file_meta.path}") elif change_type == Change.deleted: - await self.memory_store.delete_file(path, MemorySource.MEMORY) + await self.file_store.delete_file(path, MemorySource.MEMORY) logger.info(f"Deleted {path}") else: diff --git a/reme/core/flow/__init__.py b/reme/core/flow/__init__.py index 0258d63c..3b7e2753 100644 --- a/reme/core/flow/__init__.py +++ b/reme/core/flow/__init__.py @@ -3,7 +3,7 @@ from .base_flow import BaseFlow from .cmd_flow import CmdFlow from .expression_flow import ExpressionFlow -from ..context import R +from ..registry_factory import R __all__ = [ "BaseFlow", diff --git a/reme/core/flow/base_flow.py b/reme/core/flow/base_flow.py index 7e0e1f0c..68e66fd5 100644 --- a/reme/core/flow/base_flow.py +++ b/reme/core/flow/base_flow.py @@ -7,10 +7,12 @@ from abc import ABC, abstractmethod from loguru import logger -from ..context import RuntimeContext, ServiceContext, R from ..enumeration import ChunkEnum from ..op import BaseOp, SequentialOp, ParallelOp +from ..registry_factory import R +from ..runtime_context import RuntimeContext from ..schema import Response, ToolCall +from ..service_context import ServiceContext from ..utils import camel_to_snake, CacheHandler diff --git a/reme/core/flow/expression_flow.py b/reme/core/flow/expression_flow.py index b32c9257..a844c719 100644 --- a/reme/core/flow/expression_flow.py +++ b/reme/core/flow/expression_flow.py @@ -1,9 +1,9 @@ """Expression-based flow implementation driven by configuration objects.""" from .base_flow import BaseFlow -from ..context import ServiceContext from ..op import BaseOp from ..schema import FlowConfig, ToolCall +from ..service_context import ServiceContext class ExpressionFlow(BaseFlow): diff --git a/reme/core/llm/__init__.py b/reme/core/llm/__init__.py index 1cb7e03a..baa7808e 100644 --- a/reme/core/llm/__init__.py +++ b/reme/core/llm/__init__.py @@ -5,7 +5,7 @@ from .lite_llm import LiteLLM from .lite_llm_sync import LiteLLMSync from .openai_llm import OpenAILLM from .openai_llm_sync import OpenAILLMSync -from ..context import R +from ..registry_factory import R __all__ = [ "BaseLLM", diff --git a/reme/core/llm/base_llm.py b/reme/core/llm/base_llm.py index e26d3e64..856c7f3b 100644 --- a/reme/core/llm/base_llm.py +++ b/reme/core/llm/base_llm.py @@ -10,9 +10,7 @@ from typing import Callable, Generator, AsyncGenerator, Any from loguru import logger from ..enumeration import ChunkEnum, Role -from ..schema import Message -from ..schema import StreamChunk -from ..schema import ToolCall +from ..schema import Message, StreamChunk, ToolCall from ..utils import extract_content @@ -510,8 +508,14 @@ class BaseLLM(ABC): **kwargs, ) - async def close(self): - """Release async resources.""" + def start_sync(self): + """Synchronously initialize resources.""" + + async def start(self): + """Asynchronously initialize resources.""" def close_sync(self): - """Release sync resources.""" + """Synchronously release resources and close connections.""" + + async def close(self): + """Asynchronously release resources and close connections.""" diff --git a/reme/core/llm/lite_llm.py b/reme/core/llm/lite_llm.py index d3d5b044..13bc1872 100644 --- a/reme/core/llm/lite_llm.py +++ b/reme/core/llm/lite_llm.py @@ -6,9 +6,7 @@ from loguru import logger from .base_llm import BaseLLM from ..enumeration import ChunkEnum -from ..schema import Message -from ..schema import StreamChunk -from ..schema import ToolCall +from ..schema import Message, StreamChunk, ToolCall class LiteLLM(BaseLLM): diff --git a/reme/core/llm/lite_llm_sync.py b/reme/core/llm/lite_llm_sync.py index b44b9269..d074a9f3 100644 --- a/reme/core/llm/lite_llm_sync.py +++ b/reme/core/llm/lite_llm_sync.py @@ -4,9 +4,7 @@ from typing import Generator from .lite_llm import LiteLLM from ..enumeration import ChunkEnum -from ..schema import Message -from ..schema import StreamChunk -from ..schema import ToolCall +from ..schema import Message, StreamChunk, ToolCall class LiteLLMSync(LiteLLM): diff --git a/reme/core/llm/openai_llm.py b/reme/core/llm/openai_llm.py index 10030bb2..6dd90887 100644 --- a/reme/core/llm/openai_llm.py +++ b/reme/core/llm/openai_llm.py @@ -7,9 +7,7 @@ from openai import AsyncOpenAI from .base_llm import BaseLLM from ..enumeration import ChunkEnum -from ..schema import Message -from ..schema import StreamChunk -from ..schema import ToolCall +from ..schema import Message, StreamChunk, ToolCall class OpenAILLM(BaseLLM): diff --git a/reme/core/llm/openai_llm_sync.py b/reme/core/llm/openai_llm_sync.py index 8f1f4517..dadd1ca2 100644 --- a/reme/core/llm/openai_llm_sync.py +++ b/reme/core/llm/openai_llm_sync.py @@ -6,9 +6,7 @@ from openai import OpenAI from .openai_llm import OpenAILLM from ..enumeration import ChunkEnum -from ..schema import Message -from ..schema import StreamChunk -from ..schema import ToolCall +from ..schema import Message, StreamChunk, ToolCall class OpenAILLMSync(OpenAILLM): diff --git a/reme/core/memory_store/__init__.py b/reme/core/memory_store/__init__.py deleted file mode 100644 index a3079db7..00000000 --- a/reme/core/memory_store/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Memory store module for persistent memory management. - -This module provides storage backends for memory chunks and file metadata, -including SQLite-based, ChromaDB-based, and pure-Python local implementations -with vector and full-text search. -""" - -from .base_memory_store import BaseMemoryStore -from .chroma_memory_store import ChromaMemoryStore -from .local_memory_store import LocalMemoryStore -from .sqlite_memory_store import SqliteMemoryStore -from ..context import R - -__all__ = [ - "BaseMemoryStore", - "ChromaMemoryStore", - "LocalMemoryStore", - "SqliteMemoryStore", -] - -R.memory_stores.register("sqlite")(SqliteMemoryStore) -R.memory_stores.register("chroma")(ChromaMemoryStore) -R.memory_stores.register("local")(LocalMemoryStore) diff --git a/reme/core/op/__init__.py b/reme/core/op/__init__.py index c2ea5291..c99bed3a 100644 --- a/reme/core/op/__init__.py +++ b/reme/core/op/__init__.py @@ -8,7 +8,7 @@ from .base_tool import BaseTool from .mcp_tool import MCPTool from .parallel_op import ParallelOp from .sequential_op import SequentialOp -from ..context import R +from ..registry_factory import R __all__ = [ "BaseOp", diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index 8592de37..af5158c4 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -10,11 +10,14 @@ from typing import Callable, Optional, Any from loguru import logger from tqdm import tqdm -from ..context import RuntimeContext, PromptHandler, ServiceContext from ..embedding import BaseEmbeddingModel +from ..file_store import BaseFileStore from ..llm import BaseLLM -from ..memory_store import BaseMemoryStore -from ..schema import Response +from ..prompt_handler import PromptHandler +from ..runtime_context import RuntimeContext +from ..schema import Response, ServiceConfig +from ..schema.service_config import OpConfig +from ..service_context import ServiceContext from ..token_counter import BaseTokenCounter from ..utils import camel_to_snake, CacheHandler, timer from ..vector_store import BaseVectorStore @@ -42,7 +45,7 @@ class BaseOp(metaclass=ABCMeta): llm: str | BaseLLM = "default", embedding_model: str | BaseEmbeddingModel = "default", vector_store: str | BaseVectorStore = "default", - memory_store: str | BaseMemoryStore = "default", + file_store: str | BaseFileStore = "default", token_counter: str | BaseTokenCounter = "default", enable_cache: bool = False, cache_path: str = "cache/op", @@ -64,7 +67,7 @@ class BaseOp(metaclass=ABCMeta): self._llm = llm self._embedding_model = embedding_model self._vector_store = vector_store - self._memory_store = memory_store + self._file_store = file_store self._token_counter = token_counter self.enable_cache = enable_cache @@ -121,6 +124,11 @@ class BaseOp(metaclass=ABCMeta): assert self.context, "Service context is not initialized!" return self.context.service_context + @property + def service_config(self) -> ServiceConfig: + """Access the service configuration.""" + return self.service_context.service_config + @property def llm(self) -> BaseLLM: """Get the LLM instance from ServiceContext.""" @@ -143,11 +151,11 @@ class BaseOp(metaclass=ABCMeta): return self._vector_store @property - def memory_store(self) -> BaseMemoryStore: - """Lazily initialize and return the memory store instance.""" - if isinstance(self._memory_store, str): - self._memory_store = self.service_context.memory_stores[self._memory_store] - return self._memory_store + def file_store(self) -> BaseFileStore: + """Lazily initialize and return the file store instance.""" + if isinstance(self._file_store, str): + self._file_store = self.service_context.file_stores[self._file_store] + return self._file_store @property def token_counter(self) -> BaseTokenCounter: @@ -159,7 +167,7 @@ class BaseOp(metaclass=ABCMeta): @property def service_metadata(self) -> dict: """Get service configuration metadata.""" - return self.service_context.service_config.model_extra + return self.service_context.service_config.metadata @property def response(self) -> Response: @@ -167,12 +175,42 @@ class BaseOp(metaclass=ABCMeta): return self.context.response def before_execute_sync(self): - """Prepare context and validate before sync execution.""" + """Prepare context and validate before sync execution. + + This method performs the following steps: + 1. Apply input mapping to transform context variables + 2. Load operator-specific configuration from service config if available + 3. Override operator parameters and prompts based on config + """ self.context.apply_mapping(self.input_mapping) + if self.context.service_context is None: + return + + service_config = self.service_context.service_config + if self.name not in service_config.ops: + return + + op_config: OpConfig = service_config.ops[self.name] + + # Override operator parameters from config + if op_config.params: + for k, v in op_config.params.items(): + if hasattr(self, k): + setattr(self, k, v) + logger.info(f"[{self.__class__.__name__}] Set attribute '{k}' = {v}") + else: + self.op_params[k] = v + logger.info(f"[{self.__class__.__name__}] Set op_param '{k}' = {v}") + + # Load custom prompt templates from config + if op_config.prompt_dict: + self.prompt.load_prompt_dict(op_config.prompt_dict) + logger.info(f"[{self.__class__.__name__}] Loaded prompt keys={list(op_config.prompt_dict.keys())}") + async def before_execute(self): """Prepare context and validate before async execution.""" - self.context.apply_mapping(self.input_mapping) + self.before_execute_sync() def execute_sync(self): """Define core sync logic in subclasses.""" diff --git a/reme/core/op/base_ray_op.py b/reme/core/op/base_ray_op.py index 84c0ed25..f63e7b83 100644 --- a/reme/core/op/base_ray_op.py +++ b/reme/core/op/base_ray_op.py @@ -8,7 +8,7 @@ from loguru import logger from tqdm import tqdm from .base_op import BaseOp -from ..context import BaseContext +from ..base_dict import BaseDict _RAY_IMPORT_ERROR = None @@ -53,7 +53,7 @@ class BaseRayOp(BaseOp, metaclass=ABCMeta): # Put large shared objects into the Ray Object Store once optimized_kwargs = { - k: (ray.put(v) if isinstance(v, (pd.DataFrame, pd.Series, dict, list, BaseContext)) else v) + k: (ray.put(v) if isinstance(v, (pd.DataFrame, pd.Series, dict, list, BaseDict)) else v) for k, v in kwargs.items() } diff --git a/reme/core/op/mcp_tool.py b/reme/core/op/mcp_tool.py index 68e78946..bcd64ad2 100644 --- a/reme/core/op/mcp_tool.py +++ b/reme/core/op/mcp_tool.py @@ -32,7 +32,14 @@ class MCPTool(BaseTool): self.timeout: float | None = timeout # Example MCP marketplace: https://bailian.console.aliyun.com/?tab=mcp#/mcp-market - self._client = MCPClient(self.service_context.service_config.mcp_servers) + self._client: MCPClient | None = None + + @property + def client(self) -> MCPClient: + """Lazily initialize and return the MCP client.""" + if self._client is None: + self._client = MCPClient(self.service_context.service_config.mcp_servers) + return self._client def _build_tool_call(self) -> ToolCall: tool_call_dict = self.service_context.mcp_server_mapping[self.mcp_server] @@ -61,7 +68,7 @@ class MCPTool(BaseTool): return tool_call async def execute(self): - tool_result: CallToolResult = await self._client.call_tool( + tool_result: CallToolResult = await self.client.call_tool( server_name=self.mcp_server, tool_name=self.tool_name, arguments=self.input_dict, diff --git a/reme/core/context/prompt_handler.py b/reme/core/prompt_handler.py similarity index 96% rename from reme/core/context/prompt_handler.py rename to reme/core/prompt_handler.py index 34ae39fe..22e1b095 100644 --- a/reme/core/context/prompt_handler.py +++ b/reme/core/prompt_handler.py @@ -8,15 +8,16 @@ from typing import Any, Dict, Optional, Union import yaml from loguru import logger -from .base_context import BaseContext +from .base_dict import BaseDict -class PromptHandler(BaseContext): +class PromptHandler(BaseDict): """A context-aware handler for loading, retrieving, and formatting prompt templates.""" def __init__(self, language: str = "", **kwargs): super().__init__(**kwargs) - self.language: str = language.strip() + # Use object.__setattr__ to avoid storing 'language' in the dict + object.__setattr__(self, "language", language.strip()) def load_prompt_by_file( self, diff --git a/reme/core/context/registry_factory.py b/reme/core/registry_factory.py similarity index 89% rename from reme/core/context/registry_factory.py rename to reme/core/registry_factory.py index e0633cfd..b319b049 100644 --- a/reme/core/context/registry_factory.py +++ b/reme/core/registry_factory.py @@ -3,13 +3,13 @@ import inspect from typing import Callable, TypeVar -from .base_context import BaseContext -from ..utils import singleton +from .base_dict import BaseDict +from .utils import singleton T = TypeVar("T") -class Registry(BaseContext): +class Registry(BaseDict): """A registry container that uses decorators to map and store class references.""" def register(self, name: str | type = "") -> Callable[[type[T]], type[T]] | type[T]: @@ -36,7 +36,7 @@ class RegistryFactory: self.llms = Registry() self.embedding_models = Registry() self.vector_stores = Registry() - self.memory_stores = Registry() + self.file_stores = Registry() self.ops = Registry() self.flows = Registry() self.services = Registry() diff --git a/reme/core/context/runtime_context.py b/reme/core/runtime_context.py similarity index 95% rename from reme/core/context/runtime_context.py rename to reme/core/runtime_context.py index 653fc0fc..43c7c7d8 100644 --- a/reme/core/context/runtime_context.py +++ b/reme/core/runtime_context.py @@ -2,13 +2,13 @@ import asyncio -from .base_context import BaseContext +from .base_dict import BaseDict +from .enumeration import ChunkEnum +from .schema import Response, StreamChunk from .service_context import ServiceContext -from ..enumeration import ChunkEnum -from ..schema import Response, StreamChunk -class RuntimeContext(BaseContext): +class RuntimeContext(BaseDict): """Context for execution state, response metadata, and stream queues.""" def __init__( diff --git a/reme/core/schema/__init__.py b/reme/core/schema/__init__.py index 4d66089d..de167c69 100644 --- a/reme/core/schema/__init__.py +++ b/reme/core/schema/__init__.py @@ -11,10 +11,12 @@ from .response import Response from .service_config import ( CmdConfig, EmbeddingModelConfig, + FileWatcherConfig, FlowConfig, HttpConfig, LLMConfig, MCPConfig, + FileStoreConfig, ServiceConfig, TokenCounterConfig, VectorStoreConfig, @@ -30,6 +32,7 @@ __all__ = [ "ContentBlock", "EmbeddingModelConfig", "FileMetadata", + "FileWatcherConfig", "FlowConfig", "HttpConfig", "LLMConfig", @@ -37,6 +40,7 @@ __all__ = [ "MemoryChunk", "MemoryNode", "MemorySearchResult", + "FileStoreConfig", "Message", "Request", "Response", diff --git a/reme/core/schema/message.py b/reme/core/schema/message.py index d9201979..c015eb27 100644 --- a/reme/core/schema/message.py +++ b/reme/core/schema/message.py @@ -78,6 +78,14 @@ class Message(BaseModel): return self.content return [block.simple_dump() for block in self.content] + def get_text_content(self) -> str: + """Extract plain text content from message, handling both str and list[ContentBlock].""" + if isinstance(self.content, str): + return self.content + return " ".join( + block.content if isinstance(block.content, str) else str(block.content) for block in self.content + ) + def simple_dump( self, add_name: bool = False, diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py index 4b438ce3..5b89367a 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -36,6 +36,15 @@ class CmdConfig(BaseModel): flow: str = Field(default="") +class OpConfig(BaseModel): + """Configuration for op settings and parameters.""" + + model_config = ConfigDict(extra="allow") + + prompt_dict: dict[str, str] = Field(default_factory=dict) + params: dict = Field(default_factory=dict) + + class FlowConfig(ToolCall): """Configuration for workflow execution, caching, and error handling.""" @@ -77,8 +86,8 @@ class VectorStoreConfig(BaseModel): embedding_model: str = Field(default="default") -class MemoryStoreConfig(BaseModel): - """Configuration for memory database storage and associated embeddings.""" +class FileStoreConfig(BaseModel): + """Configuration for file store database storage and associated embeddings.""" model_config = ConfigDict(extra="allow") @@ -102,7 +111,7 @@ class FileWatcherConfig(BaseModel): model_config = ConfigDict(extra="allow") backend: str = Field(default="") - memory_store: str = Field(default="") + file_store: str = Field(default="") watch_paths: list[str] = Field(default_factory=list) @@ -126,11 +135,12 @@ class ServiceConfig(BaseModel): mcp: MCPConfig = Field(default_factory=MCPConfig) http: HttpConfig = Field(default_factory=HttpConfig) cmd: CmdConfig = Field(default_factory=CmdConfig) + ops: dict[str, OpConfig] = 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) + file_stores: dict[str, FileStoreConfig] = 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 4cd60264..ba4e8c40 100644 --- a/reme/core/service/__init__.py +++ b/reme/core/service/__init__.py @@ -4,7 +4,7 @@ from .base_service import BaseService from .cmd_service import CmdService from .http_service import HttpService from .mcp_service import MCPService -from ..context import R +from ..registry_factory import R __all__ = [ "BaseService", diff --git a/reme/core/service/base_service.py b/reme/core/service/base_service.py index 9496391c..65d06060 100644 --- a/reme/core/service/base_service.py +++ b/reme/core/service/base_service.py @@ -1,22 +1,26 @@ """Base service definitions for flow management.""" from abc import ABC, abstractmethod +from typing import TYPE_CHECKING from loguru import logger from pydantic import BaseModel -from ..context import ServiceContext from ..flow import BaseFlow from ..schema import ToolCall from ..utils import create_pydantic_model +if TYPE_CHECKING: + from ..application import Application + class BaseService(ABC): """Abstract base class for services that integrate and execute flows.""" - def __init__(self, service_context: ServiceContext, **kwargs): + def __init__(self, app: "Application", **kwargs): """Initialize the base service.""" - self.service_context: ServiceContext = service_context + self.app: "Application" = app + self.service_context = self.app.service_context self.service_config = self.service_context.service_config self.kwargs = kwargs diff --git a/reme/core/service/cmd_service.py b/reme/core/service/cmd_service.py index c35dd486..4d43b5c3 100644 --- a/reme/core/service/cmd_service.py +++ b/reme/core/service/cmd_service.py @@ -4,7 +4,7 @@ from loguru import logger from .base_service import BaseService from ..flow import CmdFlow, BaseFlow -from ..utils.common_utils import run_coro_safely +from ..utils import run_coro_safely class CmdService(BaseService): @@ -18,10 +18,14 @@ class CmdService(BaseService): def integrate_flow(self, flow: BaseFlow) -> str | None: """Integrate the workflow configuration into the command service.""" self._cmd_flow = CmdFlow(flow=self.service_config.cmd.flow, service_context=self.service_context) + return self._cmd_flow.tool_call.name if self._cmd_flow else None def run(self): """Execute the command flow in either asynchronous or synchronous mode.""" super().run() + if not self._cmd_flow: + logger.warning("No command flow configured, skipping execution") + return kwargs = self.service_config.cmd.model_extra if self._cmd_flow.async_mode: diff --git a/reme/core/service/http_service.py b/reme/core/service/http_service.py index ebaeba49..c634a11d 100644 --- a/reme/core/service/http_service.py +++ b/reme/core/service/http_service.py @@ -12,7 +12,7 @@ from fastapi.responses import StreamingResponse from .base_service import BaseService from ..flow import BaseFlow from ..schema import Response -from ..utils.common_utils import execute_stream_task +from ..utils import execute_stream_task class HttpService(BaseService): @@ -24,20 +24,20 @@ class HttpService(BaseService): @asynccontextmanager async def lifespan(_: FastAPI): - await self.service_context.start() + await self.app.start() yield - await self.service_context.close() + await self.app.close() - self.app = FastAPI(title=self.service_config.app_name, lifespan=lifespan) + self.http_service = FastAPI(title=self.service_config.app_name, lifespan=lifespan) - self.app.add_middleware( + self.http_service.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) - self.app.get("/health")(lambda: {"status": "healthy"}) + self.http_service.get("/health")(lambda: {"status": "healthy"}) def _integrate_flow(self, flow: BaseFlow) -> str: """Register a standard flow as a POST endpoint.""" @@ -46,7 +46,7 @@ class HttpService(BaseService): async def execute_endpoint(request: request_model) -> Response: return await flow.call(**request.model_dump(exclude_none=True)) - self.app.post( + self.http_service.post( path=f"/{tool_call.name}", response_model=Response, description=tool_call.description, @@ -72,7 +72,7 @@ class HttpService(BaseService): return StreamingResponse(generate_stream(), media_type="text/event-stream") - self.app.post(f"/{tool_call.name}")(execute_stream_endpoint) + self.http_service.post(f"/{tool_call.name}")(execute_stream_endpoint) return tool_call.name def integrate_flow(self, flow: BaseFlow) -> str | None: @@ -84,7 +84,7 @@ class HttpService(BaseService): super().run() cfg = self.service_config.http uvicorn.run( - self.app, + self.http_service, host=cfg.host, port=cfg.port, timeout_keep_alive=cfg.timeout_keep_alive, diff --git a/reme/core/service/mcp_service.py b/reme/core/service/mcp_service.py index 8b7bebc3..6ed9148c 100644 --- a/reme/core/service/mcp_service.py +++ b/reme/core/service/mcp_service.py @@ -18,11 +18,11 @@ class MCPService(BaseService): @asynccontextmanager async def lifespan(_: FastMCP): - await self.service_context.start() + await self.app.start() yield {} - await self.service_context.close() + await self.app.close() - self.mcp = FastMCP(name=self.service_config.app_name, lifespan=lifespan) + self.mcp_service = FastMCP(name=self.service_config.app_name, lifespan=lifespan) def integrate_flow(self, flow: BaseFlow) -> str | None: """Register a non-streaming flow as an MCP tool.""" @@ -37,7 +37,7 @@ class MCPService(BaseService): response = await flow.call(**request_instance.model_dump(exclude_none=True)) return response.answer - self.mcp.add_tool( + self.mcp_service.add_tool( FunctionTool( name=tool_call.name, # noqa description=tool_call.description, # noqa @@ -50,8 +50,8 @@ class MCPService(BaseService): def run(self): """Run the MCP server with specified transport protocol.""" super().run() - cfg = self.service_config.mcp + cfg = self.service_config.mcp_service run_args: dict = {"transport": cfg.transport, "show_banner": False, **cfg.model_extra} if cfg.transport != "stdio": run_args.update({"host": cfg.host, "port": cfg.port}) - self.mcp.run(**run_args) + self.mcp_service.run(**run_args) diff --git a/reme/core/context/service_context.py b/reme/core/service_context.py similarity index 81% rename from reme/core/context/service_context.py rename to reme/core/service_context.py index e5a9860d..0b9c5a1b 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/service_context.py @@ -6,21 +6,21 @@ from typing import TYPE_CHECKING from loguru import logger -from .base_context import BaseContext -from ..schema import ServiceConfig -from ..utils import load_env, PydanticConfigParser +from .base_dict import BaseDict +from .schema import ServiceConfig +from .utils import load_env, PydanticConfigParser if TYPE_CHECKING: - from ..llm import BaseLLM - from ..embedding import BaseEmbeddingModel - from ..vector_store import BaseVectorStore - from ..memory_store import BaseMemoryStore - from ..token_counter import BaseTokenCounter - from ..flow import BaseFlow - from ..file_watcher import BaseFileWatcher + from .llm import BaseLLM + from .embedding import BaseEmbeddingModel + from .vector_store import BaseVectorStore + from .file_store import BaseFileStore + from .token_counter import BaseTokenCounter + from .flow import BaseFlow + from .file_watcher import BaseFileWatcher -class ServiceContext(BaseContext): +class ServiceContext(BaseDict): """Service context.""" def __init__( @@ -39,7 +39,7 @@ class ServiceContext(BaseContext): 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_file_store_config: dict | None = None, default_token_counter_config: dict | None = None, default_file_watcher_config: dict | None = None, **kwargs, @@ -57,7 +57,7 @@ class ServiceContext(BaseContext): if service_config is None: parser_class = parser if parser is not None else PydanticConfigParser - parser = parser_class(ServiceConfig) + parser_instance = parser_class(ServiceConfig) input_args = [] if config_path: input_args.append(f"config={config_path}") @@ -72,8 +72,8 @@ class ServiceContext(BaseContext): 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_store_config: + self._update_section_config(kwargs, "file_stores", **default_file_store_config) if default_file_watcher_config: self._update_section_config(kwargs, "file_watchers", **default_file_watcher_config) @@ -85,7 +85,7 @@ class ServiceContext(BaseContext): }, ) logger.info(f"update with args: {input_args} kwargs: {kwargs}") - service_config = parser.parse_args(*input_args, **kwargs) + service_config = parser_instance.parse_args(*input_args, **kwargs) self.service_config: ServiceConfig = service_config @@ -94,7 +94,7 @@ class ServiceContext(BaseContext): self.embedding_models: dict[str, "BaseEmbeddingModel"] = {} self.token_counters: dict[str, "BaseTokenCounter"] = {} self.vector_stores: dict[str, "BaseVectorStore"] = {} - self.memory_stores: dict[str, "BaseMemoryStore"] = {} + self.file_stores: dict[str, "BaseFileStore"] = {} self.file_watchers: dict[str, "BaseFileWatcher"] = {} self.flows: dict[str, "BaseFlow"] = {} self.mcp_server_mapping: dict[str, dict] = {} diff --git a/reme/core/token_counter/__init__.py b/reme/core/token_counter/__init__.py index 7792046f..6a1bca1d 100644 --- a/reme/core/token_counter/__init__.py +++ b/reme/core/token_counter/__init__.py @@ -3,7 +3,7 @@ from .base_token_counter import BaseTokenCounter from .hf_token_counter import HFTokenCounter from .openai_token_counter import OpenAITokenCounter -from ..context import R +from ..registry_factory import R __all__ = [ "BaseTokenCounter", diff --git a/reme/core/token_counter/base_token_counter.py b/reme/core/token_counter/base_token_counter.py index 6f98cd02..4fca5015 100644 --- a/reme/core/token_counter/base_token_counter.py +++ b/reme/core/token_counter/base_token_counter.py @@ -39,10 +39,7 @@ class BaseTokenCounter: # Extract text from messages segments = [] for msg in messages: - content = msg.content - if isinstance(content, bytes): - content = content.decode("utf-8", errors="ignore") - segments.extend([content, msg.reasoning_content]) + segments.extend([msg.get_text_content(), msg.reasoning_content]) # Extract text from tools if tools: diff --git a/reme/core/token_counter/openai_token_counter.py b/reme/core/token_counter/openai_token_counter.py index 672f8a45..0c482145 100644 --- a/reme/core/token_counter/openai_token_counter.py +++ b/reme/core/token_counter/openai_token_counter.py @@ -42,7 +42,7 @@ class OpenAITokenCounter(BaseTokenCounter): # Every message has <|start|>{role/name}\n{content}<|end|>\n total_tokens += 3 # Base overhead per message if msg.content: - total_tokens += len(enc.encode(msg.content)) + total_tokens += len(enc.encode(msg.get_text_content())) if msg.tool_calls: for tc in msg.tool_calls: diff --git a/reme/core/tools/__init__.py b/reme/core/tools/__init__.py new file mode 100644 index 00000000..8005793a --- /dev/null +++ b/reme/core/tools/__init__.py @@ -0,0 +1,45 @@ +"""tools""" + +from .execute_code import ExecuteCode +from .execute_shell import ExecuteShell + +# file tools +from .file.base_file_tool import BaseFileTool +from .file.bash_tool import BashTool +from .file.edit_tool import EditTool +from .file.find_tool import FindTool +from .file.grep_tool import GrepTool +from .file.ls_tool import LsTool +from .file.read_tool import ReadTool +from .file.write_tool import WriteTool + +# search tools +from .search.dashscope_search import DashscopeSearch +from .search.mock_search import MockSearch +from .search.tavily_search import TavilySearch +from .think_tool import ThinkTool +from ..registry_factory import R + +__all__ = [ + # base tools + "ThinkTool", + "ExecuteCode", + "ExecuteShell", + # file tools + "BaseFileTool", + "BashTool", + "EditTool", + "FindTool", + "GrepTool", + "LsTool", + "ReadTool", + "WriteTool", + # search tools + "DashscopeSearch", + "TavilySearch", + "MockSearch", +] + +for name in __all__: + tool_class = globals()[name] + R.ops.register(tool_class) diff --git a/reme/tool/gallery/execute_code.py b/reme/core/tools/execute_code.py similarity index 90% rename from reme/tool/gallery/execute_code.py rename to reme/core/tools/execute_code.py index 262b7f2c..7631777d 100644 --- a/reme/tool/gallery/execute_code.py +++ b/reme/core/tools/execute_code.py @@ -4,10 +4,9 @@ This module provides an operation that can execute Python code strings and return the output or error messages. """ -from ...core.op import BaseTool -from ...core.schema import ToolCall - -from ...core.utils import exec_code, async_exec_code +from ..op import BaseTool +from ..schema import ToolCall +from ..utils import exec_code, async_exec_code class ExecuteCode(BaseTool): diff --git a/reme/tool/gallery/execute_code.yaml b/reme/core/tools/execute_code.yaml similarity index 100% rename from reme/tool/gallery/execute_code.yaml rename to reme/core/tools/execute_code.yaml diff --git a/reme/tool/gallery/execute_shell.py b/reme/core/tools/execute_shell.py similarity index 92% rename from reme/tool/gallery/execute_shell.py rename to reme/core/tools/execute_shell.py index de862602..aa77c036 100644 --- a/reme/tool/gallery/execute_shell.py +++ b/reme/core/tools/execute_shell.py @@ -4,10 +4,9 @@ This module provides an operation that can execute shell commands asynchronously and return the output, error, and exit code. """ -from ...core.op import BaseTool -from ...core.schema import ToolCall - -from ...core.utils import run_shell_command +from ..op import BaseTool +from ..schema import ToolCall +from ..utils import run_shell_command class ExecuteShell(BaseTool): diff --git a/reme/tool/gallery/execute_shell.yaml b/reme/core/tools/execute_shell.yaml similarity index 100% rename from reme/tool/gallery/execute_shell.yaml rename to reme/core/tools/execute_shell.yaml diff --git a/reme/agent/memory/personal/__init__.py b/reme/core/tools/file/__init__.py similarity index 100% rename from reme/agent/memory/personal/__init__.py rename to reme/core/tools/file/__init__.py diff --git a/reme/tool/fs/base_fs_tool.py b/reme/core/tools/file/base_file_tool.py similarity index 92% rename from reme/tool/fs/base_fs_tool.py rename to reme/core/tools/file/base_file_tool.py index 92c9092a..e7183f65 100644 --- a/reme/tool/fs/base_fs_tool.py +++ b/reme/core/tools/file/base_file_tool.py @@ -2,11 +2,11 @@ from loguru import logger -from ...core.context import RuntimeContext -from ...core.op import BaseTool +from ...op import BaseTool +from ...runtime_context import RuntimeContext -class BaseFsTool(BaseTool): +class BaseFileTool(BaseTool): """Base class for file system tools. Features: @@ -34,6 +34,7 @@ class BaseFsTool(BaseTool): 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)}" diff --git a/reme/tool/fs/bash_tool.py b/reme/core/tools/file/bash_tool.py similarity index 98% rename from reme/tool/fs/bash_tool.py rename to reme/core/tools/file/bash_tool.py index 1b084323..c40c2018 100644 --- a/reme/tool/fs/bash_tool.py +++ b/reme/core/tools/file/bash_tool.py @@ -11,9 +11,9 @@ import platform import signal from pathlib import Path -from .base_fs_tool import BaseFsTool +from .base_file_tool import BaseFileTool from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncate_tail -from ...core.schema import ToolCall, TruncationResult +from ...schema import ToolCall, TruncationResult def get_shell_config() -> tuple[str, list[str]]: @@ -53,7 +53,7 @@ def kill_process_tree(pid: int) -> None: pass # Best effort -class BashTool(BaseFsTool): +class BashTool(BaseFileTool): """Production-grade tool for executing bash commands. Features: diff --git a/reme/tool/fs/edit_diff.py b/reme/core/tools/file/edit_diff.py similarity index 100% rename from reme/tool/fs/edit_diff.py rename to reme/core/tools/file/edit_diff.py diff --git a/reme/tool/fs/edit_tool.py b/reme/core/tools/file/edit_tool.py similarity index 97% rename from reme/tool/fs/edit_tool.py rename to reme/core/tools/file/edit_tool.py index 8733c2af..59ba8632 100644 --- a/reme/tool/fs/edit_tool.py +++ b/reme/core/tools/file/edit_tool.py @@ -3,7 +3,7 @@ import os from pathlib import Path -from .base_fs_tool import BaseFsTool +from .base_file_tool import BaseFileTool from .edit_diff import ( detect_line_ending, fuzzy_find_text, @@ -13,10 +13,10 @@ from .edit_diff import ( restore_line_endings, strip_bom, ) -from ...core.schema import ToolCall +from ...schema import ToolCall -class EditTool(BaseFsTool): +class EditTool(BaseFileTool): """Edit a file by replacing exact text.""" def __init__(self, cwd: str | None = None): diff --git a/reme/tool/fs/find_tool.py b/reme/core/tools/file/find_tool.py similarity index 98% rename from reme/tool/fs/find_tool.py rename to reme/core/tools/file/find_tool.py index 731722cd..73ee285f 100644 --- a/reme/tool/fs/find_tool.py +++ b/reme/core/tools/file/find_tool.py @@ -3,12 +3,12 @@ import os from pathlib import Path -from .base_fs_tool import BaseFsTool +from .base_file_tool import BaseFileTool from .truncate import FIND_MAX_BYTES, FIND_MAX_LINES, format_size, truncate_head -from ...core.schema import ToolCall +from ...schema import ToolCall -class FindTool(BaseFsTool): +class FindTool(BaseFileTool): """Search for files by glob pattern, respecting .gitignore.""" def __init__(self, cwd: str | None = None): diff --git a/reme/tool/fs/grep_tool.py b/reme/core/tools/file/grep_tool.py similarity index 98% rename from reme/tool/fs/grep_tool.py rename to reme/core/tools/file/grep_tool.py index 0546a495..a0f1c1ec 100644 --- a/reme/tool/fs/grep_tool.py +++ b/reme/core/tools/file/grep_tool.py @@ -13,7 +13,7 @@ import os import shutil from pathlib import Path -from .base_fs_tool import BaseFsTool +from .base_file_tool import BaseFileTool from .truncate import ( DEFAULT_MAX_BYTES, GREP_MAX_LINE_LENGTH, @@ -21,13 +21,13 @@ from .truncate import ( truncate_head, truncate_line, ) -from ...core.schema import ToolCall +from ...schema import ToolCall # Default limits DEFAULT_LIMIT = 100 # Maximum number of matches -class GrepTool(BaseFsTool): +class GrepTool(BaseFileTool): """Tool for searching file contents using ripgrep. Features: diff --git a/reme/tool/fs/ls_tool.py b/reme/core/tools/file/ls_tool.py similarity index 97% rename from reme/tool/fs/ls_tool.py rename to reme/core/tools/file/ls_tool.py index f231119c..e0c6fc3c 100644 --- a/reme/tool/fs/ls_tool.py +++ b/reme/core/tools/file/ls_tool.py @@ -3,14 +3,14 @@ import os from pathlib import Path -from .base_fs_tool import BaseFsTool +from .base_file_tool import BaseFileTool from .truncate import DEFAULT_MAX_BYTES, truncate_head -from ...core.schema import ToolCall +from ...schema import ToolCall DEFAULT_LIMIT = 500 -class LsTool(BaseFsTool): +class LsTool(BaseFileTool): """List directory contents with smart truncation. Features: diff --git a/reme/tool/fs/read_tool.py b/reme/core/tools/file/read_tool.py similarity index 98% rename from reme/tool/fs/read_tool.py rename to reme/core/tools/file/read_tool.py index ace2339a..1c7e8b7e 100644 --- a/reme/tool/fs/read_tool.py +++ b/reme/core/tools/file/read_tool.py @@ -9,9 +9,9 @@ Features: import os from pathlib import Path -from .base_fs_tool import BaseFsTool +from .base_file_tool import BaseFileTool from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, format_size, truncate_head -from ...core.schema import ToolCall +from ...schema import ToolCall # Supported image extensions IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"} @@ -29,7 +29,7 @@ def is_image_file(path: str) -> bool: return Path(path).suffix.lower() in IMAGE_EXTENSIONS -class ReadTool(BaseFsTool): +class ReadTool(BaseFileTool): """Read file contents with smart truncation. Features: diff --git a/reme/tool/fs/truncate.py b/reme/core/tools/file/truncate.py similarity index 99% rename from reme/tool/fs/truncate.py rename to reme/core/tools/file/truncate.py index b4de0197..22f0c4b8 100644 --- a/reme/tool/fs/truncate.py +++ b/reme/core/tools/file/truncate.py @@ -2,7 +2,7 @@ from typing import Literal -from ...core.schema import TruncationResult +from ...schema import TruncationResult # Default limits for output truncation DEFAULT_MAX_LINES = 1000 # Maximum lines to keep for tail truncation diff --git a/reme/tool/fs/write_tool.py b/reme/core/tools/file/write_tool.py similarity index 96% rename from reme/tool/fs/write_tool.py rename to reme/core/tools/file/write_tool.py index a291e0ac..75cf6be5 100644 --- a/reme/tool/fs/write_tool.py +++ b/reme/core/tools/file/write_tool.py @@ -8,11 +8,11 @@ This module provides a tool for writing content to files with: import os -from .base_fs_tool import BaseFsTool -from ...core.schema import ToolCall +from .base_file_tool import BaseFileTool +from ...schema import ToolCall -class WriteTool(BaseFsTool): +class WriteTool(BaseFileTool): """Tool for writing content to files. Features: diff --git a/reme/agent/memory/procedural/__init__.py b/reme/core/tools/search/__init__.py similarity index 100% rename from reme/agent/memory/procedural/__init__.py rename to reme/core/tools/search/__init__.py diff --git a/reme/tool/search/dashscope_search.py b/reme/core/tools/search/dashscope_search.py similarity index 98% rename from reme/tool/search/dashscope_search.py rename to reme/core/tools/search/dashscope_search.py index f0e0ae9d..da28e251 100644 --- a/reme/tool/search/dashscope_search.py +++ b/reme/core/tools/search/dashscope_search.py @@ -9,8 +9,8 @@ from typing import Literal from loguru import logger -from ...core.op import BaseTool -from ...core.schema import ToolCall +from ...op import BaseTool +from ...schema import ToolCall class DashscopeSearch(BaseTool): diff --git a/reme/tool/search/dashscope_search.yaml b/reme/core/tools/search/dashscope_search.yaml similarity index 100% rename from reme/tool/search/dashscope_search.yaml rename to reme/core/tools/search/dashscope_search.yaml diff --git a/reme/tool/search/mock_search.py b/reme/core/tools/search/mock_search.py similarity index 92% rename from reme/tool/search/mock_search.py rename to reme/core/tools/search/mock_search.py index ed695b0c..f2bc36ca 100644 --- a/reme/tool/search/mock_search.py +++ b/reme/core/tools/search/mock_search.py @@ -9,10 +9,10 @@ import random from loguru import logger -from ...core.enumeration import Role -from ...core.op import BaseTool -from ...core.schema import ToolCall, Message -from ...core.utils import extract_content +from ...enumeration import Role +from ...op import BaseTool +from ...schema import ToolCall, Message +from ...utils import extract_content class MockSearch(BaseTool): diff --git a/reme/tool/search/mock_search.yaml b/reme/core/tools/search/mock_search.yaml similarity index 100% rename from reme/tool/search/mock_search.yaml rename to reme/core/tools/search/mock_search.yaml diff --git a/reme/tool/search/tavily_search.py b/reme/core/tools/search/tavily_search.py similarity index 98% rename from reme/tool/search/tavily_search.py rename to reme/core/tools/search/tavily_search.py index bfe29879..65b53734 100644 --- a/reme/tool/search/tavily_search.py +++ b/reme/core/tools/search/tavily_search.py @@ -9,8 +9,8 @@ import os from loguru import logger -from ...core.op import BaseTool -from ...core.schema import ToolCall +from ...op import BaseTool +from ...schema import ToolCall class TavilySearch(BaseTool): diff --git a/reme/tool/search/tavily_search.yaml b/reme/core/tools/search/tavily_search.yaml similarity index 100% rename from reme/tool/search/tavily_search.yaml rename to reme/core/tools/search/tavily_search.yaml diff --git a/reme/tool/gallery/think_tool.py b/reme/core/tools/think_tool.py similarity index 95% rename from reme/tool/gallery/think_tool.py rename to reme/core/tools/think_tool.py index f01646db..828d4043 100644 --- a/reme/tool/gallery/think_tool.py +++ b/reme/core/tools/think_tool.py @@ -4,8 +4,8 @@ This module provides a tool that prompts the model for explicit reflection before taking actions, helping agents reason about their next steps. """ -from ...core.op import BaseTool -from ...core.schema import ToolCall +from ..op import BaseTool +from ..schema import ToolCall class ThinkTool(BaseTool): diff --git a/reme/tool/gallery/think_tool.yaml b/reme/core/tools/think_tool.yaml similarity index 100% rename from reme/tool/gallery/think_tool.yaml rename to reme/core/tools/think_tool.yaml diff --git a/reme/core/utils/env_utils.py b/reme/core/utils/env_utils.py index 433039cd..b36c7bd5 100644 --- a/reme/core/utils/env_utils.py +++ b/reme/core/utils/env_utils.py @@ -54,8 +54,6 @@ def load_env(path: str | Path | None = None, enable_log: bool = True) -> None: _ENV_LOADED = True return - logger.warning(".env file not found in search path") - def reset_env_flag() -> None: """Reset the internal load state flag.""" diff --git a/reme/core/utils/singleton.py b/reme/core/utils/singleton.py index 0f6c5907..8c2c071b 100644 --- a/reme/core/utils/singleton.py +++ b/reme/core/utils/singleton.py @@ -1,17 +1,21 @@ """Module providing a decorator to implement the Singleton design pattern.""" +import threading + def singleton(cls): """A class decorator that ensures only one instance of a class exists.""" # Dictionary to cache the single instance of the class _instance = {} + _lock = threading.Lock() def _singleton(*args, **kwargs): """Return the existing instance or create a new one if it doesn't exist.""" - if cls not in _instance: - # Create and store the instance if it's the first call - _instance[cls] = cls(*args, **kwargs) + with _lock: + if cls not in _instance: + # Create and store the instance if it's the first call + _instance[cls] = cls(*args, **kwargs) return _instance[cls] return _singleton diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py index 84b68c3c..9d411d03 100644 --- a/reme/core/vector_store/__init__.py +++ b/reme/core/vector_store/__init__.py @@ -6,7 +6,7 @@ from .es_vector_store import ESVectorStore from .local_vector_store import LocalVectorStore from .pgvector_store import PGVectorStore from .qdrant_vector_store import QdrantVectorStore -from ..context import R +from ..registry_factory import R __all__ = [ "BaseVectorStore", diff --git a/reme/core/vector_store/base_vector_store.py b/reme/core/vector_store/base_vector_store.py index 62a73a3e..63791fe8 100644 --- a/reme/core/vector_store/base_vector_store.py +++ b/reme/core/vector_store/base_vector_store.py @@ -1,10 +1,7 @@ """Base vector store interface for managing vector embeddings and similarity search.""" -import asyncio from abc import ABC, abstractmethod -from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor -from functools import partial +from pathlib import Path from ..embedding import BaseEmbeddingModel from ..schema import VectorNode @@ -16,21 +13,16 @@ class BaseVectorStore(ABC): def __init__( self, collection_name: str, + db_path: str | Path, embedding_model: BaseEmbeddingModel, - thread_pool: ThreadPoolExecutor, **kwargs, ): """Initialize the vector store with a collection name and an embedding model.""" self.collection_name: str = collection_name + self.db_path: Path = Path(db_path) self.embedding_model: BaseEmbeddingModel = embedding_model - self.thread_pool: ThreadPoolExecutor = thread_pool self.kwargs: dict = kwargs - async def _run_sync_in_executor(self, sync_func: Callable, *args, **kwargs): - """Run a synchronous function in the context-defined thread pool executor.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor(self.thread_pool, partial(sync_func, *args, **kwargs)) # noqa - async def get_node_embedding(self, node: VectorNode) -> VectorNode: """Generate and assign embedding for a single vector node.""" return await self.embedding_model.get_node_embedding(node) @@ -119,5 +111,13 @@ class BaseVectorStore(ABC): reverse: If True, sort in descending order; if False, sort in ascending order """ + async def start(self) -> None: + """Initialize the vector store and ensure the collection exists. + + This method should be called after instantiation to perform async initialization. + Subclasses should call super().start() to ensure collection creation. + """ + await self.create_collection(self.collection_name) + async def close(self) -> None: """Release resources and close active connections to the vector store.""" diff --git a/reme/core/vector_store/chroma_vector_store.py b/reme/core/vector_store/chroma_vector_store.py index 5b0748f5..2bd9f4ae 100644 --- a/reme/core/vector_store/chroma_vector_store.py +++ b/reme/core/vector_store/chroma_vector_store.py @@ -1,6 +1,6 @@ """ChromaDB vector store implementation for the ReMe framework.""" -from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Any from loguru import logger @@ -26,12 +26,11 @@ class ChromaVectorStore(BaseVectorStore): def __init__( self, collection_name: str, + db_path: str | Path, embedding_model: BaseEmbeddingModel, - thread_pool: ThreadPoolExecutor, client: chromadb.ClientAPI | None = None, host: str | None = None, port: int | None = None, - path: str | None = None, api_key: str | None = None, tenant: str | None = None, database: str | None = None, @@ -45,13 +44,14 @@ class ChromaVectorStore(BaseVectorStore): super().__init__( collection_name=collection_name, + db_path=db_path, embedding_model=embedding_model, - thread_pool=thread_pool, **kwargs, ) self.client: chromadb.ClientAPI self.collection: chromadb.Collection + self.is_local = client is None and not (api_key and tenant) and not (host and port) if client: self.client = client @@ -66,18 +66,9 @@ class ChromaVectorStore(BaseVectorStore): logger.info(f"Initializing ChromaDB HTTP client at {host}:{port}") self.client = chromadb.HttpClient(host=host, port=port) else: - if path is None: - path = "./chroma_vector_store" - logger.info(f"Initializing local ChromaDB at {path}") - self.client = chromadb.PersistentClient( - path=path, - settings=Settings(anonymized_telemetry=False), - ) + self.client = None # Will be initialized in start() - self.collection = self.client.get_or_create_collection( - name=collection_name, - metadata={"hnsw:space": "cosine"}, - ) + self.collection: chromadb.Collection | None = None @staticmethod def _parse_results( @@ -213,63 +204,47 @@ class ChromaVectorStore(BaseVectorStore): async def list_collections(self) -> list[str]: """Retrieve a list of all existing collection names.""" - - def _list(): - return [col.name for col in self.client.list_collections()] - - return await self._run_sync_in_executor(_list) + return [col.name for col in self.client.list_collections()] async def create_collection(self, collection_name: str, **kwargs): """Create a new collection with specified distance metrics and metadata.""" - - def _create(): - distance_metric = kwargs.get("distance_metric", "cosine") - metadata = kwargs.get("metadata", {}) - metadata["hnsw:space"] = distance_metric - return self.client.get_or_create_collection(name=collection_name, metadata=metadata) - - new_collection = await self._run_sync_in_executor(_create) + distance_metric = kwargs.get("distance_metric", "cosine") + metadata = kwargs.get("metadata", {}) + metadata["hnsw:space"] = distance_metric + new_collection = self.client.get_or_create_collection(name=collection_name, metadata=metadata) if collection_name == self.collection_name: self.collection = new_collection logger.info(f"Created collection `{collection_name}`") async def delete_collection(self, collection_name: str, **kwargs): """Delete a specified collection from the database.""" - - def _delete(): - try: - self.client.delete_collection(name=collection_name) - return True - except Exception as _e: - logger.warning(f"Failed to delete collection {collection_name}: {_e}") - return False - - deleted = await self._run_sync_in_executor(_delete) + try: + self.client.delete_collection(name=collection_name) + deleted = True + except Exception as _e: + logger.warning(f"Failed to delete collection {collection_name}: {_e}") + deleted = False if deleted and collection_name == self.collection_name: self.collection = None logger.info(f"Deleted collection {collection_name}") async def copy_collection(self, collection_name: str, **kwargs): """Copy all data from the current collection to a new collection.""" + source_data = self.collection.get(include=["documents", "metadatas", "embeddings"]) + if not source_data["ids"]: + logger.warning(f"Source collection {self.collection_name} is empty") + return - def _copy(): - source_data = self.collection.get(include=["documents", "metadatas", "embeddings"]) - if not source_data["ids"]: - logger.warning(f"Source collection {self.collection_name} is empty") - return - - target_collection = self.client.get_or_create_collection( - name=collection_name, - metadata={"hnsw:space": "cosine"}, - ) - target_collection.add( - ids=source_data["ids"], - documents=source_data["documents"], - metadatas=source_data["metadatas"], - embeddings=source_data["embeddings"], - ) - - await self._run_sync_in_executor(_copy) + target_collection = self.client.get_or_create_collection( + name=collection_name, + metadata={"hnsw:space": "cosine"}, + ) + target_collection.add( + ids=source_data["ids"], + documents=source_data["documents"], + metadatas=source_data["metadatas"], + embeddings=source_data["embeddings"], + ) logger.info(f"Copied collection {self.collection_name} to {collection_name}") async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): @@ -291,16 +266,14 @@ class ChromaVectorStore(BaseVectorStore): batch_size = kwargs.get("batch_size", 100) - def _insert_batch(batch_nodes: list[VectorNode]): + for i in range(0, len(nodes_to_insert), batch_size): + batch_nodes = nodes_to_insert[i : i + batch_size] self.collection.add( ids=[n.vector_id for n in batch_nodes], documents=[n.content for n in batch_nodes], embeddings=[n.vector for n in batch_nodes], metadatas=[n.metadata for n in batch_nodes], ) - - for i in range(0, len(nodes_to_insert), batch_size): - await self._run_sync_in_executor(_insert_batch, nodes_to_insert[i : i + batch_size]) logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}") async def search( @@ -315,18 +288,15 @@ class ChromaVectorStore(BaseVectorStore): where_clause = self._generate_where_clause(filters) include_embeddings = kwargs.get("include_embeddings", False) - def _search(): - include: list = ["documents", "metadatas", "distances"] - if include_embeddings: - include.append("embeddings") - return self.collection.query( - query_embeddings=[query_vector], - n_results=limit, - where=where_clause, - include=include, - ) - - results = await self._run_sync_in_executor(_search) + include: list = ["documents", "metadatas", "distances"] + if include_embeddings: + include.append("embeddings") + results = self.collection.query( + query_embeddings=[query_vector], + n_results=limit, + where=where_clause, + include=include, + ) nodes = self._parse_results(results, include_score=True) score_threshold = kwargs.get("score_threshold") @@ -341,26 +311,19 @@ class ChromaVectorStore(BaseVectorStore): if not vector_ids: return - def _delete(): - self.collection.delete(ids=vector_ids) - - await self._run_sync_in_executor(_delete) + self.collection.delete(ids=vector_ids) logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}") async def delete_all(self, **kwargs): """Remove all vectors from the collection.""" - - def _delete_all(): - # Get all IDs in the collection - result = self.collection.get() - if result and result.get("ids"): - ids = result["ids"] - if ids: - self.collection.delete(ids=ids) - return len(ids) - return 0 - - count = await self._run_sync_in_executor(_delete_all) + # Get all IDs in the collection + result = self.collection.get() + count = 0 + if result and result.get("ids"): + ids = result["ids"] + if ids: + self.collection.delete(ids=ids) + count = len(ids) logger.info(f"Deleted all {count} nodes from {self.collection_name}") async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): @@ -380,15 +343,12 @@ class ChromaVectorStore(BaseVectorStore): else: nodes_to_update = nodes - def _update(): - self.collection.upsert( - ids=[n.vector_id for n in nodes_to_update], - documents=[n.content for n in nodes_to_update], - embeddings=[n.vector for n in nodes_to_update if n.vector] or None, - metadatas=[n.metadata for n in nodes_to_update], - ) - - await self._run_sync_in_executor(_update) + self.collection.upsert( + ids=[n.vector_id for n in nodes_to_update], + documents=[n.content for n in nodes_to_update], + embeddings=[n.vector for n in nodes_to_update], + metadatas=[n.metadata for n in nodes_to_update], + ) logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}") async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None: @@ -396,10 +356,7 @@ class ChromaVectorStore(BaseVectorStore): is_single = isinstance(vector_ids, str) ids = [vector_ids] if is_single else vector_ids - def _get(): - return self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"]) - - results = await self._run_sync_in_executor(_get) + results = self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"]) nodes = self._parse_results(results) return nodes[0] if is_single and nodes else (nodes if not is_single else None) @@ -423,14 +380,11 @@ class ChromaVectorStore(BaseVectorStore): # If sorting is needed, fetch all records first, then apply limit after sorting fetch_limit = None if sort_key else limit - def _list(): - return self.collection.get( - where=where_clause, - limit=fetch_limit, - include=["documents", "metadatas", "embeddings"], - ) - - results = await self._run_sync_in_executor(_list) + results = self.collection.get( + where=where_clause, + limit=fetch_limit, + include=["documents", "metadatas", "embeddings"], + ) nodes = self._parse_results(results) # Apply sorting if sort_key is provided @@ -453,22 +407,38 @@ class ChromaVectorStore(BaseVectorStore): async def count(self) -> int: """Return the total number of vectors in the current collection.""" - return await self._run_sync_in_executor(self.collection.count) + return self.collection.count() async def reset(self): """Reset the current collection by clearing all its data.""" logger.warning(f"Resetting collection {self.collection_name}...") await self.delete_collection(self.collection_name) - def _recreate(): - self.collection = self.client.get_or_create_collection( - name=self.collection_name, - metadata={"hnsw:space": "cosine"}, - ) - - await self._run_sync_in_executor(_recreate) + self.collection = self.client.get_or_create_collection( + name=self.collection_name, + metadata={"hnsw:space": "cosine"}, + ) logger.info(f"Collection {self.collection_name} has been reset") + async def start(self) -> None: + """Initialize the ChromaDB collection. + + Creates or retrieves the collection with cosine similarity metric. + For local mode, creates the db_path directory if it doesn't exist. + """ + if self.is_local: + self.db_path.mkdir(parents=True, exist_ok=True) + logger.info(f"Initializing local ChromaDB at {self.db_path}") + self.client = chromadb.PersistentClient( + path=str(self.db_path), + settings=Settings(anonymized_telemetry=False), + ) + self.collection = self.client.get_or_create_collection( + name=self.collection_name, + metadata={"hnsw:space": "cosine"}, + ) + logger.info(f"ChromaDB collection {self.collection_name} initialized") + async def close(self): """Close the vector store and log the shutdown process.""" logger.info(f"ChromaDB vector store for collection {self.collection_name} closed") diff --git a/reme/core/vector_store/es_vector_store.py b/reme/core/vector_store/es_vector_store.py index 62b28bc1..e0df83a0 100644 --- a/reme/core/vector_store/es_vector_store.py +++ b/reme/core/vector_store/es_vector_store.py @@ -4,7 +4,7 @@ This module provides an Elasticsearch-based vector store that implements the Bas interface for high-performance dense vector storage and retrieval. """ -from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Any from loguru import logger @@ -30,8 +30,8 @@ class ESVectorStore(BaseVectorStore): def __init__( self, collection_name: str, + db_path: str | Path, embedding_model: BaseEmbeddingModel, - thread_pool: ThreadPoolExecutor, hosts: str | list[str] | None = None, basic_auth: tuple[str, str] | None = None, cloud_id: str | None = None, @@ -44,8 +44,8 @@ class ESVectorStore(BaseVectorStore): Args: collection_name: Name of the Elasticsearch index (converted to lowercase). + db_path: Database path (not used for remote Elasticsearch, kept for API consistency). embedding_model: Model instance used to generate vector embeddings. - thread_pool: ThreadPoolExecutor for running synchronous operations. hosts: Connection host(s) for the Elasticsearch cluster. basic_auth: Credentials for basic authentication. cloud_id: Deployment ID for Elastic Cloud. @@ -64,8 +64,8 @@ class ESVectorStore(BaseVectorStore): super().__init__( collection_name=collection_name, + db_path=db_path, embedding_model=embedding_model, - thread_pool=thread_pool, **kwargs, ) @@ -521,6 +521,14 @@ class ESVectorStore(BaseVectorStore): await self.create_collection(collection_name) logger.info(f"Collection reset to {collection_name}") + async def start(self) -> None: + """Initialize the Elasticsearch index. + + Creates the index with dense vector mappings if it doesn't exist. + """ + await super().start() + logger.info(f"Elasticsearch index {self.collection_name} initialized") + async def close(self): """Terminate the Elasticsearch client session and release resources.""" await self.client.close() diff --git a/reme/core/vector_store/local_vector_store.py b/reme/core/vector_store/local_vector_store.py index e3af5a49..fb2a4a1b 100644 --- a/reme/core/vector_store/local_vector_store.py +++ b/reme/core/vector_store/local_vector_store.py @@ -1,85 +1,95 @@ """Local file system vector store implementation for ReMe.""" import json -from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import numpy as np from loguru import logger from .base_vector_store import BaseVectorStore from ..embedding import BaseEmbeddingModel from ..schema import VectorNode -from ..utils import cosine_similarity +from ..utils import batch_cosine_similarity class LocalVectorStore(BaseVectorStore): - """Local file system-based vector store using JSON files and manual cosine similarity.""" + """Local file system-based vector store with in-memory caching. + + All operations are performed in memory after start(). + Changes are persisted to disk on close(). + """ def __init__( self, collection_name: str, + db_path: str | Path, embedding_model: BaseEmbeddingModel, - thread_pool: ThreadPoolExecutor, - root_path: str = "./local_vector_store", **kwargs, ): - """Initialize the local vector store with a root path and collection name.""" + """Initialize the local vector store with a db_path and collection name.""" super().__init__( collection_name=collection_name, + db_path=db_path, embedding_model=embedding_model, - thread_pool=thread_pool, **kwargs, ) - self.root_path = Path(root_path) - self.collection_path = self.root_path / collection_name - self.root_path.mkdir(parents=True, exist_ok=True) + # In-memory cache: vector_id -> VectorNode + self._cache: dict[str, VectorNode] = {} + self._dirty: bool = False # Track if cache has unsaved changes def _get_collection_path(self, collection_name: str) -> Path: """Get the file system path for a specific collection.""" - return self.root_path / collection_name + return self.db_path / collection_name def _get_node_file_path(self, vector_id: str, collection_name: str | None = None) -> Path: """Get the JSON file path for a specific vector node.""" col_path = self._get_collection_path(collection_name or self.collection_name) return col_path / f"{vector_id}.json" - def _save_node(self, node: VectorNode, collection_name: str | None = None): + def _save_node_to_disk(self, node: VectorNode): """Save a vector node to a JSON file on disk.""" - file_path = self._get_node_file_path(node.vector_id, collection_name) + file_path = self._get_node_file_path(node.vector_id) file_path.parent.mkdir(parents=True, exist_ok=True) - with open(file_path, "w", encoding="utf-8") as f: json.dump(node.model_dump(), f, ensure_ascii=False, indent=2) - def _load_node(self, vector_id: str, collection_name: str | None = None) -> VectorNode | None: - """Load a vector node from a JSON file.""" - file_path = self._get_node_file_path(vector_id, collection_name) - - if not file_path.exists(): - return None - - with open(file_path, "r", encoding="utf-8") as f: - data = json.load(f) - return VectorNode(**data) - - def _load_all_nodes(self, collection_name: str | None = None) -> list[VectorNode]: - """Load all vector nodes existing in a collection.""" - col_path = self._get_collection_path(collection_name or self.collection_name) - + def _load_all_from_disk(self) -> dict[str, VectorNode]: + """Load all vector nodes from disk into a dictionary.""" + col_path = self._get_collection_path(self.collection_name) if not col_path.exists(): - return [] + return {} - nodes = [] + nodes = {} for file_path in col_path.glob("*.json"): try: with open(file_path, "r", encoding="utf-8") as f: data = json.load(f) - nodes.append(VectorNode(**data)) + node = VectorNode(**data) + nodes[node.vector_id] = node except Exception as e: logger.warning(f"Failed to load node from {file_path}: {e}") - return nodes + def _flush_to_disk(self): + """Persist all cached nodes to disk.""" + col_path = self._get_collection_path(self.collection_name) + col_path.mkdir(parents=True, exist_ok=True) + + # Remove files that are no longer in cache + existing_files = set(col_path.glob("*.json")) + cached_ids = set(self._cache.keys()) + for file_path in existing_files: + vector_id = file_path.stem + if vector_id not in cached_ids: + file_path.unlink() + + # Write all cached nodes + for node in self._cache.values(): + self._save_node_to_disk(node) + + self._dirty = False + logger.info(f"Flushed {len(self._cache)} nodes to disk") + @staticmethod def _match_filters(node: VectorNode, filters: dict | None) -> bool: """Check if a vector node matches the provided metadata filters. @@ -114,11 +124,11 @@ class LocalVectorStore(BaseVectorStore): return True async def list_collections(self) -> list[str]: - """List all collection directories in the root path.""" - if not self.root_path.exists(): + """List all collection directories in the db_path.""" + if not self.db_path.exists(): return [] - return [d.name for d in self.root_path.iterdir() if d.is_dir() and not d.name.startswith(".")] + return [d.name for d in self.db_path.iterdir() if d.is_dir() and not d.name.startswith(".")] async def create_collection(self, collection_name: str, **kwargs): """Create a new collection directory.""" @@ -158,7 +168,7 @@ class LocalVectorStore(BaseVectorStore): logger.info(f"Copied collection {self.collection_name} to {collection_name}") async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): - """Insert vector nodes into the local store, generating embeddings if necessary.""" + """Insert vector nodes into the cache, generating embeddings if necessary.""" if isinstance(nodes, VectorNode): nodes = [nodes] @@ -171,8 +181,9 @@ class LocalVectorStore(BaseVectorStore): nodes_to_insert = nodes for node in nodes_to_insert: - self._save_node(node) + self._cache[node.vector_id] = node + self._dirty = True logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}") async def search( @@ -182,73 +193,69 @@ class LocalVectorStore(BaseVectorStore): filters: dict | None = None, **kwargs, ) -> list[VectorNode]: - """Search for nodes similar to the query using brute-force cosine similarity.""" + """Search for nodes similar to the query using batch cosine similarity.""" query_vector = await self.get_embedding(query) - all_nodes = self._load_all_nodes() - filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)] - scored_nodes = [] - for node in filtered_nodes: - if node.vector is None: - logger.warning(f"Node {node.vector_id} has no vector, skipping") - continue + # Filter nodes from cache + filtered_nodes = [node for node in self._cache.values() if self._match_filters(node, filters)] - try: - score = cosine_similarity(query_vector, node.vector) - scored_nodes.append((node, score)) - except ValueError as e: - logger.warning(f"Failed to calculate similarity for node {node.vector_id}: {e}") + # Separate nodes with and without vectors + nodes_with_vectors = [node for node in filtered_nodes if node.vector is not None] + if not nodes_with_vectors: + return [] - scored_nodes.sort(key=lambda x: x[1], reverse=True) + # Build matrix for batch similarity computation + node_vectors = np.array([node.vector for node in nodes_with_vectors]) + query_matrix = np.array([query_vector]) + # Compute similarities in batch: shape (1, num_nodes) -> flatten to (num_nodes,) + similarities = batch_cosine_similarity(query_matrix, node_vectors).flatten() + + # Apply score threshold if specified score_threshold = kwargs.get("score_threshold") + + # Pair nodes with scores and filter/sort + scored_nodes = list(zip(nodes_with_vectors, similarities)) if score_threshold is not None: scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold] + scored_nodes.sort(key=lambda x: x[1], reverse=True) scored_nodes = scored_nodes[:limit] + + # Attach scores to metadata results = [] for node, score in scored_nodes: - node.metadata["score"] = score + node.metadata["score"] = float(score) results.append(node) return results async def delete(self, vector_ids: str | list[str], **kwargs): - """Delete specific vector nodes by their IDs.""" + """Delete specific vector nodes by their IDs from cache.""" if isinstance(vector_ids, str): vector_ids = [vector_ids] deleted_count = 0 for vector_id in vector_ids: - file_path = self._get_node_file_path(vector_id) - if file_path.exists(): - file_path.unlink() + if vector_id in self._cache: + del self._cache[vector_id] deleted_count += 1 else: logger.warning(f"Node {vector_id} does not exist") + if deleted_count > 0: + self._dirty = True logger.info(f"Deleted {deleted_count} nodes from {self.collection_name}") async def delete_all(self, **kwargs): - """Remove all vectors from the collection.""" - col_path = self._get_collection_path(self.collection_name) - - if not col_path.exists(): - logger.warning(f"Collection {self.collection_name} does not exist") - return - - deleted_count = 0 - for file_path in col_path.glob("*.json"): - try: - file_path.unlink() - deleted_count += 1 - except Exception as e: - logger.warning(f"Failed to delete file {file_path}: {e}") - - logger.info(f"Deleted all {deleted_count} nodes from {self.collection_name}") + """Remove all vectors from the cache.""" + count = len(self._cache) + self._cache.clear() + self._dirty = True + logger.info(f"Deleted all {count} nodes from {self.collection_name}") async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): - """Update existing vector nodes with new data or embeddings.""" + """Update existing vector nodes in the cache.""" if isinstance(nodes, VectorNode): nodes = [nodes] @@ -262,23 +269,24 @@ class LocalVectorStore(BaseVectorStore): updated_count = 0 for node in nodes_to_update: - file_path = self._get_node_file_path(node.vector_id) - if file_path.exists(): - self._save_node(node) + if node.vector_id in self._cache: + self._cache[node.vector_id] = node updated_count += 1 else: logger.warning(f"Node {node.vector_id} does not exist, skipping update") + if updated_count > 0: + self._dirty = True logger.info(f"Updated {updated_count} nodes in {self.collection_name}") async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]: - """Retrieve one or more vector nodes by their unique IDs.""" + """Retrieve one or more vector nodes from cache by their unique IDs.""" is_single = isinstance(vector_ids, str) ids = [vector_ids] if is_single else vector_ids results = [] for vector_id in ids: - node = self._load_node(vector_id) + node = self._cache.get(vector_id) if node: results.append(node) else: @@ -291,9 +299,9 @@ class LocalVectorStore(BaseVectorStore): filters: dict | None = None, limit: int | None = None, sort_key: str | None = None, - reverse: bool = False, + reverse: bool = True, ) -> list[VectorNode]: - """List vector nodes in the collection with optional filtering and limits. + """List vector nodes from cache with optional filtering and limits. Args: filters: Dictionary of filter conditions to match vectors @@ -301,16 +309,14 @@ class LocalVectorStore(BaseVectorStore): sort_key: Key to sort the results by (e.g., field name in metadata). None for no sorting reverse: If True, sort in descending order; if False, sort in ascending order """ - all_nodes = self._load_all_nodes() - filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)] + filtered_nodes = [node for node in self._cache.values() if self._match_filters(node, filters)] # Apply sorting if sort_key is provided if sort_key: - # Sort with proper handling of None and missing values + def sort_key_func(node): value = node.metadata.get(sort_key) if value is None: - # Return appropriate default based on reverse flag return float("-inf") if not reverse else float("inf") return value @@ -321,6 +327,15 @@ class LocalVectorStore(BaseVectorStore): return filtered_nodes + async def start(self) -> None: + """Initialize the local vector store and load all nodes into memory.""" + await super().start() + self._cache = self._load_all_from_disk() + self._dirty = False + logger.info(f"Local vector store loaded {len(self._cache)} nodes from {self.collection_name}") + async def close(self): - """Close the vector store (no-op for local file system).""" + """Persist all cached data to disk and close the vector store.""" + if self._dirty: + self._flush_to_disk() logger.info("Local vector store closed") diff --git a/reme/core/vector_store/pgvector_store.py b/reme/core/vector_store/pgvector_store.py index 0ec7d6a8..bc665cd9 100644 --- a/reme/core/vector_store/pgvector_store.py +++ b/reme/core/vector_store/pgvector_store.py @@ -2,7 +2,7 @@ import json import re -from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Any from loguru import logger @@ -47,8 +47,8 @@ class PGVectorStore(BaseVectorStore): def __init__( self, collection_name: str, + db_path: str | Path, embedding_model: BaseEmbeddingModel, - thread_pool: ThreadPoolExecutor, host: str = "localhost", port: int = 5432, database: str = "postgres", @@ -72,8 +72,8 @@ class PGVectorStore(BaseVectorStore): super().__init__( collection_name=collection_name, + db_path=db_path, embedding_model=embedding_model, - thread_pool=thread_pool, **kwargs, ) @@ -117,12 +117,6 @@ class PGVectorStore(BaseVectorStore): return self._pool - async def _ensure_collection_exists(self): - """Check if the collection table exists and create it if missing.""" - collections = await self.list_collections() - if self.collection_name not in collections: - await self.create_collection(self.collection_name) - async def list_collections(self) -> list[str]: """List all available table names in the current database.""" pool = await self._get_pool() @@ -219,8 +213,6 @@ class PGVectorStore(BaseVectorStore): async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): """Insert or upsert vector nodes into the PostgreSQL collection.""" - await self._ensure_collection_exists() - if isinstance(nodes, VectorNode): nodes = [nodes] @@ -337,8 +329,6 @@ class PGVectorStore(BaseVectorStore): **kwargs, ) -> list[VectorNode]: """Perform vector similarity search with optional metadata filtering.""" - await self._ensure_collection_exists() - query_vector = await self.get_embedding(query) vector_str = f"[{','.join(map(str, query_vector))}]" pool = await self._get_pool() @@ -394,8 +384,6 @@ class PGVectorStore(BaseVectorStore): async def delete(self, vector_ids: str | list[str], **kwargs): """Remove specific vector records from the collection by their IDs.""" - await self._ensure_collection_exists() - if isinstance(vector_ids, str): vector_ids = [vector_ids] @@ -414,8 +402,6 @@ class PGVectorStore(BaseVectorStore): async def delete_all(self, **kwargs): """Remove all vectors from the collection.""" - await self._ensure_collection_exists() - pool = await self._get_pool() async with pool.acquire() as conn: result = await conn.execute(f"DELETE FROM {self.collection_name}") @@ -424,8 +410,6 @@ class PGVectorStore(BaseVectorStore): async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): """Update existing vector nodes with new content, embeddings, or metadata.""" - await self._ensure_collection_exists() - if isinstance(nodes, VectorNode): nodes = [nodes] @@ -474,8 +458,6 @@ class PGVectorStore(BaseVectorStore): async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None: """Retrieve vector nodes by their unique identifiers.""" - await self._ensure_collection_exists() - single_result = isinstance(vector_ids, str) if single_result: vector_ids = [vector_ids] @@ -531,8 +513,6 @@ class PGVectorStore(BaseVectorStore): sort_key: Key to sort the results by (e.g., field name in metadata). None for no sorting reverse: If True, sort in descending order; if False, sort in ascending order """ - await self._ensure_collection_exists() - pool = await self._get_pool() filter_clause, filter_params = self._build_filter_clause(filters) @@ -613,6 +593,15 @@ class PGVectorStore(BaseVectorStore): await self.create_collection(collection_name) logger.info(f"Collection reset to {collection_name}") + async def start(self) -> None: + """Initialize the PGVector store. + + Creates the connection pool and ensures the collection table exists. + """ + await self._get_pool() + await super().start() + logger.info(f"PGVector collection {self.collection_name} initialized") + async def close(self): """Terminate the database connection pool and release associated resources.""" if self._pool is not None: diff --git a/reme/core/vector_store/qdrant_vector_store.py b/reme/core/vector_store/qdrant_vector_store.py index 93ccee70..b981a23f 100644 --- a/reme/core/vector_store/qdrant_vector_store.py +++ b/reme/core/vector_store/qdrant_vector_store.py @@ -1,6 +1,6 @@ """Qdrant vector store implementation for the ReMe project.""" -from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Any from loguru import logger @@ -42,11 +42,10 @@ class QdrantVectorStore(BaseVectorStore): def __init__( self, collection_name: str, + db_path: str | Path, embedding_model: BaseEmbeddingModel, - thread_pool: ThreadPoolExecutor, host: str | None = None, port: int = 6333, - path: str | None = None, url: str | None = None, api_key: str | None = None, https: bool | None = None, @@ -60,11 +59,10 @@ class QdrantVectorStore(BaseVectorStore): Args: collection_name: Name of the collection. + db_path: Local storage path for on-disk/in-memory mode. embedding_model: Model used for generating vector embeddings. - thread_pool: ThreadPoolExecutor for running synchronous operations. host: Server host address. port: HTTP port for the server. - path: Local storage path for on-disk/in-memory mode. url: Full connection URL. api_key: Authentication key for Qdrant Cloud. https: Use secure connection if True. @@ -81,26 +79,23 @@ class QdrantVectorStore(BaseVectorStore): super().__init__( collection_name=collection_name, + db_path=db_path, embedding_model=embedding_model, - thread_pool=thread_pool, **kwargs, ) - client_kwargs = {k: v for k, v in kwargs.items() if k != "thread_pool"} + self.is_local = host is None and url is None and api_key is None + self.client: AsyncQdrantClient - self.client = AsyncQdrantClient( - host=host, - port=port, - path=path, - url=url, - api_key=api_key, - https=https, - grpc_port=grpc_port, - prefer_grpc=prefer_grpc, - **client_kwargs, - ) + # Store connection parameters for deferred initialization in start() + self._host = host + self._port = port + self._url = url + self._api_key = api_key + self._https = https + self._grpc_port = grpc_port + self._prefer_grpc = prefer_grpc - self.is_local = path is not None distance_map = { "cosine": Distance.COSINE, "euclid": Distance.EUCLID, @@ -526,6 +521,29 @@ class QdrantVectorStore(BaseVectorStore): return results + async def start(self) -> None: + """Initialize the Qdrant collection. + + Creates the collection if it doesn't exist with configured vector parameters. + For local mode, creates the db_path directory if it doesn't exist. + """ + if self.is_local: + self.db_path.mkdir(parents=True, exist_ok=True) + self.client = AsyncQdrantClient(path=str(self.db_path)) + else: + self.client = AsyncQdrantClient( + host=self._host, + port=self._port, + url=self._url, + api_key=self._api_key, + https=self._https, + grpc_port=self._grpc_port, + prefer_grpc=self._prefer_grpc, + **self.kwargs, + ) + await super().start() + logger.info(f"Qdrant collection {self.collection_name} initialized") + async def close(self): """Close the AsyncQdrantClient connection and release resources.""" await self.client.close() diff --git a/reme/extension/__init__.py b/reme/extension/__init__.py new file mode 100644 index 00000000..84bb2ea9 --- /dev/null +++ b/reme/extension/__init__.py @@ -0,0 +1,18 @@ +"""Extension operations and tools.""" + +from .simple_chat import SimpleChat +from .stream_chat import StreamChat +from .test_op import TestOp +from .translate_ts import TranslateTs +from ..core.registry_factory import R + +__all__ = [ + "SimpleChat", + "StreamChat", + "TestOp", + "TranslateTs", +] + +for name in __all__: + op_class = globals()[name] + R.ops.register(op_class) diff --git a/reme/agent/chat/simple_chat.py b/reme/extension/simple_chat.py similarity index 95% rename from reme/agent/chat/simple_chat.py rename to reme/extension/simple_chat.py index 36181346..04a36656 100644 --- a/reme/agent/chat/simple_chat.py +++ b/reme/extension/simple_chat.py @@ -2,9 +2,9 @@ from loguru import logger -from ...core.enumeration import Role -from ...core.op import BaseTool -from ...core.schema import Message, ToolCall +from ..core.enumeration import Role +from ..core.op import BaseTool +from ..core.schema import Message, ToolCall class SimpleChat(BaseTool): diff --git a/reme/agent/chat/stream_chat.py b/reme/extension/stream_chat.py similarity index 94% rename from reme/agent/chat/stream_chat.py rename to reme/extension/stream_chat.py index f5cd7a5a..dd0aeda7 100644 --- a/reme/agent/chat/stream_chat.py +++ b/reme/extension/stream_chat.py @@ -2,9 +2,9 @@ from loguru import logger -from ...core.enumeration import Role, ChunkEnum -from ...core.op import BaseTool -from ...core.schema import Message, ToolCall +from ..core.enumeration import Role, ChunkEnum +from ..core.op import BaseTool +from ..core.schema import Message, ToolCall class StreamChat(BaseTool): diff --git a/reme/workflow/gallery/test_op.py b/reme/extension/test_op.py similarity index 91% rename from reme/workflow/gallery/test_op.py rename to reme/extension/test_op.py index 1efa53af..acc1ddcf 100644 --- a/reme/workflow/gallery/test_op.py +++ b/reme/extension/test_op.py @@ -2,7 +2,7 @@ from loguru import logger -from ...core.op import BaseOp +from ..core.op import BaseOp class TestOp(BaseOp): diff --git a/reme/workflow/gallery/translate_ts.py b/reme/extension/translate_ts.py similarity index 94% rename from reme/workflow/gallery/translate_ts.py rename to reme/extension/translate_ts.py index 8bac9166..eb032386 100644 --- a/reme/workflow/gallery/translate_ts.py +++ b/reme/extension/translate_ts.py @@ -4,9 +4,9 @@ from pathlib import Path from loguru import logger -from ...core.enumeration import Role -from ...core.op import BaseOp -from ...core.schema import Message +from ..core.enumeration import Role +from ..core.op import BaseOp +from ..core.schema import Message class TranslateTs(BaseOp): diff --git a/reme/workflow/gallery/translate_ts.yaml b/reme/extension/translate_ts.yaml similarity index 100% rename from reme/workflow/gallery/translate_ts.yaml rename to reme/extension/translate_ts.yaml diff --git a/reme/memory/__init__.py b/reme/memory/__init__.py new file mode 100644 index 00000000..ac825bc6 --- /dev/null +++ b/reme/memory/__init__.py @@ -0,0 +1,11 @@ +"""memory""" + +from . import file_based +from . import tools +from . import vector_based + +__all__ = [ + "file_based", + "tools", + "vector_based", +] diff --git a/reme/memory/file_based/__init__.py b/reme/memory/file_based/__init__.py new file mode 100644 index 00000000..2f994785 --- /dev/null +++ b/reme/memory/file_based/__init__.py @@ -0,0 +1,18 @@ +"""File-based memory operations.""" + +from .fb_cli import FbCli +from .fb_compactor import FbCompactor +from .fb_context_checker import FbContextChecker +from .fb_summarizer import FbSummarizer +from ...core.registry_factory import R + +__all__ = [ + "FbCli", + "FbCompactor", + "FbContextChecker", + "FbSummarizer", +] + +for name in __all__: + op_class = globals()[name] + R.ops.register(op_class) diff --git a/reme/agent/chat/fs_cli.py b/reme/memory/file_based/fb_cli.py similarity index 94% rename from reme/agent/chat/fs_cli.py rename to reme/memory/file_based/fb_cli.py index 2ee1f352..6299c1a4 100644 --- a/reme/agent/chat/fs_cli.py +++ b/reme/memory/file_based/fb_cli.py @@ -1,4 +1,4 @@ -"""FsCli system prompt""" +"""FbCli system prompt""" import asyncio from datetime import datetime @@ -9,12 +9,12 @@ from loguru import logger from ...core.enumeration import Role, ChunkEnum from ...core.op import BaseReactStream from ...core.schema import Message, StreamChunk +from ...core.tools import BashTool, LsTool, ReadTool, WriteTool, EditTool from ...core.utils import format_messages -from ...tool.fs import BashTool, LsTool, ReadTool, WriteTool, EditTool -class FsCli(BaseReactStream): - """FsCli agent with system prompt.""" +class FbCli(BaseReactStream): + """FbCli agent with system prompt.""" def __init__( self, @@ -50,11 +50,11 @@ class FsCli(BaseReactStream): remaining_tasks.append(task) self.summary_tasks = remaining_tasks - from ..fs import FsSummarizer + from .fb_summarizer import FbSummarizer # Summarize current conversation and save to memory files current_date = datetime.now().strftime("%Y-%m-%d") - summarizer = FsSummarizer( + summarizer = FbSummarizer( tools=[ BashTool(cwd=self.working_dir), LsTool(cwd=self.working_dir), @@ -94,10 +94,10 @@ class FsCli(BaseReactStream): async def context_check(self) -> dict: """Check if messages exceed token limits.""" # Import required modules - from ..fs import FsContextChecker + from .fb_context_checker import FbContextChecker # Step 1: Check and find cut point - checker = FsContextChecker( + checker = FbContextChecker( context_window_tokens=self.context_window_tokens, reserve_tokens=self.reserve_tokens, keep_recent_tokens=self.keep_recent_tokens, @@ -120,7 +120,7 @@ class FsCli(BaseReactStream): return "No history to compact." # Import required modules - from ..fs import FsCompactor + from .fb_compactor import FbCompactor # Step 1: Check and find cut point cut_result = await self.context_check() @@ -137,7 +137,7 @@ class FsCli(BaseReactStream): turn_prefix_messages = cut_result.get("turn_prefix_messages", []) left_messages = cut_result.get("left_messages", []) - compactor = FsCompactor(language=self.language) + compactor = FbCompactor(language=self.language) summary_content = await compactor.call( messages_to_summarize=messages_to_summarize, turn_prefix_messages=turn_prefix_messages, diff --git a/reme/agent/chat/fs_cli.yaml b/reme/memory/file_based/fb_cli.yaml similarity index 100% rename from reme/agent/chat/fs_cli.yaml rename to reme/memory/file_based/fb_cli.yaml diff --git a/reme/agent/fs/fs_compactor.py b/reme/memory/file_based/fb_compactor.py similarity index 99% rename from reme/agent/fs/fs_compactor.py rename to reme/memory/file_based/fb_compactor.py index 4ddb7dd3..75d1fd54 100644 --- a/reme/agent/fs/fs_compactor.py +++ b/reme/memory/file_based/fb_compactor.py @@ -8,7 +8,7 @@ from ...core.schema import Message from ...core.utils import format_messages -class FsCompactor(BaseOp): +class FbCompactor(BaseOp): """Generate summaries for conversation history compaction.""" def __init__(self, return_prompt: bool = False, **kwargs): diff --git a/reme/agent/fs/fs_compactor.yaml b/reme/memory/file_based/fb_compactor.yaml similarity index 100% rename from reme/agent/fs/fs_compactor.yaml rename to reme/memory/file_based/fb_compactor.yaml diff --git a/reme/agent/fs/fs_context_checker.py b/reme/memory/file_based/fb_context_checker.py similarity index 99% rename from reme/agent/fs/fs_context_checker.py rename to reme/memory/file_based/fb_context_checker.py index 6a472591..99b6a580 100644 --- a/reme/agent/fs/fs_context_checker.py +++ b/reme/memory/file_based/fb_context_checker.py @@ -7,7 +7,7 @@ from ...core.op import BaseOp from ...core.schema import CutPointResult, Message -class FsContextChecker(BaseOp): +class FbContextChecker(BaseOp): """Check if context exceeds token limits and find cut point for compaction.""" def __init__( diff --git a/reme/agent/fs/fs_summarizer.py b/reme/memory/file_based/fb_summarizer.py similarity index 98% rename from reme/agent/fs/fs_summarizer.py rename to reme/memory/file_based/fb_summarizer.py index d924983f..67583d1f 100644 --- a/reme/agent/fs/fs_summarizer.py +++ b/reme/memory/file_based/fb_summarizer.py @@ -10,7 +10,7 @@ from ...core.schema import Message from ...core.utils import format_messages -class FsSummarizer(BaseReact): +class FbSummarizer(BaseReact): """Retrieve personal memories through vector search and history reading.""" def __init__( diff --git a/reme/agent/fs/fs_summarizer.yaml b/reme/memory/file_based/fb_summarizer.yaml similarity index 100% rename from reme/agent/fs/fs_summarizer.yaml rename to reme/memory/file_based/fb_summarizer.yaml diff --git a/reme/tool/memory/__init__.py b/reme/memory/tools/__init__.py similarity index 58% rename from reme/tool/memory/__init__.py rename to reme/memory/tools/__init__.py index d7ee2e93..af9b851f 100644 --- a/reme/tool/memory/__init__.py +++ b/reme/memory/tools/__init__.py @@ -1,51 +1,60 @@ """memory tools""" from .base_memory_tool import BaseMemoryTool + +# chunk tools +from .chunk.memory_get import MemoryGet +from .chunk.memory_search import MemorySearch from .delegate_task import DelegateTask + +# history tools from .history.add_history import AddHistory from .history.read_history import ReadHistory from .history.read_history_v2 import ReadHistoryV2 + +# profiles tools from .profiles.add_draft_and_read_all_profiles import AddDraftAndReadAllProfiles from .profiles.add_profile import AddProfile from .profiles.delete_profile import DeleteProfile -from .profiles.profile_handler import ProfileHandler from .profiles.read_all_profiles import ReadAllProfiles from .profiles.update_profile import UpdateProfile from .profiles.update_profiles_v1 import UpdateProfilesV1 -from .vector.add_and_retrieve_similar_memory import AddAndRetrieveSimilarMemory -from .vector.add_draft_and_retrieve_similar_memory import AddDraftAndRetrieveSimilarMemory -from .vector.add_memory import AddMemory -from .vector.delete_memory import DeleteMemory -from .vector.memory_handler import MemoryHandler -from .vector.retrieve_memory import RetrieveMemory -from .vector.retrieve_recent_memory import RetrieveRecentMemory -from .vector.update_memory import UpdateMemory -from .vector.update_memory_v1 import UpdateMemoryV1 -from .vector.update_memory_v2 import UpdateMemoryV2 + +# record tools +from .record.add_and_retrieve_similar_memory import AddAndRetrieveSimilarMemory +from .record.add_draft_and_retrieve_similar_memory import AddDraftAndRetrieveSimilarMemory +from .record.add_memory import AddMemory +from .record.delete_memory import DeleteMemory +from .record.retrieve_memory import RetrieveMemory +from .record.retrieve_recent_memory import RetrieveRecentMemory +from .record.update_memory import UpdateMemory +from .record.update_memory_v1 import UpdateMemoryV1 +from .record.update_memory_v2 import UpdateMemoryV2 from ...core import R __all__ = [ - # Base + # base "BaseMemoryTool", "DelegateTask", - # History + # chunk tools + "MemoryGet", + "MemorySearch", + # history tools "AddHistory", "ReadHistory", "ReadHistoryV2", - # Profiles + # profiles tools "AddDraftAndReadAllProfiles", "AddProfile", - "ProfileHandler", + "DeleteProfile", "ReadAllProfiles", "UpdateProfile", - "DeleteProfile", "UpdateProfilesV1", - # Vector + # record tools "AddAndRetrieveSimilarMemory", "AddDraftAndRetrieveSimilarMemory", "AddMemory", "DeleteMemory", - "MemoryHandler", "RetrieveMemory", "RetrieveRecentMemory", "UpdateMemory", @@ -55,5 +64,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.ops.register(tool_class) + R.ops.register(tool_class) diff --git a/reme/tool/memory/base_memory_tool.py b/reme/memory/tools/base_memory_tool.py similarity index 100% rename from reme/tool/memory/base_memory_tool.py rename to reme/memory/tools/base_memory_tool.py diff --git a/reme/agent/memory/tool/__init__.py b/reme/memory/tools/chunk/__init__.py similarity index 100% rename from reme/agent/memory/tool/__init__.py rename to reme/memory/tools/chunk/__init__.py diff --git a/reme/tool/fs/fs_memory_get.py b/reme/memory/tools/chunk/memory_get.py similarity index 75% rename from reme/tool/fs/fs_memory_get.py rename to reme/memory/tools/chunk/memory_get.py index 7c19fbe7..572a4a26 100644 --- a/reme/tool/fs/fs_memory_get.py +++ b/reme/memory/tools/chunk/memory_get.py @@ -3,16 +3,20 @@ import os from pathlib import Path -from reme.core.schema import ToolCall -from .base_fs_tool import BaseFsTool +from loguru import logger + +from ....core import RuntimeContext +from ....core.op import BaseTool +from ....core.schema import ToolCall -class FsMemoryGet(BaseFsTool): +class MemoryGet(BaseTool): """Read specific snippets from memory files.""" def __init__(self, cwd: str | None = None, **kwargs): """Initialize memory get tool.""" - kwargs.setdefault("name", "memory_get") + kwargs.setdefault("max_retries", 1) + kwargs.setdefault("raise_exception", False) super().__init__(**kwargs) self.cwd = cwd or os.getcwd() @@ -87,3 +91,23 @@ class FsMemoryGet(BaseFsTool): # Extract slice (1-indexed to 0-indexed conversion) selected = lines[start - 1 : start - 1 + count] return "\n".join(selected) + + 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/fs_memory_search.py b/reme/memory/tools/chunk/memory_search.py similarity index 73% rename from reme/tool/fs/fs_memory_search.py rename to reme/memory/tools/chunk/memory_search.py index df8e561f..7423b331 100644 --- a/reme/tool/fs/fs_memory_search.py +++ b/reme/memory/tools/chunk/memory_search.py @@ -2,12 +2,15 @@ import json -from reme.core.enumeration import MemorySource -from reme.core.schema import ToolCall -from .base_fs_tool import BaseFsTool +from loguru import logger + +from ....core.enumeration import MemorySource +from ....core.op import BaseTool +from ....core.runtime_context import RuntimeContext +from ....core.schema import ToolCall -class FsMemorySearch(BaseFsTool): +class MemorySearch(BaseTool): """Semantically search MEMORY.md and memory files.""" def __init__( @@ -21,7 +24,8 @@ class FsMemorySearch(BaseFsTool): ): """Initialize memory search tool.""" assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}" - kwargs.setdefault("name", "memory_search") + kwargs.setdefault("max_retries", 1) + kwargs.setdefault("raise_exception", False) super().__init__(**kwargs) self.sources = sources or [MemorySource.MEMORY] self.min_score = min_score @@ -73,8 +77,8 @@ class FsMemorySearch(BaseFsTool): isinstance(max_results, int) and max_results > 0 ), f"max_results must be a positive integer, got {max_results}" - # Use hybrid_search from memory_store - results = await self.memory_store.hybrid_search( + # Use hybrid_search from file_store + results = await self.file_store.hybrid_search( query=query, limit=max_results, sources=self.sources, @@ -86,3 +90,23 @@ class FsMemorySearch(BaseFsTool): results = [r for r in results if r.score >= min_score] return json.dumps([result.model_dump(exclude_none=True) for result in results], indent=2, ensure_ascii=False) + + 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/memory/delegate_task.py b/reme/memory/tools/delegate_task.py similarity index 98% rename from reme/tool/memory/delegate_task.py rename to reme/memory/tools/delegate_task.py index b0d77d07..d3a0298d 100644 --- a/reme/tool/memory/delegate_task.py +++ b/reme/memory/tools/delegate_task.py @@ -3,7 +3,7 @@ from loguru import logger from .base_memory_tool import BaseMemoryTool -from ...agent.memory import BaseMemoryAgent +from ..vector_based import BaseMemoryAgent from ...core.enumeration import MemoryType from ...core.schema import ToolCall diff --git a/reme/tool/memory/history/__init__.py b/reme/memory/tools/history/__init__.py similarity index 100% rename from reme/tool/memory/history/__init__.py rename to reme/memory/tools/history/__init__.py diff --git a/reme/tool/memory/history/add_history.py b/reme/memory/tools/history/add_history.py similarity index 100% rename from reme/tool/memory/history/add_history.py rename to reme/memory/tools/history/add_history.py diff --git a/reme/tool/memory/history/read_history.py b/reme/memory/tools/history/read_history.py similarity index 100% rename from reme/tool/memory/history/read_history.py rename to reme/memory/tools/history/read_history.py diff --git a/reme/tool/memory/history/read_history_v2.py b/reme/memory/tools/history/read_history_v2.py similarity index 100% rename from reme/tool/memory/history/read_history_v2.py rename to reme/memory/tools/history/read_history_v2.py diff --git a/reme/tool/memory/profiles/__init__.py b/reme/memory/tools/profiles/__init__.py similarity index 100% rename from reme/tool/memory/profiles/__init__.py rename to reme/memory/tools/profiles/__init__.py diff --git a/reme/tool/memory/profiles/add_draft_and_read_all_profiles.py b/reme/memory/tools/profiles/add_draft_and_read_all_profiles.py similarity index 100% rename from reme/tool/memory/profiles/add_draft_and_read_all_profiles.py rename to reme/memory/tools/profiles/add_draft_and_read_all_profiles.py diff --git a/reme/tool/memory/profiles/add_profile.py b/reme/memory/tools/profiles/add_profile.py similarity index 100% rename from reme/tool/memory/profiles/add_profile.py rename to reme/memory/tools/profiles/add_profile.py diff --git a/reme/tool/memory/profiles/delete_profile.py b/reme/memory/tools/profiles/delete_profile.py similarity index 100% rename from reme/tool/memory/profiles/delete_profile.py rename to reme/memory/tools/profiles/delete_profile.py diff --git a/reme/tool/memory/profiles/profile_handler.py b/reme/memory/tools/profiles/profile_handler.py similarity index 100% rename from reme/tool/memory/profiles/profile_handler.py rename to reme/memory/tools/profiles/profile_handler.py diff --git a/reme/tool/memory/profiles/read_all_profiles.py b/reme/memory/tools/profiles/read_all_profiles.py similarity index 100% rename from reme/tool/memory/profiles/read_all_profiles.py rename to reme/memory/tools/profiles/read_all_profiles.py diff --git a/reme/tool/memory/profiles/update_profile.py b/reme/memory/tools/profiles/update_profile.py similarity index 100% rename from reme/tool/memory/profiles/update_profile.py rename to reme/memory/tools/profiles/update_profile.py diff --git a/reme/tool/memory/profiles/update_profiles_v1.py b/reme/memory/tools/profiles/update_profiles_v1.py similarity index 100% rename from reme/tool/memory/profiles/update_profiles_v1.py rename to reme/memory/tools/profiles/update_profiles_v1.py diff --git a/reme/tool/memory/vector/__init__.py b/reme/memory/tools/record/__init__.py similarity index 100% rename from reme/tool/memory/vector/__init__.py rename to reme/memory/tools/record/__init__.py diff --git a/reme/tool/memory/vector/add_and_retrieve_similar_memory.py b/reme/memory/tools/record/add_and_retrieve_similar_memory.py similarity index 100% rename from reme/tool/memory/vector/add_and_retrieve_similar_memory.py rename to reme/memory/tools/record/add_and_retrieve_similar_memory.py diff --git a/reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py b/reme/memory/tools/record/add_draft_and_retrieve_similar_memory.py similarity index 100% rename from reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py rename to reme/memory/tools/record/add_draft_and_retrieve_similar_memory.py diff --git a/reme/tool/memory/vector/add_memory.py b/reme/memory/tools/record/add_memory.py similarity index 100% rename from reme/tool/memory/vector/add_memory.py rename to reme/memory/tools/record/add_memory.py diff --git a/reme/tool/memory/vector/delete_memory.py b/reme/memory/tools/record/delete_memory.py similarity index 100% rename from reme/tool/memory/vector/delete_memory.py rename to reme/memory/tools/record/delete_memory.py diff --git a/reme/tool/memory/vector/memory_handler.py b/reme/memory/tools/record/memory_handler.py similarity index 99% rename from reme/tool/memory/vector/memory_handler.py rename to reme/memory/tools/record/memory_handler.py index 2b304df7..7997ac86 100644 --- a/reme/tool/memory/vector/memory_handler.py +++ b/reme/memory/tools/record/memory_handler.py @@ -3,7 +3,7 @@ import numpy as np from loguru import logger -from ....core.context import ServiceContext +from ....core import ServiceContext from ....core.enumeration import MemoryType from ....core.schema import MemoryNode from ....core.utils.common_utils import batch_cosine_similarity diff --git a/reme/tool/memory/vector/retrieve_memory.py b/reme/memory/tools/record/retrieve_memory.py similarity index 100% rename from reme/tool/memory/vector/retrieve_memory.py rename to reme/memory/tools/record/retrieve_memory.py diff --git a/reme/tool/memory/vector/retrieve_recent_memory.py b/reme/memory/tools/record/retrieve_recent_memory.py similarity index 100% rename from reme/tool/memory/vector/retrieve_recent_memory.py rename to reme/memory/tools/record/retrieve_recent_memory.py diff --git a/reme/tool/memory/vector/update_memory.py b/reme/memory/tools/record/update_memory.py similarity index 100% rename from reme/tool/memory/vector/update_memory.py rename to reme/memory/tools/record/update_memory.py diff --git a/reme/tool/memory/vector/update_memory_v1.py b/reme/memory/tools/record/update_memory_v1.py similarity index 100% rename from reme/tool/memory/vector/update_memory_v1.py rename to reme/memory/tools/record/update_memory_v1.py diff --git a/reme/tool/memory/vector/update_memory_v2.py b/reme/memory/tools/record/update_memory_v2.py similarity index 100% rename from reme/tool/memory/vector/update_memory_v2.py rename to reme/memory/tools/record/update_memory_v2.py diff --git a/reme/agent/memory/__init__.py b/reme/memory/vector_based/__init__.py similarity index 71% rename from reme/agent/memory/__init__.py rename to reme/memory/vector_based/__init__.py index a1e941ee..69469603 100644 --- a/reme/agent/memory/__init__.py +++ b/reme/memory/vector_based/__init__.py @@ -1,12 +1,8 @@ """memory agent""" from .base_memory_agent import BaseMemoryAgent -from .personal.personal_halumem_retriever import PersonalHalumemRetriever -from .personal.personal_halumem_summarizer import PersonalHalumemSummarizer from .personal.personal_retriever import PersonalRetriever from .personal.personal_summarizer import PersonalSummarizer -from .personal.personal_v1_retriever import PersonalV1Retriever -from .personal.personal_v1_summarizer import PersonalV1Summarizer from .procedural.procedural_retriever import ProceduralRetriever from .procedural.procedural_summarizer import ProceduralSummarizer from .reme_retriever import ReMeRetriever @@ -19,10 +15,6 @@ __all__ = [ "BaseMemoryAgent", "PersonalRetriever", "PersonalSummarizer", - "PersonalV1Retriever", - "PersonalV1Summarizer", - "PersonalHalumemRetriever", - "PersonalHalumemSummarizer", "ProceduralRetriever", "ProceduralSummarizer", "ReMeRetriever", diff --git a/reme/agent/memory/base_memory_agent.py b/reme/memory/vector_based/base_memory_agent.py similarity index 100% rename from reme/agent/memory/base_memory_agent.py rename to reme/memory/vector_based/base_memory_agent.py diff --git a/reme/memory/vector_based/personal/__init__.py b/reme/memory/vector_based/personal/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/agent/memory/personal/personal_v1_retriever.py b/reme/memory/vector_based/personal/personal_retriever.py similarity index 95% rename from reme/agent/memory/personal/personal_v1_retriever.py rename to reme/memory/vector_based/personal/personal_retriever.py index 525d597a..b8803717 100644 --- a/reme/agent/memory/personal/personal_v1_retriever.py +++ b/reme/memory/vector_based/personal/personal_retriever.py @@ -7,13 +7,15 @@ from ....core.schema import Message from ....core.utils import format_messages -class PersonalV1Retriever(BaseMemoryAgent): +class PersonalRetriever(BaseMemoryAgent): """Retrieve personal memories through vector search and history reading. + clear && python benchmark/halumem/eval_reme.py \ --data_path /Users/yuli/workspace/HaluMem/data/HaluMem-Medium.jsonl \ - --reme_model_name qwen3-30b-a3b-instruct-2507 \ - --algo_version v1 \ - --enable_thinking_params + --reme_model_name qwen3.5-plus \ + --batch_size 10000 \ + --algo_version default + 📊 Question Answering (with LLM answer): Correct (all): 0.8537 Hallucination (all): 0.1159 diff --git a/reme/agent/memory/personal/personal_v1_retriever.yaml b/reme/memory/vector_based/personal/personal_retriever.yaml similarity index 100% rename from reme/agent/memory/personal/personal_v1_retriever.yaml rename to reme/memory/vector_based/personal/personal_retriever.yaml diff --git a/reme/agent/memory/personal/personal_summarizer.py b/reme/memory/vector_based/personal/personal_summarizer.py similarity index 80% rename from reme/agent/memory/personal/personal_summarizer.py rename to reme/memory/vector_based/personal/personal_summarizer.py index 624cbcce..8e570718 100644 --- a/reme/agent/memory/personal/personal_summarizer.py +++ b/reme/memory/vector_based/personal/personal_summarizer.py @@ -16,35 +16,28 @@ class PersonalSummarizer(BaseMemoryAgent): async def _build_s1_messages(self) -> list[Message]: return [ Message( - role=Role.SYSTEM, + role=Role.USER, content=self.prompt_format( - prompt_name="system_prompt_s1", + prompt_name="user_message_s1", context=self.context.history_node.content, memory_type=self.memory_type.value, memory_target=self.memory_target, ), ), - Message( - role=Role.USER, - content=self.get_prompt("user_message_s1"), - ), ] - async def _build_s2_messages(self) -> list[Message]: + async def _build_s2_messages(self, profiles: str) -> list[Message]: return [ Message( - role=Role.SYSTEM, + role=Role.USER, content=self.prompt_format( - prompt_name="system_prompt_s2", + prompt_name="user_message_s2", + profiles=profiles, context=self.context.history_node.content, memory_type=self.memory_type.value, memory_target=self.memory_target, ), ), - Message( - role=Role.USER, - content=self.get_prompt("user_message_s2"), - ), ] async def _acting_step( @@ -72,9 +65,12 @@ class PersonalSummarizer(BaseMemoryAgent): async def execute(self): memory_tools = [] profile_tools = [] + read_all_profiles_tool: BaseTool | None = None for i, tool in enumerate(self.tools): tool_name = tool.tool_call.name - if "_memory" in tool_name: + if tool_name == "read_all_profiles": + read_all_profiles_tool = tool + elif "_memory" in tool_name: memory_tools.append(tool) elif "_profile" in tool_name: profile_tools.append(tool) @@ -89,9 +85,17 @@ class PersonalSummarizer(BaseMemoryAgent): logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage) + if read_all_profiles_tool is not None: + profiles = await read_all_profiles_tool.call( + memory_target=self.memory_target, + service_context=self.service_context, + ) + else: + profiles = "" + if profile_tools: stage = "s2-profile" - messages_s2 = await self._build_s2_messages() + messages_s2 = await self._build_s2_messages(profiles) for i, message in enumerate(messages_s2): role = message.name or message.role logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") @@ -99,7 +103,9 @@ class PersonalSummarizer(BaseMemoryAgent): else: tools_s2, messages_s2, success_s2 = [], [], True - answer = (messages_s1[-1].content if success_s1 else "") + (messages_s2[-1].content if success_s2 else "") + answer = (messages_s1[-1].content if success_s1 and messages_s1 else "") + ( + messages_s2[-1].content if success_s2 and messages_s2 else "" + ) success = success_s1 and success_s2 messages = messages_s1 + messages_s2 tools = tools_s1 + tools_s2 diff --git a/reme/agent/memory/personal/personal_v1_summarizer.yaml b/reme/memory/vector_based/personal/personal_summarizer.yaml similarity index 100% rename from reme/agent/memory/personal/personal_v1_summarizer.yaml rename to reme/memory/vector_based/personal/personal_summarizer.yaml diff --git a/reme/memory/vector_based/procedural/__init__.py b/reme/memory/vector_based/procedural/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/agent/memory/procedural/procedural_retriever.py b/reme/memory/vector_based/procedural/procedural_retriever.py similarity index 100% rename from reme/agent/memory/procedural/procedural_retriever.py rename to reme/memory/vector_based/procedural/procedural_retriever.py diff --git a/reme/agent/memory/procedural/procedural_summarizer.py b/reme/memory/vector_based/procedural/procedural_summarizer.py similarity index 100% rename from reme/agent/memory/procedural/procedural_summarizer.py rename to reme/memory/vector_based/procedural/procedural_summarizer.py diff --git a/reme/agent/memory/reme_retriever.py b/reme/memory/vector_based/reme_retriever.py similarity index 100% rename from reme/agent/memory/reme_retriever.py rename to reme/memory/vector_based/reme_retriever.py diff --git a/reme/agent/memory/reme_retriever.yaml b/reme/memory/vector_based/reme_retriever.yaml similarity index 100% rename from reme/agent/memory/reme_retriever.yaml rename to reme/memory/vector_based/reme_retriever.yaml diff --git a/reme/agent/memory/reme_summarizer.py b/reme/memory/vector_based/reme_summarizer.py similarity index 100% rename from reme/agent/memory/reme_summarizer.py rename to reme/memory/vector_based/reme_summarizer.py diff --git a/reme/agent/memory/reme_summarizer.yaml b/reme/memory/vector_based/reme_summarizer.yaml similarity index 100% rename from reme/agent/memory/reme_summarizer.yaml rename to reme/memory/vector_based/reme_summarizer.yaml diff --git a/reme/memory/vector_based/tool/__init__.py b/reme/memory/vector_based/tool/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/agent/memory/tool/tool_retriever.py b/reme/memory/vector_based/tool/tool_retriever.py similarity index 100% rename from reme/agent/memory/tool/tool_retriever.py rename to reme/memory/vector_based/tool/tool_retriever.py diff --git a/reme/agent/memory/tool/tool_summarizer.py b/reme/memory/vector_based/tool/tool_summarizer.py similarity index 100% rename from reme/agent/memory/tool/tool_summarizer.py rename to reme/memory/vector_based/tool/tool_summarizer.py diff --git a/reme/reme.py b/reme/reme.py index cab3a2c4..1c1990a3 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -3,41 +3,32 @@ import sys from pathlib import Path -from .agent.memory import ( - BaseMemoryAgent, - ReMeSummarizer, - ReMeRetriever, - PersonalV1Summarizer, - PersonalV1Retriever, - PersonalHalumemSummarizer, - PersonalHalumemRetriever, - PersonalSummarizer, - PersonalRetriever, - ProceduralSummarizer, - ProceduralRetriever, - ToolSummarizer, - ToolRetriever, -) from .config import ReMeConfigParser from .core import Application from .core.enumeration import MemoryType, Role from .core.schema import Message, MemoryNode -from .tool.memory import ( - RetrieveMemory, - DelegateTask, - ReadHistory, - ReadHistoryV2, - ProfileHandler, - MemoryHandler, - AddAndRetrieveSimilarMemory, +from .memory.tools import ( AddDraftAndRetrieveSimilarMemory, - UpdateMemoryV2, - AddDraftAndReadAllProfiles, - UpdateProfile, AddHistory, - ReadAllProfiles, - UpdateProfilesV1, AddMemory, + DelegateTask, + ReadAllProfiles, + ReadHistory, + RetrieveMemory, + UpdateProfilesV1, +) +from .memory.tools.profiles.profile_handler import ProfileHandler +from .memory.tools.record.memory_handler import MemoryHandler +from .memory.vector_based import ( + BaseMemoryAgent, + PersonalRetriever, + PersonalSummarizer, + ProceduralRetriever, + ProceduralSummarizer, + ReMeRetriever, + ReMeSummarizer, + ToolRetriever, + ToolSummarizer, ) @@ -52,7 +43,7 @@ class ReMe(Application): embedding_api_key: str | None = None, embedding_base_url: str | None = None, working_dir: str = ".reme", - config_path: str = "default", + config_path: str = "vector", enable_logo: bool = True, log_to_console: bool = True, default_llm_config: dict | None = None, @@ -182,29 +173,8 @@ class ReMe(Application): message = Message(**message) format_messages.append(message) - personal_summarizer: BaseMemoryAgent if version == "default": - personal_summarizer = PersonalSummarizer( - llm=llm_config_name, - tools=[ - AddAndRetrieveSimilarMemory( - enable_thinking_params=enable_thinking_params, - top_k=retrieve_top_k, - ), - UpdateMemoryV2(enable_thinking_params=enable_thinking_params), - AddDraftAndReadAllProfiles( - enable_thinking_params=enable_thinking_params, - profile_dir=self.profile_dir, - ), - UpdateProfile( - enable_thinking_params=enable_thinking_params, - profile_dir=self.profile_dir, - ), - ], - ) - - elif version == "v1": - personal_summarizer = PersonalV1Summarizer( + personal_summarizer: BaseMemoryAgent = PersonalSummarizer( llm=llm_config_name, tools=[ AddDraftAndRetrieveSimilarMemory( @@ -212,6 +182,7 @@ class ReMe(Application): enable_memory_target=False, enable_when_to_use=False, enable_multiple=True, + top_k=retrieve_top_k, ), AddMemory( enable_thinking_params=enable_thinking_params, @@ -232,71 +203,12 @@ class ReMe(Application): ), ], ) - elif version == "v2": - personal_summarizer = PersonalV1Summarizer( - llm=llm_config_name, - tools=[ - AddDraftAndRetrieveSimilarMemory( - enable_thinking_params=enable_thinking_params, - enable_memory_target=False, - enable_when_to_use=False, - enable_multiple=True, - ), - AddMemory( - enable_thinking_params=enable_thinking_params, - enable_memory_target=False, - enable_when_to_use=False, - enable_multiple=True, - ), - ReadAllProfiles( - enable_thinking_params=enable_thinking_params, - enable_memory_target=False, - profile_dir=self.profile_dir, - ), - UpdateProfilesV1( - enable_thinking_params=enable_thinking_params, - enable_memory_target=False, - enable_multiple=True, - profile_dir=self.profile_dir, - ), - ], - ) - elif version == "halumem": - personal_summarizer = PersonalHalumemSummarizer( - llm=llm_config_name, - tools=[ - AddAndRetrieveSimilarMemory( - enable_thinking_params=enable_thinking_params, - top_k=retrieve_top_k, - ), - UpdateMemoryV2( - enable_thinking_params=enable_thinking_params, - ), - # 处理userprofile - ReadAllProfiles( - enable_thinking_params=enable_thinking_params, - profile_dir=self.profile_dir, - ), - UpdateProfile( - enable_thinking_params=enable_thinking_params, - profile_dir=self.profile_dir, - ), - ], - ) - else: - raise NotImplementedError - procedural_summarizer: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - procedural_summarizer = ProceduralSummarizer(tools=[]) else: - raise NotImplementedError + raise NotImplementedError(f"version={version} is not supported") - tool_summarizer: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - tool_summarizer = ToolSummarizer(tools=[]) - else: - raise NotImplementedError + procedural_summarizer: BaseMemoryAgent = ProceduralSummarizer(tools=[]) + tool_summarizer: BaseMemoryAgent = ToolSummarizer(tools=[]) memory_agents = [] memory_targets = [] @@ -342,11 +254,9 @@ class ReMe(Application): if not memory_agents: memory_agents = [personal_summarizer, procedural_summarizer, tool_summarizer] - reme_summarizer: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - reme_summarizer = ReMeSummarizer(tools=[AddHistory(), DelegateTask(memory_agents=memory_agents)]) - else: - raise NotImplementedError + reme_summarizer: BaseMemoryAgent = ReMeSummarizer( + tools=[AddHistory(), DelegateTask(memory_agents=memory_agents)], + ) result = await reme_summarizer.call( messages=format_messages, @@ -379,28 +289,9 @@ class ReMe(Application): ) -> str | dict: """Retrieve relevant personal, procedural and tool memories for a query.""" - personal_retriever: BaseMemoryAgent if version == "default": - personal_retriever = PersonalRetriever( + personal_retriever: BaseMemoryAgent = PersonalRetriever( llm=llm_config_name, - tools=[ - ReadAllProfiles( - enable_thinking_params=enable_thinking_params, - profile_dir=self.profile_dir, - ), - RetrieveMemory( - enable_thinking_params=enable_thinking_params, - top_k=retrieve_top_k, - enable_time_filter=enable_time_filter, - ), - ReadHistory(enable_thinking_params=enable_thinking_params), - ], - ) - - elif version == "v1": - personal_retriever = PersonalV1Retriever( - llm=llm_config_name, - return_memory_nodes=False, tools=[ ReadAllProfiles( enable_thinking_params=enable_thinking_params, @@ -419,63 +310,11 @@ class ReMe(Application): ), ], ) - elif version == "v2": - personal_retriever = PersonalV1Retriever( - llm=llm_config_name, - return_memory_nodes=False, - tools=[ - ReadAllProfiles( - enable_thinking_params=enable_thinking_params, - enable_memory_target=False, - profile_dir=self.profile_dir, - ), - RetrieveMemory( - top_k=retrieve_top_k, - enable_thinking_params=enable_thinking_params, - enable_time_filter=enable_time_filter, - enable_multiple=True, - ), - ReadHistoryV2( - message_block_size=4, - vector_top_k=3, - enable_multiple=True, - enable_thinking_params=enable_thinking_params, - ), - ], - ) - elif version == "halumem": - personal_retriever = PersonalHalumemRetriever( - llm=llm_config_name, - tools=[ - ReadAllProfiles( - enable_thinking_params=enable_thinking_params, - profile_dir=self.profile_dir, - ), - RetrieveMemory( - enable_thinking_params=enable_thinking_params, - top_k=retrieve_top_k, - enable_time_filter=enable_time_filter, - ), - ReadHistoryV2( - message_block_size=4, - vector_top_k=3, - ), - ], - ) else: - raise NotImplementedError + raise NotImplementedError(f"version={version} is not supported") - procedural_retriever: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - procedural_retriever = ProceduralRetriever(tools=[]) - else: - raise NotImplementedError - - tool_retriever: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - tool_retriever = ToolRetriever(tools=[]) - else: - raise NotImplementedError + procedural_retriever: BaseMemoryAgent = ProceduralRetriever(tools=[]) + tool_retriever: BaseMemoryAgent = ToolRetriever(tools=[]) memory_agents = [] memory_targets = [] @@ -518,11 +357,9 @@ class ReMe(Application): if not memory_agents: memory_agents = [personal_retriever, procedural_retriever, tool_retriever] - reme_retriever: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - reme_retriever = ReMeRetriever(tools=[DelegateTask(memory_agents=memory_agents)]) - else: - raise NotImplementedError + reme_retriever: BaseMemoryAgent = ReMeRetriever( + tools=[DelegateTask(memory_agents=memory_agents)], + ) result = await reme_retriever.call( query=query, @@ -710,16 +547,13 @@ class ReMe(Application): """Get the profile handler for the specified user.""" return ProfileHandler(memory_target=user_name, profile_path=self.profile_path) - async def context_offload(self): - """working memory summary""" - - async def context_reload(self): - """working memory retrieve""" - def main(): """Main entry point for running ReMe from command line.""" - ReMe(*sys.argv[1:]).run_service() + from . import extension # noqa: F401 # pylint: disable=unused-import + from . import memory # noqa: F401 # pylint: disable=unused-import + + ReMe(*sys.argv[1:], config_path="service").run_service() if __name__ == "__main__": diff --git a/reme/reme_cli.py b/reme/reme_cli.py index c0362158..9255bef1 100644 --- a/reme/reme_cli.py +++ b/reme/reme_cli.py @@ -7,25 +7,26 @@ from typing import AsyncGenerator from prompt_toolkit import PromptSession -from reme.core.op import BaseTool -from .agent.chat import FsCli from .core.enumeration import ChunkEnum +from .core.op import BaseTool from .core.schema import StreamChunk -from .core.utils import execute_stream_task, play_horse_easter_egg -from .reme_fs import ReMeFs -from .tool.fs import ( +from .core.tools import ( BashTool, EditTool, - FsMemorySearch, LsTool, ReadTool, WriteTool, + ExecuteCode, + DashscopeSearch, + TavilySearch, ) -from .tool.gallery import ExecuteCode -from .tool.search import DashscopeSearch, TavilySearch +from .core.utils import execute_stream_task, play_horse_easter_egg +from .memory.file_based import FbCli +from .memory.tools import MemorySearch +from .reme_fb import ReMeFb -class ReMeCli(ReMeFs): +class ReMeCli(ReMeFb): """ReMe Cli""" def __init__(self, *args, config_path: str = "cli", **kwargs): @@ -46,7 +47,7 @@ class ReMeCli(ReMeFs): language = self.service_config.language print(f"ReMe language={language or 'default'}") tools: list[BaseTool] = [ - FsMemorySearch( + MemorySearch( vector_weight=self.service_config.metadata["vector_weight"], candidate_multiplier=self.service_config.metadata["candidate_multiplier"], ), @@ -68,7 +69,7 @@ class ReMeCli(ReMeFs): else: print("No Tavily or Dashscope API key found, skip Tavily and Dashscope search tool") - fs_cli = FsCli( + fb_cli = FbCli( tools=tools, context_window_tokens=self.service_config.metadata["context_window_tokens"], reserve_tokens=self.service_config.metadata["reserve_tokens"], @@ -88,7 +89,7 @@ class ReMeCli(ReMeFs): """Execute chat query and yield streaming chunks.""" stream_queue = asyncio.Queue() task = asyncio.create_task( - fs_cli.call( + fb_cli.call( query=q, stream_queue=stream_queue, service_context=self.service_context, @@ -115,22 +116,22 @@ class ReMeCli(ReMeFs): break if user_input == "/new": - result = await fs_cli.new() + result = await fb_cli.new() print(f"{result}\nConversation reset\n") continue if user_input == "/compact": - result = await fs_cli.compact(force_compact=True) + result = await fb_cli.compact(force_compact=True) print(f"{result}\nHistory compacted.\n") continue if user_input == "/history": - result = fs_cli.format_history() + result = fb_cli.format_history() print(f"Formated History:\n{result}\n") continue if user_input == "/clear": - fs_cli.messages.clear() + fb_cli.messages.clear() print("History cleared.\n") continue diff --git a/reme/reme_fs.py b/reme/reme_fb.py similarity index 91% rename from reme/reme_fs.py rename to reme/reme_fb.py index e6e31b1b..914d237b 100644 --- a/reme/reme_fs.py +++ b/reme/reme_fb.py @@ -1,30 +1,29 @@ -"""ReMe File System""" +"""ReMe File Based""" from pathlib import Path -from .agent.fs import FsCompactor, FsContextChecker, FsSummarizer from .config import ReMeConfigParser from .core import Application from .core.schema import Message -from .tool.fs import ( +from .core.tools import ( BashTool, EditTool, - FsMemoryGet, - FsMemorySearch, LsTool, ReadTool, WriteTool, ) +from .memory.file_based import FbCompactor, FbContextChecker, FbSummarizer +from .memory.tools import MemoryGet, MemorySearch -class ReMeFs(Application): - """ReMe File System""" +class ReMeFb(Application): + """ReMe File Based""" def __init__( self, *args, working_dir: str = ".reme", - config_path: str = "fs", + config_path: str = "file", enable_logo: bool = True, log_to_console: bool = True, llm_api_key: str | None = None, @@ -33,7 +32,7 @@ class ReMeFs(Application): embedding_base_url: str | None = None, default_llm_config: dict | None = None, default_embedding_model_config: dict | None = None, - default_memory_store_config: dict | None = None, + default_file_store_config: dict | None = None, default_token_counter_config: dict | None = None, default_file_watcher_config: dict | None = None, context_window_tokens: int = 128000, @@ -70,7 +69,7 @@ class ReMeFs(Application): parser=ReMeConfigParser, default_llm_config=default_llm_config, default_embedding_model_config=default_embedding_model_config, - default_memory_store_config=default_memory_store_config, + default_file_store_config=default_file_store_config, default_token_counter_config=default_token_counter_config, default_file_watcher_config=default_file_watcher_config, **kwargs, @@ -84,7 +83,7 @@ class ReMeFs(Application): async def context_check(self, messages: list[Message | dict]) -> dict: """Check if messages exceed context limits.""" - checker = FsContextChecker( + checker = FbContextChecker( context_window_tokens=self.service_config.metadata["context_window_tokens"], reserve_tokens=self.service_config.metadata["reserve_tokens"], keep_recent_tokens=self.service_config.metadata["keep_recent_tokens"], @@ -100,7 +99,7 @@ class ReMeFs(Application): **kwargs, ) -> str | dict: """Compact messages into a summary.""" - compactor = FsCompactor(language=language, **kwargs) + compactor = FbCompactor(language=language, **kwargs) return await compactor.call( messages_to_summarize=messages_to_summarize or [], turn_prefix_messages=turn_prefix_messages or [], @@ -117,7 +116,7 @@ class ReMeFs(Application): **kwargs, ) -> str | dict: """Generate a summary of the given messages.""" - summarizer = FsSummarizer( + summarizer = FbSummarizer( tools=[ BashTool(cwd=self.working_dir), LsTool(cwd=self.working_dir), @@ -146,7 +145,7 @@ class ReMeFs(Application): Returns: Search results as formatted string """ - search_tool = FsMemorySearch( + search_tool = MemorySearch( vector_weight=self.service_config.metadata["vector_weight"], candidate_multiplier=self.service_config.metadata["candidate_multiplier"], ) @@ -170,13 +169,13 @@ class ReMeFs(Application): Returns: Memory file content as string """ - get_tool = FsMemoryGet(cwd=self.working_dir) + get_tool = MemoryGet(cwd=self.working_dir) return await get_tool.call(path=path, offset=offset, limit=limit, service_context=self.service_context) async def needs_compaction(self, messages: list[Message | dict]) -> bool: """Check if messages need compaction based on context window limits.""" messages = [Message(**message) if isinstance(message, dict) else message for message in messages] - checker = FsContextChecker( + checker = FbContextChecker( context_window_tokens=self.service_config.metadata["context_window_tokens"], reserve_tokens=self.service_config.metadata["reserve_tokens"], ) diff --git a/reme/tool/__init__.py b/reme/tool/__init__.py deleted file mode 100644 index 5c0d019b..00000000 --- a/reme/tool/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Tool""" - -from . import gallery -from . import memory -from . import search - -__all__ = [ - "gallery", - "memory", - "search", -] diff --git a/reme/tool/fs/__init__.py b/reme/tool/fs/__init__.py deleted file mode 100644 index a9f003ae..00000000 --- a/reme/tool/fs/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -"""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 -from .write_tool import WriteTool -from ...core import R - -__all__ = [ - "BaseFsTool", - "BashTool", - "EditTool", - "FindTool", - "FsMemoryGet", - "FsMemorySearch", - "GrepTool", - "LsTool", - "ReadTool", - "WriteTool", -] - -for name in __all__: - tool_class = globals()[name] - R.ops.register(tool_class) diff --git a/reme/tool/gallery/__init__.py b/reme/tool/gallery/__init__.py deleted file mode 100644 index 2f3ab74f..00000000 --- a/reme/tool/gallery/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -"""execute tool""" - -from .execute_code import ExecuteCode -from .execute_shell import ExecuteShell -from .think_tool import ThinkTool -from ...core import R - -__all__ = [ - "ExecuteCode", - "ExecuteShell", - "ThinkTool", -] - -for name in __all__: - tool_class = globals()[name] - R.ops.register(tool_class) diff --git a/reme/tool/search/__init__.py b/reme/tool/search/__init__.py deleted file mode 100644 index 69995c37..00000000 --- a/reme/tool/search/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -"""search tool""" - -from .dashscope_search import DashscopeSearch -from .mock_search import MockSearch -from .tavily_search import TavilySearch -from ...core import R - -__all__ = [ - "DashscopeSearch", - "MockSearch", - "TavilySearch", -] - -for name in __all__: - tool_class = globals()[name] - R.ops.register(tool_class) diff --git a/reme/workflow/__init__.py b/reme/workflow/__init__.py index 485fda51..e69de29b 100644 --- a/reme/workflow/__init__.py +++ b/reme/workflow/__init__.py @@ -1,9 +0,0 @@ -"""workflow""" - -from . import gallery -from . import procedural_memory - -__all__ = [ - "gallery", - "procedural_memory", -] diff --git a/reme/workflow/gallery/__init__.py b/reme/workflow/gallery/__init__.py deleted file mode 100644 index 6203089d..00000000 --- a/reme/workflow/gallery/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""test""" - -from .test_op import TestOp -from .translate_ts import TranslateTs -from ...core import R - -__all__ = [ - "TestOp", - "TranslateTs", -] - -for name in __all__: - agent_class = globals()[name] - R.ops.register(agent_class) diff --git a/tests/demo_memory_search.py b/tests/demo_memory_search.py index a1d37091..95c04eef 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_store import SqliteMemoryStore +from reme.core.file_store import SqliteFileStore from reme.core.utils import load_env load_env() @@ -60,7 +60,7 @@ def append_to_file(file_path: str, content: str): print(f"✓ 追加内容到: {file_path}") -async def verify_database(store: SqliteMemoryStore, title: str): +async def verify_database(store: SqliteFileStore, title: str): """验证数据库内容""" print_separator(title) @@ -110,7 +110,7 @@ async def test_full_file_watcher(): dimensions=EMBEDDING_DIMENSIONS, ) - store = SqliteMemoryStore( + store = SqliteFileStore( db_path=DB_PATH_FULL, vec_ext_path="", embedding_model=embedding_model, @@ -122,7 +122,7 @@ async def test_full_file_watcher(): # 创建 FullFileWatcher watcher = FullFileWatcher( watch_paths=temp_workspace, - memory_store=store, + file_store=store, chunk_tokens=200, chunk_overlap=20, recursive=True, @@ -211,7 +211,7 @@ async def test_delta_file_watcher(): dimensions=EMBEDDING_DIMENSIONS, ) - store = SqliteMemoryStore( + store = SqliteFileStore( db_path=DB_PATH_DELTA, vec_ext_path="", embedding_model=embedding_model, @@ -223,7 +223,7 @@ async def test_delta_file_watcher(): # 创建 DeltaFileWatcher watcher = DeltaFileWatcher( watch_paths=temp_workspace, - memory_store=store, + file_store=store, chunk_tokens=200, chunk_overlap=20, overlap_lines=2, diff --git a/tests/test_base_context.py b/tests/test_base_context.py index 61a355c0..a9d614bb 100644 --- a/tests/test_base_context.py +++ b/tests/test_base_context.py @@ -1,16 +1,16 @@ """ -Unit tests for the BaseContext class in reme_ai.core.context. +Unit tests for the BaseDict class in reme.core.base_dict. Ensures attribute-style and dict-style access work interchangeably. """ import pickle -from reme.core.context import BaseContext +from reme.core.base_dict import BaseDict def test_attribute_access(): """Test setting values via attributes and retrieving via items.""" - context = BaseContext() + context = BaseDict() context.xxx = 123 assert context.xxx == 123 assert context["xxx"] == 123 @@ -18,7 +18,7 @@ def test_attribute_access(): def test_dict_access(): """Test setting values via items and retrieving via attributes.""" - context = BaseContext() + context = BaseDict() context["yyy"] = 456 assert context.yyy == 456 assert context["yyy"] == 456 @@ -26,7 +26,7 @@ def test_dict_access(): def test_delete_attribute(): """Test that deleting an attribute removes it from the internal state.""" - context = BaseContext() + context = BaseDict() context.zzz = 789 del context.zzz assert "zzz" not in context @@ -34,7 +34,7 @@ def test_delete_attribute(): def test_attribute_error(): """Test that accessing non-existent attributes raises the correct error.""" - context = BaseContext() + context = BaseDict() try: _ = context.nonexistent assert False, "Should raise AttributeError" @@ -43,8 +43,8 @@ def test_attribute_error(): def test_pickling(): - """Test that BaseContext instances can be serialized and deserialized.""" - context = BaseContext() + """Test that BaseDict instances can be serialized and deserialized.""" + context = BaseDict() context.test_value = "bar" context.num = 42 @@ -53,12 +53,12 @@ def test_pickling(): assert restored.test_value == "bar" assert restored.num == 42 - assert isinstance(restored, BaseContext) + assert isinstance(restored, BaseDict) def test_init_with_data(): """Test that the constructor correctly handles initial dictionary data.""" - context = BaseContext({"a": 1, "b": 2}) + context = BaseDict({"a": 1, "b": 2}) assert context.a == 1 assert context.b == 2 diff --git a/tests/test_chunking_utils.py b/tests/test_chunking_utils.py index c0937021..0ac44d0a 100644 --- a/tests/test_chunking_utils.py +++ b/tests/test_chunking_utils.py @@ -42,11 +42,8 @@ def test_chunk_markdown_empty(): overlap=10, ) - # Empty string splits to [""] which creates one chunk with empty text - assert len(chunks) == 1 - assert chunks[0].text == "" - assert chunks[0].start_line == 1 - assert chunks[0].end_line == 1 + # Empty text is filtered out (no meaningful content to store) + assert len(chunks) == 0 def test_chunk_markdown_single_line(): diff --git a/tests/test_memory_store.py b/tests/test_file_store.py similarity index 86% rename from tests/test_memory_store.py rename to tests/test_file_store.py index b8d6e0d3..fd44332d 100644 --- a/tests/test_memory_store.py +++ b/tests/test_file_store.py @@ -1,15 +1,15 @@ # pylint: disable=too-many-lines -"""Unified test suite for memory store implementations. +"""Unified test suite for file store implementations. -This module provides comprehensive test coverage for SqliteMemoryStore, ChromaMemoryStore, -LocalMemoryStore and future memory store implementations. Tests can be run for specific stores +This module provides comprehensive test coverage for SqliteFileStore, ChromaFileStore, +LocalFileStore and future file store implementations. Tests can be run for specific stores or all implementations. Usage: - python test_memory_store.py --sqlite # Test SqliteMemoryStore only - python test_memory_store.py --chroma # Test ChromaMemoryStore only - python test_memory_store.py --local # Test LocalMemoryStore only - python test_memory_store.py --all # Test all memory stores + python test_file_store.py --sqlite # Test SqliteFileStore only + python test_file_store.py --chroma # Test ChromaFileStore only + python test_file_store.py --local # Test LocalFileStore only + python test_file_store.py --all # Test all file stores """ import argparse @@ -17,7 +17,6 @@ import asyncio import hashlib import shutil import time -from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import List @@ -25,10 +24,10 @@ from loguru import logger from reme.core.embedding import OpenAIEmbeddingModel from reme.core.enumeration.memory_source import MemorySource -from reme.core.memory_store.base_memory_store import BaseMemoryStore -from reme.core.memory_store.chroma_memory_store import ChromaMemoryStore -from reme.core.memory_store.local_memory_store import LocalMemoryStore -from reme.core.memory_store.sqlite_memory_store import SqliteMemoryStore +from reme.core.file_store.base_file_store import BaseFileStore +from reme.core.file_store.chroma_file_store import ChromaFileStore +from reme.core.file_store.local_file_store import LocalFileStore +from reme.core.file_store.sqlite_file_store import SqliteFileStore from reme.core.schema.file_metadata import FileMetadata from reme.core.schema.memory_chunk import MemoryChunk from reme.core.utils import load_env @@ -44,18 +43,18 @@ load_env() class TestConfig: """Configuration for test execution.""" - # SqliteMemoryStore settings + # SqliteFileStore settings NAME = "test" - SQLITE_DB_PATH = "./test_memory_store_sqlite/memory.db" + SQLITE_DB_PATH = "./test_file_store_sqlite/memory.db" SQLITE_VEC_EXT_PATH = "" # Empty string to use default vec0/sqlite_vec/vector0 SQLITE_FTS_ENABLED = True - # ChromaMemoryStore settings - CHROMA_DB_PATH = "./test_memory_store_chroma" + # ChromaFileStore settings + CHROMA_DB_PATH = "./test_file_store_chroma" CHROMA_FTS_ENABLED = True - # LocalMemoryStore settings - LOCAL_DB_PATH = "./test_memory_store_local" + # LocalFileStore settings + LOCAL_DB_PATH = "./test_file_store_local" LOCAL_FTS_ENABLED = True # Embedding model settings @@ -182,36 +181,36 @@ class SampleDataGenerator: ) -# ==================== Memory Store Factory ==================== +# ==================== File Store Factory ==================== -def get_store_type(store: BaseMemoryStore) -> str: - """Get the type identifier of a memory store instance. +def get_store_type(store: BaseFileStore) -> str: + """Get the type identifier of a file store instance. Args: - store: Memory store instance + store: File store instance Returns: str: Type identifier ("sqlite", "chroma", etc.) """ - if isinstance(store, SqliteMemoryStore): + if isinstance(store, SqliteFileStore): return "sqlite" - elif isinstance(store, ChromaMemoryStore): + elif isinstance(store, ChromaFileStore): return "chroma" - elif isinstance(store, LocalMemoryStore): + elif isinstance(store, LocalFileStore): return "local" else: - raise ValueError(f"Unknown memory store type: {type(store)}") + raise ValueError(f"Unknown file store type: {type(store)}") -def create_memory_store(store_type: str) -> BaseMemoryStore: - """Create a memory store instance based on type. +def create_file_store(store_type: str) -> BaseFileStore: + """Create a file store instance based on type. Args: - store_type: Type of memory store ("sqlite", "chroma", etc.) + store_type: Type of file store ("sqlite", "chroma", etc.) Returns: - BaseMemoryStore: Initialized memory store instance + BaseFileStore: Initialized file store instance """ config = TestConfig() @@ -221,32 +220,27 @@ def create_memory_store(store_type: str) -> BaseMemoryStore: dimensions=config.EMBEDDING_DIMENSIONS, ) - thread_pool = ThreadPoolExecutor() - if store_type == "sqlite": - return SqliteMemoryStore( + return SqliteFileStore( store_name=config.NAME, db_path=config.SQLITE_DB_PATH, embedding_model=embedding_model, vec_ext_path=config.SQLITE_VEC_EXT_PATH, fts_enabled=config.SQLITE_FTS_ENABLED, - thread_pool=thread_pool, ) elif store_type == "chroma": - return ChromaMemoryStore( + return ChromaFileStore( store_name=config.NAME, db_path=config.CHROMA_DB_PATH, embedding_model=embedding_model, fts_enabled=config.CHROMA_FTS_ENABLED, - thread_pool=thread_pool, ) elif store_type == "local": - return LocalMemoryStore( + return LocalFileStore( store_name=config.NAME, db_path=config.LOCAL_DB_PATH, embedding_model=embedding_model, fts_enabled=config.LOCAL_FTS_ENABLED, - thread_pool=thread_pool, ) else: raise ValueError(f"Unknown store type: {store_type}") @@ -255,7 +249,7 @@ def create_memory_store(store_type: str) -> BaseMemoryStore: # ==================== Test Functions ==================== -async def test_start_store(store: BaseMemoryStore, _store_name: str): +async def test_start_store(store: BaseFileStore, _store_name: str): """Test store initialization.""" logger.info("=" * 20 + " START STORE TEST " + "=" * 20) @@ -263,7 +257,7 @@ async def test_start_store(store: BaseMemoryStore, _store_name: str): logger.info("✓ Store initialized successfully") # Verify tables created (SQLite specific) - if isinstance(store, SqliteMemoryStore): + if isinstance(store, SqliteFileStore): cursor = store.conn.cursor() cursor.execute( "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name", @@ -277,21 +271,21 @@ async def test_start_store(store: BaseMemoryStore, _store_name: str): logger.info("✓ Required tables created") # Verify ChromaDB collection created - if isinstance(store, ChromaMemoryStore): + if isinstance(store, ChromaFileStore): assert store.client is not None, "ChromaDB client should be initialized" assert store.chunks_collection is not None, "ChromaDB collection should exist" logger.info(f"✓ ChromaDB collection created: {store.collection_name}") - # Verify LocalMemoryStore initialized (access internals for test assertions) - if isinstance(store, LocalMemoryStore): + # Verify LocalFileStore initialized (access internals for test assertions) + if isinstance(store, LocalFileStore): # pylint: disable=protected-access - assert store._started, "LocalMemoryStore should be marked as started" + assert store._started, "LocalFileStore should be marked as started" assert isinstance(store._chunks, dict), "Chunks index should be a dict" assert isinstance(store._files, dict), "Files index should be a dict" - logger.info(f"✓ LocalMemoryStore ready (chunks file: {store._chunks_file})") + logger.info(f"✓ LocalFileStore ready (chunks file: {store._chunks_file})") -async def test_upsert_file(store: BaseMemoryStore, _store_name: str) -> tuple[FileMetadata, List[MemoryChunk]]: +async def test_upsert_file(store: BaseFileStore, _store_name: str) -> tuple[FileMetadata, List[MemoryChunk]]: """Test file and chunks insertion.""" logger.info("=" * 20 + " UPSERT FILE TEST " + "=" * 20) @@ -322,7 +316,7 @@ async def test_upsert_file(store: BaseMemoryStore, _store_name: str) -> tuple[Fi return file_meta, chunks -async def test_upsert_multiple_sources(store: BaseMemoryStore, _store_name: str): +async def test_upsert_multiple_sources(store: BaseFileStore, _store_name: str): """Test upserting files from different sources.""" logger.info("=" * 20 + " UPSERT MULTIPLE SOURCES TEST " + "=" * 20) @@ -358,7 +352,7 @@ async def test_upsert_multiple_sources(store: BaseMemoryStore, _store_name: str) logger.info("✓ Multiple sources test passed") -async def test_update_file(store: BaseMemoryStore, _store_name: str): +async def test_update_file(store: BaseFileStore, _store_name: str): """Test updating an existing file.""" logger.info("=" * 20 + " UPDATE FILE TEST " + "=" * 20) @@ -405,7 +399,7 @@ async def test_update_file(store: BaseMemoryStore, _store_name: str): logger.info(f"✓ Verified update: new chunk count = {new_meta.chunk_count}") -async def test_get_file_metadata(store: BaseMemoryStore, _store_name: str): +async def test_get_file_metadata(store: BaseFileStore, _store_name: str): """Test retrieving file metadata.""" logger.info("=" * 20 + " GET FILE METADATA TEST " + "=" * 20) @@ -422,7 +416,7 @@ async def test_get_file_metadata(store: BaseMemoryStore, _store_name: str): logger.info("✓ Get file metadata test passed") -async def test_list_files(store: BaseMemoryStore, _store_name: str): +async def test_list_files(store: BaseFileStore, _store_name: str): """Test listing files by source.""" logger.info("=" * 20 + " LIST FILES TEST " + "=" * 20) @@ -441,7 +435,7 @@ async def test_list_files(store: BaseMemoryStore, _store_name: str): logger.info("✓ List files test passed") -async def test_get_file_chunks(store: BaseMemoryStore, _store_name: str): +async def test_get_file_chunks(store: BaseFileStore, _store_name: str): """Test retrieving chunks for a file.""" logger.info("=" * 20 + " GET FILE CHUNKS TEST " + "=" * 20) @@ -461,7 +455,7 @@ async def test_get_file_chunks(store: BaseMemoryStore, _store_name: str): logger.info("✓ Get file chunks test passed") -async def test_vector_search(store: BaseMemoryStore, _store_name: str): +async def test_vector_search(store: BaseFileStore, _store_name: str): """Test vector similarity search.""" logger.info("=" * 20 + " VECTOR SEARCH TEST " + "=" * 20) @@ -492,7 +486,7 @@ async def test_vector_search(store: BaseMemoryStore, _store_name: str): logger.info("\n✓ Vector search test passed") -async def test_vector_search_with_source_filter(store: BaseMemoryStore, _store_name: str): +async def test_vector_search_with_source_filter(store: BaseFileStore, _store_name: str): """Test vector search with source filtering.""" logger.info("=" * 20 + " VECTOR SEARCH WITH SOURCE FILTER TEST " + "=" * 20) @@ -548,7 +542,7 @@ async def test_vector_search_with_source_filter(store: BaseMemoryStore, _store_n logger.info("\n✓ Vector search with source filter test passed") -async def test_keyword_search(store: BaseMemoryStore, _store_name: str): +async def test_keyword_search(store: BaseFileStore, _store_name: str): """Test full-text keyword search.""" logger.info("=" * 20 + " KEYWORD SEARCH TEST " + "=" * 20) @@ -580,7 +574,7 @@ async def test_keyword_search(store: BaseMemoryStore, _store_name: str): logger.info("\n⊘ No results found (may be expected depending on data)") -async def test_keyword_search_with_source_filter(store: BaseMemoryStore, _store_name: str): +async def test_keyword_search_with_source_filter(store: BaseFileStore, _store_name: str): """Test keyword search with source filtering.""" logger.info("=" * 20 + " KEYWORD SEARCH WITH SOURCE FILTER TEST " + "=" * 20) @@ -626,7 +620,7 @@ async def test_keyword_search_with_source_filter(store: BaseMemoryStore, _store_ logger.info("\n✓ Keyword search with source filter test passed") -async def test_keyword_search_special_chars(store: BaseMemoryStore, _store_name: str): +async def test_keyword_search_special_chars(store: BaseFileStore, _store_name: str): """Test keyword search with special characters like ?, *, etc.""" logger.info("=" * 20 + " KEYWORD SEARCH SPECIAL CHARS TEST " + "=" * 20) @@ -662,7 +656,7 @@ async def test_keyword_search_special_chars(store: BaseMemoryStore, _store_name: logger.info("\n✓ Keyword search with special characters test passed") -async def test_delete_file(store: BaseMemoryStore, _store_name: str): +async def test_delete_file(store: BaseFileStore, _store_name: str): """Test file deletion.""" logger.info("=" * 20 + " DELETE FILE TEST " + "=" * 20) @@ -693,7 +687,7 @@ async def test_delete_file(store: BaseMemoryStore, _store_name: str): logger.info("✓ Verified deletion") -async def test_batch_upsert(store: BaseMemoryStore, _store_name: str): +async def test_batch_upsert(store: BaseFileStore, _store_name: str): """Test batch file upsertion.""" logger.info("=" * 20 + " BATCH UPSERT TEST " + "=" * 20) @@ -718,7 +712,7 @@ async def test_batch_upsert(store: BaseMemoryStore, _store_name: str): logger.info(f"✓ Verified {len(batch_files)} batch files") -async def test_concurrent_searches(store: BaseMemoryStore, _store_name: str): +async def test_concurrent_searches(store: BaseFileStore, _store_name: str): """Test concurrent search operations.""" logger.info("=" * 20 + " CONCURRENT SEARCHES TEST " + "=" * 20) @@ -748,7 +742,7 @@ async def test_concurrent_searches(store: BaseMemoryStore, _store_name: str): logger.info("✓ Concurrent searches test passed") -async def test_edge_cases(store: BaseMemoryStore, _store_name: str): +async def test_edge_cases(store: BaseFileStore, _store_name: str): """Test edge cases and boundary conditions.""" logger.info("=" * 20 + " EDGE CASES TEST " + "=" * 20) @@ -866,7 +860,7 @@ async def test_edge_cases(store: BaseMemoryStore, _store_name: str): logger.info("✓ Edge cases test passed") -async def test_clear_all(store: BaseMemoryStore, _store_name: str): +async def test_clear_all(store: BaseFileStore, _store_name: str): """Test clearing all data.""" logger.info("=" * 20 + " CLEAR ALL TEST " + "=" * 20) @@ -904,18 +898,18 @@ async def test_clear_all(store: BaseMemoryStore, _store_name: str): async def run_all_tests_for_store(store_type: str, store_name: str): - """Run all tests for a specific memory store type. + """Run all tests for a specific file store type. Args: - store_type: Type of memory store ("sqlite", etc.) - store_name: Display name for the memory store + store_type: Type of file store ("sqlite", etc.) + store_name: Display name for the file store """ logger.info(f"\n\n{'#' * 60}") logger.info(f"# Running all tests for: {store_name}") logger.info(f"{'#' * 60}") - # Create memory store instance - store = create_memory_store(store_type) + # Create file store instance + store = create_file_store(store_type) try: # ========== Basic Tests ========== @@ -971,12 +965,12 @@ async def run_all_tests_for_store(store_type: str, store_name: str): await cleanup_store(store, store_type) -async def cleanup_store(store: BaseMemoryStore, store_type: str): - """Clean up test resources for a memory store. +async def cleanup_store(store: BaseFileStore, store_type: str): + """Clean up test resources for a file store. Args: - store: Memory store instance - store_type: Type of memory store ("sqlite", "chroma", etc.) + store: File store instance + store_type: Type of file store ("sqlite", "chroma", etc.) """ logger.info("=" * 20 + " CLEANUP " + "=" * 20) @@ -985,7 +979,7 @@ async def cleanup_store(store: BaseMemoryStore, store_type: str): await store.close() logger.info("✓ Closed store connections") - # Clean up local directory if SqliteMemoryStore + # Clean up local directory if SqliteFileStore if store_type == "sqlite": config = TestConfig() db_dir = Path(config.SQLITE_DB_PATH).parent @@ -993,7 +987,7 @@ async def cleanup_store(store: BaseMemoryStore, store_type: str): shutil.rmtree(db_dir) logger.info(f"✓ Cleaned up directory: {db_dir}") - # Clean up local directory if ChromaMemoryStore + # Clean up local directory if ChromaFileStore if store_type == "chroma": config = TestConfig() db_dir = Path(config.CHROMA_DB_PATH) @@ -1006,7 +1000,7 @@ async def cleanup_store(store: BaseMemoryStore, store_type: str): metadata_file.unlink() logger.info(f"✓ Cleaned up metadata file: {metadata_file}") - # Clean up LocalMemoryStore JSON persistence files + # Clean up LocalFileStore JSON persistence files if store_type == "local": config = TestConfig() db_dir = Path(config.LOCAL_DB_PATH) @@ -1030,68 +1024,68 @@ async def cleanup_store(store: BaseMemoryStore, store_type: str): async def main(): """Main entry point for running tests.""" parser = argparse.ArgumentParser( - description="Run memory store tests", + description="Run file store tests", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: - python test_memory_store.py --sqlite # Test SqliteMemoryStore only - python test_memory_store.py --chroma # Test ChromaMemoryStore only - python test_memory_store.py --local # Test LocalMemoryStore only - python test_memory_store.py --all # Test all memory stores + python test_file_store.py --sqlite # Test SqliteFileStore only + python test_file_store.py --chroma # Test ChromaFileStore only + python test_file_store.py --local # Test LocalFileStore only + python test_file_store.py --all # Test all file stores """, ) parser.add_argument( "--sqlite", action="store_true", - help="Test SqliteMemoryStore", + help="Test SqliteFileStore", ) parser.add_argument( "--chroma", action="store_true", - help="Test ChromaMemoryStore", + help="Test ChromaFileStore", ) parser.add_argument( "--local", action="store_true", - help="Test LocalMemoryStore", + help="Test LocalFileStore", ) parser.add_argument( "--all", action="store_true", - help="Run tests for all available memory stores", + help="Run tests for all available file stores", ) args = parser.parse_args() - # Determine which memory stores to test + # Determine which file stores to test stores_to_test = [] if args.all: stores_to_test = [ - ("sqlite", "SqliteMemoryStore"), - ("chroma", "ChromaMemoryStore"), - ("local", "LocalMemoryStore"), + ("sqlite", "SqliteFileStore"), + ("chroma", "ChromaFileStore"), + ("local", "LocalFileStore"), ] else: # Build list based on individual flags if args.sqlite: - stores_to_test.append(("sqlite", "SqliteMemoryStore")) + stores_to_test.append(("sqlite", "SqliteFileStore")) if args.chroma: - stores_to_test.append(("chroma", "ChromaMemoryStore")) + stores_to_test.append(("chroma", "ChromaFileStore")) if args.local: - stores_to_test.append(("local", "LocalMemoryStore")) + stores_to_test.append(("local", "LocalFileStore")) if not stores_to_test: - # Default to all memory stores if no argument provided + # Default to all file stores if no argument provided stores_to_test = [ - ("sqlite", "SqliteMemoryStore"), - ("chroma", "ChromaMemoryStore"), - ("local", "LocalMemoryStore"), + ("sqlite", "SqliteFileStore"), + ("chroma", "ChromaFileStore"), + ("local", "LocalFileStore"), ] - print("No memory store specified, defaulting to test all memory stores") + print("No file store specified, defaulting to test all file stores") print("Use --sqlite, --chroma, or --local to test specific ones\n") - # Run tests for each memory store + # Run tests for each file store for store_type, store_name in stores_to_test: try: await run_all_tests_for_store(store_type, store_name) @@ -1104,7 +1098,7 @@ Examples: print(f"\n\n{'#' * 60}") print("# TEST SUMMARY") print(f"{'#' * 60}") - print(f"✓ All tests passed for {len(stores_to_test)} memory store(s):") + print(f"✓ All tests passed for {len(stores_to_test)} file store(s):") for _, store_name in stores_to_test: print(f" - {store_name}") print(f"{'#' * 60}\n") diff --git a/tests/test_fs_compactor.py b/tests/test_fs_compactor.py index f239ea57..00be5320 100644 --- a/tests/test_fs_compactor.py +++ b/tests/test_fs_compactor.py @@ -1,12 +1,12 @@ -"""Tests for FsCompactor - conversation history summarization. +"""Tests for FbCompactor - conversation history summarization. -This module tests the summary generation logic of FsCompactor class, +This module tests the summary generation logic of FbCompactor class, which creates compact summaries of conversation history using LLM. """ import asyncio -from reme import ReMeFs +from reme import ReMeFb from reme.core.enumeration import Role from reme.core.schema import Message @@ -560,7 +560,7 @@ async def test_full_compact_with_summary(): print("TEST: Full Compaction with LLM Summary Generation") print("=" * 80) - reme_fs = ReMeFs( + reme_fs = ReMeFb( enable_logo=False, vector_store=None, compact_params={ @@ -607,7 +607,7 @@ async def test_realistic_personal_conversation_compact(): print("TEST: Realistic Personal Conversation Compaction") print("=" * 80) - reme_fs = ReMeFs( + reme_fs = ReMeFb( enable_logo=False, vector_store=None, compact_params={ @@ -638,7 +638,7 @@ async def test_realistic_personal_conversation_compact(): async def main(): """Run compactor tests.""" print("\n" + "=" * 80) - print("FsCompactor - Summary Generation Test Suite") + print("FbCompactor - Summary Generation Test Suite") print("=" * 80) print("\nThis test suite validates the LLM-based summarization:") print(" - Full compaction flow (context check + summary generation)") diff --git a/tests/test_fs_context_checker.py b/tests/test_fs_context_checker.py index bff7fbcf..14a7fe9f 100644 --- a/tests/test_fs_context_checker.py +++ b/tests/test_fs_context_checker.py @@ -1,12 +1,12 @@ -"""Tests for FsContextChecker - context window limit checking and cut point finding. +"""Tests for FbContextChecker - context window limit checking and cut point finding. -This module tests the cut point finding logic of FsContextChecker class, +This module tests the cut point finding logic of FbContextChecker class, which determines where to split conversation history when token limits are exceeded. """ import asyncio -from reme import ReMeFs +from reme import ReMeFb from reme.core.enumeration import Role from reme.core.schema import Message @@ -67,7 +67,7 @@ async def test_no_compaction_needed(): print("TEST 1: Below Threshold - No Cut Point Needed") print("=" * 80) - reme_fs = ReMeFs( + reme_fs = ReMeFb( "vector_stores={}", # Override config to disable vector stores enable_logo=False, context_window_tokens=5000, @@ -112,7 +112,7 @@ async def test_compaction_needed_above_threshold(): print("TEST 2: Compaction Needed Above Threshold") print("=" * 80) - reme_fs = ReMeFs( + reme_fs = ReMeFb( "vector_stores={}", # Override config to disable vector stores enable_logo=False, context_window_tokens=1500, @@ -181,7 +181,7 @@ async def test_split_turn_scenario(): print("TEST 3: Split Turn - Cut in Middle of Assistant Response") print("=" * 80) - reme_fs = ReMeFs( + reme_fs = ReMeFb( "vector_stores={}", # Override config to disable vector stores enable_logo=False, context_window_tokens=2000, @@ -265,7 +265,7 @@ async def test_split_turn_scenario(): async def main(): """Run context checker tests.""" print("\n" + "=" * 80) - print("FsContextChecker - Cut Point Finding Test Suite") + print("FbContextChecker - Cut Point Finding Test Suite") print("=" * 80) print("\nThis test suite validates the cut point finding logic:") print(" 1. Below threshold - no compaction needed") diff --git a/tests/test_fs_file_watch_integration.py b/tests/test_fs_file_watch_integration.py index ac40d918..8f47b472 100644 --- a/tests/test_fs_file_watch_integration.py +++ b/tests/test_fs_file_watch_integration.py @@ -1,8 +1,8 @@ -"""Integration test for ReMeFs file watching with memory_search and memory_get. +"""Integration test for ReMeFb file watching with memory_search and memory_get. This test demonstrates the complete workflow: 1. Create markdown files with personal information in test_reme folder -2. Initialize ReMeFs with file watching enabled +2. Initialize ReMeFb with file watching enabled 3. Start file watching to automatically index files into the database 4. Use memory_search and memory_get to retrieve the indexed content 5. Modify the markdown files @@ -19,7 +19,7 @@ import json import shutil from pathlib import Path -from reme import ReMeFs +from reme import ReMeFb # ==================== Test Configuration ==================== @@ -278,13 +278,13 @@ async def test_file_watch_integration(): test_files = create_test_markdown_files(TestConfig.WORKING_DIR) print(f"\n✓ Created {len(test_files)} markdown files in {TestConfig.WORKING_DIR}") - # ==================== STEP 2: Initialize ReMeFs ==================== - print_separator("STEP 2: Initializing ReMeFs with File Watching") + # ==================== STEP 2: Initialize ReMeFb ==================== + print_separator("STEP 2: Initializing ReMeFb with File Watching") - reme_fs = ReMeFs( + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_integration", "embedding_model": "default", @@ -299,7 +299,7 @@ async def test_file_watch_integration(): }, ) - print("✓ ReMeFs instance created") + print("✓ ReMeFb instance created") print(f" Working directory: {TestConfig.WORKING_DIR}") print(f" Watch paths: {TestConfig.WORKING_DIR}, {TestConfig.WORKING_DIR}/memory") print(" File filters: .md files") @@ -469,7 +469,7 @@ async def test_file_watch_integration(): print_separator("STEP 10: Cleanup") await reme_fs.close() - print("✓ ReMeFs closed") + print("✓ ReMeFb closed") # Clean up test directory if test_dir.exists(): @@ -489,11 +489,11 @@ async def test_file_watch_integration(): async def main(): """Run the file watch integration test.""" print("\n" + "=" * 80) - print(" ReMeFs File Watch Integration Test") + print(" ReMeFb File Watch Integration Test") print("=" * 80) print("\nThis test validates the complete file watching workflow:") print(" 1. Create markdown files with personal information") - print(" 2. Initialize ReMeFs and start file watching") + print(" 2. Initialize ReMeFb and start file watching") print(" 3. Verify automatic indexing into database") print(" 4. Search and retrieve initial content") print(" 5. Modify files and verify re-indexing") diff --git a/tests/test_fs_memory_get.py b/tests/test_fs_memory_get.py index 246a695e..d353751d 100644 --- a/tests/test_fs_memory_get.py +++ b/tests/test_fs_memory_get.py @@ -1,6 +1,6 @@ -"""Tests for ReMeFs memory_get interface. +"""Tests for ReMeFb memory_get interface. -This module tests the memory_get() method of ReMeFs class which provides +This module tests the memory_get() method of ReMeFb class which provides a high-level interface for reading specific snippets from memory files. The memory_get function should enable the LLM to: @@ -13,7 +13,7 @@ import asyncio import os from pathlib import Path -from reme import ReMeFs +from reme import ReMeFb def print_result(content: str, title: str = "RESULT", max_len: int = 300): @@ -105,7 +105,7 @@ async def test_memory_get_full_file(): print("=" * 80) workspace_dir = ".reme_test_get" - reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + reme_fs = ReMeFb(enable_logo=False, working_dir=workspace_dir) await reme_fs.start() # Create test file @@ -144,7 +144,7 @@ async def test_memory_get_with_offset(): print("=" * 80) workspace_dir = ".reme_test_get" - reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + reme_fs = ReMeFb(enable_logo=False, working_dir=workspace_dir) await reme_fs.start() test_file_path = "memory/test_profile.md" @@ -180,7 +180,7 @@ async def test_memory_get_with_offset_and_limit(): print("=" * 80) workspace_dir = ".reme_test_get" - reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + reme_fs = ReMeFb(enable_logo=False, working_dir=workspace_dir) await reme_fs.start() test_file_path = "memory/test_profile.md" @@ -219,7 +219,7 @@ async def test_memory_get_beginning_lines(): print("=" * 80) workspace_dir = ".reme_test_get" - reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + reme_fs = ReMeFb(enable_logo=False, working_dir=workspace_dir) await reme_fs.start() test_file_path = "memory/test_profile.md" @@ -257,7 +257,7 @@ async def test_memory_get_single_line(): print("=" * 80) workspace_dir = ".reme_test_get" - reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + reme_fs = ReMeFb(enable_logo=False, working_dir=workspace_dir) await reme_fs.start() test_file_path = "memory/test_profile.md" @@ -294,7 +294,7 @@ async def test_memory_get_with_absolute_path(): print("=" * 80) workspace_dir = ".reme_test_get" - reme_fs = ReMeFs(enable_logo=False, working_dir=workspace_dir) + reme_fs = ReMeFb(enable_logo=False, working_dir=workspace_dir) await reme_fs.start() # Get absolute path @@ -324,7 +324,7 @@ async def test_memory_get_with_absolute_path(): async def main(): """Run core memory_get interface tests.""" print("\n" + "=" * 80) - print("ReMeFs Memory Get Interface - Tests") + print("ReMeFb Memory Get Interface - Tests") print("=" * 80) print("\nThis test suite validates that the memory_get() function:") print(" 1. Reads entire memory files without parameters") diff --git a/tests/test_fs_memory_search.py b/tests/test_fs_memory_search.py index c6de7e44..2a65fd5a 100644 --- a/tests/test_fs_memory_search.py +++ b/tests/test_fs_memory_search.py @@ -1,6 +1,6 @@ -"""Tests for ReMeFs memory_search interface. +"""Tests for ReMeFb memory_search interface. -This module tests the memory_search() method of ReMeFs class which provides +This module tests the memory_search() method of ReMeFb class which provides a high-level interface for searching personal information stored in memory files. The memory_search function should enable: @@ -16,7 +16,7 @@ import hashlib import shutil from pathlib import Path -from reme import ReMeFs +from reme import ReMeFb from reme.core.enumeration import MemorySource from reme.core.schema import FileMetadata, MemoryChunk @@ -207,11 +207,11 @@ async def test_memory_search_basic(): print("TEST 1: Basic Memory Search") print("=" * 80) - # Initialize ReMeFs with unique store name - reme_fs = ReMeFs( + # Initialize ReMeFb with unique store name + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_basic", "embedding_model": "default", @@ -222,13 +222,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.default_memory_store.get_chunk_embeddings(personal_chunks) + personal_chunks = await reme_fs.default_file_store.get_chunk_embeddings(personal_chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/personal_info.md", len(personal_chunks), ) - await reme_fs.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( file_meta, MemorySource.MEMORY, personal_chunks, @@ -268,11 +268,11 @@ async def test_memory_search_technical_content(): print("TEST 2: Technical Content Search") print("=" * 80) - # Initialize ReMeFs with unique store name - reme_fs = ReMeFs( + # Initialize ReMeFb with unique store name + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_technical", "embedding_model": "default", @@ -283,13 +283,13 @@ async def test_memory_search_technical_content(): # Insert technical chunks tech_chunks = SampleDataGenerator.create_technical_chunks("test_technical") - tech_chunks = await reme_fs.default_memory_store.get_chunk_embeddings(tech_chunks) + tech_chunks = await reme_fs.default_file_store.get_chunk_embeddings(tech_chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/technical_notes.md", len(tech_chunks), ) - await reme_fs.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( file_meta, MemorySource.MEMORY, tech_chunks, @@ -334,8 +334,8 @@ async def test_memory_search_with_source_filter(): print("TEST 3: Memory Search with Source Filter") print("=" * 80) - # Initialize ReMeFs with unique store name - reme_fs = ReMeFs( + # Initialize ReMeFb with unique store name + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, ) @@ -343,12 +343,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.default_memory_store.get_chunk_embeddings(personal_chunks) + personal_chunks = await reme_fs.default_file_store.get_chunk_embeddings(personal_chunks) personal_meta = SampleDataGenerator.create_file_metadata( "memory/personal_info.md", len(personal_chunks), ) - await reme_fs.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( personal_meta, MemorySource.MEMORY, personal_chunks, @@ -357,12 +357,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.default_memory_store.get_chunk_embeddings(session_chunks) + session_chunks = await reme_fs.default_file_store.get_chunk_embeddings(session_chunks) session_meta = SampleDataGenerator.create_file_metadata( "sessions/2024-01-15.jsonl", len(session_chunks), ) - await reme_fs.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( session_meta, MemorySource.SESSIONS, session_chunks, @@ -374,7 +374,7 @@ async def test_memory_search_with_source_filter(): # Search only MEMORY source print(f"\n--- Searching MEMORY source for: '{query}' ---") # Create a new instance with MEMORY source filter - reme_fs_memory = ReMeFs( + reme_fs_memory = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, search_params={"sources": [MemorySource.MEMORY]}, @@ -396,7 +396,7 @@ async def test_memory_search_with_source_filter(): # Search only SESSIONS source print(f"\n--- Searching SESSIONS source for: '{query}' ---") # Create a new instance with SESSIONS source filter - reme_fs_sessions = ReMeFs( + reme_fs_sessions = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, search_params={"sources": [MemorySource.SESSIONS]}, @@ -437,11 +437,11 @@ async def test_memory_search_score_filtering(): print("TEST 4: Memory Search with Score Filtering") print("=" * 80) - # Initialize ReMeFs with unique store name - reme_fs = ReMeFs( + # Initialize ReMeFb with unique store name + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_score_filter", "embedding_model": "default", @@ -452,12 +452,12 @@ async def test_memory_search_score_filtering(): # Insert test data chunks = SampleDataGenerator.create_technical_chunks("test_score") - chunks = await reme_fs.default_memory_store.get_chunk_embeddings(chunks) + chunks = await reme_fs.default_file_store.get_chunk_embeddings(chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/technical_notes.md", len(chunks), ) - await reme_fs.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( file_meta, MemorySource.MEMORY, chunks, @@ -503,11 +503,11 @@ async def test_memory_search_max_results(): print("TEST 5: Memory Search with Result Limiting") print("=" * 80) - # Initialize ReMeFs with unique store name - reme_fs = ReMeFs( + # Initialize ReMeFb with unique store name + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_max_results", "embedding_model": "default", @@ -521,14 +521,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.default_memory_store.get_chunk_embeddings(all_chunks) + all_chunks = await reme_fs.default_file_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.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( combined_meta, MemorySource.MEMORY, all_chunks, @@ -569,11 +569,11 @@ async def test_memory_search_hybrid_mode(): print("TEST 6: Memory Search with Hybrid Mode") print("=" * 80) - # Initialize ReMeFs with unique store name - reme_fs = ReMeFs( + # Initialize ReMeFb with unique store name + reme_fs = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_hybrid", "embedding_model": "default", @@ -584,12 +584,12 @@ async def test_memory_search_hybrid_mode(): # Insert test data chunks = SampleDataGenerator.create_technical_chunks("test_hybrid") - chunks = await reme_fs.default_memory_store.get_chunk_embeddings(chunks) + chunks = await reme_fs.default_file_store.get_chunk_embeddings(chunks) file_meta = SampleDataGenerator.create_file_metadata( "memory/technical_notes.md", len(chunks), ) - await reme_fs.default_memory_store.upsert_file( + await reme_fs.default_file_store.upsert_file( file_meta, MemorySource.MEMORY, chunks, @@ -601,10 +601,10 @@ async def test_memory_search_hybrid_mode(): # Test with hybrid enabled print(f"\n--- Hybrid search (enabled) for: '{query}' ---") # Create instance with hybrid enabled - reme_fs_hybrid = ReMeFs( + reme_fs_hybrid = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_hybrid", "embedding_model": "default", @@ -630,10 +630,10 @@ async def test_memory_search_hybrid_mode(): # Test with hybrid disabled (vector only) print(f"\n--- Vector-only search for: '{query}' ---") # Create instance with hybrid disabled - reme_fs_vector = ReMeFs( + reme_fs_vector = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_hybrid", "embedding_model": "default", @@ -662,10 +662,10 @@ async def test_memory_search_hybrid_mode(): for vec_weight, text_weight in weight_configs: # Create instance with specific weights - reme_fs_weights = ReMeFs( + reme_fs_weights = ReMeFb( enable_logo=False, working_dir=TestConfig.WORKING_DIR, - default_memory_store_config={ + default_file_store_config={ "backend": "sqlite", "store_name": "test_hybrid", "embedding_model": "default", @@ -708,7 +708,7 @@ async def cleanup_test_data(): async def main(): """Run all memory search tests.""" print("\n" + "=" * 80) - print("ReMeFs Memory Search Interface Tests") + print("ReMeFb Memory Search Interface Tests") print("=" * 80) print("\nThis test suite validates the memory_search() function:") print(" 1. Basic semantic search functionality") diff --git a/tests/test_fs_summary.py b/tests/test_fs_summary.py index 7c19055a..b2701344 100644 --- a/tests/test_fs_summary.py +++ b/tests/test_fs_summary.py @@ -1,6 +1,6 @@ -"""ReMeFs summary接口测试。 +"""ReMeFb summary接口测试。 -本模块测试ReMeFs类的summary()方法，该方法提供了 +本模块测试ReMeFb类的summary()方法，该方法提供了 将用户个人信息存储到记忆文件的高级接口。 summary函数应该能够让LLM： @@ -13,7 +13,7 @@ import asyncio import shutil from pathlib import Path -from reme import ReMeFs +from reme import ReMeFb from reme.core.enumeration import Role from reme.core.schema import Message @@ -114,7 +114,7 @@ async def test_summary_first_write(): if Path(working_dir).exists(): shutil.rmtree(working_dir) - reme_fs = ReMeFs(enable_logo=False, working_dir=working_dir, vector_store=None) + reme_fs = ReMeFb(enable_logo=False, working_dir=working_dir, vector_store=None) await reme_fs.start() # 确保记忆文件已删除 @@ -177,7 +177,7 @@ async def test_summary_complementary_info(): if Path(working_dir).exists(): shutil.rmtree(working_dir) - reme_fs = ReMeFs(enable_logo=False, working_dir=working_dir, vector_store=None) + reme_fs = ReMeFb(enable_logo=False, working_dir=working_dir, vector_store=None) await reme_fs.start() # 确保记忆文件已删除 @@ -271,7 +271,7 @@ async def test_summary_conflicting_info(): if Path(working_dir).exists(): shutil.rmtree(working_dir) - reme_fs = ReMeFs(enable_logo=False, working_dir=working_dir, vector_store=None) + reme_fs = ReMeFb(enable_logo=False, working_dir=working_dir, vector_store=None) await reme_fs.start() # 确保记忆文件已删除 @@ -353,7 +353,7 @@ async def test_summary_conflicting_info(): async def main(): """运行时间对齐的summary接口测试。""" print("\n" + "=" * 80) - print("ReMeFs Summary接口 - 时间对齐的记忆存储测试") + print("ReMeFb Summary接口 - 时间对齐的记忆存储测试") print("=" * 80) print("\n本测试套件验证summary()函数：") print(" 1. 正确处理消息中的time_created字段（%Y-%m-%d %H:%M:%S）") diff --git a/tests/test_fs_tool.py b/tests/test_fs_tool.py index 15f2d858..d5ae5f01 100644 --- a/tests/test_fs_tool.py +++ b/tests/test_fs_tool.py @@ -8,7 +8,7 @@ from pathlib import Path async def test_bash_tool(): """Test BashTool.""" - from reme.tool.fs import BashTool + from reme.core.tools import BashTool print("=== Testing BashTool ===") bash_tool = BashTool() @@ -20,7 +20,7 @@ async def test_bash_tool(): async def test_edit_tool(): """Test EditTool.""" - from reme.tool.fs import EditTool + from reme.core.tools import EditTool print("=== Testing EditTool ===") @@ -77,7 +77,7 @@ async def test_edit_tool(): async def test_find_tool(): """Test FindTool.""" - from reme.tool.fs import FindTool + from reme.core.tools import FindTool print("=== Testing FindTool ===") @@ -131,7 +131,7 @@ async def test_find_tool(): async def test_grep_tool(): """Test GrepTool.""" - from reme.tool.fs import GrepTool + from reme.core.tools import GrepTool print("=== Testing GrepTool ===") @@ -206,7 +206,7 @@ async def test_grep_tool(): async def test_ls_tool(): """Test LsTool.""" - from reme.tool.fs import LsTool + from reme.core.tools import LsTool print("=== Testing LsTool ===") @@ -277,7 +277,7 @@ async def test_ls_tool(): async def test_read_tool(): """Test ReadTool.""" - from reme.tool.fs import ReadTool + from reme.core.tools import ReadTool print("=== Testing ReadTool ===") @@ -357,7 +357,7 @@ async def test_read_tool(): async def test_write_tool(): """Test WriteTool.""" - from reme.tool.fs import WriteTool + from reme.core.tools import WriteTool print("=== Testing WriteTool ===") diff --git a/tests/test_keyword_search_performance.py b/tests/test_keyword_search_performance.py new file mode 100644 index 00000000..345d5f7e --- /dev/null +++ b/tests/test_keyword_search_performance.py @@ -0,0 +1,129 @@ +"""Performance test for LocalFileStore keyword_search. + +Tests keyword_search efficiency with: +- Query length: 20 characters +- Chunk count: 1000 chunks +""" + +import asyncio +import hashlib +import random +import shutil +import time +from pathlib import Path + +from reme.core.enumeration.memory_source import MemorySource +from reme.core.file_store.local_file_store import LocalFileStore +from reme.core.schema.memory_chunk import MemoryChunk + + +def generate_random_text(length: int = 200) -> str: + """Generate random text content.""" + words = [ + "python", "function", "class", "memory", "search", "algorithm", + "database", "vector", "embedding", "chunk", "file", "store", + "query", "result", "performance", "test", "data", "index", + "keyword", "text", "content", "process", "system", "module", + "import", "return", "value", "parameter", "method", "object", + "instance", "variable", "constant", "string", "integer", "float", + "list", "dictionary", "tuple", "set", "array", "matrix", + ] + text_words = [] + current_length = 0 + while current_length < length: + word = random.choice(words) + text_words.append(word) + current_length += len(word) + 1 # +1 for space + return " ".join(text_words)[:length] + + +def create_test_chunks(count: int = 1000, text_length: int = 10000) -> list[MemoryChunk]: + """Create test chunks for performance testing.""" + chunks = [] + for i in range(count): + text = generate_random_text(text_length) + chunk = MemoryChunk( + id=f"perf_test_chunk_{i}", + path=f"/test/file_{i % 100}.py", + source=MemorySource.MEMORY, + start_line=i * 10 + 1, + end_line=(i + 1) * 10, + text=text, + hash=hashlib.md5(text.encode()).hexdigest(), + embedding=None, + metadata={"index": i}, + ) + chunks.append(chunk) + return chunks + + +async def run_performance_test(): + """Run keyword_search performance test.""" + # Setup + test_db_path = Path("./test_keyword_perf") + test_db_path.mkdir(exist_ok=True) + + store = LocalFileStore( + db_path=test_db_path, + store_name="perf_test", + vector_enabled=False, # Disable vector search for this test + fts_enabled=True, + ) + await store.start() + + # Create test data + print("Creating 1000 test chunks...") + chunks = create_test_chunks(1000) + + # Manually add chunks to store (bypass embedding) + for chunk in chunks: + store._chunks[chunk.id] = chunk + + print(f"Loaded {len(store._chunks)} chunks into memory") + + # Create a 20-character query + query = "python function data" # 20 characters including spaces + print(f"Query: '{query}' (length: {len(query)})") + + # Warmup + await store.keyword_search(query, limit=10) + + # Performance test - multiple runs + num_runs = 100 + times = [] + + print(f"\nRunning {num_runs} iterations...") + + for _ in range(num_runs): + start = time.perf_counter() + results = await store.keyword_search(query, limit=10) + elapsed = time.perf_counter() - start + times.append(elapsed) + + # Statistics + avg_time = sum(times) / len(times) + min_time = min(times) + max_time = max(times) + + print("\n" + "=" * 50) + print("Performance Results (keyword_search)") + print("=" * 50) + print(f"Query length: {len(query)} characters") + print(f"Chunk count: {len(store._chunks)}") + print(f"Iterations: {num_runs}") + print("-" * 50) + print(f"Average time: {avg_time * 1000:.4f} ms") + print(f"Min time: {min_time * 1000:.4f} ms") + print(f"Max time: {max_time * 1000:.4f} ms") + print(f"Total time: {sum(times) * 1000:.2f} ms") + print("=" * 50) + + # Cleanup + await store.close() + + # Remove test directory + shutil.rmtree(test_db_path, ignore_errors=True) + + +if __name__ == "__main__": + asyncio.run(run_performance_test()) diff --git a/tests/test_tool.py b/tests/test_tool.py index b4639776..cf07d6c4 100644 --- a/tests/test_tool.py +++ b/tests/test_tool.py @@ -17,7 +17,7 @@ async def test_search(_app): Tests DashscopeSearch, MockSearch, and TavilySearch operations with a sample query to verify they work correctly. """ - from reme.tool.search import DashscopeSearch, MockSearch, TavilySearch + from reme.core.tools import DashscopeSearch, MockSearch, TavilySearch query = "美股DFDV是做什么的？" @@ -41,7 +41,7 @@ async def test_execute(_app): including successful execution, syntax errors, runtime errors, and invalid commands to verify error handling. """ - from reme.tool.gallery import ExecuteCode, ExecuteShell + from reme.core.tools import ExecuteCode, ExecuteShell # Test ExecuteCode print("\n" + "=" * 60) @@ -153,7 +153,7 @@ async def test_simple_chat(app): Tests the SimpleChat agent with a basic query to verify it can process and respond to user input. """ - from reme.agent.chat import SimpleChat + from reme.extension import SimpleChat op = SimpleChat() output = await op.call(query="你好", service_context=app.service_context) @@ -166,9 +166,9 @@ async def test_stream_chat(app): Tests the StreamChat agent with a query to verify it can process and stream responses in real-time using async operations. """ - from reme.agent.chat import StreamChat + from reme.extension import StreamChat from reme.core.utils import execute_stream_task - from reme.core.context import RuntimeContext + from reme.core import RuntimeContext from asyncio import Queue op = StreamChat() diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index a5c2a849..ccdd9c28 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -18,7 +18,7 @@ Usage: import argparse import asyncio import shutil -from concurrent.futures import ThreadPoolExecutor +import tempfile from pathlib import Path from typing import List @@ -216,21 +216,17 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor dimensions=config.EMBEDDING_DIMENSIONS, ) - # Create thread pool executor for vector stores - thread_pool = ThreadPoolExecutor(max_workers=4) - if store_type == "local": return LocalVectorStore( collection_name=collection_name, embedding_model=embedding_model, - thread_pool=thread_pool, - root_path=config.LOCAL_ROOT_PATH, + db_path=config.LOCAL_ROOT_PATH, ) elif store_type == "es": return ESVectorStore( collection_name=collection_name, embedding_model=embedding_model, - thread_pool=thread_pool, + db_path=tempfile.mkdtemp(prefix="test_es_"), hosts=config.ES_HOSTS, basic_auth=config.ES_BASIC_AUTH, ) @@ -238,8 +234,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor return QdrantVectorStore( collection_name=collection_name, embedding_model=embedding_model, - thread_pool=thread_pool, - path=config.QDRANT_PATH, + db_path=config.QDRANT_PATH or tempfile.mkdtemp(prefix="test_qdrant_"), host=config.QDRANT_HOST, port=config.QDRANT_PORT, url=config.QDRANT_URL, @@ -251,7 +246,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor return PGVectorStore( collection_name=collection_name, embedding_model=embedding_model, - thread_pool=thread_pool, + db_path=tempfile.mkdtemp(prefix="test_pgvector_"), dsn=config.PG_DSN, min_size=config.PG_MIN_SIZE, max_size=config.PG_MAX_SIZE, @@ -262,8 +257,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor return ChromaVectorStore( collection_name=collection_name, embedding_model=embedding_model, - thread_pool=thread_pool, - path=config.CHROMA_PATH, + db_path=config.CHROMA_PATH, host=config.CHROMA_HOST, port=config.CHROMA_PORT, api_key=config.CHROMA_API_KEY, @@ -607,6 +601,7 @@ async def test_copy_collection(store: BaseVectorStore, store_name: str): # Verify content in copied collection copied_store = create_vector_store(store_type, copy_collection_name) + await copied_store.start() copied_nodes = await copied_store.list() logger.info(f"✓ Copied collection has {len(copied_nodes)} nodes") await copied_store.close() @@ -1464,14 +1459,13 @@ async def test_sql_injection_protection(store: BaseVectorStore, store_name: str) # Test 1: Invalid collection name (SQL injection attempt) try: embedding_model = OpenAIEmbeddingModel() - thread_pool = ThreadPoolExecutor(max_workers=4) # This should raise ValueError due to invalid table name try: _ = PGVectorStore( collection_name="test'; DROP TABLE users; --", + db_path=".", embedding_model=embedding_model, - thread_pool=thread_pool, ) logger.error("❌ FAILED: Invalid collection name was accepted (SQL injection risk!)") assert False, "Should have raised ValueError for invalid collection name" @@ -1657,6 +1651,9 @@ async def run_all_tests_for_store(store_type: str, store_name: str): store = create_vector_store(store_type, collection_name) try: + # Initialize the store (connect to database, create client, etc.) + await store.start() + # Run cosine similarity test first (only for LocalVectorStore) await test_cosine_similarity(store_name)