From e9757bccf609ec44d5217fd695c79bde61de6cd4 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 24 Feb 2026 19:45:47 +0800 Subject: [PATCH] refactor(core): restructure application initialization and component management --- reme/__init__.py | 3 +- reme/agent/fs/fs_compactor.py | 62 ++-- reme/agent/fs/fs_summarizer.py | 52 ++-- reme/agent/fs/fs_summarizer.yaml | 111 ++----- reme/agent/memory/__init__.py | 4 +- reme/config/cli.yaml | 6 +- reme/config/default.yaml | 6 +- reme/config/fs.yaml | 12 +- reme/core/application.py | 267 ++++++++++++----- reme/core/context/prompt_handler.py | 280 ++---------------- reme/core/context/runtime_context.py | 6 +- reme/core/context/service_context.py | 216 +------------- reme/core/embedding/base_embedding_model.py | 34 ++- reme/core/file_watcher/base_file_watcher.py | 4 +- reme/core/memory_store/base_memory_store.py | 1 + reme/core/memory_store/chroma_memory_store.py | 2 - reme/core/memory_store/local_memory_store.py | 5 +- reme/core/memory_store/sqlite_memory_store.py | 5 +- reme/core/schema/service_config.py | 5 +- reme/core/utils/__init__.py | 2 + reme/core/utils/execute_utils.py | 2 +- reme/{ => core/utils}/horse.py | 2 +- reme/reme_cli.py | 5 +- reme/reme_fs.py | 4 +- .../procedural_memory/summarizer/__init__.py | 4 +- tests/test_horse.py | 2 +- 26 files changed, 406 insertions(+), 696 deletions(-) rename reme/{ => core/utils}/horse.py (99%) diff --git a/reme/__init__.py b/reme/__init__.py index 537864d0..b02aaa94 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -20,8 +20,7 @@ __all__ = [ "ReMeFs", ] -__version__ = "0.3.0.0b3" - +__version__ = "0.3.0.0b4" """ conda create -n fl_test2 python=3.10 diff --git a/reme/agent/fs/fs_compactor.py b/reme/agent/fs/fs_compactor.py index 50e3a53a..a2dd8a17 100644 --- a/reme/agent/fs/fs_compactor.py +++ b/reme/agent/fs/fs_compactor.py @@ -11,6 +11,10 @@ from ...core.utils import format_messages class FsCompactor(BaseOp): """Generate summaries for conversation history compaction.""" + def __init__(self, return_prompt: bool = False, **kwargs): + super().__init__(**kwargs) + self.return_prompt = return_prompt + @staticmethod def _normalize_messages(messages: list[Message | dict]) -> list[Message]: """Convert dict messages to Message objects.""" @@ -27,7 +31,7 @@ class FsCompactor(BaseOp): return format_messages( messages=messages, add_index=False, - add_time=False, + add_time=True, use_name=True, add_reasoning=False, add_tools=True, @@ -58,33 +62,57 @@ class FsCompactor(BaseOp): system_prompt = self.get_prompt("system_prompt") conversation_text = self._serialize_conversation(turn_prefix_messages) - turn_prefix_prompt = self.prompt_format("turn_prefix_summarization", conversation_text=conversation_text) + turn_prefix_prompt = self.prompt_format( + "turn_prefix_summarization", + conversation_text=conversation_text, + ) return [ Message(role=Role.SYSTEM, content=system_prompt), Message(role=Role.USER, content=turn_prefix_prompt), ] - async def execute(self) -> str: + async def execute(self) -> str | dict: """Generate summary for conversation history.""" messages_to_summarize = self.context.get("messages_to_summarize", []) turn_prefix_messages = self.context.get("turn_prefix_messages", []) previous_summary = self.context.get("previous_summary", "") messages_to_summarize = self._normalize_messages(messages_to_summarize) - if messages_to_summarize: - history_prompt_messages = self._build_history_prompt(messages_to_summarize, previous_summary) - history_summary = "**History Summary**:\n\n" + await self._generate_summary(history_prompt_messages) - else: - history_summary = "" - turn_prefix_messages = self._normalize_messages(turn_prefix_messages) - if turn_prefix_messages: - turn_prefix_prompt_messages = self._build_turn_prefix_prompt(turn_prefix_messages) - turn_prefix_summary = "**Turn Context**:\n\n" + await self._generate_summary(turn_prefix_prompt_messages) - else: - turn_prefix_summary = "" - summary = "\n\n---".join([history_summary, turn_prefix_summary]) - logger.info(f"Generated summary: {summary}") - return summary + if self.return_prompt: + result = { + "system": self.get_prompt("system_prompt"), + } + + if messages_to_summarize: + history_prompt_messages = self._build_history_prompt(messages_to_summarize, previous_summary) + if len(history_prompt_messages) == 2: + result["history_user"] = history_prompt_messages[-1].content + + if turn_prefix_messages: + turn_prefix_prompt_messages = self._build_turn_prefix_prompt(turn_prefix_messages) + if len(turn_prefix_prompt_messages) == 2: + result["turn_prefix_user"] = turn_prefix_prompt_messages[-1].content + + return result + + else: + if messages_to_summarize: + history_prompt_messages = self._build_history_prompt(messages_to_summarize, previous_summary) + history_summary = "**History Summary**:\n\n" + await self._generate_summary(history_prompt_messages) + else: + history_summary = "" + + if turn_prefix_messages: + turn_prefix_prompt_messages = self._build_turn_prefix_prompt(turn_prefix_messages) + turn_prefix_summary = "**Turn Context**:\n\n" + await self._generate_summary( + turn_prefix_prompt_messages, + ) + else: + turn_prefix_summary = "" + + summary = "\n\n---".join([history_summary, turn_prefix_summary]) + logger.info(f"Generated summary: {summary}") + return summary diff --git a/reme/agent/fs/fs_summarizer.py b/reme/agent/fs/fs_summarizer.py index 38df4fb9..1fe9ddb3 100644 --- a/reme/agent/fs/fs_summarizer.py +++ b/reme/agent/fs/fs_summarizer.py @@ -13,30 +13,25 @@ from ...core.utils import format_messages class FsSummarizer(BaseReact): """Retrieve personal memories through vector search and history reading.""" - def __init__(self, working_dir: str, memory_dir: str = "memory", version: str = "v1", **kwargs): + def __init__( + self, + working_dir: str, + memory_dir: str = "memory", + version: str = "default", + return_prompt: bool = False, + **kwargs, + ): super().__init__(**kwargs) self.working_dir: str = working_dir self.memory_dir: str = memory_dir self.version: str = version + self.return_prompt = return_prompt async def build_messages(self) -> list[Message]: messages: list[Message] = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] date_str: str = self.context.get("date", datetime.datetime.now().strftime("%Y-%m-%d")) if self.version == "default": - messages.append( - Message( - role=Role.USER, - content=self.prompt_format( - "user_message_default", - working_dir=self.working_dir, - date=date_str, - memory_dir=self.memory_dir, - ), - ), - ) - - elif self.version == "v1": conversation = format_messages(messages, add_index=False) messages = [ Message( @@ -52,14 +47,27 @@ class FsSummarizer(BaseReact): ), ] + elif self.version == "v1": + messages.append( + Message( + role=Role.USER, + content=self.prompt_format( + "user_message_default", + working_dir=self.working_dir, + date=date_str, + memory_dir=self.memory_dir, + ), + ), + ) + else: messages.extend( [ - Message(role=Role.SYSTEM, content=self.get_prompt("system_prompt")), + Message(role=Role.SYSTEM, content=self.get_prompt("system_prompt_deprecated")), Message( role=Role.USER, content=self.prompt_format( - "user_message", + "user_message_deprecated", date=date_str, memory_dir=self.memory_dir, ), @@ -69,7 +77,13 @@ class FsSummarizer(BaseReact): return messages async def execute(self): - result = await super().execute() - answer = result["answer"] - logger.info(f"[{self.__class__.__name__}] answer={answer}") + if self.return_prompt: + result = {} + messages: list[Message] = await self.build_messages() + result["prompt"] = messages[-1].content + return result + else: + result = await super().execute() + answer = str(result["answer"]) + logger.info(f"[{self.__class__.__name__}] answer={answer}") return result diff --git a/reme/agent/fs/fs_summarizer.yaml b/reme/agent/fs/fs_summarizer.yaml index 477a79e8..e17eacf7 100644 --- a/reme/agent/fs/fs_summarizer.yaml +++ b/reme/agent/fs/fs_summarizer.yaml @@ -1,39 +1,38 @@ -system_prompt: | +system_prompt_deprecated: | Pre-compaction memory flush turn. The session is near auto-compaction; capture durable memories to disk. You may reply, but usually [SILENT] is correct. -user_message: | +user_message_deprecated: | Pre-compaction memory flush. Current Date: {date} Store durable memories now (use {memory_dir}/YYYY-MM-DD.md; create {memory_dir}/ if needed). If nothing to store, reply with [SILENT]. user_message_default: | - Pre-compaction memory flush turn. - The session is near auto-compaction; capture durable memories to disk. + Memory Pre-compression Flush Cycle Initiated + The current session is about to enter the automatic compression phase. Please capture persistent memory and write it to disk. - Current Date: {date} - Working Dir: {working_dir} + Current date: {date} + Working directory: {working_dir} - Store durable memories now (use {memory_dir}/YYYY-MM-DD.md). + Immediately store persistent memory to: {memory_dir}/YYYY-MM-DD.md Workflow: - 1. Use read_tool to read {memory_dir}/YYYY-MM-DD.md (if file doesn't exist, read_tool tool will return an error) - 2. Intelligently merge new information with existing content (skip if file doesn't exist): - - Avoid duplicating information that's already recorded - - Enrich existing entries with new details when relevant - - Maintain chronological order when applicable + 1. First, `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned). + 2. Intelligently merge new information with existing content (skip merging if the file doesn’t exist): + - Avoid duplicating already recorded information + - Enrich existing entries with new details where relevant + - Maintain chronological order wherever applicable 3. Write the updated content: - - Use edit_tool to update specific sections when possible - - Use write_tool to overwrite the entire file if major restructuring is needed - 4. Create {memory_dir}/ if it doesn't exist + - Prefer using `edit` to update specific sections when possible + - Use `write` to overwrite the entire file only if substantial restructuring is required Principles: - - Always preserve timestamps, dates, and time-related context - - Only add truly new or enriching information - - Keep entries concise but complete - - If nothing meaningful to store, reply with [SILENT] + - Always preserve timestamps and any date/time-related context + - Add only genuinely new or meaningfully enriching information + - Keep entries concise yet complete + - If there’s nothing to store, respond with [SILENT] user_message_default_zh: | @@ -46,83 +45,17 @@ user_message_default_zh: | 立即存储持久化记忆（使用路径 {memory_dir}/YYYY-MM-DD.md）。 工作流程： - 1. 使用 read_tool 读取 {memory_dir}/YYYY-MM-DD.md（如文件不存在，read_tool 会返回错误提示） + 1. 先 `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示） 2. 智能合并新信息与现有内容（若文件不存在则跳过合并）： - 避免重复已记录的信息 - 在相关时丰富现有条目的新细节 - 在适用时保持时间顺序 3. 写入更新后的内容： - - 尽可能使用 edit_tool 更新特定部分 - - 如需大幅重构则使用 write_tool 覆盖整个文件 - 4. 如 {memory_dir}/ 不存在则创建 + - 尽可能使用 `edit` 更新特定部分 + - 如需大幅重构则使用 `write` 覆盖整个文件 原则： - 始终保留时间戳、日期和时间相关上下文 - 仅添加真正新的或有丰富价值的信息 - 保持条目简洁但完整 - - 若无有意义的内容可存储，请回复 [SILENT] - - -user_message_v1: | - - {conversation} - - - The conversation is about to be compacted. Please extract persistent memories to disk. - Current Date: {date} - Working Dir: {working_dir} - - Execution Flow: - 1. Determine if the conversation contains information worth storing - - If no: Reply with reason + [SILENT] - - If yes: Continue to step 2 - - 2. Use Read tool to read {memory_dir}/YYYY-MM-DD.md (use actual date) - - If file doesn't exist (Read returns error): Write new memories directly - - If file exists: - a) Compare and identify new/updated information from the read content - c) Prefer `edit_tool` for precise additions (preserves existing content); `write_tool` overwrites entire file - d) If no new information: Reply with explanation + [SILENT] - - Update Principles: - - Only add unrecorded information, preserve all existing content - - Intelligently merge duplicate information, retain key details like timestamps - - Example (information merging): - Existing: "Alice joined Company A as Software Engineer on 2023-01-01" - New: "Alice joined Company B as Senior Engineer on 2024-01-01" - Result: "Alice joined Company A as Software Engineer on 2023-01-01, moved to Company B as Senior Engineer on 2024-01-01" - - Please store persistent memories, keeping entries concise and well-structured. - -user_message_v1_zh: | - - {conversation} - - - conversation即将压缩，请提取持久性记忆存储至磁盘。 - 当前日期：{date} - 工作目录: {working_dir} - - 执行流程： - 1. 判断对话是否包含值得存储的信息 - - 若无：回复原因 + [SILENT] - - 若有：继续步骤 2 - - 2. 使用 Read 工具读取 {memory_dir}/YYYY-MM-DD.md（使用实际日期） - - 文件不存在（Read 返回错误）：直接使用 `write_tool` 写入新记忆 - - 文件已存在： - a) 从读取的内容中对比识别新增/更新信息 - c) 优先使用 `edit_tool` 精准添加新信息（保留已有内容） - d) 若无新信息：回复说明 + [SILENT] - - 更新原则： - - 仅添加未记录的信息，保留所有已有内容 - - 智能合并重复信息，保留时间等关键细节 - - 示例（信息合并）： - 已有："Alice 于 2023-01-01 加入 A 公司任软件工程师" - 新增："Alice 于 2024-01-01 加入 B 公司任高级工程师" - 结果："Alice 于 2023-01-01 加入 A 公司任软件工程师，2024-01-01 转入 B 公司任高级工程师" - - 请存储持久性记忆，保持条目简洁、结构清晰。 + - 若无内容可存储，请回复 [SILENT] diff --git a/reme/agent/memory/__init__.py b/reme/agent/memory/__init__.py index ae158f8b..a1e941ee 100644 --- a/reme/agent/memory/__init__.py +++ b/reme/agent/memory/__init__.py @@ -1,12 +1,12 @@ """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 .personal.personal_halumem_retriever import PersonalHalumemRetriever -from .personal.personal_halumem_summarizer import PersonalHalumemSummarizer from .procedural.procedural_retriever import ProceduralRetriever from .procedural.procedural_summarizer import ProceduralSummarizer from .reme_retriever import ReMeRetriever diff --git a/reme/config/cli.yaml b/reme/config/cli.yaml index b71988c9..71f05647 100644 --- a/reme/config/cli.yaml +++ b/reme/config/cli.yaml @@ -23,7 +23,7 @@ embedding_models: memory_stores: default: backend: chroma -# backend: local + # backend: local db_name: reme.db store_name: reme embedding_model: default @@ -34,8 +34,8 @@ file_watchers: default: backend: full memory_store: default - watch_paths: [".reme", ".reme/memory"] - suffix_filters: [".md"] + watch_paths: [ ".reme", ".reme/memory" ] + suffix_filters: [ ".md" ] recursive: false scan_on_start: true diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 3a7e7073..36fd24b4 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -21,9 +21,9 @@ llms: default: backend: openai model_name: qwen3-30b-a3b-instruct-2507 -# model_name: qwen3-30b-a3b-thinking-2507 + # model_name: qwen3-30b-a3b-thinking-2507 request_interval: 1 -# temperature: 0.0001 + # temperature: 0.0001 qwen3_max_instruct: backend: openai @@ -46,7 +46,7 @@ embedding_models: vector_stores: default: backend: chroma -# backend: local + # backend: local embedding_model: default collection_name: reme diff --git a/reme/config/fs.yaml b/reme/config/fs.yaml index d9da8f1f..e4c191bb 100644 --- a/reme/config/fs.yaml +++ b/reme/config/fs.yaml @@ -3,8 +3,8 @@ backend: cmd llms: default: backend: openai -# model_name: qwen3-30b-a3b-instruct-2507 -# model_name: qwen3-30b-a3b-thinking-2507 + # model_name: qwen3-30b-a3b-instruct-2507 + # model_name: qwen3-30b-a3b-thinking-2507 model_name: qwen3-235b-a22b-thinking-2507 request_interval: 1 # temperature: 0.0001 @@ -17,9 +17,9 @@ embedding_models: memory_stores: default: -# backend: sqlite + # backend: sqlite backend: chroma -# backend: local + # backend: local db_name: reme.db store_name: reme embedding_model: default @@ -30,8 +30,8 @@ file_watchers: default: backend: full memory_store: default - watch_paths: [".reme", ".reme/memory"] - suffix_filters: [".md"] + watch_paths: [ ".reme", ".reme/memory" ] + suffix_filters: [ ".md" ] recursive: false scan_on_start: true diff --git a/reme/core/application.py b/reme/core/application.py index c1792f5a..0d85fc29 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -1,8 +1,12 @@ """High-level entry point for configuring and running ReMe services and flows.""" import asyncio +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path -from .context import PromptHandler, ServiceContext +from loguru import logger + +from .context import PromptHandler, ServiceContext, R from .embedding import BaseEmbeddingModel from .file_watcher import BaseFileWatcher from .flow import BaseFlow @@ -10,7 +14,7 @@ from .llm import BaseLLM from .memory_store import BaseMemoryStore from .schema import Response, ServiceConfig from .token_counter import BaseTokenCounter -from .utils import execute_stream_task, PydanticConfigParser +from .utils import execute_stream_task, PydanticConfigParser, init_logger, print_logo, MCPClient from .vector_store import BaseVectorStore @@ -57,80 +61,207 @@ class Application: default_file_watcher_config=default_file_watcher_config, **kwargs, ) - self.prompt_handler = PromptHandler(language=self.service_context.language) + self.prompt_handler = PromptHandler(language=self.service_config.language) self._started: bool = False - def update_api_envs( - self, - llm_api_key: str | None = None, - llm_base_url: str | None = None, - embedding_api_key: str | None = None, - embedding_base_url: str | None = None, - ): - """Update the API environment variables.""" - self.service_context.update_api_envs( - llm_api_key=llm_api_key, - llm_base_url=llm_base_url, - embedding_api_key=embedding_api_key, - embedding_base_url=embedding_base_url, - ) - @classmethod - async def create( - cls, - *args, - llm_api_key: str | None = None, - llm_base_url: str | None = None, - embedding_api_key: str | None = None, - embedding_base_url: str | None = None, - enable_logo: bool = True, - parser: type[PydanticConfigParser] | None = None, - llm: dict | None = None, - embedding_model: dict | None = None, - vector_store: dict | None = None, - memory_store: dict | None = None, - token_counter: dict | None = None, - file_watcher: dict | None = None, - **kwargs, - ) -> "Application": + async def create(cls, *args, **kwargs) -> "Application": """Create and start an Application instance asynchronously.""" - instance = cls( - *args, - llm_api_key=llm_api_key, - llm_base_url=llm_base_url, - embedding_api_key=embedding_api_key, - embedding_base_url=embedding_base_url, - enable_logo=enable_logo, - parser=parser, - default_llm_config=llm, - default_embedding_model_config=embedding_model, - default_vector_store_config=vector_store, - default_memory_store_config=memory_store, - default_token_counter_config=token_counter, - default_file_watcher_config=file_watcher, - **kwargs, - ) + instance = cls(*args, **kwargs) 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 application.""" + """Start the service context by initializing all configured components.""" if self._started: - return self - else: - await self.service_context.start() - self._started = True + logger.warning("Application has already started.") return self - async def close(self): - """Close the application.""" - if self._started: - await self.service_context.close() - self._started = False + 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, + ) + + expression_flow_cls = None + for name, flow_cls in R.flows.items(): + if not self._filter_flows(name): + continue + + if name == "ExpressionFlow": + expression_flow_cls = flow_cls + else: + flow: "BaseFlow" = flow_cls(name=name, service_context=self.service_context) + self.service_context.flows[flow.name] = flow + + if expression_flow_cls is not None: + for name, flow_config in self.service_config.flows.items(): + if not self._filter_flows(name): + continue + flow_config.name = name + flow: BaseFlow = expression_flow_cls( # noqa + flow_config=flow_config, + service_context=self.service_context, + ) + self.service_context.flows[flow.name] = flow else: - raise RuntimeError("Application is not started") + logger.info("No expression flow found, please check your configuration.") + + 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) + + for name, config in self.service_config.embedding_models.items(): + if config.backend not in R.embedding_models: + logger.warning(f"Embedding model backend {config.backend} is not supported.") + else: + 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) + + for name, config in self.service_config.token_counters.items(): + if config.backend not in R.token_counters: + logger.warning(f"Token counter backend {config.backend} is not supported.") + else: + config_dict = config.model_dump(exclude={"backend"}) + self.service_context.token_counters[name] = R.token_counters[config.backend](**config_dict) + + for name, config in self.service_config.vector_stores.items(): + if config.backend not in R.vector_stores: + logger.warning(f"Vector 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, + }, + ) + 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) + + 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.") + 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", + }, + ) + self.service_context.memory_stores[name] = R.memory_stores[config.backend](**config_dict) + await self.service_context.memory_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] + self.service_context.file_watchers[name] = R.file_watchers[config.backend](**config_dict) + await self.service_context.file_watchers[name].start() + + if self.service_config.mcp_servers: + await self.prepare_mcp_servers() + + 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}) + for server_name in self.service_config.mcp_servers.keys(): + try: + tool_calls = await mcp_client.list_tool_calls(server_name=server_name, return_dict=False) + self.service_context.mcp_server_mapping[server_name] = { + tool_call.name: tool_call for tool_call in tool_calls + } + for tool_call in tool_calls: + logger.info(f"list_tool_calls: {server_name}@{tool_call.name} {tool_call.simple_input_dump()}") + except Exception as e: + logger.exception(f"list_tool_calls: {server_name} error: {e}") + + async def close(self) -> bool: + """Close all service components asynchronously.""" + if not self._started: + logger.warning("Application is not started") + return True + + for name, vector_store in self.service_context.vector_stores.items(): + 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_watcher in self.service_context.file_watchers.items(): + logger.info(f"Closing file watcher: {name}") + await file_watcher.close() + + for name, llm in self.service_context.llms.items(): + logger.info(f"Closing LLM: {name}") + await llm.close() + + for name, embedding_model in self.service_context.embedding_models.items(): + logger.info(f"Closing embedding model: {name}") + await embedding_model.close() + + self.shutdown_thread_pool() + self.shutdown_ray() + + self._started = False return False + def shutdown_thread_pool(self, wait: bool = True): + """Shutdown the thread pool executor.""" + if self.service_context.thread_pool: + self.service_context.thread_pool.shutdown(wait=wait) + + def shutdown_ray(self, wait: bool = True): + """Shutdown Ray cluster if it was initialized.""" + if self.service_config and self.service_config.ray_max_workers > 1: + import ray + + ray.shutdown(_exiting_interpreter=not wait) + async def __aenter__(self): """Async context manager entry.""" return await self.start() @@ -218,11 +349,6 @@ class Application: """Get the default token counter instance.""" return self.service_context.token_counters.get("default") - @property - def service_config(self) -> ServiceConfig: - """Get the service configuration.""" - return self.service_context.service_config - def get_token_counter(self, name: str): """Get a token counter instance by name.""" return self.service_context.token_counters.get(name) @@ -232,4 +358,9 @@ class Application: import warnings warnings.filterwarnings("ignore", category=DeprecationWarning) - self.service_context.service.run() + service = R.services[self.service_config.backend](service_context=self.service_context) + service.run() + + async def reset_default_collection(self, collection_name: str): + """Reset the default vector store.""" + await self.service_context.vector_stores["default"].reset_collection(collection_name) diff --git a/reme/core/context/prompt_handler.py b/reme/core/context/prompt_handler.py index e93ee4c2..34ae39fe 100644 --- a/reme/core/context/prompt_handler.py +++ b/reme/core/context/prompt_handler.py @@ -1,12 +1,4 @@ -"""Module for managing and formatting prompt templates from files or dictionaries. - -This module provides a PromptHandler class that: -- Loads prompts from YAML/JSON files or dictionaries -- Supports multi-language prompts with automatic suffix handling -- Provides conditional line filtering using boolean flags -- Formats prompts with template variable substitution -- Validates format strings and provides helpful error messages -""" +"""Module for managing and formatting prompt templates from files or dictionaries.""" import json from pathlib import Path @@ -19,54 +11,10 @@ from loguru import logger from .base_context import BaseContext -class PromptNotFoundError(KeyError): - """Exception raised when a requested prompt template is not found.""" - - def __init__(self, prompt_name: str, available_prompts: list[str]): - self.prompt_name = prompt_name - self.available_prompts = available_prompts - super().__init__( - f"Prompt '{prompt_name}' not found. " - f"Available prompts: {', '.join(available_prompts[:10])}" - f"{'...' if len(available_prompts) > 10 else ''}", - ) - - -class PromptFormattingError(ValueError): - """Exception raised when prompt formatting fails.""" - - class PromptHandler(BaseContext): - """A context-aware handler for loading, retrieving, and formatting prompt templates. - - This handler supports: - - Loading prompts from YAML/JSON files or dictionaries - - Multi-language prompt support with automatic language suffix - - Conditional line filtering using boolean flags (e.g., [debug], [verbose]) - - Template variable substitution with validation - - Method chaining for fluent API - - Examples: - >>> handler = PromptHandler(language="en") - >>> handler.load_prompt_dict({ - ... "greeting_en": "Hello, {name}!", - ... "farewell_en": "[debug]Debug mode\\nGoodbye, {name}!" - ... }) - >>> handler.prompt_format("greeting", name="Alice") - 'Hello, Alice!' - >>> handler.prompt_format("farewell", name="Bob", debug=False) - 'Goodbye, Bob!' - """ + """A context-aware handler for loading, retrieving, and formatting prompt templates.""" def __init__(self, language: str = "", **kwargs): - """Initialize the PromptHandler with optional language configuration. - - Args: - language: Language code to append as suffix (e.g., "en", "zh", "ja"). - If provided, get_prompt will automatically try to find - prompts with this suffix (e.g., "greeting" -> "greeting_en"). - **kwargs: Additional key-value pairs to initialize the context. - """ super().__init__(**kwargs) self.language: str = language.strip() @@ -75,25 +23,7 @@ class PromptHandler(BaseContext): prompt_file_path: Optional[Union[Path, str]] = None, overwrite: bool = True, ) -> "PromptHandler": - """Load prompt configurations from a YAML or JSON file into the context. - - Supports both YAML (.yaml, .yml) and JSON (.json) file formats. - Non-existent files are silently skipped. - - Args: - prompt_file_path: Path to the prompt configuration file. - If None, returns self without changes. - overwrite: If True, allows overwriting existing prompts with warnings. - If False, skips existing prompts without overwriting. - - Returns: - Self for method chaining. - - Raises: - ValueError: If file format is not supported. - yaml.YAMLError: If YAML parsing fails. - json.JSONDecodeError: If JSON parsing fails. - """ + """Load prompt configurations from a YAML or JSON file.""" if prompt_file_path is None: return self @@ -105,23 +35,15 @@ class PromptHandler(BaseContext): suffix = prompt_file_path.suffix.lower() - try: - with prompt_file_path.open(encoding="utf-8") as f: - if suffix in [".yaml", ".yml"]: - prompt_dict = yaml.safe_load(f) - elif suffix == ".json": - prompt_dict = json.load(f) - else: - raise ValueError( - f"Unsupported file format: {suffix}. " f"Supported formats: .yaml, .yml, .json", - ) - - self.load_prompt_dict(prompt_dict, overwrite=overwrite) - - except (yaml.YAMLError, json.JSONDecodeError) as e: - logger.error(f"Failed to parse prompt file {prompt_file_path}: {e}") - raise + with prompt_file_path.open(encoding="utf-8") as f: + if suffix in [".yaml", ".yml"]: + prompt_dict = yaml.safe_load(f) + elif suffix == ".json": + prompt_dict = json.load(f) + else: + raise ValueError(f"Unsupported file format: {suffix}") + self.load_prompt_dict(prompt_dict, overwrite=overwrite) return self def load_prompt_dict( @@ -129,233 +51,95 @@ class PromptHandler(BaseContext): prompt_dict: Optional[Dict[str, Any]] = None, overwrite: bool = True, ) -> "PromptHandler": - """Merge a dictionary of prompt strings into the current context. - - Only string values are stored as prompts. Non-string values are skipped. - - Args: - prompt_dict: Dictionary mapping prompt names to prompt template strings. - overwrite: If True, allows overwriting existing prompts with warnings. - If False, skips existing prompts without overwriting. - - Returns: - Self for method chaining. - """ + """Merge a dictionary of prompt strings into the current context.""" if not prompt_dict: return self for key, value in prompt_dict.items(): if not isinstance(value, str): - logger.debug(f"Skipping non-string prompt: key={key}, type={type(value)}") continue - if key in self: if overwrite: - logger.warning( - f"Overwriting prompt '{key}': " f"old length={len(self[key])}, new length={len(value)}", - ) + logger.warning(f"Overwriting prompt '{key}'") self[key] = value - else: - logger.debug(f"Skipping existing prompt: key={key}") else: - logger.debug(f"Adding new prompt: key={key}, length={len(value)}") self[key] = value return self def get_prompt(self, prompt_name: str, fallback_to_base: bool = True) -> str: - """Retrieve a prompt by name with automatic language suffix handling. - - If a language is configured, this method will: - 1. First try to find the prompt with language suffix (e.g., "greeting_en") - 2. If not found and fallback_to_base is True, try the base name (e.g., "greeting") - 3. Otherwise, raise PromptNotFoundError - - Args: - prompt_name: Name of the prompt to retrieve. - fallback_to_base: If True and language-specific prompt not found, - fallback to prompt without language suffix. - - Returns: - The prompt template string, stripped of leading/trailing whitespace. - - Raises: - PromptNotFoundError: If the prompt is not found. - """ - # Try with language suffix first + """Retrieve a prompt by name with automatic language suffix handling.""" if self.language and not prompt_name.endswith(f"_{self.language}"): key_with_lang = f"{prompt_name}_{self.language}" if key_with_lang in self: return self[key_with_lang].strip() - # Try base name if prompt_name in self: return self[prompt_name].strip() - # Try fallback if enabled - if fallback_to_base and self.language: - # Check if prompt_name already has language suffix, try without it - if prompt_name.endswith(f"_{self.language}"): - base_name = prompt_name[: -(len(self.language) + 1)] - if base_name in self: - return self[base_name].strip() + if fallback_to_base and self.language and prompt_name.endswith(f"_{self.language}"): + base_name = prompt_name[: -(len(self.language) + 1)] + if base_name in self: + return self[base_name].strip() - # Not found, raise error with helpful message - available = list(self.keys()) - raise PromptNotFoundError(prompt_name, available) + raise KeyError(f"Prompt '{prompt_name}' not found. Available: {list(self.keys())[:10]}") def has_prompt(self, prompt_name: str) -> bool: - """Check if a prompt exists (with or without language suffix). - - Args: - prompt_name: Name of the prompt to check. - - Returns: - True if the prompt exists, False otherwise. - """ + """Check if a prompt exists.""" try: self.get_prompt(prompt_name) return True - except PromptNotFoundError: + except KeyError: return False def list_prompts(self, language_filter: Optional[str] = None) -> list[str]: - """List all available prompt names. - - Args: - language_filter: If provided, only return prompts for this language. - If None, return all prompts. - - Returns: - List of prompt names. - """ + """List all available prompt names.""" if language_filter is None: return list(self.keys()) - suffix = f"_{language_filter.strip()}" return [key for key in self.keys() if key.endswith(suffix)] @staticmethod def _extract_format_fields(template: str) -> set[str]: - """Extract all format field names from a template string. - - Args: - template: Template string with {variable} placeholders. - - Returns: - Set of field names used in the template. - """ + """Extract all format field names from a template string.""" return {field_name for _, field_name, _, _ in Formatter().parse(template) if field_name is not None} @staticmethod def _filter_conditional_lines(prompt: str, flags: Dict[str, bool]) -> str: - """Filter lines based on boolean flags. - - Lines starting with [flag_name] are conditionally included based on - the value of flags[flag_name]. If True, the line is included (without - the flag marker). If False, the line is excluded. - - Args: - prompt: The prompt text with conditional markers. - flags: Dictionary of flag names to boolean values. - - Returns: - Filtered prompt text. - """ + """Filter lines based on boolean flags.""" filtered_lines = [] - for line in prompt.split("\n"): - # Check each flag matched_flag = None for flag_name in flags: - marker = f"[{flag_name}]" - if line.startswith(marker): + if line.startswith(f"[{flag_name}]"): matched_flag = flag_name break - if matched_flag is None: - # No flag marker, always include filtered_lines.append(line) elif flags[matched_flag]: - # Flag is True, include without marker - marker = f"[{matched_flag}]" - filtered_lines.append(line[len(marker) :]) - # else: Flag is False, skip this line - + filtered_lines.append(line[len(f"[{matched_flag}]") :]) return "\n".join(filtered_lines) - def prompt_format( - self, - prompt_name: str, - validate: bool = True, - **kwargs, - ) -> str: - """Format a prompt with conditional line filtering and variable substitution. - - This method performs two-stage formatting: - 1. Conditional line filtering: Lines marked with [flag] are included only - if the corresponding boolean kwarg is True. - 2. Variable substitution: Template variables {var} are replaced with - provided values. - - Args: - prompt_name: Name of the prompt to format. - validate: If True, check that all required template variables are provided. - **kwargs: Keyword arguments for formatting. Boolean values are treated as - conditional flags, other values are used for template substitution. - - Returns: - Formatted prompt string. - - Raises: - PromptNotFoundError: If the prompt is not found. - PromptFormattingError: If validation fails or formatting errors occur. - - Examples: - >>> handler = PromptHandler() - >>> handler["test"] = "[debug]Debug: {info}\\nResult: {value}" - >>> handler.prompt_format("test", debug=False, info="test", value=42) - 'Result: 42' - >>> handler.prompt_format("test", debug=True, info="test", value=42) - 'Debug: test\\nResult: 42' - """ - # Get the prompt template + def prompt_format(self, prompt_name: str, validate: bool = True, **kwargs) -> str: + """Format a prompt with conditional line filtering and variable substitution.""" prompt = self.get_prompt(prompt_name) - # Separate boolean flags from format variables flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} format_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} - # Step 1: Filter conditional lines if flag_kwargs: prompt = self._filter_conditional_lines(prompt, flag_kwargs) - # Step 2: Validate required fields if requested if validate: required_fields = self._extract_format_fields(prompt) missing_fields = required_fields - set(format_kwargs.keys()) - if missing_fields: - raise PromptFormattingError( - f"Missing required format variables for prompt '{prompt_name}': " - f"{', '.join(sorted(missing_fields))}", - ) + raise ValueError(f"Missing format variables for '{prompt_name}': {sorted(missing_fields)}") - # Step 3: Format with variables - try: - if format_kwargs: - prompt = prompt.format(**format_kwargs) - except KeyError as e: - raise PromptFormattingError( - f"Format error in prompt '{prompt_name}': missing variable {e}", - ) from e - except (ValueError, IndexError) as e: - raise PromptFormattingError( - f"Format error in prompt '{prompt_name}': {e}", - ) from e + if format_kwargs: + prompt = prompt.format(**format_kwargs) return prompt.strip() def __repr__(self) -> str: - """Return a string representation of the PromptHandler.""" - return f"PromptHandler(language='{self.language}', " f"num_prompts={len(self)})" + return f"PromptHandler(language='{self.language}', num_prompts={len(self)})" diff --git a/reme/core/context/runtime_context.py b/reme/core/context/runtime_context.py index d4c02050..653fc0fc 100644 --- a/reme/core/context/runtime_context.py +++ b/reme/core/context/runtime_context.py @@ -70,7 +70,11 @@ class RuntimeContext(BaseContext): self[target] = self[source] return self - def validate_required_keys(self, required_keys: dict[str, bool], context_name: str = "context") -> "RuntimeContext": + def validate_required_keys( + self, + required_keys: dict[str, bool], + context_name: str = "context", + ) -> "RuntimeContext": """Ensure all required keys are present in the context. Args: diff --git a/reme/core/context/service_context.py b/reme/core/context/service_context.py index febf956d..e5a9860d 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/context/service_context.py @@ -2,15 +2,13 @@ import os from concurrent.futures import ThreadPoolExecutor -from pathlib import Path from typing import TYPE_CHECKING from loguru import logger from .base_context import BaseContext -from .registry_factory import R from ..schema import ServiceConfig -from ..utils import load_env, MCPClient, print_logo, PydanticConfigParser, init_logger +from ..utils import load_env, PydanticConfigParser if TYPE_CHECKING: from ..llm import BaseLLM @@ -19,7 +17,6 @@ if TYPE_CHECKING: from ..memory_store import BaseMemoryStore from ..token_counter import BaseTokenCounter from ..flow import BaseFlow - from ..service import BaseService from ..file_watcher import BaseFileWatcher @@ -49,8 +46,14 @@ class ServiceContext(BaseContext): ): super().__init__() + # Load environment variables load_env() - self.update_api_envs(llm_api_key, llm_base_url, embedding_api_key, embedding_base_url) + + # Update common environment variables for LLM and embedding services. + self.update_env("REME_LLM_API_KEY", llm_api_key) + self.update_env("REME_LLM_BASE_URL", llm_base_url) + self.update_env("REME_EMBEDDING_API_KEY", embedding_api_key) + self.update_env("REME_EMBEDDING_BASE_URL", embedding_base_url) if service_config is None: parser_class = parser if parser is not None else PydanticConfigParser @@ -85,37 +88,16 @@ class ServiceContext(BaseContext): service_config = parser.parse_args(*input_args, **kwargs) self.service_config: ServiceConfig = service_config - init_logger(log_to_console=self.service_config.log_to_console) - logger.info(f"ReMe Config: {service_config.model_dump_json()}") - - if self.service_config.working_dir: - self.working_path = Path(self.service_config.working_dir) - self.working_path.mkdir(parents=True, exist_ok=True) - - if self.service_config.enable_logo: - print_logo(service_config=self.service_config) - - self.language: str = self.service_config.language - self.thread_pool: ThreadPoolExecutor = ThreadPoolExecutor( - max_workers=self.service_config.thread_pool_max_workers, - ) - if self.service_config.ray_max_workers > 1: - import ray - - ray.init(num_cpus=self.service_config.ray_max_workers) + self.thread_pool: ThreadPoolExecutor | None = None self.llms: dict[str, "BaseLLM"] = {} 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_watchers: dict[str, "BaseFileWatcher"] = {} - self.flows: dict[str, "BaseFlow"] = {} self.mcp_server_mapping: dict[str, dict] = {} - self.service: "BaseService" = R.services[self.service_config.backend](service_context=self) - - self._build_flows() @staticmethod def update_env(key: str, value: str | None): @@ -123,19 +105,6 @@ class ServiceContext(BaseContext): if value: os.environ[key] = value - def update_api_envs( - self, - llm_api_key: str | None = None, - llm_base_url: str | None = None, - embedding_api_key: str | None = None, - embedding_base_url: str | None = None, - ): - """Update common environment variables for LLM and embedding services.""" - self.update_env("REME_LLM_API_KEY", llm_api_key) - self.update_env("REME_LLM_BASE_URL", llm_base_url) - self.update_env("REME_EMBEDDING_API_KEY", embedding_api_key) - self.update_env("REME_EMBEDDING_BASE_URL", embedding_base_url) - @staticmethod def _update_section_config(config: dict, section_name: str, **kwargs): """Update a specific section of the service config with new values.""" @@ -144,170 +113,3 @@ class ServiceContext(BaseContext): if "default" not in config[section_name]: config[section_name]["default"] = {} config[section_name]["default"].update(kwargs) - - def _build_flows(self): - expression_flow_cls = None - for name, flow_cls in R.flows.items(): - if not self._filter_flows(name): - continue - - if name == "ExpressionFlow": - expression_flow_cls = flow_cls - else: - flow: "BaseFlow" = flow_cls(name=name, service_context=self) - self.flows[flow.name] = flow - - if expression_flow_cls is not None: - for name, flow_config in self.service_config.flows.items(): - if not self._filter_flows(name): - continue - flow_config.name = name - flow: BaseFlow = expression_flow_cls(flow_config=flow_config, service_context=self) # noqa - self.flows[flow.name] = flow - else: - logger.info("No expression flow found, please check your configuration.") - - async def start(self): - """Start the service context by initializing all configured components.""" - # Recreate thread pool if it was shut down - if self.thread_pool is None or self.thread_pool._shutdown: # pylint: disable=protected-access - self.thread_pool = ThreadPoolExecutor( - max_workers=self.service_config.thread_pool_max_workers, - ) - - # Re-initialize Ray if it was shut down - 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) - - 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: - self.llms[name] = R.llms[config.backend](model_name=config.model_name, **config.model_extra) - - for name, config in self.service_config.embedding_models.items(): - if config.backend not in R.embedding_models: - logger.warning(f"Embedding model backend {config.backend} is not supported.") - else: - self.embedding_models[name] = R.embedding_models[config.backend]( - cache_dir=self.working_path / "embedding_cache", - model_name=config.model_name, - **config.model_extra, - ) - - for name, config in self.service_config.token_counters.items(): - if config.backend not in R.token_counters: - logger.warning(f"Token counter backend {config.backend} is not supported.") - else: - self.token_counters[name] = R.token_counters[config.backend]( - model_name=config.model_name, - **config.model_extra, - ) - - for name, config in self.service_config.vector_stores.items(): - if config.backend not in R.vector_stores: - logger.warning(f"Vector store backend {config.backend} is not supported.") - else: - # Extract config dict and replace special fields with actual instances - config_dict = config.model_dump(exclude={"backend", "embedding_model"}) - config_dict.update( - { - "embedding_model": self.embedding_models[config.embedding_model], - "thread_pool": self.thread_pool, - }, - ) - self.vector_stores[name] = R.vector_stores[config.backend](**config_dict) - await self.vector_stores[name].create_collection(config.collection_name) - - 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.") - else: - # Extract config dict and replace embedding_model string with actual instance - config_dict = config.model_dump(exclude={"backend", "embedding_model"}) - config_dict.update( - { - "embedding_model": self.embedding_models[config.embedding_model], - "thread_pool": self.thread_pool, - "db_path": self.working_path / config.db_name, - }, - ) - self.memory_stores[name] = R.memory_stores[config.backend](**config_dict) - await self.memory_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.memory_stores[config.memory_store] - self.file_watchers[name] = R.file_watchers[config.backend](**config_dict) - await self.file_watchers[name].start() - - if self.service_config.mcp_servers: - await self.prepare_mcp_servers() - - 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}) - for server_name in self.service_config.mcp_servers.keys(): - try: - tool_calls = await mcp_client.list_tool_calls(server_name=server_name, return_dict=False) - self.mcp_server_mapping[server_name] = {tool_call.name: tool_call for tool_call in tool_calls} - for tool_call in tool_calls: - logger.info(f"list_tool_calls: {server_name}@{tool_call.name} {tool_call.simple_input_dump()}") - except Exception as e: - logger.exception(f"list_tool_calls: {server_name} error: {e}") - - async def reset_default_collection(self, collection_name: str): - """Reset the default vector store.""" - await self.vector_stores["default"].reset_collection(collection_name) - - async def close(self): - """Close all service components asynchronously.""" - for name, vector_store in self.vector_stores.items(): - logger.info(f"Closing vector store: {name}") - await vector_store.close() - - for name, memory_store in self.memory_stores.items(): - logger.info(f"Closing memory store: {name}") - await memory_store.close() - - for name, file_watcher in self.file_watchers.items(): - logger.info(f"Closing file watcher: {name}") - await file_watcher.close() - - for name, llm in self.llms.items(): - logger.info(f"Closing LLM: {name}") - await llm.close() - - for name, embedding_model in self.embedding_models.items(): - logger.info(f"Closing embedding model: {name}") - await embedding_model.close() - - self.shutdown_thread_pool() - self.shutdown_ray() - - def shutdown_thread_pool(self, wait: bool = True): - """Shutdown the thread pool executor.""" - if self.thread_pool: - self.thread_pool.shutdown(wait=wait) - - def shutdown_ray(self, wait: bool = True): - """Shutdown Ray cluster if it was initialized.""" - if self.service_config and self.service_config.ray_max_workers > 1: - import ray - - ray.shutdown(_exiting_interpreter=not wait) diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index 7d467f17..b474fd29 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -37,6 +37,7 @@ class BaseEmbeddingModel(ABC): max_input_length: int = 8192, cache_dir: str | Path = ".reme", max_cache_size: int = 2000, + enable_cache: bool = True, **kwargs, ): """Initialize model configuration and parameters. @@ -51,6 +52,7 @@ class BaseEmbeddingModel(ABC): raise_exception: Whether to raise exceptions on failure max_input_length: Maximum input text length max_cache_size: Maximum number of embeddings to cache in memory (LRU) + enable_cache: Whether to enable embedding cache **kwargs: Additional model-specific parameters """ self._api_key: str = api_key @@ -63,6 +65,7 @@ class BaseEmbeddingModel(ABC): self.max_input_length = max_input_length self.cache_dir = cache_dir self.max_cache_size = max_cache_size + self.enable_cache = enable_cache self.kwargs = kwargs # Initialize LRU cache for embeddings @@ -89,9 +92,7 @@ class BaseEmbeddingModel(ABC): def _truncate_text(self, text: str) -> str: """Truncate text to max_input_length if it exceeds the limit.""" if len(text) > self.max_input_length: - logger.warning( - f"Text length {len(text)} exceeds max_input_length {self.max_input_length}, truncating", - ) + logger.warning(f"Text length {len(text)} exceeds {self.max_input_length}, truncating") return text[: self.max_input_length] return text @@ -133,6 +134,9 @@ class BaseEmbeddingModel(ABC): Loads in reverse order (newest first) to prioritize recent embeddings when max_cache_size is smaller than the file content. """ + if not self.enable_cache: + return + cache_file = self._get_cache_file_path() if not cache_file.exists(): logger.info(f"No cache file found at {cache_file}, starting with empty cache") @@ -151,8 +155,10 @@ class BaseEmbeddingModel(ABC): continue try: data = json.loads(line) - cache_key = data.get("key") - embedding = data.get("embedding") + if not data: + continue + # Each line is {cache_key: embedding} + cache_key, embedding = next(iter(data.items())) if cache_key and embedding: # Skip if already loaded (keep the newest) @@ -174,7 +180,12 @@ class BaseEmbeddingModel(ABC): logger.info(f"Loaded {loaded_count} embeddings from cache file: {cache_file}") except Exception as e: - logger.error(f"Failed to load cache from {cache_file}: {e}") + logger.error(f"Failed to load cache from {cache_file}: {e}, deleting cache file") + try: + cache_file.unlink() + logger.info(f"Deleted corrupted cache file: {cache_file}") + except Exception as del_e: + logger.error(f"Failed to delete cache file {cache_file}: {del_e}") def _save_cache(self) -> None: """Save embedding cache to disk (JSONL format). @@ -182,6 +193,9 @@ class BaseEmbeddingModel(ABC): Each line contains a JSON object with the cache key and embedding vector. Only saves if cache is non-empty. """ + if not self.enable_cache: + return + logger.info(f"Attempting to save cache, current size: {len(self._embedding_cache)}") if not self._embedding_cache: logger.info("Cache is empty, skipping save") @@ -191,7 +205,7 @@ class BaseEmbeddingModel(ABC): try: with open(cache_file, "w", encoding="utf-8") as f: for cache_key, embedding in self._embedding_cache.items(): - cache_entry = {"key": cache_key, "embedding": embedding} + cache_entry = {cache_key: embedding} f.write(json.dumps(cache_entry, ensure_ascii=False) + "\n") logger.info(f"Saved {len(self._embedding_cache)} embeddings to cache file: {cache_file}") @@ -207,6 +221,9 @@ class BaseEmbeddingModel(ABC): Returns: Cached embedding vector or None if not found """ + if not self.enable_cache: + return None + cache_key = self._get_cache_key(text) if cache_key in self._embedding_cache: # Move to end (most recently used) @@ -227,6 +244,9 @@ class BaseEmbeddingModel(ABC): text: Input text used as cache key embedding: Embedding vector to cache """ + if not self.enable_cache: + return + if self.max_cache_size <= 0: return diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index c5313a7b..98c0325a 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -134,10 +134,10 @@ class BaseFileWatcher: existing_files.add((Change.added, str(file_path))) if existing_files: - logger.info(f"Found {len(existing_files)} existing files to process") + logger.info(f"[SCAN_ON_START] Found {len(existing_files)} existing files matching watch criteria") await self.on_changes(existing_files) else: - logger.info("No existing files found matching watch criteria") + logger.info("[SCAN_ON_START] No existing files found matching watch criteria") files: list[str] = await self.memory_store.list_files(MemorySource.MEMORY) for file_path in files: diff --git a/reme/core/memory_store/base_memory_store.py b/reme/core/memory_store/base_memory_store.py index 01121114..c9083a3c 100644 --- a/reme/core/memory_store/base_memory_store.py +++ b/reme/core/memory_store/base_memory_store.py @@ -38,6 +38,7 @@ class BaseMemoryStore(ABC): 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.vector_enabled: bool = vector_enabled diff --git a/reme/core/memory_store/chroma_memory_store.py b/reme/core/memory_store/chroma_memory_store.py index d5ea1737..ee7c4911 100644 --- a/reme/core/memory_store/chroma_memory_store.py +++ b/reme/core/memory_store/chroma_memory_store.py @@ -112,8 +112,6 @@ class ChromaMemoryStore(BaseMemoryStore): if self.client is not None: return - self.db_path.mkdir(parents=True, exist_ok=True) - # Initialize persistent ChromaDB client self.client = chromadb.PersistentClient( path=str(self.db_path), diff --git a/reme/core/memory_store/local_memory_store.py b/reme/core/memory_store/local_memory_store.py index b7ecd51b..bee6882b 100644 --- a/reme/core/memory_store/local_memory_store.py +++ b/reme/core/memory_store/local_memory_store.py @@ -51,8 +51,8 @@ class LocalMemoryStore(BaseMemoryStore): self._chunks: dict[str, _ChunkRecord] = {} self._files: dict[str, dict[str, FileMetadata]] = {} # source -> path -> meta # Persistence paths (mirror ChromaMemoryStore convention) - self._chunks_file: Path = self.db_path.parent / f"{self.store_name}_chunks.jsonl" - self._metadata_file: Path = self.db_path.parent / f"{self.store_name}_file_metadata.json" + 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" # ------------------------------------------------------------------ # Persistence helpers @@ -142,7 +142,6 @@ class LocalMemoryStore(BaseMemoryStore): if self._started: return self._started = True - self.db_path.mkdir(parents=True, exist_ok=True) await self._load_metadata() await self._load_chunks() logger.info( diff --git a/reme/core/memory_store/sqlite_memory_store.py b/reme/core/memory_store/sqlite_memory_store.py index b3ee67b9..e05c6adc 100644 --- a/reme/core/memory_store/sqlite_memory_store.py +++ b/reme/core/memory_store/sqlite_memory_store.py @@ -4,7 +4,6 @@ import json import sqlite3 import struct import time -from pathlib import Path from loguru import logger @@ -63,9 +62,7 @@ class SqliteMemoryStore(BaseMemoryStore): if self.conn is not None: return - Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) - - self.conn = sqlite3.connect(self.db_path, check_same_thread=False) + self.conn = sqlite3.connect(self.db_path / "reme.db", check_same_thread=False) # Only load sqlite-vec extension if vector search is enabled if self.vector_enabled: diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py index 84c4e056..4b438ce3 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -83,7 +83,6 @@ class MemoryStoreConfig(BaseModel): model_config = ConfigDict(extra="allow") backend: str = Field(default="sqlite") - db_name: str = Field(default="reme.db") store_name: str = Field(default="reme") embedding_model: str = Field(default="default") @@ -114,7 +113,7 @@ class ServiceConfig(BaseModel): backend: str = Field(default="") app_name: str = Field(default=os.getenv("APP_NAME", "ReMe")) - working_dir: str | None = Field(default=None) + working_dir: str = Field(default=".reme") enable_logo: bool = Field(default=True) language: str = Field(default="") thread_pool_max_workers: int = Field(default=16) @@ -122,8 +121,8 @@ class ServiceConfig(BaseModel): log_to_console: bool = Field(default=True) disabled_flows: list[str] = Field(default_factory=list) enabled_flows: list[str] = Field(default_factory=list) - mcp_servers: dict[str, dict] = Field(default_factory=dict) + mcp_servers: dict[str, dict] = Field(default_factory=dict) mcp: MCPConfig = Field(default_factory=MCPConfig) http: HttpConfig = Field(default_factory=HttpConfig) cmd: CmdConfig = Field(default_factory=CmdConfig) diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index f8be9d4f..b43784c8 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -7,6 +7,7 @@ from .chunking_utils import chunk_markdown from .common_utils import run_coro_safely, execute_stream_task, hash_text, cosine_similarity, batch_cosine_similarity from .env_utils import load_env from .execute_utils import exec_code, run_shell_command, async_exec_code +from .horse import play_horse_easter_egg from .http_client import HttpClient from .llm_utils import extract_content, format_messages, deduplicate_memories from .logger_utils import init_logger @@ -32,6 +33,7 @@ __all__ = [ "exec_code", "async_exec_code", "run_shell_command", + "play_horse_easter_egg", "HttpClient", "extract_content", "format_messages", diff --git a/reme/core/utils/execute_utils.py b/reme/core/utils/execute_utils.py index 8ea00998..d5333e2f 100644 --- a/reme/core/utils/execute_utils.py +++ b/reme/core/utils/execute_utils.py @@ -5,8 +5,8 @@ with support for async execution and output capture. """ import asyncio -import contextlib import concurrent.futures +import contextlib from io import StringIO diff --git a/reme/horse.py b/reme/core/utils/horse.py similarity index 99% rename from reme/horse.py rename to reme/core/utils/horse.py index a11ac43c..3231c8c7 100644 --- a/reme/horse.py +++ b/reme/core/utils/horse.py @@ -20,7 +20,7 @@ def _mirror_frame(frame: str) -> str: return "\n".join(mirrored) -def _play_horse_easter_egg() -> None: +def play_horse_easter_egg() -> None: """Play the /horse Easter egg: fireworks, galloping horse, and a blessing.""" cols = shutil.get_terminal_size((80, 24)).columns rows = shutil.get_terminal_size((80, 24)).lines diff --git a/reme/reme_cli.py b/reme/reme_cli.py index ee6021f0..c0362158 100644 --- a/reme/reme_cli.py +++ b/reme/reme_cli.py @@ -11,7 +11,7 @@ from reme.core.op import BaseTool from .agent.chat import FsCli from .core.enumeration import ChunkEnum from .core.schema import StreamChunk -from .core.utils import execute_stream_task +from .core.utils import execute_stream_task, play_horse_easter_egg from .reme_fs import ReMeFs from .tool.fs import ( BashTool, @@ -23,7 +23,6 @@ from .tool.fs import ( ) from .tool.gallery import ExecuteCode from .tool.search import DashscopeSearch, TavilySearch -from .horse import _play_horse_easter_egg class ReMeCli(ReMeFs): @@ -142,7 +141,7 @@ class ReMeCli(ReMeFs): continue if user_input == "/horse": - _play_horse_easter_egg() + play_horse_easter_egg() continue # Stream processing state diff --git a/reme/reme_fs.py b/reme/reme_fs.py index 7a90adaa..5e83d14f 100644 --- a/reme/reme_fs.py +++ b/reme/reme_fs.py @@ -101,7 +101,7 @@ class ReMeFs(Application): previous_summary: str = "", language: str = "zh", **kwargs, - ) -> str: + ) -> str | dict: """Compact messages into a summary.""" compactor = FsCompactor(language=language, **kwargs) return await compactor.call( @@ -118,7 +118,7 @@ class ReMeFs(Application): version: str = "default", language: str = "zh", **kwargs, - ): + ) -> str | dict: """Generate a summary of the given messages.""" summarizer = FsSummarizer( tools=[ diff --git a/reme/workflow/procedural_memory/summarizer/__init__.py b/reme/workflow/procedural_memory/summarizer/__init__.py index 34644cd0..b83a7e2a 100644 --- a/reme/workflow/procedural_memory/summarizer/__init__.py +++ b/reme/workflow/procedural_memory/summarizer/__init__.py @@ -4,9 +4,9 @@ This package exposes and registers summarization-related operators such as `TrajectoryPreprocess` and `SuccessExtraction` to the global operator registry. """ -from ....core import R -from .trajectory_preprocess import TrajectoryPreprocess from .success_extraction import SuccessExtraction +from .trajectory_preprocess import TrajectoryPreprocess +from ....core import R __all__ = ["TrajectoryPreprocess", "SuccessExtraction"] diff --git a/tests/test_horse.py b/tests/test_horse.py index 86f6c989..5cd8a50e 100644 --- a/tests/test_horse.py +++ b/tests/test_horse.py @@ -3,7 +3,7 @@ import os import time -from reme.horse import _mirror_frame +from reme.core.utils.horse import _mirror_frame def clear_screen():