diff --git a/reme_cli/component/component_registry.py b/reme_cli/component/component_registry.py new file mode 100644 index 00000000..ab0cbd9c --- /dev/null +++ b/reme_cli/component/component_registry.py @@ -0,0 +1,128 @@ +""" +Component registry module. + +Provides a global registry for managing component class registration and lookup. +Supports two registration methods: +1. Direct registration: R.register(MyClass, "name") +2. Decorator registration: @R.register("name") +""" + +from typing import Callable, TypeVar, cast + +from .base_component import BaseComponent +from ..enumeration import ComponentEnum + +T = TypeVar("T", bound=BaseComponent) + + +class ComponentRegistry: + """Registry for managing component class registration and lookup. + + Components are organized by type (ComponentEnum), with each type + containing multiple named component implementations. + + Attributes: + _registry: Internal storage structure, format: {ComponentEnum: {name: component_class}} + + Usage: + # Direct registration + R.register(OpenAIChatModel, "openai") + + # Decorator registration + @R.register("openai") + class OpenAIChatModel(BaseComponent): + ... + + # Get registered component + cls = R.get(ComponentEnum.LLM, "openai") + """ + + def __init__(self) -> None: + """Initialize the component registry.""" + self._registry: dict[ComponentEnum, dict[str, type[BaseComponent]]] = {} + + def _do_register(self, cls: type[T], name: str) -> type[T]: + """Internal method to register a class. + + Args: + cls: The component class to register. + name: The name identifier for the component. + + Returns: + The registered class. + """ + component_type = cls.component_type + if component_type not in self._registry: + self._registry[component_type] = {} + self._registry[component_type][name] = cls + return cls + + def register( + self, cls_or_name: type[T] | str, name: str | None = None + ) -> Callable[[type[T]], type[T]] | type[T]: + """Register a component class. + + Supports two calling patterns: + - Direct: R.register(MyClass, "name") -> returns MyClass + - Decorator: @R.register("name") -> returns decorator function + + Args: + cls_or_name: Either a class to register (direct mode), or a name string + for decorator mode. + name: Optional name when using direct registration with a class. + If not provided, the class name will be used. + + Returns: + Either the registered class (direct mode) or a decorator function. + + Example: + # Direct registration + R.register(OpenAIChatModel, "openai") + + # Decorator registration + @R.register("openai") + class OpenAIChatModel(BaseComponent): + ... + """ + # Direct registration: R.register(MyClass, "name") + if isinstance(cls_or_name, type): + cls = cast(type[T], cls_or_name) + _key = name or cls.__name__ + return self._do_register(cls, _key) + + # Decorator mode: @R.register("name") + _decorator_name = cls_or_name + + def decorator(decorated_cls: type[T]) -> type[T]: + key = _decorator_name or decorated_cls.__name__ + return self._do_register(decorated_cls, key) + + return decorator + + def get(self, component_type: ComponentEnum, name: str) -> type[BaseComponent] | None: + """Get a registered component by type and name. + + Args: + component_type: The component type enum value. + name: The registered name of the component. + + Returns: + The component class if found, None otherwise. + """ + return self._registry.get(component_type, {}).get(name) + + def get_all(self, component_type: ComponentEnum) -> dict[str, type[BaseComponent]]: + """Get all registered components for a given type. + + Args: + component_type: The component type enum value. + + Returns: + Dictionary mapping component names to their classes. + Returns empty dict if no components registered for the type. + """ + return self._registry.get(component_type, {}) + + +# Global registry instance +R = ComponentRegistry() \ No newline at end of file diff --git a/reme_cli/component/file_store/__init__.py b/reme_cli/component/file_store/__init__.py new file mode 100644 index 00000000..a84d7b1a --- /dev/null +++ b/reme_cli/component/file_store/__init__.py @@ -0,0 +1,13 @@ +"""File store module for persistent memory management. + +This module provides storage backends for memory chunks and file metadata, +including pure-Python local implementations with vector and full-text search. +""" + +from .base_file_store import BaseFileStore +from .local_file_store import LocalFileStore + +__all__ = [ + "BaseFileStore", + "LocalFileStore", +] \ No newline at end of file diff --git a/reme_cli/component/file_store/base_file_store.py b/reme_cli/component/file_store/base_file_store.py new file mode 100644 index 00000000..0d0f06f5 --- /dev/null +++ b/reme_cli/component/file_store/base_file_store.py @@ -0,0 +1,168 @@ +"""Base storage interface for file store.""" + +import re +from abc import abstractmethod +from pathlib import Path +from typing import TYPE_CHECKING + +from ..base_component import BaseComponent +from ..embedding import BaseEmbeddingModel +from ...enumeration import ComponentEnum +from ...schema import FileChunk, FileMetadata + +if TYPE_CHECKING: + from ..application_context import ApplicationContext + + +class BaseFileStore(BaseComponent): + """Abstract base class for file storage backends.""" + + def __init__( + 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 + self.embedding_model: BaseEmbeddingModel | None = None + self.store_name: str = store_name + self.db_path: Path = Path(db_path) + self.db_path.mkdir(parents=True, exist_ok=True) + self.vector_enabled: bool = bool(embedding_model) + self.fts_enabled: bool = fts_enabled + + if not re.match(r"^[a-zA-Z0-9_]+$", store_name): + raise ValueError(f"Invalid '{store_name}'. Only alphanumeric characters and underscores are allowed.") + if not self.vector_enabled and not self.fts_enabled: + raise ValueError("At least one of embedding_model or fts_enabled must be set.") + + async def _start(self, app_context: ApplicationContext | None = None): + """Initialize the storage backend and resolve embedding_model from app_context.""" + if not self._embedding_model_name: + return + models = app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {}) + if self._embedding_model_name not in models: + raise ValueError(f"Embedding model '{self._embedding_model_name}' not found.") + model = models[self._embedding_model_name] + if not isinstance(model, BaseEmbeddingModel): + raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}") + self.embedding_model = model + + async def _close(self): + """Close storage and release resources.""" + self.embedding_model = None + + @property + def embedding_dim(self) -> int: + """Get the embedding model's dimensionality.""" + return self.embedding_model.dimensions if self.embedding_model else 1024 + + def _disable_vector_search(self, reason: str = "embedding API error") -> None: + """Disable vector search and log a warning.""" + if self.vector_enabled: + self.logger.warning(f"[{self.store_name}] Disabling vector search: {reason}") + self.vector_enabled = False + + async def _get_embeddings_safe(self, texts: list[str], **kwargs) -> list[list[float]] | None: + """Get embeddings, return None if vector search is disabled or failed.""" + if not self.vector_enabled: + return None + try: + return await self.embedding_model.get_embeddings(texts, **kwargs) + except Exception as e: + self._disable_vector_search(str(e)) + return None + + async def get_embedding(self, query: str, **kwargs) -> list[float] | None: + """Get embedding for a single query string.""" + result = await self._get_embeddings_safe([query], **kwargs) + return result[0] if result else None + + async def get_embeddings(self, queries: list[str], **kwargs) -> list[list[float]] | None: + """Get embeddings for a batch of query strings.""" + return await self._get_embeddings_safe(queries, **kwargs) + + async def get_chunk_embedding(self, chunk: FileChunk, **kwargs) -> FileChunk: + """Generate embedding for a single FileChunk.""" + chunk.embedding = await self.get_embedding(chunk.text, **kwargs) + return chunk + + async def get_chunk_embeddings(self, chunks: list[FileChunk], **kwargs) -> list[FileChunk]: + """Generate embeddings for a batch of FileChunk.""" + if not chunks: + return chunks + embeddings = await self.get_embeddings([c.text for c in chunks], **kwargs) + if embeddings and len(embeddings) == len(chunks): + for chunk, emb in zip(chunks, embeddings): + chunk.embedding = emb + else: + for chunk in chunks: + chunk.embedding = None + return chunks + + @abstractmethod + async def clear_all(self): + """Clear all indexed data.""" + + @abstractmethod + async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]): + """Insert or update a file and its chunks.""" + + @abstractmethod + async def delete_file(self, path: str): + """Delete a file and all its chunks.""" + + @abstractmethod + async def delete_file_chunks(self, path: str, chunk_ids: list[str]): + """Delete chunks for a file.""" + + @abstractmethod + async def upsert_chunks(self, chunks: list[FileChunk]): + """Insert or update specific chunks without affecting other chunks.""" + + @abstractmethod + async def list_files(self) -> list[str]: + """List all indexed file paths.""" + + @abstractmethod + async def get_file_metadata(self, path: str) -> FileMetadata | None: + """Get full file metadata with statistics.""" + + @abstractmethod + async def update_file_metadata(self, file_meta: FileMetadata) -> None: + """Update file metadata without affecting chunks.""" + + @abstractmethod + async def get_file_chunks(self, path: str) -> list[FileChunk]: + """Get all chunks for a file.""" + + @abstractmethod + async def vector_search(self, query: str, limit: int) -> list[FileChunk]: + """Perform vector similarity search.""" + + @abstractmethod + async def keyword_search(self, query: str, limit: int) -> list[FileChunk]: + """Perform keyword/full-text search.""" + + @abstractmethod + async def hybrid_search( + self, + query: str, + limit: int, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[FileChunk]: + """Perform hybrid search combining vector and keyword search. + + Args: + query: Search query text + limit: Maximum number of results + vector_weight: Weight for vector search results (0.0-1.0) + candidate_multiplier: Multiplier for candidate pool size + + Returns: + List of FileChunk with score populated, sorted by combined relevance + """ diff --git a/reme_cli/component/file_store/local_file_store.py b/reme_cli/component/file_store/local_file_store.py new file mode 100644 index 00000000..713a4cc8 --- /dev/null +++ b/reme_cli/component/file_store/local_file_store.py @@ -0,0 +1,392 @@ +"""Pure-Python in-memory storage backend for file store, with JSON file persistence.""" + +import json +from pathlib import Path + +import numpy as np + +from .base_file_store import BaseFileStore +from ...schema import FileChunk, FileMetadata +from ...utils import batch_cosine_similarity + + +class LocalFileStore(BaseFileStore): + """Pure-Python in-memory file storage with JSONL file persistence. + + No external dependencies required. All data lives in Python dicts; + writes are persisted to JSONL files on disk so state survives restarts. + + Inherits embedding methods from BaseFileStore: + - get_chunk_embedding / get_chunk_embeddings (async) + - get_embedding / get_embeddings (async) + + Provides: + - Vector similarity search (cosine similarity, pure Python) + - Full-text / keyword search (Python substring matching) + - Efficient chunk and file metadata management + """ + + def __init__(self, encoding: str = "utf-8", **kwargs): + super().__init__(**kwargs) + self._encoding: str = encoding + self._chunks: dict[str, FileChunk] = {} + self._files: dict[str, FileMetadata] = {} + 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 + # ------------------------------------------------------------------ + + async def _load_chunks(self) -> None: + """Load chunks from JSONL file into memory.""" + if not self._chunks_file.exists(): + return + try: + data = self._chunks_file.read_text(encoding=self._encoding) + self._chunks = {} + for line in data.strip().split("\n"): + if not line: + continue + chunk = FileChunk.model_validate(json.loads(line)) + self._chunks[chunk.id] = chunk + except Exception as e: + self.logger.warning(f"Failed to load chunks: {e}") + + async def _save_chunks(self) -> None: + """Persist chunks to JSONL file with atomic write.""" + lines = [json.dumps(c.model_dump(mode="json"), ensure_ascii=False) for c in self._chunks.values()] + content = "\n".join(lines) + temp_path = self._chunks_file.with_suffix('.tmp') + try: + temp_path.write_text(content, encoding=self._encoding) + temp_path.replace(self._chunks_file) + except Exception as e: + self.logger.error(f"Failed to save chunks: {e}") + raise + finally: + if temp_path.exists(): + temp_path.unlink() + + async def _load_metadata(self) -> None: + """Load file metadata from JSON file into memory.""" + if not self._metadata_file.exists(): + return + try: + data = self._metadata_file.read_text(encoding=self._encoding) + raw: dict = json.loads(data) + self._files = {path: FileMetadata(**meta) for path, meta in raw.items()} + except Exception as e: + self.logger.warning(f"Failed to load metadata: {e}") + + async def _save_metadata(self) -> None: + """Persist file metadata to JSON file with atomic write.""" + raw = {path: meta.model_dump(exclude={"content", "metadata"}, mode="json") + for path, meta in self._files.items()} + content = json.dumps(raw, indent=2, ensure_ascii=False) + temp_path = self._metadata_file.with_suffix('.tmp') + try: + temp_path.write_text(content, encoding=self._encoding) + temp_path.replace(self._metadata_file) + except Exception as e: + self.logger.error(f"Failed to save metadata: {e}") + raise + finally: + if temp_path.exists(): + temp_path.unlink() + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def _start(self, app_context=None) -> None: + """Load persisted data into memory.""" + await self._load_metadata() + await self._load_chunks() + self.logger.info( + f"LocalFileStore '{self.store_name}' ready: " + f"{len(self._chunks)} chunks, metadata at {self._metadata_file}", + ) + await super()._start() + + async def _close(self) -> None: + """Flush state to disk and release memory.""" + await self._save_metadata() + await self._save_chunks() + self._chunks.clear() + self._files.clear() + await super()._close() + + # ------------------------------------------------------------------ + # Write operations + # ------------------------------------------------------------------ + + async def upsert_file( + self, + file_meta: FileMetadata, + chunks: list[FileChunk], + ) -> None: + """Insert or update file and its chunks.""" + if not chunks: + return + + # Remove existing chunks for this file first + await self.delete_file(file_meta.path) + + # Batch generate embeddings (base class returns mock embeddings when vector_enabled=False) + chunks = await self.get_chunk_embeddings(chunks) + + for chunk in chunks: + self._chunks[chunk.id] = chunk + + self._files[file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=len(chunks), + ) + + async def delete_file(self, path: str) -> None: + """Delete file and all its chunks.""" + to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path] + for cid in to_delete: + del self._chunks[cid] + + self._files.pop(path, None) + + async def delete_file_chunks(self, path: str, chunk_ids: list[str]) -> None: + """Delete specific chunks for a file.""" + if not chunk_ids: + return + + for cid in chunk_ids: + self._chunks.pop(cid, None) + + # Recalculate chunk_count in file metadata + if path in self._files: + self._files[path].chunk_count = sum( + 1 for chunk in self._chunks.values() if chunk.path == path + ) + + async def upsert_chunks( + self, + chunks: list[FileChunk], + ) -> None: + """Insert or update specific chunks without affecting other chunks.""" + if not chunks: + return + + chunks = await self.get_chunk_embeddings(chunks) + + for chunk in chunks: + self._chunks[chunk.id] = chunk + + # ------------------------------------------------------------------ + # Read operations + # ------------------------------------------------------------------ + + async def list_files(self) -> list[str]: + """List all indexed files.""" + return list(self._files.keys()) + + async def get_file_metadata( + self, + path: str, + ) -> FileMetadata | None: + """Get file metadata.""" + return self._files.get(path) + + async def update_file_metadata(self, file_meta: FileMetadata) -> None: + """Update file metadata without affecting chunks.""" + self._files[file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + ) + + async def get_file_chunks( + self, + path: str, + ) -> list[FileChunk]: + """Get all chunks for a file, sorted by start_line.""" + chunks = [chunk for chunk in self._chunks.values() if chunk.path == path] + chunks.sort(key=lambda c: c.start_line) + return chunks + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + async def vector_search( + self, + query: str, + limit: int, + ) -> list[FileChunk]: + """Perform cosine-similarity vector search over in-memory embeddings.""" + if not self.vector_enabled or not query: + return [] + + query_embedding = await self.get_embedding(query) + if not query_embedding: + return [] + + expected_dim = self.embedding_dim + + # Collect candidate chunks with embeddings + candidates = [c for c in self._chunks.values() if c.embedding] + if not candidates: + return [] + + # Validate and fix chunk embedding dimensions + valid_embeddings = [] + for chunk in candidates: + emb = chunk.embedding + emb_len = len(emb) + if emb_len != expected_dim: + if emb_len < expected_dim: + emb = emb + [0.0] * (expected_dim - emb_len) + else: + emb = emb[:expected_dim] + valid_embeddings.append(emb) + + # Compute similarities + query_array = np.array([query_embedding]) + chunk_embeddings = np.array(valid_embeddings) + similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] + + # Build results + results = [] + for chunk, sim in zip(candidates, similarities): + results.append( + FileChunk( + id=chunk.id, + path=chunk.path, + start_line=chunk.start_line, + end_line=chunk.end_line, + hash=chunk.hash, + text=chunk.text, + embedding=chunk.embedding, + scores={"vector": float(sim), "score": float(sim)}, + ) + ) + + results.sort(key=lambda r: r.score, reverse=True) + return results[:limit] + + async def keyword_search( + self, + query: str, + limit: int, + ) -> list[FileChunk]: + """Perform keyword/full-text search via Python substring matching.""" + if not self.fts_enabled or not query: + return [] + + words = query.split() + if not words: + return [] + + query_lower = query.lower() + words_lower = [w.lower() for w in words] + n_words = len(words) + + results = [] + for chunk in self._chunks.values(): + text_lower = chunk.text.lower() + match_count = sum(1 for w in words_lower if w in text_lower) + if match_count == 0: + continue + + base_score = match_count / n_words + phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0 + score = min(1.0, base_score + phrase_bonus) + + results.append( + FileChunk( + id=chunk.id, + path=chunk.path, + start_line=chunk.start_line, + end_line=chunk.end_line, + hash=chunk.hash, + text=chunk.text, + scores={"keyword": score, "score": score}, + ) + ) + + results.sort(key=lambda r: r.score, reverse=True) + return results[:limit] + + async def hybrid_search( + self, + query: str, + limit: int, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[FileChunk]: + """Perform hybrid search combining vector and keyword search.""" + assert 0.0 <= vector_weight <= 1.0 + + candidates = min(200, max(1, int(limit * candidate_multiplier))) + text_weight = 1.0 - vector_weight + + if self.vector_enabled and self.fts_enabled: + keyword_results = await self.keyword_search(query, candidates) + vector_results = await self.vector_search(query, candidates) + + if not keyword_results: + return vector_results[:limit] + elif not vector_results: + return keyword_results[:limit] + + merged = self._merge_hybrid_results( + vector=vector_results, + keyword=keyword_results, + vector_weight=vector_weight, + text_weight=text_weight, + ) + return merged[:limit] + elif self.vector_enabled: + return await self.vector_search(query, limit) + elif self.fts_enabled: + return await self.keyword_search(query, limit) + else: + return [] + + @staticmethod + def _merge_hybrid_results( + vector: list[FileChunk], + keyword: list[FileChunk], + vector_weight: float, + text_weight: float, + ) -> list[FileChunk]: + """Merge vector and keyword search results with weighted scoring.""" + merged: dict[str, FileChunk] = {} + + for result in vector: + v_score = result.scores.get("vector", 0) + result.scores["score"] = v_score * vector_weight + merged[result.merge_key] = result + + for result in keyword: + key = result.merge_key + k_score = result.scores.get("keyword", 0) + if key in merged: + merged[key].scores["score"] += k_score * text_weight + else: + result.scores["score"] = k_score * text_weight + merged[key] = result + + results = list(merged.values()) + results.sort(key=lambda r: r.score, reverse=True) + return results + + async def clear_all(self) -> None: + """Clear all indexed data from memory and disk.""" + self._chunks.clear() + self._files.clear() + await self._save_chunks() + await self._save_metadata() + self.logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'") diff --git a/reme_cli/component/file_watcher/__init__.py b/reme_cli/component/file_watcher/__init__.py new file mode 100644 index 00000000..c0485c71 --- /dev/null +++ b/reme_cli/component/file_watcher/__init__.py @@ -0,0 +1,13 @@ +"""File watcher module for monitoring file system changes. + +This module provides file watcher implementations for monitoring file changes +and updating memory stores accordingly. +""" + +from .base_file_watcher import BaseFileWatcher +from .md_file_watcher import MdFileWatcher + +__all__ = [ + "BaseFileWatcher", + "MdFileWatcher", +] \ No newline at end of file diff --git a/reme_cli/component/file_watcher/base_file_watcher.py b/reme_cli/component/file_watcher/base_file_watcher.py new file mode 100644 index 00000000..0919517a --- /dev/null +++ b/reme_cli/component/file_watcher/base_file_watcher.py @@ -0,0 +1,223 @@ +"""Base file watcher implementation. + +This module provides the base class for file watcher implementations +that monitor file system changes and trigger callbacks. +""" + +import asyncio +from abc import abstractmethod +from pathlib import Path +from typing import TYPE_CHECKING + +from watchfiles import Change, awatch + +from ..base_component import BaseComponent +from ..file_store import BaseFileStore +from ...enumeration import ComponentEnum + +if TYPE_CHECKING: + from ..application_context import ApplicationContext + + +class BaseFileWatcher(BaseComponent): + """Abstract base class for file watcher implementations. + + This base class provides file monitoring functionality that can be extended + to implement specific file monitoring requirements. + """ + + def __init__( + self, + watch_paths: list[str] | str, + suffix_filters: list[str] | None = None, + recursive: bool = False, + debounce: int = 2000, + chunk_tokens: int = 400, + chunk_overlap: int = 80, + file_store: str = "default", + rebuild_index_on_start: bool = True, + poll_delay_ms: int = 2000, + **kwargs, + ): + """Initialize the file watcher. + + Args: + watch_paths: Paths to watch for changes + suffix_filters: File suffix filters (e.g., ['.py', '.txt']) + recursive: Whether to watch directories recursively + debounce: Debounce time in milliseconds + chunk_tokens: Token size for chunking + chunk_overlap: Overlap size for chunks + file_store: Name of the file store component to use + rebuild_index_on_start: If True, clear all indexed data on start and rescan existing files. + If False, only monitor new changes without initialization. + poll_delay_ms: Polling delay in milliseconds. + **kwargs: Additional keyword arguments + """ + super().__init__(**kwargs) + self._file_store_name: str = file_store + self.file_store: BaseFileStore | None = None + self.watch_paths: list[str] = [watch_paths] if isinstance(watch_paths, str) else watch_paths + self.suffix_filters: list[str] = suffix_filters or [] + self.recursive: bool = recursive + self.debounce: int = debounce + self.chunk_tokens: int = chunk_tokens + self.chunk_overlap: int = chunk_overlap + self.rebuild_index_on_start: bool = rebuild_index_on_start + self.poll_delay_ms: int = poll_delay_ms + + self._stop_event = asyncio.Event() + self._watch_task: asyncio.Task | None = None + + async def _start(self, app_context: ApplicationContext | None = None): + """Initialize the file watcher and resolve file_store from app_context.""" + # Resolve file_store from app_context + if self._file_store_name: + stores = app_context.components.get(ComponentEnum.FILE_STORE, {}) + if self._file_store_name not in stores: + raise ValueError(f"File store '{self._file_store_name}' not found.") + store = stores[self._file_store_name] + if not isinstance(store, BaseFileStore): + raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}") + self.file_store = store + + # Start watching task + async def _initialize_and_watch(): + if self.rebuild_index_on_start and self.file_store: + await self.file_store.clear_all() + self.logger.info("Cleared all indexed data on start") + await self._scan_existing_files() + await self._watch_loop() + + self._stop_event.clear() + self._watch_task = asyncio.create_task(_initialize_and_watch()) + self.logger.info(f"Started watching: {self.watch_paths}") + + async def _close(self): + """Stop the file watcher and release resources.""" + # Signal stop and cancel watch task + self._stop_event.set() + if self._watch_task and not self._watch_task.done(): + self._watch_task.cancel() + try: + await self._watch_task + except asyncio.CancelledError: + pass + + self._watch_task = None + self._stop_event.clear() + self.file_store = None + self.logger.info("Stopped watching") + + def watch_filter(self, _change: Change, path: str) -> bool: + """Filter function for file watching.""" + if not self.suffix_filters: + return True + + for suffix in self.suffix_filters: + if path.endswith("." + suffix.strip(".")): + return True + + return False + + async def _scan_existing_files(self): + """Scan existing files matching watch criteria and trigger on_changes with Change.added.""" + if not self.file_store: + return + + existing_files: set[tuple[Change, str]] = set() + + for watch_path_str in self.watch_paths: + watch_path = Path(watch_path_str) + + if not watch_path.exists(): + self.logger.warning(f"Watch path does not exist: {watch_path}") + continue + + if watch_path.is_file(): + if self.watch_filter(Change.added, str(watch_path)): + existing_files.add((Change.added, str(watch_path))) + elif watch_path.is_dir(): + if self.recursive: + for file_path in watch_path.rglob("*"): + if file_path.is_file() and self.watch_filter(Change.added, str(file_path)): + existing_files.add((Change.added, str(file_path))) + else: + for file_path in watch_path.iterdir(): + if file_path.is_file() and self.watch_filter(Change.added, str(file_path)): + existing_files.add((Change.added, str(file_path))) + + if existing_files: + self.logger.info(f"[SCAN_ON_START] Found {len(existing_files)} existing files matching watch criteria") + await self.on_changes(existing_files) + self.logger.info(f"[SCAN_ON_START] Added {len(existing_files)} files to memory store") + else: + self.logger.info("[SCAN_ON_START] No existing files found matching watch criteria") + + files: list[str] = await self.file_store.list_files() + for file_path in files: + chunks = await self.file_store.get_file_chunks(file_path) + self.logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks") + + async def _interruptible_sleep(self, seconds: float): + """Sleep that can be interrupted by stop_event.""" + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=seconds) + except asyncio.TimeoutError: + pass + + async def _watch_loop(self): + """Core monitoring loop with auto-restart on failure.""" + if not self.watch_paths: + self.logger.warning("No watch paths specified") + return + + while not self._stop_event.is_set(): + valid_paths = [p for p in self.watch_paths if Path(p).exists()] + + if not valid_paths: + self.logger.warning("No valid watch paths exist, waiting 10 seconds before retry...") + await self._interruptible_sleep(10) + continue + + invalid_paths = set(self.watch_paths) - set(valid_paths) + if invalid_paths: + self.logger.warning(f"Skipping non-existent paths: {invalid_paths}") + + try: + self.logger.info(f"Starting watch on valid paths: {valid_paths}") + async for changes in awatch( + *valid_paths, + watch_filter=self.watch_filter, + 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 + + await self.on_changes(changes) + + except FileNotFoundError as e: + self.logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...") + if not self._stop_event.is_set(): + await self._interruptible_sleep(10) + + except Exception as e: + self.logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True) + if not self._stop_event.is_set(): + await self._interruptible_sleep(10) + + async def on_changes(self, changes: set[tuple[Change, str]]): + """Hook method to handle file changes.""" + await self._on_changes(changes) + self.logger.info(f"[{self.__class__.__name__}] on_changes: {changes}") + + @abstractmethod + async def _on_changes(self, changes: set[tuple[Change, str]]): + """Callback method to handle file changes. + + Args: + changes: Set of (Change, path) tuples representing file changes + """ diff --git a/reme_cli/component/file_watcher/md_file_watcher.py b/reme_cli/component/file_watcher/md_file_watcher.py new file mode 100644 index 00000000..7ee0fb18 --- /dev/null +++ b/reme_cli/component/file_watcher/md_file_watcher.py @@ -0,0 +1,89 @@ +"""Markdown file watcher for Markdown file synchronization. + +This module provides a file watcher that processes Markdown files +on any change, ensuring complete synchronization. +""" + +import asyncio +from pathlib import Path + +from watchfiles import Change + +from .base_file_watcher import BaseFileWatcher +from ...schema import FileMetadata +from ...utils import hash_text, chunk_markdown + + +class MdFileWatcher(BaseFileWatcher): + """Markdown file watcher implementation for Markdown file synchronization.""" + + def __init__(self, encoding: str = "utf-8", **kwargs): + """Initialize Markdown file watcher. + + Args: + encoding: File encoding (default: "utf-8") + **kwargs: Additional keyword arguments passed to BaseFileWatcher + """ + super().__init__(**kwargs) + self.encoding = encoding + + async def _on_changes(self, changes: set[tuple[Change, str]]): + """Handle file changes with full synchronization.""" + if not self.file_store: + self.logger.warning("File store not initialized, skipping changes") + return + + for change_type, path in changes: + try: + if change_type in [Change.added, Change.modified]: + file_meta = await self._build_file_metadata(path) + chunks = ( + chunk_markdown( + file_meta.content, + file_meta.path, + self.chunk_tokens, + self.chunk_overlap, + ) + or [] + ) + if chunks: + chunks = await self.file_store.get_chunk_embeddings(chunks) + file_meta.chunk_count = len(chunks) + + await self.file_store.delete_file(file_meta.path) + self.logger.info(f"delete_file {file_meta.path}") + + await self.file_store.upsert_file(file_meta, chunks) + self.logger.info(f"Upserted {file_meta.chunk_count} chunks for {file_meta.path}") + + elif change_type == Change.deleted: + await self.file_store.delete_file(path) + self.logger.info(f"Deleted {path}") + + else: + self.logger.warning(f"Unknown change type: {change_type}") + + self.logger.info(f"File {change_type} changed: {path}") + + except FileNotFoundError: + self.logger.warning(f"File not found: {path}, skipping") + except PermissionError: + self.logger.warning(f"Permission denied: {path}, skipping") + except Exception as e: + self.logger.error(f"Error processing {path}: {e}", exc_info=True) + + async def _build_file_metadata(self, path: str) -> FileMetadata: + """Build file metadata from path.""" + file_path = Path(path) + + def _read_file_sync(): + return file_path.stat(), file_path.read_text(encoding=self.encoding) + + stat, content = await asyncio.to_thread(_read_file_sync) + return FileMetadata( + hash=hash_text(content), + mtime_ms=stat.st_mtime * 1000, + size=stat.st_size, + path=str(file_path.absolute()), + content=content, + ) diff --git a/reme_cli/enumeration/component_enum.py b/reme_cli/enumeration/component_enum.py index 46f7ee3f..a5eded0c 100644 --- a/reme_cli/enumeration/component_enum.py +++ b/reme_cli/enumeration/component_enum.py @@ -6,6 +6,8 @@ class ComponentEnum(str, Enum): BASE = "base" + EMBEDDING_MODEL = "embedding_model" + AS_LLM = "as_llm" AS_LLM_FORMATTER = "as_llm_formatter" diff --git a/reme_cli/schema/__init__.py b/reme_cli/schema/__init__.py index ed4b2c12..3bab477e 100644 --- a/reme_cli/schema/__init__.py +++ b/reme_cli/schema/__init__.py @@ -1,7 +1,11 @@ from .application_config import ApplicationConfig from .base_node import BaseNode +from .file_chunk import FileChunk +from .file_metadata import FileMetadata __all__ = [ "ApplicationConfig", "BaseNode", + "FileChunk", + "FileMetadata", ] diff --git a/reme_cli/schema/file_chunk.py b/reme_cli/schema/file_chunk.py new file mode 100644 index 00000000..d5195c76 --- /dev/null +++ b/reme_cli/schema/file_chunk.py @@ -0,0 +1,24 @@ +"""File chunk schema.""" + +from pydantic import Field +from .base_node import BaseNode + + +class FileChunk(BaseNode): + """A chunk of file content with metadata.""" + + path: str = Field(..., description="File path relative to workspace") + start_line: int = Field(..., description="Starting line number in the source file") + end_line: int = Field(..., description="Ending line number in the source file") + hash: str = Field(..., description="Hash of the chunk content") + scores: dict[str, float] = Field(default_factory=dict, description="Search scores by type") + + @property + def score(self) -> float: + """Final score for search result.""" + return self.scores.get("score", 0.0) + + @property + def merge_key(self) -> str: + """Key for merging search results.""" + return f"{self.path}:{self.start_line}:{self.end_line}" \ No newline at end of file diff --git a/reme_cli/schema/file_metadata.py b/reme_cli/schema/file_metadata.py new file mode 100644 index 00000000..672e9c27 --- /dev/null +++ b/reme_cli/schema/file_metadata.py @@ -0,0 +1,15 @@ +"""File metadata schema.""" + +from pydantic import BaseModel, Field + + +class FileMetadata(BaseModel): + """File metadata with optional extended fields for various use cases.""" + + hash: str = Field(default=..., description="Hash of the file content") + mtime_ms: float = Field(default=..., description="Last modification time in milliseconds") + size: int = Field(default=..., description="File size in bytes") + path: str | None = Field(default=None, description="Relative path to the session file") + content: str | None = Field(default=None, description="Parsed content from the session file") + chunk_count: int | None = Field(default=None, description="Number of chunks in the file") + metadata: dict = Field(default_factory=dict, description="Additional metadata") diff --git a/reme_cli/utils/__init__.py b/reme_cli/utils/__init__.py index f2fdaa87..b0ca8dbe 100644 --- a/reme_cli/utils/__init__.py +++ b/reme_cli/utils/__init__.py @@ -1,9 +1,15 @@ +from .chunking_utils import chunk_markdown +from .common_utils import hash_text from .logger_utils import get_logger from .pydantic_config_parser import PydanticConfigParser +from .similarity_utils import batch_cosine_similarity from .singleton import singleton __all__ = [ + "chunk_markdown", + "hash_text", "get_logger", "PydanticConfigParser", + "batch_cosine_similarity", "singleton", ] diff --git a/reme_cli/utils/chunking_utils.py b/reme_cli/utils/chunking_utils.py new file mode 100644 index 00000000..bfd1f556 --- /dev/null +++ b/reme_cli/utils/chunking_utils.py @@ -0,0 +1,121 @@ +"""Chunking logic for Markdown files.""" + +from .common_utils import hash_text +from ..schema import FileChunk + + +def chunk_markdown( + text: str, + path: str, + chunk_tokens: int, + overlap: int, +) -> list[FileChunk]: + """ + Markdown chunking logic implemented based on the TypeScript version. + + Args: + text: Input text + path: File path + chunk_tokens: Maximum tokens per chunk + overlap: Overlap tokens between chunks + + Returns: + List of FileChunk objects + """ + if not text.strip(): + return [] + + lines = text.split("\n") + + # Convert tokens to characters (~1 token = 4 chars) + max_chars = max(32, chunk_tokens * 4) + overlap_chars = max(0, overlap * 4) + + chunks: list[FileChunk] = [] + + # Currently building chunk + current: list[dict] = [] # [{'line': str, 'line_no': int}] + current_chars = 0 + + def flush(): + """Add current chunk to results list""" + if not current: + return + + first_entry = current[0] + last_entry = current[-1] + + if not first_entry or not last_entry: + return + + chunk_text = "\n".join([entry["line"] for entry in current]) + start_line = first_entry["line_no"] + end_line = last_entry["line_no"] + + chunk_hash = hash_text(chunk_text) + + chunks.append( + FileChunk( + id=hash_text(f"{path}:{start_line}:{end_line}:{chunk_hash}:{len(chunks)}"), + path=path, + start_line=start_line, + end_line=end_line, + text=chunk_text, + hash=chunk_hash, + ), + ) + + def carry_overlap(): + """Keep overlapping part and clear the rest""" + nonlocal current, current_chars + + if overlap_chars <= 0 or not current: + current = [] + current_chars = 0 + return + + acc = 0 + kept = [] + + # Collect lines from the end until reaching overlap size + for j in range(len(current) - 1, -1, -1): + entry = current[j] + if not entry: + continue + + acc += len(entry["line"]) + 1 # +1 for newline + kept.insert(0, entry) # Insert at the beginning to maintain order + + if acc >= overlap_chars: + break + + current = kept + current_chars = sum(len(entry["line"]) + 1 for entry in kept) + + for i, line in enumerate(lines): + line_no = i + 1 + + # Split long lines into multiple segments + segments = [] + if not line: # Empty line + segments.append("") + 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]) + + for segment in segments: + line_size = len(segment) + 1 # +1 for newline + + # If adding current segment would exceed the limit, flush current chunk + if current_chars + line_size > max_chars and current: + flush() + carry_overlap() + + current.append({"line": segment, "line_no": line_no}) + current_chars += line_size + + # Process the final chunk + flush() + + return [c for c in chunks if c.text.strip()] diff --git a/reme_cli/utils/common_utils.py b/reme_cli/utils/common_utils.py new file mode 100644 index 00000000..c4b1dce3 --- /dev/null +++ b/reme_cli/utils/common_utils.py @@ -0,0 +1,14 @@ +import hashlib + + +def hash_text(text: str, encoding: str = "utf-8") -> str: + """Generate SHA-256 hash of text content. + + Args: + text: Input text to hash + encoding: Encoding of the text (default: "utf-8") + + Returns: + Hexadecimal representation of the SHA-256 hash + """ + return hashlib.sha256(text.encode(encoding)).hexdigest() diff --git a/reme_cli/utils/similarity_utils.py b/reme_cli/utils/similarity_utils.py new file mode 100644 index 00000000..4550404b --- /dev/null +++ b/reme_cli/utils/similarity_utils.py @@ -0,0 +1,52 @@ +import numpy as np + + +def cosine_similarity(vec1: list[float], vec2: list[float]) -> float: + """Calculate the cosine similarity between two numeric vectors.""" + if len(vec1) != len(vec2): + raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}") + + dot_product = sum(a * b for a, b in zip(vec1, vec2)) + magnitude1 = sum(a * a for a in vec1) ** 0.5 + magnitude2 = sum(b * b for b in vec2) ** 0.5 + + if magnitude1 == 0 or magnitude2 == 0: + return 0.0 + + return dot_product / (magnitude1 * magnitude2) + + +def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np.ndarray: + """Calculate cosine similarity matrix between two batches of vectors. + + Args: + nd_array1: Matrix of shape (batch_size1, emb_size) + nd_array2: Matrix of shape (batch_size2, emb_size) + + Returns: + Similarity matrix of shape (batch_size1, batch_size2) where + result[i, j] is the cosine similarity between nd_array1[i] and nd_array2[j] + + Raises: + ValueError: If embedding dimensions don't match + """ + if nd_array1.shape[1] != nd_array2.shape[1]: + raise ValueError(f"Embedding dimensions must match: {nd_array1.shape[1]} != {nd_array2.shape[1]}") + + # Compute dot products: (batch_size1, emb_size) @ (emb_size, batch_size2) + # Result shape: (batch_size1, batch_size2) + dot_products = np.dot(nd_array1, nd_array2.T) + + # Compute L2 norms for each vector + norms1 = np.linalg.norm(nd_array1, axis=1) # Shape: (batch_size1,) + norms2 = np.linalg.norm(nd_array2, axis=1) # Shape: (batch_size2,) + + # Compute outer product of norms: (batch_size1, 1) @ (1, batch_size2) + # Result shape: (batch_size1, batch_size2) + norm_products = np.outer(norms1, norms2) + + # Avoid division by zero + norm_products = np.where(norm_products == 0, 1e-10, norm_products) + + # Compute cosine similarities + return dot_products / norm_products