diff --git a/reme2/application.py b/reme2/application.py index c9f52e39..89c4f667 100644 --- a/reme2/application.py +++ b/reme2/application.py @@ -76,7 +76,6 @@ class Application(BaseComponent): f"of type '{ComponentEnum.JOB}'", ) params = job_config.model_dump() - params.setdefault("name", job_config.name) params["app_context"] = self.context self.context.jobs[job_config.name] = job_cls(**params) diff --git a/reme2/component/__init__.py b/reme2/component/__init__.py index 1e6c6bfb..b7dd0e95 100644 --- a/reme2/component/__init__.py +++ b/reme2/component/__init__.py @@ -1,12 +1,5 @@ """Components""" -from .application_context import ApplicationContext -from .base_component import BaseComponent -from .base_step import BaseStep -from .component_registry import ComponentRegistry, R -from .prompt_handler import PromptHandler -from .runtime_context import RuntimeContext - from . import as_llm from . import as_llm_formatter from . import client @@ -16,6 +9,12 @@ from . import file_store from . import file_watcher from . import job from . import service +from .application_context import ApplicationContext +from .base_component import BaseComponent +from .base_step import BaseStep +from .component_registry import ComponentRegistry, R +from .prompt_handler import PromptHandler +from .runtime_context import RuntimeContext __all__ = [ "ApplicationContext", diff --git a/reme2/component/as_llm_formatter/__init__.py b/reme2/component/as_llm_formatter/__init__.py index 3e8eb5b0..39b00334 100644 --- a/reme2/component/as_llm_formatter/__init__.py +++ b/reme2/component/as_llm_formatter/__init__.py @@ -18,6 +18,7 @@ class BaseAsLLMFormatter(BaseComponent): self.formatter: FormatterBase | None = None async def _start(self) -> None: + """Initialize the formatter.""" async def _close(self) -> None: self.formatter = None @@ -28,6 +29,7 @@ class AsOpenAIChatFormatter(BaseAsLLMFormatter): """Wrapper for OpenAI chat completion formatter.""" async def _start(self) -> None: + """Initialize the OpenAI chat formatter.""" self.formatter = ReMeOpenAIChatFormatter(**self.kwargs) diff --git a/reme2/component/as_llm_formatter/reme_openai_chat_formatter.py b/reme2/component/as_llm_formatter/reme_openai_chat_formatter.py index a501da26..66dbfd51 100644 --- a/reme2/component/as_llm_formatter/reme_openai_chat_formatter.py +++ b/reme2/component/as_llm_formatter/reme_openai_chat_formatter.py @@ -35,8 +35,8 @@ class ReMeOpenAIChatFormatter(OpenAIChatFormatter): """Extends OpenAIChatFormatter with tool result image promotion and reasoning content support.""" async def _format( - self, - msgs: list[Msg], + self, + msgs: list[Msg], ) -> list[dict[str, Any]]: """Format messages into OpenAI API format. @@ -108,7 +108,7 @@ class ReMeOpenAIChatFormatter(OpenAIChatFormatter): TextBlock( type="text", text="The following are the image contents from the tool " - f"result of '{block['name']}':", + f"result of '{block['name']}':", ), *promoted_blocks, TextBlock(type="text", text=""), diff --git a/reme2/component/as_token_counter/__init__.py b/reme2/component/as_token_counter/__init__.py index 8c102add..1b75c2ba 100644 --- a/reme2/component/as_token_counter/__init__.py +++ b/reme2/component/as_token_counter/__init__.py @@ -22,6 +22,7 @@ class BaseAsTokenCounter(BaseComponent): self.token_counter: TokenCounterBase | None = None async def _start(self) -> None: + """Initialize the token counter.""" async def _close(self) -> None: """Release token counter resources.""" diff --git a/reme2/component/as_token_counter/estimate_token_counter.py b/reme2/component/as_token_counter/estimate_token_counter.py index eca3f446..01f106c7 100644 --- a/reme2/component/as_token_counter/estimate_token_counter.py +++ b/reme2/component/as_token_counter/estimate_token_counter.py @@ -12,9 +12,9 @@ class EstimatedTokenCounter(TokenCounterBase): """ def __init__( - self, - estimate_divisor: float = 4, - encoding: str = "utf-8", + self, + estimate_divisor: float = 4, + encoding: str = "utf-8", ): """Initialize the estimated token counter. diff --git a/reme2/component/base_component.py b/reme2/component/base_component.py index 0900f8da..bea33830 100644 --- a/reme2/component/base_component.py +++ b/reme2/component/base_component.py @@ -1,10 +1,11 @@ """Base class for components.""" +import asyncio from abc import ABC, abstractmethod from typing import TYPE_CHECKING from ..enumeration import ComponentEnum -from ..utils.logger_utils import get_logger +from ..utils import get_logger if TYPE_CHECKING: from .application_context import ApplicationContext @@ -19,41 +20,45 @@ class BaseComponent(ABC): component_type = ComponentEnum.BASE def __init__( - self, - name: str | None = None, - backend: str | None = None, - app_context: "ApplicationContext | None" = None, - **kwargs, + self, + name: str | None = None, + backend: str = "", + app_context: "ApplicationContext | None" = None, + **kwargs, ) -> None: self.name: str = name or self.__class__.__name__ - self.backend: str | None = backend + self.backend: str = backend self.app_context: "ApplicationContext | None" = app_context self.kwargs: dict = dict(kwargs) self.logger = get_logger() if hasattr(self.logger, "bind"): self.logger = self.logger.bind(component=self.name) + self._is_started: bool = False + self._lock: asyncio.Lock = asyncio.Lock() @abstractmethod - async def _start(self) -> None: ... + async def _start(self) -> None: + """Start the component.""" @abstractmethod - async def _close(self) -> None: ... + async def _close(self) -> None: + """Close the component.""" async def start(self) -> None: """Start the component. No-op if already started.""" - if self._is_started: - return - await self._start() - self._is_started = True + async with self._lock: + if self._is_started: + return + await self._start() + self._is_started = True async def close(self) -> None: """Close the component. No-op if not started.""" - if not self._is_started: - return - try: + async with self._lock: + if not self._is_started: + return await self._close() - finally: self._is_started = False async def restart(self) -> None: @@ -65,26 +70,12 @@ class BaseComponent(ABC): def is_started(self) -> bool: return self._is_started - async def __call__(self, **kwargs): ... + async def __call__(self, **kwargs): + raise NotImplementedError async def __aenter__(self) -> "BaseComponent": await self.start() return self - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb, - ) -> bool: - if self._is_started: - if exc_val is not None: - try: - await self._close() - except BaseException as close_exc: - raise close_exc from exc_val - finally: - self._is_started = False - else: - await self.close() - return False + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + await self.close() diff --git a/reme2/component/base_step.py b/reme2/component/base_step.py index e4ec87ac..3d7bb173 100644 --- a/reme2/component/base_step.py +++ b/reme2/component/base_step.py @@ -7,15 +7,12 @@ from agentscope.formatter import FormatterBase from agentscope.model import ChatModelBase from agentscope.token import TokenCounterBase -from .application_context import ApplicationContext from .base_component import BaseComponent from .embedding import BaseEmbeddingModel from .file_store import BaseFileStore from .prompt_handler import PromptHandler from .runtime_context import RuntimeContext from ..enumeration import ComponentEnum -from ..schema import ApplicationConfig -from ..utils import camel_to_snake class BaseStep(BaseComponent): @@ -24,24 +21,20 @@ class BaseStep(BaseComponent): component_type = ComponentEnum.STEP def __new__(cls, *args, **kwargs): - """Capture init args for object cloning.""" instance = super().__new__(cls) instance._init_args = copy.copy(args) instance._init_kwargs = copy.copy(kwargs) return instance def __init__( - self, - name: str = "", - language: str = "", - prompt_dict: dict[str, str] | None = None, - input_mapping: dict[str, str] | None = None, - output_mapping: dict[str, str] | None = None, - **kwargs, + self, + language: str = "", + prompt_dict: dict[str, str] | None = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + **kwargs, ): - """Initialize step configurations.""" super().__init__(**kwargs) - self.name = name or camel_to_snake(self.__class__.__name__) self.language = language self.prompt = PromptHandler(language=self.language) self.prompt.load_prompt_by_class(self.__class__).load_prompt_dict(prompt_dict) @@ -49,115 +42,62 @@ class BaseStep(BaseComponent): self.output_mapping = output_mapping self.context: RuntimeContext | None = None - async def _start(self) -> None: - """Apply input mapping before execution.""" - if self.input_mapping and self.context: - self.context.apply_mapping(self.input_mapping) - - async def _close(self) -> None: - """Apply output mapping after execution.""" - if self.output_mapping and self.context: - self.context.apply_mapping(self.output_mapping) - @abstractmethod async def execute(self): """Execute the step logic.""" async def __call__(self, context: RuntimeContext | None = None, **kwargs): - """Execute the step with lifecycle management.""" self.context = RuntimeContext.from_context(context, **kwargs) - await self.start() - try: - response = await self.execute() - return response - finally: - await self.close() + assert self.context is not None - @property - def application_context(self) -> ApplicationContext: - """Get the application context from runtime context.""" - assert self.context is not None, "Runtime context not set." - return self.context.application_context + if self.input_mapping: + self.context.apply_mapping(self.input_mapping) - @property - def app_config(self) -> ApplicationConfig: - """Get the application configuration.""" - return self.application_context.app_config + result = await self.execute() + + if self.output_mapping: + self.context.apply_mapping(self.output_mapping) + + return result + + def _get_component(self, key: ComponentEnum, name: str, attr: str | None = None): + assert self.app_context is not None + comp = self.app_context.components[key][name] + return getattr(comp, attr) if attr else comp @property def as_llm(self) -> ChatModelBase: - """Get the AsLLM instance by name.""" - name_or_instance = self.kwargs.get("as_llm", "default") - if isinstance(name_or_instance, ChatModelBase): - return name_or_instance - - name = name_or_instance - as_llm_dict = self.application_context.components[ComponentEnum.AS_LLM] - if name not in as_llm_dict: - raise ValueError(f"AsLLM '{name}' not found.") - wrapper = as_llm_dict[name] - return wrapper.model + name = self.kwargs.get("as_llm", "default") + return name if isinstance(name, ChatModelBase) else self._get_component(ComponentEnum.AS_LLM, name, "model") @property def as_llm_formatter(self) -> FormatterBase: - """Get the AsLLMFormatter instance by name.""" - name_or_instance = self.kwargs.get("as_llm_formatter", "default") - if isinstance(name_or_instance, FormatterBase): - return name_or_instance - - name = name_or_instance - formatter_dict = self.application_context.components[ComponentEnum.AS_LLM_FORMATTER] - if name not in formatter_dict: - raise ValueError(f"AsLLMFormatter '{name}' not found.") - wrapper = formatter_dict[name] - return wrapper.formatter + name = self.kwargs.get("as_llm_formatter", "default") + return name if isinstance(name, FormatterBase) else self._get_component(ComponentEnum.AS_LLM_FORMATTER, name, + "formatter") @property def as_token_counter(self) -> TokenCounterBase: - """Get the TokenCounter instance by name.""" - name_or_instance = self.kwargs.get("as_token_counter", "default") - if isinstance(name_or_instance, TokenCounterBase): - return name_or_instance - - name = name_or_instance - counter_dict = self.application_context.components[ComponentEnum.AS_TOKEN_COUNTER] - if name not in counter_dict: - raise ValueError(f"AsTokenCounter '{name}' not found.") - wrapper = counter_dict[name] - return wrapper.token_counter + name = self.kwargs.get("as_token_counter", "default") + return name if isinstance(name, TokenCounterBase) else self._get_component(ComponentEnum.AS_TOKEN_COUNTER, name, + "token_counter") @property def file_store(self) -> BaseFileStore: - """Get the FileStore instance by name.""" - name: str = self.kwargs.get("file_store", "default") - stores = self.application_context.components[ComponentEnum.FILE_STORE] - if name not in stores: - raise ValueError(f"FileStore {name} not found.") - store = stores[name] - if not isinstance(store, BaseFileStore): - raise TypeError(f"{name} is not a BaseFileStore instance.") - return store + name = self.kwargs.get("file_store", "default") + return name if isinstance(name, BaseFileStore) else self._get_component(ComponentEnum.FILE_STORE, name) @property def embedding(self) -> BaseEmbeddingModel: - """Get the EmbeddingModel instance by name.""" - name: str = self.kwargs.get("embedding", "default") - models = self.application_context.components[ComponentEnum.EMBEDDING_MODEL] - if name not in models: - raise ValueError(f"EmbeddingModel {name} not found.") - model = models[name] - if not isinstance(model, BaseEmbeddingModel): - raise TypeError(f"{name} is not a BaseEmbeddingModel instance.") - return model + name = self.kwargs.get("embedding", "default") + return name if isinstance(name, BaseEmbeddingModel) else self._get_component(ComponentEnum.EMBEDDING_MODEL, + name) def prompt_format(self, prompt_name: str, **kwargs) -> str: - """Format a prompt template.""" return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs) def get_prompt(self, prompt_name: str) -> str: - """Get a prompt template by name.""" return self.prompt.get_prompt(prompt_name=prompt_name) def copy(self, **kwargs) -> "BaseStep": - """Create a copy with optional parameter overrides.""" return self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs}) diff --git a/reme2/component/client/base_client.py b/reme2/component/client/base_client.py index d92b7d50..b37e550c 100644 --- a/reme2/component/client/base_client.py +++ b/reme2/component/client/base_client.py @@ -16,6 +16,7 @@ class BaseClient(BaseComponent): self.client = None async def _start(self) -> None: + """Initialize the client.""" async def _close(self) -> None: """Close the client.""" diff --git a/reme2/component/client/http_client.py b/reme2/component/client/http_client.py index d01b81a0..d200877b 100644 --- a/reme2/component/client/http_client.py +++ b/reme2/component/client/http_client.py @@ -15,12 +15,12 @@ class HttpClient(BaseClient): """HTTP client for ReMe service.""" def __init__( - self, - action: str, - host: str | None = None, - port: int | None = None, - timeout: float = 30.0, - **kwargs, + self, + action: str, + host: str | None = None, + port: int | None = None, + timeout: float = 30.0, + **kwargs, ): super().__init__(**kwargs) diff --git a/reme2/component/component_registry.py b/reme2/component/component_registry.py index e5d5c52a..16ead1e0 100644 --- a/reme2/component/component_registry.py +++ b/reme2/component/component_registry.py @@ -25,35 +25,34 @@ class ComponentRegistry: def _do_register(self, cls: type[T], name: str) -> type[T]: """Register a component class with the given name.""" - if not hasattr(cls, "component_type"): - raise TypeError(f"{cls.__name__} must have 'component_type' attribute") + component_type = getattr(cls, "component_type", None) + if not isinstance(component_type, ComponentEnum): + raise TypeError(f"{cls.__name__} must have a ComponentEnum 'component_type' attribute") if not name: raise ValueError("Component name cannot be empty") - component_type = cls.component_type - if component_type not in self._registry: - self._registry[component_type] = {} - if name in self._registry[component_type]: + group = self._registry.setdefault(component_type, {}) + if name in group: self.logger.warning(f"Component '{name}' already registered for {component_type}, overwriting") - - self._registry[component_type][name] = cls + group[name] = cls return cls def register( - self, - cls_or_name: type[T] | str, - name: str | None = None, + self, + cls_or_name: type[T] | str, + name: str | None = None, ) -> Callable[[type[T]], type[T]] | type[T]: """Register a component class. Supports direct and decorator modes.""" # Direct registration: R.register(MyClass, "name") if isinstance(cls_or_name, type): - return self._do_register(cast(type[T], cls_or_name), name or cls_or_name.__name__) + return self._do_register(cast(type[T], cls_or_name), name if name is not None else cls_or_name.__name__) # Decorator mode: @R.register("name") - decorator_name = cls_or_name + if not isinstance(cls_or_name, str): + raise TypeError(f"Expected a class or string, got {type(cls_or_name).__name__}") def decorator(decorated_cls: type[T]) -> type[T]: - return self._do_register(decorated_cls, decorator_name or decorated_cls.__name__) + return self._do_register(decorated_cls, cls_or_name) return decorator @@ -67,8 +66,8 @@ class ComponentRegistry: def unregister(self, component_type: ComponentEnum, name: str) -> bool: """Remove a component from the registry. Returns True if found.""" - if name in self._registry.get(component_type, {}): - del self._registry[component_type][name] + if (group := self._registry.get(component_type)) and name in group: + del group[name] return True return False diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index bbc4ed3e..35f4ce97 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/reme2/component/embedding/base_embedding_model.py @@ -3,12 +3,12 @@ import hashlib import os import time - -import numpy as np from abc import abstractmethod from collections import OrderedDict from pathlib import Path +import numpy as np + from ..base_component import BaseComponent from ...enumeration import ComponentEnum from ...schema import BaseNode @@ -34,27 +34,10 @@ class BaseEmbeddingModel(BaseComponent): pass_dimensions: bool = False, max_batch_size: int = 10, max_input_length: int = 8192, - cache_dir: str | Path = ".reme", max_cache_size: int = 2000, enable_cache: bool = True, - encoding: str = "utf-8", **kwargs, ): - """Initialize embedding model configuration. - - Args: - api_key: API key for the embedding service. - base_url: Base URL for the embedding service. - model_name: Name of the embedding model. - dimensions: Vector dimensions. - pass_dimensions: Whether to pass dimensions parameter to API. - max_batch_size: Maximum batch size for embedding requests. - max_input_length: Maximum input text length. - cache_dir: Directory for cache storage. - max_cache_size: Maximum LRU cache size. - enable_cache: Whether to enable caching. - encoding: Text encoding for cache file operations. - """ super().__init__(**kwargs) self.api_key: str = api_key or os.environ.get("EMBEDDING_API_KEY", "") self.base_url: str = base_url or os.environ.get("EMBEDDING_BASE_URL", "") @@ -63,122 +46,96 @@ class BaseEmbeddingModel(BaseComponent): self.pass_dimensions = pass_dimensions self.max_batch_size = max_batch_size 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.encoding = encoding self._embedding_cache: OrderedDict[str, list[float]] = OrderedDict() self._cache_hits = 0 self._cache_misses = 0 - self.cache_path: Path = Path(self.cache_dir) + self.cache_path: Path = Path() + + def clear_cache(self) -> None: + """Clear in-memory cache and reset statistics.""" + self._embedding_cache.clear() + self._cache_hits = 0 + self._cache_misses = 0 async def _start(self) -> None: """Load cache on start.""" assert self.app_context is not None, "app_context must be provided" + self.clear_cache() working_path = Path(self.app_context.app_config.working_dir) - working_path.embedding_cache = self.cache_path + self.cache_path = working_path / "embedding_cache" / f"{self.name}.npz" self._load_cache() async def _close(self) -> None: """Save cache on close.""" self._save_cache() - def _truncate_text(self, text: str) -> str: - """Truncate text to max_input_length.""" - return text[: self.max_input_length] if len(text) > self.max_input_length else text - def _validate_and_adjust_embedding(self, embedding: list[float]) -> list[float]: """Adjust embedding dimensions to match expected dimensions.""" actual_len = len(embedding) if actual_len == self.dimensions: return embedding - if actual_len < self.dimensions: - self.logger.warning( - f"[ACTUAL_EMB_LENGTH] Embedding {actual_len} < expected {self.dimensions}, padding with zeros", - ) + self.logger.warning(f"Embedding dim {actual_len} < expected {self.dimensions}, padding") return embedding + [0.0] * (self.dimensions - actual_len) - - self.logger.warning( - f"[ACTUAL_EMB_LENGTH] Embedding {actual_len} > expected {self.dimensions}, truncating", - ) - return embedding[: self.dimensions] + self.logger.warning(f"Embedding dim {actual_len} > expected {self.dimensions}, truncating") + return embedding[:self.dimensions] def _get_cache_key(self, text: str) -> str: """Generate cache key from text + model_name + dimensions.""" - cache_string = f"{text}|{self.model_name}|{self.dimensions}" - return hashlib.sha256(cache_string.encode(self.encoding)).hexdigest() - - def _get_cache_file_path(self) -> Path: - """Return path to the cache npz file.""" - return self.cache_path / "embedding_cache.npz" + return hashlib.sha256(f"{text}|{self.model_name}|{self.dimensions}".encode()).hexdigest() def _load_cache(self) -> None: """Load embedding cache from disk (npz format).""" if not self.enable_cache: return - self.cache_path.mkdir(parents=True, exist_ok=True) - cache_file = self._get_cache_file_path() - if not cache_file.exists(): - self.logger.info(f"No cache file at {cache_file}, starting empty") + self.cache_path.parent.mkdir(parents=True, exist_ok=True) + if not self.cache_path.exists(): + self.logger.info(f"No cache file at {self.cache_path}, starting empty") return + load_start = time.time() try: - load_start = time.time() - data = np.load(cache_file) - keys = data["keys"] - embeddings = data["embeddings"] + data = np.load(self.cache_path) + except Exception: + self.logger.exception(f"Failed to load cache from {self.cache_path}, deleting file") + self.cache_path.unlink(missing_ok=True) + return - loaded_count = 0 - for key, emb in zip(keys, embeddings): - key_str = str(key) - emb_list = emb.tolist() - if len(emb_list) != self.dimensions: - self.logger.warning( - f"Cache dimension mismatch for {key_str}: " - f"expected {self.dimensions}, got {len(emb_list)}", - ) - continue - if len(self._embedding_cache) >= self.max_cache_size: - self.logger.info(f"Cache limit reached ({self.max_cache_size}), loaded {loaded_count}") - break - self._embedding_cache[key_str] = emb_list - loaded_count += 1 + loaded_count = 0 + for key, emb in zip(data["keys"], data["embeddings"]): + emb_list = emb.tolist() + if len(emb_list) != self.dimensions: + self.logger.warning(f"Cache dimension mismatch for {key}: expected {self.dimensions}, got {len(emb_list)}") + continue + if len(self._embedding_cache) >= self.max_cache_size: + self.logger.info(f"Cache limit reached ({self.max_cache_size}), loaded {loaded_count}") + break + self._embedding_cache[str(key)] = emb_list + loaded_count += 1 - self.logger.info(f"Loaded {loaded_count} embeddings from {cache_file} in {time.time() - load_start:.2f}s") - except Exception as e: - self.logger.error(f"Failed to load cache from {cache_file}: {e}, deleting file") - try: - cache_file.unlink() - except Exception as del_e: - self.logger.error(f"Failed to delete cache file: {del_e}") + self.logger.info(f"Loaded {loaded_count} embeddings from {self.cache_path} in {time.time() - load_start:.2f}s") def _save_cache(self) -> None: """Save embedding cache to disk (npz format).""" if not self.enable_cache or not self._embedding_cache: return - cache_file = self._get_cache_file_path() - try: - keys = [] - embeddings = [] - for cache_key, embedding in self._embedding_cache.items(): - if len(embedding) != self.dimensions: - self.logger.warning(f"Cache dimension mismatch for {cache_key}") - continue - keys.append(cache_key) - embeddings.append(embedding) + keys, embeddings = [], [] + for cache_key, embedding in self._embedding_cache.items(): + keys.append(cache_key) + embeddings.append(embedding) - np.savez( - cache_file, - keys=np.array(keys, dtype=str), - embeddings=np.array(embeddings, dtype=np.float32), - ) - self.logger.info(f"Saved {len(keys)} embeddings to {cache_file}") + try: + np.savez(self.cache_path, keys=np.array(keys, dtype=str), embeddings=np.array(embeddings, dtype=np.float32)) except Exception as e: - self.logger.error(f"Failed to save cache to {cache_file}: {e}") + self.logger.error(f"Failed to save cache to {self.cache_path}: {e}") + return + + self.logger.info(f"Saved {len(keys)} embeddings to {self.cache_path}") def _get_from_cache(self, text: str) -> list[float] | None: """Retrieve embedding from cache if available.""" @@ -190,29 +147,18 @@ class BaseEmbeddingModel(BaseComponent): self._cache_misses += 1 return None - embeddings = self._embedding_cache[cache_key] - if len(embeddings) != self.dimensions: - self.logger.warning("Cached embedding dimension mismatch, removing entry") - del self._embedding_cache[cache_key] - self._cache_misses += 1 - return None - self._embedding_cache.move_to_end(cache_key) self._cache_hits += 1 - preview = text[:50] + "..." if len(text) > 50 else text - self.logger.info(f"Cache hit: {preview} (hits: {self._cache_hits}, misses: {self._cache_misses})") - return embeddings + return self._embedding_cache[cache_key] def _put_to_cache(self, text: str, embedding: list[float]) -> None: """Store embedding in cache with LRU eviction.""" if not self.enable_cache or self.max_cache_size <= 0: return - - cache_key = self._get_cache_key(text) if len(embedding) != self.dimensions: - self.logger.warning(f"[PUT_TO_CACHE] Dimension mismatch for {cache_key}") return + cache_key = self._get_cache_key(text) if len(self._embedding_cache) >= self.max_cache_size and cache_key not in self._embedding_cache: self._embedding_cache.popitem(last=False) @@ -222,48 +168,26 @@ class BaseEmbeddingModel(BaseComponent): def get_cache_stats(self) -> dict[str, int | float]: """Return cache statistics: size, hits, misses, hit_rate.""" total = self._cache_hits + self._cache_misses - hit_rate = self._cache_hits / total if total > 0 else 0.0 return { "cache_size": len(self._embedding_cache), "max_cache_size": self.max_cache_size, "cache_hits": self._cache_hits, "cache_misses": self._cache_misses, - "hit_rate": hit_rate, + "hit_rate": self._cache_hits / total if total > 0 else 0.0, } - def clear_cache(self) -> None: - """Clear in-memory cache and reset statistics.""" - self._embedding_cache.clear() - self._cache_hits = 0 - self._cache_misses = 0 - @abstractmethod - async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float]]: + async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: """Fetch embeddings for a batch of texts. Override in subclasses.""" - async def get_embedding(self, input_text: str, **kwargs) -> list[float]: - """Get embedding for a single text with cache.""" - truncated_text = self._truncate_text(input_text) - cached = self._get_from_cache(truncated_text) - if cached is not None: - return cached + async def get_embedding(self, input_text: str, **kwargs) -> list[float] | None: + """Get embedding for a single text with cache. Returns None on failure.""" + results = await self.get_embeddings([input_text], **kwargs) + return results[0] if results else None - try: - result = await self._get_embeddings([truncated_text], **kwargs) - if result and len(result) == 1: - embedding = self._validate_and_adjust_embedding(result[0]) - self._put_to_cache(truncated_text, embedding) - return embedding - self.logger.warning( - f"Model {self.model_name} returned {len(result) if result else 0} results, expected 1", - ) - except Exception as e: - self.logger.error(f"Model {self.model_name} failed: {e}") - return [] - - async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float]]: + async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: """Get embeddings for multiple texts with cache and batching.""" - truncated_texts = [self._truncate_text(t) for t in input_text] + truncated_texts = [t[:self.max_input_length] for t in input_text] results: list[list[float] | None] = [None] * len(truncated_texts) texts_to_compute: list[tuple[int, str]] = [] @@ -277,8 +201,9 @@ class BaseEmbeddingModel(BaseComponent): if texts_to_compute: uncached_texts = [text for _, text in texts_to_compute] for i in range(0, len(uncached_texts), self.max_batch_size): - batch_texts = uncached_texts[i: i + self.max_batch_size] - batch_indices = [idx for idx, _ in texts_to_compute[i: i + self.max_batch_size]] + batch = texts_to_compute[i:i + self.max_batch_size] + batch_indices = [idx for idx, _ in batch] + batch_texts = [text for _, text in batch] try: batch_embeddings = await self._get_embeddings(batch_texts, **kwargs) @@ -292,27 +217,21 @@ class BaseEmbeddingModel(BaseComponent): f"Batch returned {len(batch_embeddings) if batch_embeddings else 0} " f"results for {len(batch_texts)} inputs", ) - for orig_idx in batch_indices: - if results[orig_idx] is None: - results[orig_idx] = [] except Exception as e: self.logger.error(f"Model {self.model_name} batch failed: {e}") - for orig_idx in batch_indices: - if results[orig_idx] is None: - results[orig_idx] = [] - return [r if r is not None else [] for r in results] + return results async def get_node_embeddings(self, nodes: list[BaseNode], **kwargs) -> list[BaseNode]: """Get embeddings for a list of nodes and assign to node.embedding.""" texts = [node.text for node in nodes] embeddings = await self.get_embeddings(texts, **kwargs) - if len(embeddings) == len(nodes): for node, vec in zip(nodes, embeddings): - node.embedding = vec + if vec is not None: + node.embedding = vec + else: + self.logger.warning(f"Embedding failed for node, skipping assignment") else: self.logger.warning(f"Mismatch: {len(embeddings)} vectors for {len(nodes)} nodes, skipping assignment") return nodes - - diff --git a/reme2/component/embedding/openai_embedding_model.py b/reme2/component/embedding/openai_embedding_model.py index e2405b41..6de89b33 100644 --- a/reme2/component/embedding/openai_embedding_model.py +++ b/reme2/component/embedding/openai_embedding_model.py @@ -8,49 +8,40 @@ from ..component_registry import R @R.register("openai") class OpenAIEmbeddingModel(BaseEmbeddingModel): - """Async embedding model compatible with OpenAI-style APIs.""" def __init__(self, **kwargs): - """Initialize OpenAI embedding model.""" super().__init__(**kwargs) self._client: AsyncOpenAI | None = None async def _start(self) -> None: - """Initialize the AsyncOpenAI client.""" self._client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, **self.kwargs) await super()._start() async def _close(self) -> None: - """Close the AsyncOpenAI client.""" - if self._client is not None: + if self._client: await self._client.close() self._client = None await super()._close() - async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float]]: - """Fetch embeddings for a batch of texts.""" + async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: if self._client is None: raise RuntimeError("Client not initialized. Call _start() first.") - create_kwargs: dict = { - "model": self.model_name, - "input": input_text, - **kwargs, - } + create_kwargs: dict = {"model": self.model_name, "input": input_text, **kwargs} if self.pass_dimensions: create_kwargs["dimensions"] = self.dimensions completion = await self._client.embeddings.create(**create_kwargs) - result_emb: list[list[float] | None] = [None] * len(input_text) + result: list[list[float] | None] = [None] * len(input_text) for emb in completion.data: vec = getattr(emb, "embedding", None) or getattr(emb, "dense_embedding", None) if 0 <= emb.index < len(input_text): if vec is not None: - result_emb[emb.index] = list(vec) + result[emb.index] = list(vec) else: self.logger.warning(f"Empty embedding for index {emb.index}") else: self.logger.warning(f"Invalid index {emb.index} for input length {len(input_text)}") - return [r if r is not None else [] for r in result_emb] + return result diff --git a/reme2/component/file_parser/default_file_parser.py b/reme2/component/file_parser/default_file_parser.py index 53e82629..04f7b151 100644 --- a/reme2/component/file_parser/default_file_parser.py +++ b/reme2/component/file_parser/default_file_parser.py @@ -51,13 +51,13 @@ class DefaultFileParser(BaseFileParser): chunks: list[FileChunk] = [] if content: chunks = ( - chunk_markdown( - content, - file_meta.path, - self.chunk_tokens, - self.chunk_overlap, - ) - or [] + chunk_markdown( + content, + file_meta.path, + self.chunk_tokens, + self.chunk_overlap, + ) + or [] ) file_meta.chunk_count = len(chunks) diff --git a/reme2/component/file_parser/md_file_parser.py b/reme2/component/file_parser/md_file_parser.py index 688a8c04..13878412 100644 --- a/reme2/component/file_parser/md_file_parser.py +++ b/reme2/component/file_parser/md_file_parser.py @@ -42,13 +42,13 @@ class MdFileParser(BaseFileParser): ) chunks = ( - chunk_markdown( - content, - file_meta.path, - self.chunk_tokens, - self.chunk_overlap, - ) - or [] + chunk_markdown( + content, + file_meta.path, + self.chunk_tokens, + self.chunk_overlap, + ) + or [] ) file_meta.chunk_count = len(chunks) diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index 9167426d..d1981fb1 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/file_store/base_file_store.py @@ -21,12 +21,12 @@ class BaseFileStore(BaseComponent): component_type = ComponentEnum.FILE_STORE def __init__( - self, - store_name: str, - db_path: str | Path, - embedding_model: str = "default", - fts_enabled: bool = True, - **kwargs, + self, + store_name: str, + db_path: str | Path, + embedding_model: str = "default", + fts_enabled: bool = True, + **kwargs, ): super().__init__(**kwargs) self._embedding_model_name: str = embedding_model @@ -134,12 +134,12 @@ class BaseFileStore(BaseComponent): # -- Hybrid search (concrete, delegates to abstract vector/keyword) ----- async def hybrid_search( - self, - query: str, - limit: int, - vector_weight: float = 0.7, - candidate_multiplier: float = 3.0, - search_filter: SearchFilter | None = None, + self, + query: str, + limit: int, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + search_filter: SearchFilter | None = None, ) -> list[FileChunk]: """Perform hybrid search combining vector and keyword results.""" assert 0.0 <= vector_weight <= 1.0 @@ -171,10 +171,10 @@ class BaseFileStore(BaseComponent): @staticmethod def _merge_hybrid_results( - vector: list[FileChunk], - keyword: list[FileChunk], - vector_weight: float, - text_weight: float, + vector: list[FileChunk], + keyword: list[FileChunk], + vector_weight: float, + text_weight: float, ) -> list[FileChunk]: """Merge vector and keyword results with weighted scoring.""" merged: dict[str, FileChunk] = {} @@ -182,10 +182,10 @@ class BaseFileStore(BaseComponent): for result in vector: v_score = result.scores.get("vector", 0) result.scores["score"] = v_score * vector_weight - merged[result.merge_key] = result + merged[result.unique_key] = result for result in keyword: - key = result.merge_key + key = result.unique_key k_score = result.scores.get("keyword", 0) if key in merged: merged[key].scores["score"] += k_score * text_weight @@ -200,10 +200,10 @@ class BaseFileStore(BaseComponent): # -- Filter utility ----------------------------------------------------- def _apply_filter( - self, - chunks: list[FileChunk], - search_filter: SearchFilter | None, - file_metadata: dict[str, FileMetadata] | None = None, + self, + chunks: list[FileChunk], + search_filter: SearchFilter | None, + file_metadata: dict[str, FileMetadata] | None = None, ) -> list[FileChunk]: """Apply search filter to a list of chunks. @@ -250,9 +250,9 @@ class BaseFileStore(BaseComponent): @abstractmethod async def keyword_search( - self, - query: str, - limit: int, - search_filter: SearchFilter | None = None, + self, + query: str, + limit: int, + search_filter: SearchFilter | None = None, ) -> list[FileChunk]: """Perform full-text/keyword search.""" diff --git a/reme2/component/file_store/chroma_file_store.py b/reme2/component/file_store/chroma_file_store.py index cb709be2..c4598617 100644 --- a/reme2/component/file_store/chroma_file_store.py +++ b/reme2/component/file_store/chroma_file_store.py @@ -201,10 +201,10 @@ class ChromaFileStore(BaseFileStore): return chunks[:limit] async def keyword_search( - self, - query: str, - limit: int, - search_filter: SearchFilter | None = None, + self, + query: str, + limit: int, + search_filter: SearchFilter | None = None, ) -> list[FileChunk]: """Keyword search via ChromaDB $contains with case variants.""" if not self.fts_enabled or not query: diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index fb0df8c6..004c270a 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/file_store/local_file_store.py @@ -202,10 +202,10 @@ class LocalFileStore(BaseFileStore): return results[:limit] async def keyword_search( - self, - query: str, - limit: int, - search_filter: SearchFilter | None = None, + self, + query: str, + limit: int, + search_filter: SearchFilter | None = None, ) -> list[FileChunk]: """Keyword search via substring matching.""" if not self.fts_enabled or not query: diff --git a/reme2/component/file_store/sqlite_file_store.py b/reme2/component/file_store/sqlite_file_store.py index 1e97a7d6..cd8b7773 100644 --- a/reme2/component/file_store/sqlite_file_store.py +++ b/reme2/component/file_store/sqlite_file_store.py @@ -1,11 +1,10 @@ """SQLite storage backend for file store.""" import json +import sqlite3 import struct import time -import sqlite3 - from .base_file_store import BaseFileStore from ..component_registry import R from ...schema import FileChunk, FileMetadata, SearchFilter @@ -426,10 +425,10 @@ class SqliteFileStore(BaseFileStore): cursor.close() async def keyword_search( - self, - query: str, - limit: int, - search_filter: SearchFilter | None = None, + self, + query: str, + limit: int, + search_filter: SearchFilter | None = None, ) -> list[FileChunk]: if not self.fts_enabled or not query: return [] diff --git a/reme2/component/file_watcher/base_file_watcher.py b/reme2/component/file_watcher/base_file_watcher.py index 512c65f4..9e2a8d48 100644 --- a/reme2/component/file_watcher/base_file_watcher.py +++ b/reme2/component/file_watcher/base_file_watcher.py @@ -25,17 +25,17 @@ class BaseFileWatcher(BaseComponent): component_type = ComponentEnum.FILE_WATCHER def __init__( - self, - watch_paths: list[str] | str, - recursive: bool = False, - debounce: int = 2000, - chunk_tokens: int = 400, - chunk_overlap: int = 80, - file_store: str = "default", - default_parser: str | None = None, - rebuild_index_on_start: bool = True, - poll_delay_ms: int = 2000, - **kwargs, + self, + watch_paths: list[str] | str, + recursive: bool = False, + debounce: int = 2000, + chunk_tokens: int = 400, + chunk_overlap: int = 80, + file_store: str = "default", + default_parser: str | None = None, + rebuild_index_on_start: bool = True, + poll_delay_ms: int = 2000, + **kwargs, ): super().__init__(**kwargs) self._file_store_name: str = file_store @@ -171,11 +171,11 @@ class BaseFileWatcher(BaseComponent): try: self.logger.info(f"Starting watch on: {valid_paths}") async for changes in awatch( - *valid_paths, - recursive=self.recursive, - debounce=self.debounce, - poll_delay_ms=self.poll_delay_ms, - stop_event=self._stop_event, + *valid_paths, + recursive=self.recursive, + debounce=self.debounce, + poll_delay_ms=self.poll_delay_ms, + stop_event=self._stop_event, ): if self._stop_event.is_set(): break diff --git a/reme2/component/job/base_job.py b/reme2/component/job/base_job.py index 88c41c58..62b7fc74 100644 --- a/reme2/component/job/base_job.py +++ b/reme2/component/job/base_job.py @@ -19,12 +19,12 @@ class BaseJob(BaseComponent): component_type = ComponentEnum.JOB def __init__( - self, - name: str = "", - description: str = "", - parameters: dict | None = None, - steps: list[ComponentConfig] | None = None, - **kwargs, + self, + name: str = "", + description: str = "", + parameters: dict | None = None, + steps: list[ComponentConfig] | None = None, + **kwargs, ): """Initialize the job. diff --git a/reme2/component/prompt_handler.py b/reme2/component/prompt_handler.py index 6ebd727c..891ce242 100644 --- a/reme2/component/prompt_handler.py +++ b/reme2/component/prompt_handler.py @@ -2,11 +2,14 @@ import inspect import json +import re from pathlib import Path from string import Formatter import yaml +_FLAG_PATTERN = re.compile(r"^\[(\w+)\]") + class PromptHandler: """A handler for loading, retrieving, and formatting prompt templates.""" @@ -18,9 +21,9 @@ class PromptHandler: self.language: str = language.strip() def load_prompt_by_file( - self, - prompt_file_path: str | Path | None = None, - overwrite: bool = True, + self, + prompt_file_path: str | Path | None = None, + overwrite: bool = True, ) -> "PromptHandler": """Load prompts from a YAML or JSON file.""" if prompt_file_path is None: @@ -32,7 +35,7 @@ class PromptHandler: try: with path.open(encoding="utf-8") as f: - prompt_dict = yaml.safe_load(f) if path.suffix in (".yaml", ".yml") else json.load(f) + prompt_dict = yaml.safe_load(f) if path.suffix.lower() in (".yaml", ".yml") else json.load(f) except (json.JSONDecodeError, yaml.YAMLError, OSError): return self @@ -53,7 +56,7 @@ class PromptHandler: def load_prompt_dict(self, prompt_dict: dict | None = None, overwrite: bool = True) -> "PromptHandler": """Merge prompts from a dictionary.""" - if not prompt_dict: + if not isinstance(prompt_dict, dict): return self for key, value in prompt_dict.items(): @@ -72,11 +75,12 @@ class PromptHandler: def has_prompt(self, prompt_name: str) -> bool: """Check if a prompt exists.""" - return prompt_name in self.data or f"{prompt_name}_{self.language}" in self.data + keys = (f"{prompt_name}_{self.language}", prompt_name) if self.language else (prompt_name,) + return any(k in self.data for k in keys) def list_prompts(self, language_filter: str | None = None) -> list[str]: """List all available prompt names.""" - if not language_filter: + if language_filter is None: return list(self.data.keys()) suffix = f"_{language_filter.strip()}" return [k for k in self.data if k.endswith(suffix)] @@ -90,16 +94,10 @@ class PromptHandler: if flags: lines = [] for line in prompt.split("\n"): - remaining = line - should_include = False - for flag, enabled in flags.items(): - prefix = f"[{flag}]" - while remaining.startswith(prefix): - remaining = remaining[len(prefix) :] - if enabled: - should_include = True - if should_include or not any(line.startswith(f"[{f}]") for f in flags): - lines.append(remaining) + active_flags = _FLAG_PATTERN.findall(line) + cleaned = _FLAG_PATTERN.sub("", line).lstrip() + if not active_flags or any(flags.get(f, False) for f in active_flags): + lines.append(cleaned) prompt = "\n".join(lines) if validate: @@ -107,7 +105,7 @@ class PromptHandler: if missing := required - set(formats.keys()): raise ValueError(f"Missing format variables for '{prompt_name}': {sorted(missing)}") - return prompt.format(**formats).strip() if formats else prompt.strip() + return prompt.format(**formats).strip() if formats else prompt def __repr__(self) -> str: return f"PromptHandler(language='{self.language}', num_prompts={len(self.data)})" diff --git a/reme2/component/runtime_context.py b/reme2/component/runtime_context.py index af455536..f0cd12b0 100644 --- a/reme2/component/runtime_context.py +++ b/reme2/component/runtime_context.py @@ -2,7 +2,6 @@ import asyncio -from .application_context import ApplicationContext from ..enumeration import ChunkEnum from ..schema import Response, StreamChunk @@ -11,132 +10,58 @@ class RuntimeContext: """Context for execution state, response metadata, and stream queues.""" def __init__(self, **kwargs): - """Initialize the context with all keyword arguments stored in data.""" self.data: dict = kwargs def get(self, key: str, default=None): - """Get a value from data by key, with optional default.""" return self.data.get(key, default) - def set(self, key: str, value) -> "RuntimeContext": - """Set a value in data by key.""" - self.data[key] = value - return self - - def delete(self, key: str) -> "RuntimeContext": - """Delete a key from data.""" - if key in self.data: - del self.data[key] - return self - - def contains(self, key: str) -> bool: - """Check if a key exists in data.""" - return key in self.data - def update(self, data: dict) -> "RuntimeContext": - """Update data with a dictionary.""" self.data.update(data) return self - def keys(self) -> list[str]: - """Get all keys in data.""" - return list(self.data.keys()) - - def values(self) -> list: - """Get all values in data.""" - return list(self.data.values()) - - def items(self) -> list[tuple]: - """Get all key-value pairs in data.""" - return list(self.data.items()) - def __getitem__(self, key: str): - """Get a value using bracket syntax.""" return self.data[key] def __setitem__(self, key: str, value): - """Set a value using bracket syntax.""" self.data[key] = value def __delitem__(self, key: str): - """Delete a key using bracket syntax.""" del self.data[key] def __contains__(self, key: str) -> bool: - """Check if a key exists using 'in' operator.""" return key in self.data @property def response(self) -> Response: - """Get or create the response object.""" return self.data.setdefault("response", Response()) @property def stream_queue(self) -> asyncio.Queue: - """Get the stream queue.""" return self.data["stream_queue"] - @property - def application_context(self) -> ApplicationContext: - """Get the application context.""" - return self.data["application_context"] - @classmethod def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext": - """Create a new context from an existing instance or keywords.""" if context is None: return cls(**kwargs) - context.data.update(kwargs) + context.update(kwargs) return context async def _enqueue(self, chunk: StreamChunk) -> None: - """Internal helper to put a chunk into the queue if it exists.""" if self.stream_queue: await self.stream_queue.put(chunk) async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext": - """Enqueue a stream chunk from a raw string and type.""" await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk)) return self - async def add_stream_chunk(self, stream_chunk: StreamChunk) -> "RuntimeContext": - """Enqueue an existing stream chunk.""" - await self._enqueue(stream_chunk) - return self - async def add_stream_done(self) -> "RuntimeContext": - """Enqueue a termination chunk to signal the end of the stream.""" await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True)) return self - def add_response_error(self, e: Exception) -> "RuntimeContext": - """Record an exception into the response object.""" - self.response.success = False - self.response.answer = str(e) - return self - def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext": - """Copy internal values based on a source-to-target key map.""" if not mapping: return self - for source, target in mapping.items(): if source in self.data: self.data[target] = self.data[source] return self - - def validate_required_keys( - self, - required_keys: dict[str, bool], - context_name: str = "context", - ) -> "RuntimeContext": - """Ensure all required keys are present in the context. - - Args: - required_keys: Dictionary mapping key names to boolean indicating if required - context_name: Name of the context for error messages (e.g., operator name) - """ - for key, is_required in required_keys.items(): - if is_required and key not in self.data: - raise ValueError(f"{context_name}: missing required input '{key}'") - return self diff --git a/reme2/component/service/http_service.py b/reme2/component/service/http_service.py index d264f87c..51eef25e 100644 --- a/reme2/component/service/http_service.py +++ b/reme2/component/service/http_service.py @@ -77,10 +77,10 @@ class HttpService(BaseService): async def generate_stream() -> AsyncGenerator[bytes, None]: async for chunk in execute_stream_task( - stream_queue=stream_queue, - task=task, - task_name=job.name, - output_format="bytes", + stream_queue=stream_queue, + task=task, + task_name=job.name, + output_format="bytes", ): assert isinstance(chunk, bytes) yield chunk diff --git a/reme2/config/obsidian.md b/reme2/config/obsidian.md index ca98a27e..cda53c55 100644 --- a/reme2/config/obsidian.md +++ b/reme2/config/obsidian.md @@ -1,18 +1,18 @@ # all @jinli + reme help reme start reme restart reme version reme vault="My Vault" - # daily + daily:path reme daily:xxx - - # crud + reme create file="New Note" content="# Hello" title="xxx" tags="[]" status="" reme read file=Recipe/path="Templates/Recipe.md" reme edit file=Recipe/path="Templates/Recipe.md" old="xxx" new="xxx" @@ -21,57 +21,66 @@ reme prepend file="My Note" content="New line" reme delete file="My Note"/path # reme + reme stat file/path reme list file/path # search + reme search query="search term" limit=10 tag="[]" score=0.1 copy=true # property + reme property:read reme property:update file="My Note" status=done xx=xxx reme property:delete keys="[xxxx, xxxx]" # 全局所有标签 + reme tags # show link + reme backlinks file="My Note" reme links file="My Note" - [[Algorithm Notes#Sorting]] - - - 记忆的类型 + 1. daily -> daily/xxxx-mm-dd/overview.md -> xxxx.md 2. topic -> topic/personal(agent)/xxxx.md 记忆的被动总结 + 1. 是否需要:是需要,不会主动触发 2. 什么时候调用: - - freq (every_n_turn、compact) -> daily_summarizer - - topic (/dream ) -> topic_summarizer(daily_xx -> topic_xx) - - proactive -> proactive_summarizer(personal_xxx -> proactive_query) + +- freq (every_n_turn、compact) -> daily_summarizer +- topic (/dream ) -> topic_summarizer(daily_xx -> topic_xx) +- proactive -> proactive_summarizer(personal_xxx -> proactive_query) - pre_query - - 记忆的搜索 file watch + - start 全量扫描 -> 文件变化结果 - 增量扫描 -> 文件变化结果 如果有变化,检测 增加、删除、修改 cud: + 1. file -> metadata 存json - - path + +- path - mtime - header: tags/kv/title/desc + 2. file -> link 存json + - 读全量修改的文件的全量,识别link + 3. file -> chunk 存db/json + - 切片 -> file+start+end: preview(100) diff --git a/reme2/config/paw.yaml b/reme2/config/paw.yaml index 1ef43102..7d604ddd 100644 --- a/reme2/config/paw.yaml +++ b/reme2/config/paw.yaml @@ -60,5 +60,5 @@ components: backend: full file_store: default default_parser: default - watch_paths: ["./test_data"] + watch_paths: [ "./test_data" ] recursive: true diff --git a/reme2/file_based/file_io.py b/reme2/file_based/file_io.py index f86b564a..22399015 100644 --- a/reme2/file_based/file_io.py +++ b/reme2/file_based/file_io.py @@ -40,10 +40,10 @@ class FileIO: return str(self.working_dir / file_path) async def read_file( # pylint: disable=too-many-return-statements - self, - file_path: str, - start_line: int | None = None, - end_line: int | None = None, + self, + file_path: str, + start_line: int | None = None, + end_line: int | None = None, ) -> ToolResponse: """Read a file. Relative paths resolve from WORKING_DIR. @@ -138,7 +138,7 @@ class FileIO: ) # Extract selected lines - selected_content = "\n".join(all_lines[s - 1 : e]) + selected_content = "\n".join(all_lines[s - 1: e]) # Apply smart truncation (consistent with shell output format) text = truncate_text_output( @@ -154,13 +154,13 @@ class FileIO: if text == selected_content and e < total: content_bytes = len(text.encode("utf-8")) notice = ( - TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated." - f"\nThe full content is saved to the file " - f"and contains {total} lines in total." - f"\nThis excerpt starts at line {s} and " - f"covers the next {content_bytes} bytes." - "\nIf the current content is not enough, " - f"call `read_file` with file_path={file_path} start_line={e + 1} to read more." + TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated." + f"\nThe full content is saved to the file " + f"and contains {total} lines in total." + f"\nThis excerpt starts at line {s} and " + f"covers the next {content_bytes} bytes." + "\nIf the current content is not enough, " + f"call `read_file` with file_path={file_path} start_line={e + 1} to read more." ) text = text + notice @@ -179,9 +179,9 @@ class FileIO: ) async def write_file( - self, - file_path: str, - content: str, + self, + file_path: str, + content: str, ) -> ToolResponse: """Create or overwrite a file. Relative paths resolve from working_dir. @@ -226,10 +226,10 @@ class FileIO: # pylint: disable=too-many-return-statements async def edit_file( - self, - file_path: str, - old_text: str, - new_text: str, + self, + file_path: str, + old_text: str, + new_text: str, ) -> ToolResponse: """Find-and-replace text in a file. All occurrences of old_text are replaced with new_text. Relative paths resolve from working_dir. @@ -314,9 +314,9 @@ class FileIO: ) async def append_file( - self, - file_path: str, - content: str, + self, + file_path: str, + content: str, ) -> ToolResponse: """Append content to the end of a file. Relative paths resolve from working_dir. diff --git a/reme2/file_based/file_utils.py b/reme2/file_based/file_utils.py index bb775291..e0875679 100644 --- a/reme2/file_based/file_utils.py +++ b/reme2/file_based/file_utils.py @@ -12,12 +12,12 @@ from ..constants import ( def _truncate_fresh( - text: str, - start_line: int, - total_lines: int, - max_bytes: int, - file_path: str | None, - encoding: str, + text: str, + start_line: int, + total_lines: int, + max_bytes: int, + file_path: str | None, + encoding: str, ) -> str: """Truncate fresh text (no prior truncation marker) by bytes with line integrity. @@ -66,20 +66,20 @@ def _truncate_fresh( return result notice = ( - TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated." - f"\nThe full content is saved to the file and contains {total_lines} lines in total." - f"\nThis excerpt starts at line {start_line} and covers the next {max_bytes} bytes." - f"\nIf the current content is not enough, call `read_file` with file_path={file_path or ''} " - f"start_line={read_from} to read more." + TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated." + f"\nThe full content is saved to the file and contains {total_lines} lines in total." + f"\nThis excerpt starts at line {start_line} and covers the next {max_bytes} bytes." + f"\nIf the current content is not enough, call `read_file` with file_path={file_path or ''} " + f"start_line={read_from} to read more." ) return result + notice def _retruncate( - text: str, - max_bytes: int, - encoding: str, + text: str, + max_bytes: int, + encoding: str, ) -> str: """Re-truncate text that was previously truncated (contains TRUNCATION_NOTICE_MARKER). @@ -132,12 +132,12 @@ def _retruncate( def truncate_text_output( - text: str, - start_line: int = 1, - total_lines: int = 0, - max_bytes: int = DEFAULT_MAX_BYTES, - file_path: str | None = None, - encoding: str = "utf-8", + text: str, + start_line: int = 1, + total_lines: int = 0, + max_bytes: int = DEFAULT_MAX_BYTES, + file_path: str | None = None, + encoding: str = "utf-8", ) -> str: """Truncate file output by bytes with line integrity. diff --git a/reme2/file_based/memory_search.py b/reme2/file_based/memory_search.py index 647090a4..84650ab7 100644 --- a/reme2/file_based/memory_search.py +++ b/reme2/file_based/memory_search.py @@ -35,10 +35,10 @@ class MemorySearch(BaseStep): assert query, "Query cannot be empty" assert ( - isinstance(min_score, float | int) and 0.0 <= min_score <= 1.0 + isinstance(min_score, float | int) and 0.0 <= min_score <= 1.0 ), f"min_score must be between 0 and 1, got {min_score}" assert ( - isinstance(max_results, int) and max_results > 0 + isinstance(max_results, int) and max_results > 0 ), f"max_results must be a positive integer, got {max_results}" filter_paths: list[str] | None = self.context.get("paths") or None diff --git a/reme2/file_based/summarizer.py b/reme2/file_based/summarizer.py index 62335c9b..c5ba5f21 100644 --- a/reme2/file_based/summarizer.py +++ b/reme2/file_based/summarizer.py @@ -18,16 +18,16 @@ class Summarizer(BaseStep): """Summarizer step for summarizing memory messages.""" def __init__( - self, - working_dir: str, - memory_dir: str, - memory_compact_threshold: int, - toolkit: Toolkit | None = None, - console_enabled: bool = False, - timezone: str | None = None, - add_thinking_block: bool = True, - as_token_counter: HuggingFaceTokenCounter | None = None, - **kwargs, + self, + working_dir: str, + memory_dir: str, + memory_compact_threshold: int, + toolkit: Toolkit | None = None, + console_enabled: bool = False, + timezone: str | None = None, + add_thinking_block: bool = True, + as_token_counter: HuggingFaceTokenCounter | None = None, + **kwargs, ): """Initialize the summarizer step. @@ -227,10 +227,10 @@ class Summarizer(BaseStep): return total async def _format_msgs_to_str( - self, - messages: list[Msg], - memory_compact_threshold: int, - include_thinking: bool = True, + self, + messages: list[Msg], + memory_compact_threshold: int, + include_thinking: bool = True, ) -> str: """Format list of messages to a single formatted string. diff --git a/reme2/file_based/summarizer.yaml b/reme2/file_based/summarizer.yaml index 14599f71..a9cb59e6 100644 --- a/reme2/file_based/summarizer.yaml +++ b/reme2/file_based/summarizer.yaml @@ -1,62 +1,62 @@ user_message: | - Memory Pre-compression Flush Cycle. - - The current session is about to enter the automatic compression phase. Please capture persistent memory AND session reflections, then write them to disk. - - Current date: {date} - Working directory: {working_dir} - - # Task - Immediately store persistent memory and reflections to: {memory_dir}/YYYY-MM-DD.md - - # Workflow - 1. Extract and synthesize content from the current session: - - Persistent Memory: Facts, user profile updates, project states, and important events. - - Experience Reflection: Reusable thinking logic derived from user feedback, successful problem-solving strategies, mistakes made/pitfalls to avoid, and actionable insights for future interactions. - 2. `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned) - - If the file doesn’t exist, use `write` tool directly. - - If the file exists, intelligently merge new information with existing content, prefer using `edit` to update specific sections. - - Use `write` to overwrite the entire file only if substantial restructuring is required. - - # Principles - - Intelligently merge new information with existing content: - - Categorize clearly (e.g., separate "Factual Memory" from "Reflections & Logic"). - - Avoid duplicating already recorded information. - - Enrich existing entries with new details where relevant. - - Maintain chronological order wherever applicable. - - Always preserve timestamps and any date/time-related context. - - Add only genuinely new or meaningfully enriching information. - - Reflections MUST focus on forming reusable cognitive frameworks based on user feedback, aiming to improve future task execution. - - Keep entries concise yet complete. - - If there’s nothing to store or reflect on, respond with [SILENT]. + Memory Pre-compression Flush Cycle. + + The current session is about to enter the automatic compression phase. Please capture persistent memory AND session reflections, then write them to disk. + + Current date: {date} + Working directory: {working_dir} + + # Task + Immediately store persistent memory and reflections to: {memory_dir}/YYYY-MM-DD.md + + # Workflow + 1. Extract and synthesize content from the current session: + - Persistent Memory: Facts, user profile updates, project states, and important events. + - Experience Reflection: Reusable thinking logic derived from user feedback, successful problem-solving strategies, mistakes made/pitfalls to avoid, and actionable insights for future interactions. + 2. `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned) + - If the file doesn’t exist, use `write` tool directly. + - If the file exists, intelligently merge new information with existing content, prefer using `edit` to update specific sections. + - Use `write` to overwrite the entire file only if substantial restructuring is required. + + # Principles + - Intelligently merge new information with existing content: + - Categorize clearly (e.g., separate "Factual Memory" from "Reflections & Logic"). + - Avoid duplicating already recorded information. + - Enrich existing entries with new details where relevant. + - Maintain chronological order wherever applicable. + - Always preserve timestamps and any date/time-related context. + - Add only genuinely new or meaningfully enriching information. + - Reflections MUST focus on forming reusable cognitive frameworks based on user feedback, aiming to improve future task execution. + - Keep entries concise yet complete. + - If there’s nothing to store or reflect on, respond with [SILENT]. user_message_zh: | - 预压缩内存刷新轮次。 - - 当前会话即将进入自动压缩阶段;请将持久化记忆与经验反思捕获并写入磁盘。 - - 当前日期:{date} - 工作目录:{working_dir} - - # 任务 - 立即存储持久化记忆与反思(使用路径 {memory_dir}/YYYY-MM-DD.md)。 - - # 工作流程 - 1. 从当前会话中提取并综合两类内容: - - 持久化记忆:客观事实、用户信息更新、项目状态及重要事件。 - - 经验反思:基于用户反馈形成的可复用思考逻辑、成功的问题解决策略、犯下的错误/应避免的陷阱,以及对未来交互有帮助的行动指南。 - 2. `read` {memory_dir}/YYYY-MM-DD.md(如文件不存在,会返回错误提示) - - 若文件不存在,直接使用 `write` 工具写入。 - - 若文件已存在,智能合并新信息与现有内容,尽可能使用 `edit` 更新特定部分,仅在需要大幅重构时使用 `write` 覆盖整个文件。 - - # 原则 - - 智能合并新信息与现有内容: - - 将内容进行清晰的分类(例如明确区分“事实记忆”与“反思与逻辑”)。 - - 避免重复已记录的信息。 - - 在相关时丰富现有条目的新细节。 - - 在适用时保持时间顺序。 - - 始终保留时间戳、日期和时间相关上下文。 - - 仅添加真正新的或有丰富价值的信息。 - - 反思内容必须侧重于根据用户反馈构建可复用的思维逻辑,以改善未来的任务执行。 - - 保持条目简洁但完整。 - - 若无任何新内容可存储或反思,请回复 [SILENT]。 \ No newline at end of file + 预压缩内存刷新轮次。 + + 当前会话即将进入自动压缩阶段;请将持久化记忆与经验反思捕获并写入磁盘。 + + 当前日期:{date} + 工作目录:{working_dir} + + # 任务 + 立即存储持久化记忆与反思(使用路径 {memory_dir}/YYYY-MM-DD.md)。 + + # 工作流程 + 1. 从当前会话中提取并综合两类内容: + - 持久化记忆:客观事实、用户信息更新、项目状态及重要事件。 + - 经验反思:基于用户反馈形成的可复用思考逻辑、成功的问题解决策略、犯下的错误/应避免的陷阱,以及对未来交互有帮助的行动指南。 + 2. `read` {memory_dir}/YYYY-MM-DD.md(如文件不存在,会返回错误提示) + - 若文件不存在,直接使用 `write` 工具写入。 + - 若文件已存在,智能合并新信息与现有内容,尽可能使用 `edit` 更新特定部分,仅在需要大幅重构时使用 `write` 覆盖整个文件。 + + # 原则 + - 智能合并新信息与现有内容: + - 将内容进行清晰的分类(例如明确区分“事实记忆”与“反思与逻辑”)。 + - 避免重复已记录的信息。 + - 在相关时丰富现有条目的新细节。 + - 在适用时保持时间顺序。 + - 始终保留时间戳、日期和时间相关上下文。 + - 仅添加真正新的或有丰富价值的信息。 + - 反思内容必须侧重于根据用户反馈构建可复用的思维逻辑,以改善未来的任务执行。 + - 保持条目简洁但完整。 + - 若无任何新内容可存储或反思,请回复 [SILENT]。 \ No newline at end of file diff --git a/reme2/reme.py b/reme2/reme.py index 1cf55749..99fd369f 100644 --- a/reme2/reme.py +++ b/reme2/reme.py @@ -1,6 +1,5 @@ """ReMe CLI application entry point.""" -import asyncio import sys from pathlib import Path @@ -22,17 +21,17 @@ class ReMe(Application): """ReMe memory management application.""" async def summarize( - self, - messages: list[Msg], - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase | HuggingFaceTokenCounter = "default", - toolkit: Toolkit | None = None, - language: str = "zh", - max_input_length: float = 128 * 1024, - compact_ratio: float = 0.7, - timezone: str | None = None, - add_thinking_block: bool = True, + self, + messages: list[Msg], + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase | HuggingFaceTokenCounter = "default", + toolkit: Toolkit | None = None, + language: str = "zh", + max_input_length: float = 128 * 1024, + compact_ratio: float = 0.7, + timezone: str | None = None, + add_thinking_block: bool = True, ) -> str: """Summarize and compact memory messages. @@ -95,25 +94,25 @@ class ReMe(Application): return str(e) async def dream( - self, - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase = "default", - toolkit: Toolkit | None = None, - language: str = "zh", - timezone: str | None = None, + self, + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase = "default", + toolkit: Toolkit | None = None, + language: str = "zh", + timezone: str | None = None, ) -> str: """Process and consolidate memories in background.""" return "" async def proactive( - self, - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase = "default", - toolkit: Toolkit | None = None, - language: str = "zh", - timezone: str | None = None, + self, + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase = "default", + toolkit: Toolkit | None = None, + language: str = "zh", + timezone: str | None = None, ) -> str: """Generate proactive memory insights.""" return "" diff --git a/reme2/schema/file_chunk.py b/reme2/schema/file_chunk.py index b08a95f3..3d4ff6cc 100644 --- a/reme2/schema/file_chunk.py +++ b/reme2/schema/file_chunk.py @@ -1,53 +1,21 @@ -"""File chunk schema module. - -This module defines the FileChunk model for representing chunks of file content -in the document processing and retrieval pipeline. -""" - from pydantic import Field from .base_node import BaseNode class FileChunk(BaseNode): - """A chunk of file content with positional and scoring metadata. + """文件内容分块,包含位置和评分元数据。""" - Represents a contiguous section of a file that has been extracted - for processing, embedding, or retrieval. Inherits text and embedding - capabilities from BaseNode. - - Attributes: - path: File path relative to workspace root. - start_line: Starting line number (1-indexed) in the source file. - end_line: Ending line number (1-indexed) in the source file. - hash: Hash of the chunk content for deduplication. - scores: Search relevance scores indexed by score type. - - Properties: - score: Final combined score for search result ranking. - merge_key: Unique key for merging duplicate search results. - """ - - path: str = Field(..., description="File path relative to workspace") - start_line: int = Field(..., description="Starting line number (1-indexed)") - end_line: int = Field(..., description="Ending line number (1-indexed)") - hash: str = Field(..., description="Hash of chunk content for deduplication") - scores: dict[str, float] = Field(default_factory=dict, description="Search scores by type") + path: str = Field(...) + start_line: int = Field(...) + end_line: int = Field(...) + hash: str = Field(...) + scores: dict[str, float] = Field(default_factory=dict) @property def score(self) -> float: - """Get the final score for search result ranking. - - Returns: - The combined score, or 0.0 if not set. - """ return self.scores.get("score", 0.0) @property - def merge_key(self) -> str: - """Generate a unique key for merging search results. - - Returns: - A string key in format "path:start_line:end_line". - """ + def unique_key(self) -> str: return f"{self.path}:{self.start_line}:{self.end_line}" diff --git a/reme2/utils/chunking_utils.py b/reme2/utils/chunking_utils.py index 87ee0df7..9e94d5d3 100644 --- a/reme2/utils/chunking_utils.py +++ b/reme2/utils/chunking_utils.py @@ -9,10 +9,10 @@ from ..schema import FileChunk def chunk_markdown( - text: str, - path: str, - chunk_tokens: int, - overlap: int, + text: str, + path: str, + chunk_tokens: int, + overlap: int, ) -> list[FileChunk]: """Split Markdown text into chunks with configurable size and overlap. @@ -125,7 +125,7 @@ def chunk_markdown( else: # If line is too long, split by maximum character count for start in range(0, len(line), max_chars): - segments.append(line[start : start + max_chars]) + segments.append(line[start: start + max_chars]) for segment in segments: line_size = len(segment) + 1 # +1 for newline diff --git a/reme2/utils/common_utils.py b/reme2/utils/common_utils.py index 6c86d487..513cb8ec 100644 --- a/reme2/utils/common_utils.py +++ b/reme2/utils/common_utils.py @@ -47,10 +47,10 @@ def hash_text(text: str, encoding: str = "utf-8") -> str: async def execute_stream_task( - stream_queue: asyncio.Queue, - task: asyncio.Task, - task_name: str | None = None, - output_format: Literal["str", "bytes", "chunk"] = "str", + stream_queue: asyncio.Queue, + task: asyncio.Task, + task_name: str | None = None, + output_format: Literal["str", "bytes", "chunk"] = "str", ) -> AsyncGenerator[str | bytes | StreamChunk, None]: """Core stream flow execution logic. diff --git a/reme2/utils/logger_utils.py b/reme2/utils/logger_utils.py index 0c2eeb70..b9648152 100644 --- a/reme2/utils/logger_utils.py +++ b/reme2/utils/logger_utils.py @@ -14,11 +14,11 @@ _initialized = False def get_logger( - log_dir: str = "logs", - level: str = "INFO", - log_to_console: bool = True, - log_to_file: bool = True, - force_init: bool = False, + log_dir: str = "logs", + level: str = "INFO", + log_to_console: bool = True, + log_to_file: bool = True, + force_init: bool = False, ): """Get a configured logger instance.