From 16d139583fbfd98aef27d3a4bf170263daa4631a Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 8 Apr 2026 16:19:27 +0800 Subject: [PATCH] init --- reme_cli/__init__.py | 6 + reme_cli/application.py | 20 + reme_cli/component/__init__.py | 5 + reme_cli/component/as_llm/__init__.py | 9 + .../component/as_llm_formatter/__init__.py | 9 + .../reme_openai_chat_formatter.py | 215 ++++ reme_cli/component/base_component.py | 25 + reme_cli/component/base_dict.py | 41 + reme_cli/component/file_store/__init__.py | 23 + .../component/file_store/base_file_store.py | 227 ++++ .../component/file_store/chroma_file_store.py | 633 ++++++++++++ .../component/file_store/local_file_store.py | 461 +++++++++ .../component/file_store/sqlite_file_store.py | 978 ++++++++++++++++++ reme_cli/component/file_watcher/__init__.py | 19 + .../file_watcher/base_file_watcher.py | 240 +++++ .../file_watcher/delta_file_watcher.py | 280 +++++ .../file_watcher/full_file_watcher.py | 79 ++ reme_cli/component/prompt_handler.py | 146 +++ reme_cli/component/registry_factory.py | 50 + reme_cli/config/__init__.py | 0 reme_cli/op/__init__.py | 0 reme_cli/op/base_op.py | 427 ++++++++ reme_cli/reme_cli.py | 78 ++ reme_cli/schema/__init__.py | 0 reme_cli/schema/service_config.py | 144 +++ reme_cli/utils/__init__.py | 0 reme_cli/utils/logger_utils.py | 59 ++ reme_cli/utils/singleton.py | 21 + 28 files changed, 4195 insertions(+) create mode 100644 reme_cli/__init__.py create mode 100644 reme_cli/application.py create mode 100644 reme_cli/component/__init__.py create mode 100644 reme_cli/component/as_llm/__init__.py create mode 100644 reme_cli/component/as_llm_formatter/__init__.py create mode 100644 reme_cli/component/as_llm_formatter/reme_openai_chat_formatter.py create mode 100644 reme_cli/component/base_component.py create mode 100644 reme_cli/component/base_dict.py create mode 100644 reme_cli/component/file_store/__init__.py create mode 100644 reme_cli/component/file_store/base_file_store.py create mode 100644 reme_cli/component/file_store/chroma_file_store.py create mode 100644 reme_cli/component/file_store/local_file_store.py create mode 100644 reme_cli/component/file_store/sqlite_file_store.py create mode 100644 reme_cli/component/file_watcher/__init__.py create mode 100644 reme_cli/component/file_watcher/base_file_watcher.py create mode 100644 reme_cli/component/file_watcher/delta_file_watcher.py create mode 100644 reme_cli/component/file_watcher/full_file_watcher.py create mode 100644 reme_cli/component/prompt_handler.py create mode 100644 reme_cli/component/registry_factory.py create mode 100644 reme_cli/config/__init__.py create mode 100644 reme_cli/op/__init__.py create mode 100644 reme_cli/op/base_op.py create mode 100644 reme_cli/reme_cli.py create mode 100644 reme_cli/schema/__init__.py create mode 100644 reme_cli/schema/service_config.py create mode 100644 reme_cli/utils/__init__.py create mode 100644 reme_cli/utils/logger_utils.py create mode 100644 reme_cli/utils/singleton.py diff --git a/reme_cli/__init__.py b/reme_cli/__init__.py new file mode 100644 index 00000000..bcde3ea8 --- /dev/null +++ b/reme_cli/__init__.py @@ -0,0 +1,6 @@ +"""ReMe CLI package.""" + +from reme_cli.component import BaseComponent +from reme_cli.application import Application + +__all__ = ["BaseComponent", "Application"] \ No newline at end of file diff --git a/reme_cli/application.py b/reme_cli/application.py new file mode 100644 index 00000000..a7522647 --- /dev/null +++ b/reme_cli/application.py @@ -0,0 +1,20 @@ +from reme_cli.component import BaseComponent + + +class Application(BaseComponent): + """Application component for managing the main application.""" + + def __init__(self) -> None: + super().__init__() + ... + + async def start(self) -> None: + """Start the application.""" + # 初始化llm formater + # + pass + + async def close(self) -> None: + """Close the application.""" + pass + diff --git a/reme_cli/component/__init__.py b/reme_cli/component/__init__.py new file mode 100644 index 00000000..695c24f7 --- /dev/null +++ b/reme_cli/component/__init__.py @@ -0,0 +1,5 @@ +from .base_component import BaseComponent + +__all__ = [ + "BaseComponent", +] \ No newline at end of file diff --git a/reme_cli/component/as_llm/__init__.py b/reme_cli/component/as_llm/__init__.py new file mode 100644 index 00000000..9cf527af --- /dev/null +++ b/reme_cli/component/as_llm/__init__.py @@ -0,0 +1,9 @@ +"""Module for registering AgentScope LLM models.""" + +from agentscope.model import DashScopeChatModel +from agentscope.model import OpenAIChatModel + +from ..registry_factory import R + +R.as_llms.register("openai")(OpenAIChatModel) +R.as_llms.register("dashscope")(DashScopeChatModel) diff --git a/reme_cli/component/as_llm_formatter/__init__.py b/reme_cli/component/as_llm_formatter/__init__.py new file mode 100644 index 00000000..1c7eee46 --- /dev/null +++ b/reme_cli/component/as_llm_formatter/__init__.py @@ -0,0 +1,9 @@ +"""Module for registering AgentScope LLM formatters.""" + +from agentscope.formatter import DashScopeChatFormatter + +from .reme_openai_chat_formatter import ReMeOpenAIChatFormatter +from ..registry_factory import R + +R.as_llm_formatters.register("openai")(ReMeOpenAIChatFormatter) +R.as_llm_formatters.register("dashscope")(DashScopeChatFormatter) diff --git a/reme_cli/component/as_llm_formatter/reme_openai_chat_formatter.py b/reme_cli/component/as_llm_formatter/reme_openai_chat_formatter.py new file mode 100644 index 00000000..d6807dfb --- /dev/null +++ b/reme_cli/component/as_llm_formatter/reme_openai_chat_formatter.py @@ -0,0 +1,215 @@ +"""ReMeOpenAIChatFormatter""" + +import json +from typing import Any + +from agentscope.formatter import OpenAIChatFormatter +from agentscope.formatter._openai_formatter import ( + _format_openai_image_block, + _to_openai_audio_data, +) +from agentscope.message import Msg, TextBlock, ImageBlock, URLSource +from loguru import logger + + +def _format_openai_video_block(video_block: dict) -> dict[str, Any]: + """Format a video block for OpenAI API. + + Args: + video_block: The video block to format. + + Returns: + A dictionary with video content in OpenAI format. + """ + source = video_block["source"] + if source["type"] == "url": + url = source["url"] + elif source["type"] == "base64": + data = source["data"] + media_type = source["media_type"] + url = f"data:{media_type};base64,{data}" + else: + raise ValueError(f"Unsupported video source type: {source['type']}") + + return { + "type": "video_url", + "video_url": { + "url": url, + }, + } + + +class ReMeOpenAIChatFormatter(OpenAIChatFormatter): + """ReMeOpenAIChatFormatter""" + + async def _format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into OpenAI API required format. + + Args: + msgs (`list[Msg]`): + The list of Msg objects to format. + + Returns: + `list[dict[str, Any]]`: + A list of dictionaries, where each dictionary has "name", + "role", and "content" keys. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + i = 0 + while i < len(msgs): + msg = msgs[i] + content_blocks = [] + tool_calls = [] + reasoning_content_blocks = [] + + for block in msg.get_content_blocks(): + typ = block.get("type") + if typ == "text": + content_blocks.append({**block}) + + elif typ == "thinking": + # Collect thinking blocks for reasoning_content field + # This is compatible with models like DeepSeek that support + # extended thinking via reasoning_content field + reasoning_content_blocks.append({**block}) + + elif typ == "tool_use": + tool_calls.append( + { + "id": block.get("id"), + "type": "function", + "function": { + "name": block.get("name"), + "arguments": json.dumps( + block.get("input", {}), + ensure_ascii=False, + ), + }, + }, + ) + + elif typ == "tool_result": + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block["output"]) + + messages.append( + { + "role": "tool", + "tool_call_id": block.get("id"), + "content": (textual_output), # type: ignore[arg-type] + "name": block.get("name"), + }, + ) + + # Then, handle the multimodal data if any + promoted_blocks: list = [] + for url, multimodal_block in multimodal_data: + if multimodal_block["type"] == "image" and self.promote_tool_result_images: + promoted_blocks.extend( + [ + TextBlock( + type="text", + text=f"\n- The image from '{url}': ", + ), + ImageBlock( + type="image", + source=URLSource( + type="url", + url=url, + ), + ), + ], + ) + + if promoted_blocks: + # Insert promoted blocks as new user message(s) + promoted_blocks = [ + TextBlock( + type="text", + text="The following are " + "the image contents from the tool " + f"result of '{block['name']}':", + ), + *promoted_blocks, + TextBlock( + type="text", + text="", + ), + ] + + msgs.insert( + i + 1, + Msg( + name="user", + content=promoted_blocks, + role="user", + ), + ) + + elif typ == "image": + content_blocks.append( + _format_openai_image_block( + block, # type: ignore[arg-type] + ), + ) + + elif typ == "audio": + # Filter out audio content when the multimodal model + # outputs both text and audio, to prevent errors in + # subsequent model calls + if msg.role == "assistant": + continue + input_audio = _to_openai_audio_data(block["source"]) + content_blocks.append( + { + "type": "input_audio", + "input_audio": input_audio, + }, + ) + + elif typ == "video": + # Filter out video content when the multimodal model + # outputs both text and video, to prevent errors in + # subsequent model calls + if msg.role == "assistant": + continue + content_blocks.append( + _format_openai_video_block(block), + ) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + typ, + ) + + msg_openai = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + + if tool_calls: + msg_openai["tool_calls"] = tool_calls + + # Add reasoning_content for thinking blocks (compatible with DeepSeek, etc.) + if reasoning_content_blocks: + reasoning_msg = "\n".join(reasoning.get("thinking", "") for reasoning in reasoning_content_blocks) + if reasoning_msg: + msg_openai["reasoning_content"] = reasoning_msg + + # When both content and tool_calls are None, skipped + if msg_openai["content"] or msg_openai.get("tool_calls"): + messages.append(msg_openai) + + # Move to next message + i += 1 + + return messages diff --git a/reme_cli/component/base_component.py b/reme_cli/component/base_component.py new file mode 100644 index 00000000..a21a3f80 --- /dev/null +++ b/reme_cli/component/base_component.py @@ -0,0 +1,25 @@ +"""Base class for components.""" + +from abc import ABC, abstractmethod + + +class BaseComponent(ABC): + """Base class supporting async start/close and async context management.""" + + @abstractmethod + async def start(self) -> None: + """Start the component asynchronously.""" + + @abstractmethod + async def close(self) -> None: + """Close the component asynchronously.""" + ... + + async def __aenter__(self) -> "BaseComponent": + """Enter async context manager.""" + await self.start() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + """Exit async context manager.""" + await self.close() diff --git a/reme_cli/component/base_dict.py b/reme_cli/component/base_dict.py new file mode 100644 index 00000000..b96b4c24 --- /dev/null +++ b/reme_cli/component/base_dict.py @@ -0,0 +1,41 @@ +"""Module providing a dictionary subclass with attribute-style access and pickling support.""" + +from typing import Generic, TypeVar + +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") + + +class BaseDict(dict, Generic[_KT, _VT]): + """A dictionary subclass that enables accessing and modifying keys as attributes.""" + + def __getattr__(self, name: str) -> _VT: + """Retrieve a dictionary item as an attribute.""" + try: + return self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __setattr__(self, name: str, value: _VT) -> None: + """Assign a value to a dictionary item using attribute syntax.""" + self[name] = value + + def __delattr__(self, name: str) -> None: + """Remove a dictionary item using attribute syntax.""" + try: + # Delete item from dict via key + del self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __getstate__(self) -> dict: + """Return the dictionary representation for pickling.""" + return dict(self) + + def __setstate__(self, state: dict) -> None: + """Restore the dictionary state from a pickled object.""" + self.update(state) + + def __reduce__(self): + """Define the reconstruction logic for pickling processes.""" + return self.__class__, (), self.__getstate__() diff --git a/reme_cli/component/file_store/__init__.py b/reme_cli/component/file_store/__init__.py new file mode 100644 index 00000000..1358df52 --- /dev/null +++ b/reme_cli/component/file_store/__init__.py @@ -0,0 +1,23 @@ +"""File store module for persistent memory management. + +This module provides storage backends for memory chunks and file metadata, +including SQLite-based, ChromaDB-based, and pure-Python local implementations +with vector and full-text search. +""" + +from .base_file_store import BaseFileStore +from .chroma_file_store import ChromaFileStore +from .local_file_store import LocalFileStore +from .sqlite_file_store import SqliteFileStore +from ..registry_factory import R + +__all__ = [ + "BaseFileStore", + "ChromaFileStore", + "LocalFileStore", + "SqliteFileStore", +] + +R.file_stores.register("sqlite")(SqliteFileStore) +R.file_stores.register("chroma")(ChromaFileStore) +R.file_stores.register("local")(LocalFileStore) diff --git a/reme_cli/component/file_store/base_file_store.py b/reme_cli/component/file_store/base_file_store.py new file mode 100644 index 00000000..6a6b0901 --- /dev/null +++ b/reme_cli/component/file_store/base_file_store.py @@ -0,0 +1,227 @@ +"""Base storage interface for file store.""" + +import re +from abc import ABC, abstractmethod +from pathlib import Path + +from ..embedding import BaseEmbeddingModel +from ..enumeration import MemorySource +from ..schema import FileMetadata, MemoryChunk, MemorySearchResult +from ..utils import get_logger + +logger = get_logger() + + +class BaseFileStore(ABC): + """Abstract base class for file storage backends.""" + + def __init__( + self, + store_name: str, + db_path: str | Path, + embedding_model: BaseEmbeddingModel | None = None, + vector_enabled: bool = False, + fts_enabled: bool = True, + **kwargs, + ): + """Initialize""" + # Validate store_name to prevent SQL injection + # Only allow alphanumeric characters and underscores + if not re.match(r"^[a-zA-Z0-9_]+$", store_name): + raise ValueError(f"Invalid '{store_name}'. Only alphanumeric characters and underscores are allowed.") + + # Ensure at least one search method is enabled + if not vector_enabled and not fts_enabled: + raise ValueError("At least one of vector_enabled or fts_enabled must be True.") + + # Ensure embedding_model is provided when vector search is enabled + if vector_enabled and embedding_model is None: + raise ValueError("embedding_model is required when vector_enabled is True.") + + self.store_name: str = store_name + self.db_path: Path = Path(db_path) + self.db_path.mkdir(parents=True, exist_ok=True) + self.embedding_model: BaseEmbeddingModel | None = embedding_model + self.vector_enabled: bool = vector_enabled + self.fts_enabled: bool = fts_enabled + self.kwargs: dict = kwargs + + @property + def embedding_dim(self) -> int: + """Get the embedding model's dimensionality.""" + if self.embedding_model is None: + return 1024 + return self.embedding_model.dimensions + + def _get_mock_embedding(self) -> list[float]: + """Generate a zero vector based on embedding model dimensions.""" + return [0.0] * self.embedding_dim + + def _disable_vector_search(self, reason: str = "embedding API error") -> None: + """Disable vector search and log a warning.""" + if self.vector_enabled: + logger.warning( + f"[{self.store_name}] Disabling vector search due to {reason}. " + "Falling back to full-text search only.", + ) + self.vector_enabled = False + + async def get_embedding(self, query: str, **kwargs) -> list[float]: + """Get embedding for a single query string.""" + if not self.vector_enabled: + return self._get_mock_embedding() + try: + return await self.embedding_model.get_embedding(query, **kwargs) + except Exception as e: + self._disable_vector_search(str(e)) + return self._get_mock_embedding() + + async def get_embeddings(self, queries: list[str], **kwargs) -> list[list[float]]: + """Get embeddings for a batch of query strings.""" + if not self.vector_enabled: + return [self._get_mock_embedding() for _ in queries] + try: + return await self.embedding_model.get_embeddings(queries, **kwargs) + except Exception as e: + self._disable_vector_search(str(e)) + return [self._get_mock_embedding() for _ in queries] + + async def get_chunk_embedding(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk: + """Generate and populate embedding field for a single MemoryChunk object.""" + if not self.vector_enabled: + chunk.embedding = self._get_mock_embedding() + return chunk + try: + return await self.embedding_model.get_chunk_embedding(chunk, **kwargs) + except Exception as e: + self._disable_vector_search(str(e)) + chunk.embedding = self._get_mock_embedding() + return chunk + + async def get_chunk_embeddings(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]: + """Generate and populate embedding fields for a batch of MemoryChunk objects.""" + if not self.vector_enabled: + mock_embedding = self._get_mock_embedding() + for chunk in chunks: + chunk.embedding = mock_embedding.copy() + return chunks + try: + return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs) + except Exception as e: + self._disable_vector_search(str(e)) + mock_embedding = self._get_mock_embedding() + for chunk in chunks: + chunk.embedding = mock_embedding.copy() + return chunks + + @abstractmethod + async def start(self): + """Initialize the storage backend.""" + + @abstractmethod + async def upsert_file(self, file_meta: FileMetadata, source: MemorySource, chunks: list[MemoryChunk]): + """Insert or update a file and its chunks.""" + + @abstractmethod + async def delete_file(self, path: str, source: MemorySource): + """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[MemoryChunk], source: MemorySource): + """Insert or update specific chunks without affecting other chunks.""" + + @abstractmethod + async def list_files(self, source: MemorySource) -> list[str]: + """List all indexed file paths for a source.""" + + @abstractmethod + async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None: + """Get full file metadata with statistics.""" + + @abstractmethod + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks. + + This is useful for incremental updates where only metadata needs to be updated + (e.g., after adding/removing chunks in delta file watcher). + + Args: + file_meta: Updated file metadata (hash, mtime_ms, size, chunk_count) + source: Memory source + """ + + @abstractmethod + async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]: + """Get all chunks for a file.""" + + @abstractmethod + async def vector_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform vector similarity search. + + Args: + query: Query embedding vector + limit: Maximum number of results + sources: Optional list of sources to filter + + Returns: + List of search results sorted by similarity + """ + + @abstractmethod + async def keyword_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform keyword/full-text search. + + Args: + query: Search query text + limit: Maximum number of results + sources: Optional list of sources to filter + + Returns: + List of search results sorted by relevance + """ + + @abstractmethod + async def hybrid_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[MemorySearchResult]: + """Perform hybrid search combining vector and keyword search. + + Args: + query: Search query text + limit: Maximum number of results + sources: Optional list of sources to filter + vector_weight: Weight for vector search results (0.0-1.0). + Keyword weight = 1.0 - vector_weight. + candidate_multiplier: Multiplier for candidate pool size. + candidates = limit * candidate_multiplier + + Returns: + List of search results sorted by combined relevance score + """ + + @abstractmethod + async def clear_all(self): + """Clear all indexed data.""" + + @abstractmethod + async def close(self): + """Close storage and release resources.""" diff --git a/reme_cli/component/file_store/chroma_file_store.py b/reme_cli/component/file_store/chroma_file_store.py new file mode 100644 index 00000000..6dda41b8 --- /dev/null +++ b/reme_cli/component/file_store/chroma_file_store.py @@ -0,0 +1,633 @@ +"""ChromaDB storage backend for file store.""" + +import json +import random +import time +from pathlib import Path + +from .base_file_store import BaseFileStore +from ..enumeration import MemorySource +from ..schema import FileMetadata, MemoryChunk, MemorySearchResult +from ..utils import get_logger + +logger = get_logger() + +try: + import chromadb + from chromadb.config import Settings + + _CHROMADB_IMPORT_ERROR: Exception | None = None +except Exception as e: + _CHROMADB_IMPORT_ERROR = e + chromadb = None + Settings = None + + +class ChromaFileStore(BaseFileStore): + """ChromaDB file storage with vector and full-text search. + + Inherits embedding methods from BaseFileStore: + - get_chunk_embedding / get_chunk_embeddings (async) + - get_embedding / get_embeddings (async) + + Provides ChromaDB-backed persistent storage with: + - Vector similarity search (native ChromaDB) + - Full-text search (via ChromaDB where_document filter) + - Efficient chunk and file metadata management + """ + + def __init__( + self, + **kwargs, + ): + if _CHROMADB_IMPORT_ERROR is not None: + raise _CHROMADB_IMPORT_ERROR + + super().__init__(**kwargs) + self.client: "chromadb.ClientAPI | None" = None + self.chunks_collection: "chromadb.Collection | None" = None + # Initialize metadata file path (db_path and store_name are set by base class) + self._metadata_file: Path = self.db_path.parent / f"{self.store_name}_file_metadata.json" + self._metadata_cache: dict[str, dict[str, FileMetadata]] = {} + + @property + def collection_name(self) -> str: + """Get the name of the ChromaDB collection for this store.""" + return f"chunks_{self.store_name}" + + async def _load_metadata(self) -> dict[str, dict[str, FileMetadata]]: + """Load file metadata from disk. + + Returns: + Dictionary mapping source -> path -> FileMetadata + """ + if not self._metadata_file.exists(): + return {} + + try: + data = self._metadata_file.read_text(encoding="utf-8") + metadata_dict = json.loads(data) + + # Convert dict to FileMetadata objects + result = {} + for source, files in metadata_dict.items(): + result[source] = {} + for path, meta in files.items(): + result[source][path] = FileMetadata(**meta) + + logger.debug(f"Loaded file metadata from {self._metadata_file}") + return result + except Exception as e: + logger.warning(f"Failed to load file metadata from {self._metadata_file}: {e}") + return {} + + async def _save_metadata(self, metadata: dict[str, dict[str, FileMetadata]]) -> None: + """Save file metadata to disk. + + Args: + metadata: Dictionary mapping source -> path -> FileMetadata + """ + try: + # Convert FileMetadata objects to dict for JSON serialization + metadata_dict = {} + for source, files in metadata.items(): + metadata_dict[source] = {} + for path, meta in files.items(): + metadata_dict[source][path] = { + "path": meta.path, + "hash": meta.hash, + "mtime_ms": meta.mtime_ms, + "size": meta.size, + "chunk_count": meta.chunk_count, + } + + data = json.dumps(metadata_dict, indent=2, ensure_ascii=False) + self._metadata_file.write_text(data, encoding="utf-8") + logger.debug(f"Saved file metadata to {self._metadata_file}") + except Exception as e: + logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}") + + async def start(self) -> None: + """Initialize ChromaDB client and collection.""" + if self.client is not None: + return + + # Initialize persistent ChromaDB client + self.client = chromadb.PersistentClient( + path=str(self.db_path), + settings=Settings( + anonymized_telemetry=False, + allow_reset=True, + ), + ) + + # Get or create the chunks collection + # ChromaDB uses cosine distance by default for similarity + self.chunks_collection = self.client.get_or_create_collection( + name=self.collection_name, + metadata={"hnsw:space": "cosine"}, + ) + + # Load metadata into cache + self._metadata_cache = await self._load_metadata() + + logger.info(f"ChromaDB initialized with collection: {self.collection_name}") + logger.info(f"File metadata will be persisted to: {self._metadata_file}") + + async def upsert_file( + self, + file_meta: FileMetadata, + source: MemorySource, + chunks: list[MemoryChunk], + ) -> None: + """Insert or update file and its chunks.""" + if not chunks: + return + + # Delete existing chunks for this file first + await self.delete_file(file_meta.path, source) + + # Batch generate embeddings for all chunks + # (base class returns mock embeddings when vector_enabled=False) + chunks = await self.get_chunk_embeddings(chunks) + + # Prepare data for ChromaDB batch upsert + ids = [] + documents = [] + embeddings = [] + metadatas = [] + + now = int(time.time() * 1000) + for chunk in chunks: + ids.append(chunk.id) + documents.append(chunk.text) + embeddings.append(chunk.embedding) + metadatas.append( + { + "path": file_meta.path, + "source": source.value, + "start_line": chunk.start_line, + "end_line": chunk.end_line, + "hash": chunk.hash, + "updated_at": now, + }, + ) + + # Batch upsert to ChromaDB (always pass embeddings to prevent default embedding function) + self.chunks_collection.upsert( + ids=ids, + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ) + + # Update file metadata in cache + if source.value not in self._metadata_cache: + self._metadata_cache[source.value] = {} + self._metadata_cache[source.value][file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=len(chunks), + ) + + async def delete_file(self, path: str, source: MemorySource) -> None: + """Delete file and all its chunks.""" + # Query for all chunks with this path and source + results = self.chunks_collection.get( + where={"$and": [{"path": path}, {"source": source.value}]}, + include=[], + ) + + if results["ids"]: + self.chunks_collection.delete( + ids=results["ids"], + ) + + # Remove from file metadata cache + if source.value in self._metadata_cache: + self._metadata_cache[source.value].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 + + self.chunks_collection.delete( + ids=chunk_ids, + ) + + # Update chunk count in file metadata cache + for source_meta in self._metadata_cache.values(): + if path in source_meta: + # Recalculate chunk count + results = self.chunks_collection.get( + where={"path": path}, + include=[], + ) + source_meta[path].chunk_count = len(results["ids"]) + break + + async def upsert_chunks( + self, + chunks: list[MemoryChunk], + source: MemorySource, + ) -> None: + """Insert or update specific chunks without affecting other chunks.""" + if not chunks: + return + + # Batch generate embeddings for all chunks + # (base class returns mock embeddings when vector_enabled=False) + chunks = await self.get_chunk_embeddings(chunks) + + ids = [] + documents = [] + embeddings = [] + metadatas = [] + + now = int(time.time() * 1000) + for chunk in chunks: + ids.append(chunk.id) + documents.append(chunk.text) + embeddings.append(chunk.embedding) + metadatas.append( + { + "path": chunk.path, + "source": source.value, + "start_line": chunk.start_line, + "end_line": chunk.end_line, + "hash": chunk.hash, + "updated_at": now, + }, + ) + + # Always pass embeddings to prevent default embedding function + self.chunks_collection.upsert( + ids=ids, + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ) + + async def list_files(self, source: MemorySource) -> list[str]: + """List all indexed files for a source.""" + if source.value not in self._metadata_cache: + return [] + return list(self._metadata_cache[source.value].keys()) + + async def get_file_metadata( + self, + path: str, + source: MemorySource, + ) -> FileMetadata | None: + """Get file metadata with chunk count.""" + if source.value not in self._metadata_cache: + return None + return self._metadata_cache[source.value].get(path) + + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks.""" + if source.value not in self._metadata_cache: + self._metadata_cache[source.value] = {} + + self._metadata_cache[source.value][file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + ) + + async def get_file_chunks( + self, + path: str, + source: MemorySource, + ) -> list[MemoryChunk]: + """Get all chunks for a file.""" + results = self.chunks_collection.get( + where={"$and": [{"path": path}, {"source": source.value}]}, + include=["documents", "embeddings", "metadatas"], + ) + + chunks = [] + for i, chunk_id in enumerate(results["ids"]): + metadata = results["metadatas"][i] + chunks.append( + MemoryChunk( + id=chunk_id, + path=metadata["path"], + source=MemorySource(metadata["source"]), + start_line=metadata["start_line"], + end_line=metadata["end_line"], + text=results["documents"][i], + hash=metadata["hash"], + embedding=results["embeddings"][i] if results["embeddings"] is not None else None, + ), + ) + + # Sort by start_line + chunks.sort(key=lambda c: c.start_line) + return chunks + + async def vector_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform vector similarity search.""" + if not self.vector_enabled or not query: + return [] + + # Get query embedding + query_embedding = await self.get_embedding(query) + if not query_embedding: + return [] + + # Build where filter for sources + where_filter = None + if sources: + if len(sources) == 1: + where_filter = {"source": sources[0].value} + else: + where_filter = {"source": {"$in": [s.value for s in sources]}} + + # Perform vector search + try: + results = self.chunks_collection.query( + query_embeddings=[query_embedding], + n_results=limit, + where=where_filter, + include=["documents", "metadatas", "distances"], + ) + except Exception as e: + logger.error(f"Vector search failed: {e}, falling back to random results") + # Fallback: get some documents without vector search and assign random scores + try: + fallback_results = self.chunks_collection.get( + where=where_filter, + limit=limit, + include=["documents", "metadatas"], + ) + search_results = [] + if fallback_results["ids"]: + for i, _ in enumerate(fallback_results["ids"]): + metadata = fallback_results["metadatas"][i] + search_results.append( + MemorySearchResult( + path=metadata["path"], + start_line=metadata["start_line"], + end_line=metadata["end_line"], + score=random.uniform(0.3, 0.7), # Random score in middle range + snippet=fallback_results["documents"][i], + source=MemorySource(metadata["source"]), + raw_metric=None, + ), + ) + return search_results + except Exception as fallback_e: + logger.error(f"Fallback search also failed: {fallback_e}") + return [] + + search_results = [] + if results["ids"] and results["ids"][0]: + for i, _ in enumerate(results["ids"][0]): + metadata = results["metadatas"][0][i] + distance = results["distances"][0][i] + + # Convert cosine distance to similarity score + # Cosine distance range is [0, 2], convert to [1, 0] score + score = max(0.0, 1.0 - distance / 2.0) + + search_results.append( + MemorySearchResult( + path=metadata["path"], + start_line=metadata["start_line"], + end_line=metadata["end_line"], + score=score, + snippet=results["documents"][0][i], + source=MemorySource(metadata["source"]), + raw_metric=distance, + ), + ) + + # Sort by score descending + search_results.sort(key=lambda r: r.score, reverse=True) + return search_results + + async def keyword_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform keyword/full-text search. + + ChromaDB supports where_document filter for text matching. + Note: ChromaDB's $contains is case-sensitive, so we generate multiple + case variants (original, lowercase, capitalized) for each word to + improve recall while maintaining case-insensitive scoring. + """ + if not self.fts_enabled or not query: + return [] + + # Normalize whitespace and split into words + words = query.split() + if not words: + return [] + + # Generate case variants for each word to handle case-sensitive $contains + # Include: original, lowercase, and capitalized forms + word_variants = set() + for word in words: + word_variants.add(word) # original + word_variants.add(word.lower()) # lowercase + word_variants.add(word.capitalize()) # Capitalized + word_variants.add(word.upper()) # UPPERCASE + word_variants_list = list(word_variants) + + # Build where filter for sources + where_filter = None + if sources: + if len(sources) == 1: + where_filter = {"source": sources[0].value} + else: + where_filter = {"source": {"$in": [s.value for s in sources]}} + + # ChromaDB where_document uses $contains for substring matching (case-sensitive) + # Use multiple case variants to improve recall + if len(word_variants_list) == 1: + where_document: dict = {"$contains": word_variants_list[0]} + else: + where_document = {"$or": [{"$contains": w} for w in word_variants_list]} + + # Get all matching documents + results = self.chunks_collection.get( + where=where_filter, + where_document=where_document, + include=["documents", "metadatas"], + ) + + search_results = [] + query_lower = query.lower() + words_lower = [w.lower() for w in words] # lowercase words for scoring + n_words = len(words) + + for i, _ in enumerate(results["ids"]): + metadata = results["metadatas"][i] + text = results["documents"][i] + text_lower = text.lower() + + # Calculate relevance score based on word matches + match_count = sum(1 for w in words_lower if w in text_lower) + base_score = match_count / n_words + + # Bonus for full phrase match (only applies to multi-word queries) + phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0 + # Scale base_score and add phrase bonus, max score is 1.0 + score = min(1.0, base_score + phrase_bonus) + + search_results.append( + MemorySearchResult( + path=metadata["path"], + start_line=metadata["start_line"], + end_line=metadata["end_line"], + score=score, + snippet=text, + source=MemorySource(metadata["source"]), + ), + ) + + # Sort by score descending and limit results + search_results.sort(key=lambda r: r.score, reverse=True) + return search_results[:limit] + + async def hybrid_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[MemorySearchResult]: + """Perform hybrid search combining vector and keyword search. + + Args: + query: Search query text + limit: Maximum number of results + sources: Optional list of sources to filter + vector_weight: Weight for vector search results (0.0-1.0). + Keyword weight = 1.0 - vector_weight. + candidate_multiplier: Multiplier for candidate pool size. + + Returns: + List of search results sorted by combined relevance score + """ + assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}" + + candidates = min(200, max(1, int(limit * candidate_multiplier))) + text_weight = 1.0 - vector_weight + + # Perform search based on enabled backends + if self.vector_enabled and self.fts_enabled: + keyword_results = await self.keyword_search(query, candidates, sources) + vector_results = await self.vector_search(query, candidates, sources) + + # Log original vector results + logger.info("\n=== Vector Search Results ===") + for i, r in enumerate(vector_results[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + # Log original keyword results + logger.info("\n=== Keyword Search Results ===") + for i, r in enumerate(keyword_results[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + if not keyword_results: + return vector_results[:limit] + elif not vector_results: + return keyword_results[:limit] + else: + merged = self._merge_hybrid_results( + vector=vector_results, + keyword=keyword_results, + vector_weight=vector_weight, + text_weight=text_weight, + ) + + # Log merged results + logger.info("\n=== Merged Hybrid Results ===") + for i, r in enumerate(merged[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + return merged[:limit] + elif self.vector_enabled: + vector_results = await self.vector_search(query, limit, sources) + return vector_results + elif self.fts_enabled: + keyword_results = await self.keyword_search(query, limit, sources) + return keyword_results + else: + return [] + + @staticmethod + def _merge_hybrid_results( + vector: list[MemorySearchResult], + keyword: list[MemorySearchResult], + vector_weight: float, + text_weight: float, + ) -> list[MemorySearchResult]: + """Merge vector and keyword search results with weighted scoring.""" + merged: dict[str, MemorySearchResult] = {} + + # Process vector results + for result in vector: + result.score = result.score * vector_weight + merged[result.merge_key] = result + + # Process keyword results + for result in keyword: + key = result.merge_key + if key in merged: + merged[key].score += result.score * text_weight + else: + result.score = result.score * text_weight + merged[key] = result + + # Sort by score and return + 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.""" + # Delete and recreate the collection + self.client.delete_collection( + name=self.collection_name, + ) + self.chunks_collection = self.client.get_or_create_collection( + name=self.collection_name, + metadata={"hnsw:space": "cosine"}, + ) + + # Clear file metadata cache and disk + self._metadata_cache = {} + await self._save_metadata({}) + + logger.info(f"Cleared all data from ChromaDB collection: {self.collection_name}") + + async def close(self) -> None: + """Close ChromaDB client and release resources.""" + # Persist metadata cache to disk before closing + if self._metadata_cache: + await self._save_metadata(self._metadata_cache) + + # ChromaDB PersistentClient handles persistence automatically + self.client = None + self.chunks_collection = None + await super().close() 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..38757d7a --- /dev/null +++ b/reme_cli/component/file_store/local_file_store.py @@ -0,0 +1,461 @@ +"""Pure-Python in-memory storage backend for file store, with JSON file persistence.""" + +import json +from pathlib import Path + +import numpy as np +from loguru import logger + +from .base_file_store import BaseFileStore +from ..enumeration import MemorySource +from ..schema import FileMetadata, MemoryChunk, MemorySearchResult +from ..utils.common_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, **kwargs): + super().__init__(**kwargs) + self._started: bool = False + # In-memory indexes + self._chunks: dict[str, MemoryChunk] = {} + self._files: dict[str, dict[str, FileMetadata]] = {} # source -> path -> meta + # Persistence paths (mirror ChromaFileStore convention) + self._chunks_file: Path = self.db_path / f"{self.store_name}_chunks.jsonl" + self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json" + + # ------------------------------------------------------------------ + # 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="utf-8") + self._chunks = {} + for line in data.strip().split("\n"): + if not line: + continue + rec = json.loads(line) + chunk = MemoryChunk.model_validate(rec) + self._chunks[chunk.id] = chunk + logger.debug(f"Loaded {len(self._chunks)} chunks from {self._chunks_file}") + except Exception as e: + logger.warning(f"Failed to load chunks from {self._chunks_file}: {e}") + + async def _save_chunks(self) -> None: + """Persist chunks to JSONL file.""" + try: + lines = [] + for chunk in self._chunks.values(): + chunk_dict = chunk.model_dump(mode="json") + lines.append(json.dumps(chunk_dict, ensure_ascii=False)) + data = "\n".join(lines) + self._chunks_file.write_text(data, encoding="utf-8") + logger.debug(f"Saved {len(self._chunks)} chunks to {self._chunks_file}") + except Exception as e: + logger.error(f"Failed to save chunks to {self._chunks_file}: {e}") + + 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="utf-8") + raw: dict = json.loads(data) + self._files = { + source: {path: FileMetadata(**meta) for path, meta in files.items()} for source, files in raw.items() + } + logger.debug(f"Loaded file metadata from {self._metadata_file}") + except Exception as e: + logger.warning(f"Failed to load file metadata from {self._metadata_file}: {e}") + + async def _save_metadata(self) -> None: + """Persist file metadata to JSON file.""" + try: + raw: dict = {} + for source, files in self._files.items(): + raw[source] = { + path: { + "path": meta.path, + "hash": meta.hash, + "mtime_ms": meta.mtime_ms, + "size": meta.size, + "chunk_count": meta.chunk_count, + } + for path, meta in files.items() + } + data = json.dumps(raw, indent=2, ensure_ascii=False) + self._metadata_file.write_text(data, encoding="utf-8") + logger.debug(f"Saved file metadata to {self._metadata_file}") + except Exception as e: + logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}") + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def start(self) -> None: + """Load persisted data into memory.""" + if self._started: + return + self._started = True + await self._load_metadata() + await self._load_chunks() + logger.info( + f"LocalFileStore '{self.store_name}' ready: " + f"{len(self._chunks)} chunks, metadata at {self._metadata_file}", + ) + + 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() + self._started = False + + # ------------------------------------------------------------------ + # Write operations + # ------------------------------------------------------------------ + + async def upsert_file( + self, + file_meta: FileMetadata, + source: MemorySource, + chunks: list[MemoryChunk], + ) -> None: + """Insert or update file and its chunks.""" + if not chunks: + return + + # Remove existing chunks for this file/source first + await self.delete_file(file_meta.path, source) + + # 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 + + if source.value not in self._files: + self._files[source.value] = {} + self._files[source.value][file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=len(chunks), + ) + + async def delete_file(self, path: str, source: MemorySource) -> None: + """Delete file and all its chunks.""" + to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path and chunk.source == source] + for cid in to_delete: + del self._chunks[cid] + + if source.value in self._files: + self._files[source.value].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 (per source) + for source_key, source_meta in self._files.items(): + if path in source_meta: + source_meta[path].chunk_count = sum( + 1 for chunk in self._chunks.values() if chunk.path == path and chunk.source.value == source_key + ) + + async def upsert_chunks( + self, + chunks: list[MemoryChunk], + source: MemorySource, + ) -> 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, source: MemorySource) -> list[str]: + """List all indexed files for a source.""" + return list(self._files.get(source.value, {}).keys()) + + async def get_file_metadata( + self, + path: str, + source: MemorySource, + ) -> FileMetadata | None: + """Get file metadata.""" + return self._files.get(source.value, {}).get(path) + + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks.""" + if source.value not in self._files: + self._files[source.value] = {} + + self._files[source.value][file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + ) + + async def get_file_chunks( + self, + path: str, + source: MemorySource, + ) -> list[MemoryChunk]: + """Get all chunks for a file, sorted by start_line.""" + chunks = [chunk for chunk in self._chunks.values() if chunk.path == path and chunk.source == source] + chunks.sort(key=lambda c: c.start_line) + return chunks + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + async def vector_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """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 = [ + chunk for chunk in self._chunks.values() if (not sources or chunk.source in sources) and chunk.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) + logger.warning( + f"Chunk embedding dimension {emb_len} < expected {expected_dim}, " + f"padded with zeros (chunk_id={chunk.id})", + ) + else: + emb = emb[:expected_dim] + logger.warning( + f"Chunk embedding dimension {emb_len} > expected {expected_dim}, " + f"truncated to {expected_dim} (chunk_id={chunk.id})", + ) + valid_embeddings.append(emb) + + # Build embedding matrix and compute similarities in batch + query_array = np.array([query_embedding]) # Shape: (1, emb_size) + chunk_embeddings = np.array(valid_embeddings) # Shape: (n, emb_size) + similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] # Shape: (n,) + + # Build results + results = [ + MemorySearchResult( + path=chunk.path, + start_line=chunk.start_line, + end_line=chunk.end_line, + score=float(similarity), + snippet=chunk.text, + source=chunk.source, + raw_metric=1.0 - float(similarity), + ) + for chunk, similarity in zip(candidates, similarities) + ] + + results.sort(key=lambda r: r.score, reverse=True) + return results[:limit] + + async def keyword_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """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(): + if sources and chunk.source not in sources: + continue + + 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 + # Bonus for full phrase match (multi-word queries only) + 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( + MemorySearchResult( + path=chunk.path, + start_line=chunk.start_line, + end_line=chunk.end_line, + score=score, + snippet=chunk.text, + source=chunk.source, + ), + ) + + results.sort(key=lambda r: r.score, reverse=True) + return results[:limit] + + async def hybrid_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[MemorySearchResult]: + """Perform hybrid search combining vector and keyword search. + + Args: + query: Search query text + limit: Maximum number of results + sources: Optional list of sources to filter + vector_weight: Weight for vector search results (0.0-1.0). + Keyword weight = 1.0 - vector_weight. + candidate_multiplier: Multiplier for candidate pool size. + + Returns: + List of search results sorted by combined relevance score + """ + assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}" + + 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, sources) + vector_results = await self.vector_search(query, candidates, sources) + + logger.info("\n=== Vector Search Results ===") + for i, r in enumerate(vector_results[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + logger.info("\n=== Keyword Search Results ===") + for i, r in enumerate(keyword_results[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + if not keyword_results: + return vector_results[:limit] + elif not vector_results: + return keyword_results[:limit] + else: + merged = self._merge_hybrid_results( + vector=vector_results, + keyword=keyword_results, + vector_weight=vector_weight, + text_weight=text_weight, + ) + + logger.info("\n=== Merged Hybrid Results ===") + for i, r in enumerate(merged[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + return merged[:limit] + elif self.vector_enabled: + return await self.vector_search(query, limit, sources) + elif self.fts_enabled: + return await self.keyword_search(query, limit, sources) + else: + return [] + + @staticmethod + def _merge_hybrid_results( + vector: list[MemorySearchResult], + keyword: list[MemorySearchResult], + vector_weight: float, + text_weight: float, + ) -> list[MemorySearchResult]: + """Merge vector and keyword search results with weighted scoring.""" + merged: dict[str, MemorySearchResult] = {} + + for result in vector: + result.metadata["_weighted_score"] = result.score * vector_weight + merged[result.merge_key] = result + + for result in keyword: + key = result.merge_key + if key in merged: + merged[key].metadata["_weighted_score"] += result.score * text_weight + else: + result.metadata["_weighted_score"] = result.score * text_weight + merged[key] = result + + results = list(merged.values()) + for r in results: + r.score = r.metadata.pop("_weighted_score") + + results.sort(key=lambda r: r.score, reverse=True) + return results + + 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() + logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'") diff --git a/reme_cli/component/file_store/sqlite_file_store.py b/reme_cli/component/file_store/sqlite_file_store.py new file mode 100644 index 00000000..0a494d2c --- /dev/null +++ b/reme_cli/component/file_store/sqlite_file_store.py @@ -0,0 +1,978 @@ +"""SQLite storage backend for file store.""" + +import json + +import struct +import time + +from loguru import logger + +from .base_file_store import BaseFileStore +from ..enumeration import MemorySource +from ..schema import FileMetadata, MemoryChunk, MemorySearchResult + + +class SqliteFileStore(BaseFileStore): + """SQLite file storage with vector and full-text search. + + Inherits embedding methods from BaseFileStore: + - get_chunk_embedding / get_chunk_embeddings (async) + - get_chunk_embedding_sync / get_chunk_embeddings_sync (sync) + - get_embedding / get_embeddings (async) + + Provides SQLite-backed persistent storage with: + - Vector similarity search (via sqlite-vec extension) + - Full-text search (via FTS5) + - Efficient chunk and file metadata management + """ + + def __init__(self, vec_ext_path: str = "", **kwargs): + super().__init__(**kwargs) + self.vec_ext_path = vec_ext_path + import sqlite3 + + self.conn: sqlite3.Connection | None = None + + @property + def vector_table_name(self) -> str: + """Get the name of the vector table for this store.""" + return f"chunks_vec_{self.store_name}" + + @property + def fts_table_name(self) -> str: + """Get the name of the FTS table for this store.""" + return f"chunks_fts_{self.store_name}" + + @property + def chunks_table_name(self) -> str: + """Get the name of the chunks table for this store.""" + return f"chunks_{self.store_name}" + + @property + def files_table_name(self) -> str: + """Get the name of the files table for this store.""" + return f"files_{self.store_name}" + + @staticmethod + def vector_to_blob(embedding: list[float]) -> bytes: + """Convert vector to binary blob for sqlite-vec.""" + return struct.pack(f"{len(embedding)}f", *embedding) + + async def start(self) -> None: + """Initialize database and load extensions.""" + if self.conn is not None: + return + import sqlite3 + + self.conn = sqlite3.connect(self.db_path / "reme.db", check_same_thread=False) + + # Only load sqlite-vec extension if vector search is enabled + if self.vector_enabled: + logger.warning( + "On macOS systems with version 14 or earlier, " + "loading the sqlite-vec vector extension carries a risk of crashes or hangs.", + ) + + self.conn.enable_load_extension(True) + + # Load sqlite-vec extension + if self.vec_ext_path: + try: + self.conn.load_extension(self.vec_ext_path) + logger.info(f"Loaded sqlite-vec: {self.vec_ext_path}") + except Exception as e: + logger.warning(f"Failed to load sqlite-vec: {e}") + + else: + try: + import sqlite_vec + + ext_path = sqlite_vec.loadable_path() + self.conn.load_extension(ext_path) + logger.info(f"Loaded sqlite-vec from package: {ext_path}") + + except Exception as e: + logger.warning(f"Failed to load sqlite-vec from package: {e}") + # Fallback: try common extension names + for name in ["vec0", "sqlite_vec", "vector0"]: + try: + self.conn.load_extension(name) + logger.info(f"Loaded sqlite-vec: {name}") + break + except Exception: + pass + + self.conn.enable_load_extension(False) + else: + logger.info("Vector search disabled, skipping sqlite-vec extension loading") + + await self._create_tables() + + async def _create_tables(self) -> None: + """Create database schema.""" + cursor = self.conn.cursor() + try: + # Files + cursor.execute( + f""" + CREATE TABLE IF NOT EXISTS {self.files_table_name} ( + path TEXT, + source TEXT, + hash TEXT, + mtime REAL, + size INTEGER, + PRIMARY KEY (path, source) + ) + """, + ) + + # Chunks + cursor.execute( + f""" + CREATE TABLE IF NOT EXISTS {self.chunks_table_name} ( + id TEXT PRIMARY KEY, + path TEXT, + source TEXT, + start_line INTEGER, + end_line INTEGER, + hash TEXT, + text TEXT, + embedding TEXT, + updated_at INTEGER + ) + """, + ) + + # Vector table (sqlite-vec) + if self.vector_enabled: + cursor.execute( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self.vector_table_name} USING vec0( + id TEXT PRIMARY KEY, + embedding FLOAT[{self.embedding_dim}] + ) + """, + ) + logger.info(f"Created vector table (dims={self.embedding_dim})") + + # FTS table + if self.fts_enabled: + cursor.execute( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self.fts_table_name} USING fts5( + text, + id UNINDEXED, + path UNINDEXED, + source UNINDEXED, + start_line UNINDEXED, + end_line UNINDEXED, + tokenize='trigram' + ) + """, + ) + logger.info("Created FTS5 table with trigram tokenizer") + + self.conn.commit() + except Exception as e: + logger.error(f"Failed to create tables: {e}") + raise + finally: + cursor.close() + + async def upsert_file(self, file_meta: FileMetadata, source: MemorySource, chunks: list[MemoryChunk]): + """Insert or update file and its chunks.""" + cursor = self.conn.cursor() + + try: + cursor.execute("BEGIN") + + # Insert file + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.files_table_name} (path, source, hash, mtime, size) + VALUES (?, ?, ?, ?, ?) + """, + (file_meta.path, source.value, file_meta.hash, file_meta.mtime_ms, file_meta.size), + ) + + # Insert chunks + now = int(time.time() * 1000) + for chunk in chunks: + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.chunks_table_name} ( + id, path, source, start_line, end_line, + hash, text, embedding, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + chunk.id, + file_meta.path, + source.value, + chunk.start_line, + chunk.end_line, + chunk.hash, + chunk.text, + json.dumps(chunk.embedding) if chunk.embedding else None, + now, + ), + ) + + # Insert vector (vec0 doesn't support OR REPLACE, use DELETE + INSERT) + if self.vector_enabled: + if not chunk.embedding: + logger.warning( + f"Chunk {chunk.id} missing embedding for vector insert, skipping vector indexing", + ) + else: + # Delete existing vector first + cursor.execute( + f"DELETE FROM {self.vector_table_name} WHERE id = ?", + (chunk.id,), + ) + # Then insert new vector + cursor.execute( + f"INSERT INTO {self.vector_table_name} (id, embedding) VALUES (?, ?)", + (chunk.id, self.vector_to_blob(chunk.embedding)), + ) + + # Insert FTS + if self.fts_enabled: + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.fts_table_name} ( + text, id, path, source, start_line, end_line + ) VALUES (?, ?, ?, ?, ?, ?) + """, + ( + chunk.text, + chunk.id, + file_meta.path, + source.value, + chunk.start_line, + chunk.end_line, + ), + ) + + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + logger.error(f"Failed to upsert file {file_meta.path}: {e}") + raise + finally: + cursor.close() + + async def delete_file(self, path: str, source: MemorySource): + """Delete file and all its chunks.""" + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + + # Get chunk IDs for vector deletion + cursor.execute( + f"SELECT id FROM {self.chunks_table_name} WHERE path = ? AND source = ?", + (path, source.value), + ) + chunk_ids = [row[0] for row in cursor.fetchall()] + + # Delete vectors + if self.vector_enabled and chunk_ids: + for chunk_id in chunk_ids: + cursor.execute( + f"DELETE FROM {self.vector_table_name} WHERE id = ?", + (chunk_id,), + ) + + # Delete FTS entries + if self.fts_enabled: + cursor.execute( + f"DELETE FROM {self.fts_table_name} WHERE path = ? AND source = ?", + (path, source.value), + ) + + # Delete chunks and file + cursor.execute( + f"DELETE FROM {self.chunks_table_name} WHERE path = ? AND source = ?", + (path, source.value), + ) + cursor.execute( + f"DELETE FROM {self.files_table_name} WHERE path = ? AND source = ?", + (path, source.value), + ) + + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + logger.error(f"Failed to delete file {path}: {e}") + raise + finally: + cursor.close() + + async def delete_file_chunks(self, path: str, chunk_ids: list[str]): + """Delete specific chunks for a file.""" + if not chunk_ids: + return + + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + + # Delete vectors + if self.vector_enabled: + for chunk_id in chunk_ids: + cursor.execute( + f"DELETE FROM {self.vector_table_name} WHERE id = ?", + (chunk_id,), + ) + + # Delete FTS entries + if self.fts_enabled: + placeholders = ",".join("?" * len(chunk_ids)) + cursor.execute( + f"DELETE FROM {self.fts_table_name} WHERE id IN ({placeholders})", + chunk_ids, + ) + + # Delete chunks + placeholders = ",".join("?" * len(chunk_ids)) + cursor.execute( + f"DELETE FROM {self.chunks_table_name} WHERE id IN ({placeholders})", + chunk_ids, + ) + + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + logger.error(f"Failed to delete chunks for {path}: {e}") + raise + finally: + cursor.close() + + async def upsert_chunks(self, chunks: list[MemoryChunk], source: MemorySource): + """Insert or update specific chunks without affecting other chunks.""" + if not chunks: + return + + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + + now = int(time.time() * 1000) + for chunk in chunks: + # Insert/update chunk + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.chunks_table_name} ( + id, path, source, start_line, end_line, + hash, text, embedding, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + chunk.id, + chunk.path, + source.value, + chunk.start_line, + chunk.end_line, + chunk.hash, + chunk.text, + json.dumps(chunk.embedding) if chunk.embedding else None, + now, + ), + ) + + # Insert/update vector (vec0 doesn't support OR REPLACE, use DELETE + INSERT) + if self.vector_enabled: + if not chunk.embedding: + logger.warning( + f"Chunk {chunk.id} missing embedding for vector insert, skipping vector indexing", + ) + else: + # Delete existing vector first + cursor.execute( + f"DELETE FROM {self.vector_table_name} WHERE id = ?", + (chunk.id,), + ) + # Then insert new vector + cursor.execute( + f"INSERT INTO {self.vector_table_name} (id, embedding) VALUES (?, ?)", + (chunk.id, self.vector_to_blob(chunk.embedding)), + ) + + # Insert/update FTS + if self.fts_enabled: + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.fts_table_name} ( + text, id, path, source, start_line, end_line + ) VALUES (?, ?, ?, ?, ?, ?) + """, + ( + chunk.text, + chunk.id, + chunk.path, + source.value, + chunk.start_line, + chunk.end_line, + ), + ) + + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + logger.error(f"Failed to upsert chunks: {e}") + raise + finally: + cursor.close() + + async def list_files(self, source: MemorySource) -> list[str]: + """List all indexed files.""" + cursor = self.conn.cursor() + try: + cursor.execute(f"SELECT path FROM {self.files_table_name} WHERE source = ?", (source.value,)) + paths = [row[0] for row in cursor.fetchall()] + return paths + except Exception as e: + logger.error(f"Failed to list files: {e}") + raise + finally: + cursor.close() + + async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None: + """Get file metadata with chunk count.""" + cursor = self.conn.cursor() + try: + cursor.execute( + f"SELECT hash, mtime, size FROM {self.files_table_name} WHERE path = ? AND source = ?", + (path, source.value), + ) + row = cursor.fetchone() + if not row: + return None + + hash_val, mtime, size = row + cursor.execute( + f"SELECT COUNT(*) FROM {self.chunks_table_name} WHERE path = ? AND source = ?", + (path, source.value), + ) + chunk_count = cursor.fetchone()[0] + + return FileMetadata( + hash=hash_val, + mtime_ms=mtime, + size=size, + path=path, + chunk_count=chunk_count, + ) + except Exception as e: + logger.error(f"Failed to get file metadata for {path}: {e}") + raise + finally: + cursor.close() + + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata without affecting chunks.""" + cursor = self.conn.cursor() + try: + cursor.execute( + f""" + INSERT OR REPLACE INTO {self.files_table_name} (path, source, hash, mtime, size) + VALUES (?, ?, ?, ?, ?) + """, + (file_meta.path, source.value, file_meta.hash, file_meta.mtime_ms, file_meta.size), + ) + self.conn.commit() + except Exception as e: + logger.error(f"Failed to update file metadata for {file_meta.path}: {e}") + raise + finally: + cursor.close() + + async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]: + """Get all chunks for a file.""" + cursor = self.conn.cursor() + try: + cursor.execute( + f""" + SELECT id, path, source, start_line, end_line, text, hash, embedding + FROM {self.chunks_table_name} WHERE path = ? AND source = ? + ORDER BY start_line + """, + (path, source.value), + ) + + chunks = [] + for row in cursor.fetchall(): + chunk_id, path_val, source_val, start, end, text, hash_val, emb_str = row + # Parse embedding from JSON string + embedding = None + if emb_str: + try: + embedding = json.loads(emb_str) + except (json.JSONDecodeError, TypeError): + embedding = None + + chunks.append( + MemoryChunk( + id=chunk_id, + path=path_val, + source=MemorySource(source_val), + start_line=start, + end_line=end, + text=text, + hash=hash_val, + embedding=embedding, + ), + ) + + return chunks + except Exception as e: + logger.error(f"Failed to get file chunks for {path}: {e}") + raise + finally: + cursor.close() + + async def vector_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform vector similarity search.""" + if not self.vector_enabled or not query: + return [] + + # Get query embedding + query_embedding = await self.get_embedding(query) + if not query_embedding: + return [] + + cursor = self.conn.cursor() + source_filter = "" + params: list = [] + if sources: + placeholders = ",".join("?" * len(sources)) + source_filter = f" AND c.source IN ({placeholders})" + params = [s.value for s in sources] + + try: + query_blob = self.vector_to_blob(query_embedding) + + # Correct SQLite-vec syntax for vector search with limit + # vec0 requires 'k = ?' constraint for knn queries + query_sql = f""" + SELECT c.id, c.path, c.start_line, c.end_line, c.source, c.text, v.distance + FROM {self.vector_table_name} v + JOIN {self.chunks_table_name} c ON v.id = c.id + WHERE v.embedding MATCH ? + AND k = ? + """ + query_params: list = [query_blob, limit] + + # Add source filter if specified + if source_filter: + query_sql += source_filter + query_params.extend(params) + + # Order by distance (k constraint already limits results) + query_sql += " ORDER BY v.distance" + + cursor.execute(query_sql, query_params) + + results = [] + for _, path, start, end, src, text, dist in cursor.fetchall(): + # Convert L2 distance to similarity score + # For normalized vectors, L2 distance range is [0, 2] + # Map to [1, 0] score range (higher score = more similar) + score = max(0.0, 1.0 - dist / 2.0) + snippet = text + results.append( + MemorySearchResult( + path=path, + start_line=start, + end_line=end, + score=score, + snippet=snippet, + source=MemorySource(src), + raw_metric=dist, + ), + ) + + results.sort(key=lambda r: r.score, reverse=True) + return results + except Exception as e: + logger.error(f"Vector search failed: {e}") + return [] + finally: + cursor.close() + + @staticmethod + def _sanitize_fts_query(query: str) -> str: + """Sanitize query string for FTS5 search. + + Removes or escapes special characters that have special meaning in FTS5: + - * (prefix match) + - ? (not used in FTS5, but can cause issues) + - " (phrase search, needs escaping) + - : (column filter) + - ^ (start of line anchor, not standard FTS5) + - ' (single quote, causes syntax errors) + - ` (backtick, can cause issues) + - | (pipe, OR operator) + - + (plus, can be used for required terms) + - - (minus, NOT operator) + - = (equals, can cause issues) + - < > (angle brackets, comparison operators) + - ! (exclamation, NOT operator variant) + - @ # $ % & (other special chars) + - "\" + - / (slash, can interfere) + - ; (semicolon, statement separator) + - , (comma, can interfere with phrase parsing) + + Args: + query: Raw query string + + Returns: + Sanitized query string safe for FTS5 + """ + if not query: + return "" + + # Remove FTS5 special characters that we don't want users to use + # Keep only alphanumeric, spaces, periods, and underscores + special_chars = [ + "*", + "?", + ":", + "^", + "(", + ")", + "[", + "]", + "{", + "}", + "'", + '"', + "`", + "|", + "+", + "-", + "=", + "<", + ">", + "!", + "@", + "#", + "$", + "%", + "&", + "\\", + "/", + ";", + ",", + ] + cleaned = query + for char in special_chars: + cleaned = cleaned.replace(char, " ") + + # Normalize whitespace + cleaned = " ".join(cleaned.split()) + + return cleaned + + async def keyword_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform keyword search. + + Strategy: + - FTS5 trigram (fast path): used when ALL terms >= 3 chars (trigram minimum). + - LIKE (universal fallback): used when any term < 3 chars, covering CJK + short words, single/double-char queries, and mixed-length queries. + """ + if not self.fts_enabled: + return [] + + cleaned = self._sanitize_fts_query(query) + if not cleaned: + return [] + + words = cleaned.split() + if not words: + return [] + + # FTS5 trigram requires every term >= 3 characters + if all(len(w) >= 3 for w in words): + results = await self._fts_trigram_search(words, limit, sources) + if results: + return results + + # Universal fallback: LIKE-based substring search + return await self._like_search(cleaned, words, limit, sources) + + async def _fts_trigram_search( + self, + words: list[str], + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """FTS5 trigram search. All terms must be >= 3 characters.""" + escaped_words = [w.replace('"', '""') for w in words] + fts_query = " OR ".join(escaped_words) + + cursor = self.conn.cursor() + source_filter = "" + params: list = [fts_query] + if sources: + placeholders = ",".join("?" * len(sources)) + source_filter = f" AND fts.source IN ({placeholders})" + params.extend([s.value for s in sources]) + params.append(limit) + + try: + cursor.execute( + f""" + SELECT fts.id, fts.path, fts.start_line, fts.end_line, + fts.source, fts.text, rank + FROM {self.fts_table_name} fts + WHERE fts.text MATCH ?{source_filter} + ORDER BY rank + LIMIT ? + """, + params, + ) + + results = [] + for _, path, start, end, src, text, rank in cursor.fetchall(): + score = max(0.0, 1.0 / (1.0 + abs(rank))) + results.append( + MemorySearchResult( + path=path, + start_line=start, + end_line=end, + score=score, + snippet=text, + source=MemorySource(src), + raw_metric=rank, + ), + ) + results.sort(key=lambda r: r.score, reverse=True) + return results + except Exception as e: + logger.error(f"FTS trigram search failed: {e}") + return [] + finally: + cursor.close() + + async def _like_search( + self, + phrase: str, + words: list[str], + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """LIKE-based substring search with Python-side relevance scoring. + + Handles any term length and all languages (CJK, Latin, etc.). + Scores results by: word-match ratio + full-phrase bonus. + """ + cursor = self.conn.cursor() + try: + # Build OR conditions: match any individual word + like_clauses = [] + params: list = [] + for word in words: + like_clauses.append("c.text LIKE ?") + params.append(f"%{word}%") + + where_clause = " OR ".join(like_clauses) + + source_filter = "" + if sources: + placeholders = ",".join("?" * len(sources)) + source_filter = f" AND c.source IN ({placeholders})" + params.extend([s.value for s in sources]) + + # Fetch extra candidates for re-ranking in Python + fetch_limit = min(limit * 3, 200) + params.append(fetch_limit) + + cursor.execute( + f""" + SELECT c.id, c.path, c.start_line, c.end_line, c.source, c.text + FROM {self.chunks_table_name} c + WHERE ({where_clause}){source_filter} + LIMIT ? + """, + params, + ) + + results = [] + phrase_lower = phrase.lower() + words_lower = [w.lower() for w in words] + n_words = len(words) + + for _, path, start, end, src, text in cursor.fetchall(): + text_lower = text.lower() + + # Base score: proportion of query words found in text + match_count = sum(1 for w in words_lower if w in text_lower) + base_score = match_count / n_words + + # Bonus: full phrase appears as contiguous substring + phrase_bonus = 0.2 if n_words > 1 and phrase_lower in text_lower else 0.0 + + score = min(1.0, base_score * 0.8 + phrase_bonus) + + results.append( + MemorySearchResult( + path=path, + start_line=start, + end_line=end, + score=score, + snippet=text, + source=MemorySource(src), + ), + ) + + # Sort by score descending, return top `limit` + results.sort(key=lambda r: r.score, reverse=True) + return results[:limit] + except Exception as e: + logger.error(f"LIKE search failed: {e}") + return [] + finally: + cursor.close() + + async def hybrid_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[MemorySearchResult]: + """Perform hybrid search combining vector and keyword search. + + Args: + query: Search query text + limit: Maximum number of results + sources: Optional list of sources to filter + vector_weight: Weight for vector search results (0.0-1.0). + Keyword weight = 1.0 - vector_weight. + candidate_multiplier: Multiplier for candidate pool size. + + Returns: + List of search results sorted by combined relevance score + """ + assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}" + + candidates = min(200, max(1, int(limit * candidate_multiplier))) + text_weight = 1.0 - vector_weight + + # Perform search based on enabled backends + if self.vector_enabled and self.fts_enabled: + keyword_results = await self.keyword_search(query, candidates, sources) + vector_results = await self.vector_search(query, candidates, sources) + + # Log original vector results + logger.info("\n=== Vector Search Results ===") + for i, r in enumerate(vector_results[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + # Log original keyword results + logger.info("\n=== Keyword Search Results ===") + for i, r in enumerate(keyword_results[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + if not keyword_results: + return vector_results[:limit] + elif not vector_results: + return keyword_results[:limit] + else: + merged = self._merge_hybrid_results( + vector=vector_results, + keyword=keyword_results, + vector_weight=vector_weight, + text_weight=text_weight, + ) + + # Log merged results + logger.info("\n=== Merged Hybrid Results ===") + for i, r in enumerate(merged[:10], 1): + snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet + logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}") + + return merged[:limit] + elif self.vector_enabled: + vector_results = await self.vector_search(query, limit, sources) + return vector_results + elif self.fts_enabled: + keyword_results = await self.keyword_search(query, limit, sources) + return keyword_results + else: + return [] + + @staticmethod + def _merge_hybrid_results( + vector: list[MemorySearchResult], + keyword: list[MemorySearchResult], + vector_weight: float, + text_weight: float, + ) -> list[MemorySearchResult]: + """Merge vector and keyword search results with weighted scoring.""" + merged: dict[str, MemorySearchResult] = {} + + # Process vector results + for result in vector: + result.score = result.score * vector_weight + merged[result.merge_key] = result + + # Process keyword results + for result in keyword: + key = result.merge_key + if key in merged: + merged[key].score += result.score * text_weight + else: + result.score = result.score * text_weight + merged[key] = result + + # Sort by score and return + results = list(merged.values()) + results.sort(key=lambda r: r.score, reverse=True) + return results + + async def clear_all(self): + """Clear all indexed data.""" + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + + cursor.execute(f"DELETE FROM {self.files_table_name}") + cursor.execute(f"DELETE FROM {self.chunks_table_name}") + + if self.vector_enabled: + cursor.execute(f"DELETE FROM {self.vector_table_name}") + + if self.fts_enabled: + cursor.execute(f"DELETE FROM {self.fts_table_name}") + + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + logger.error(f"Failed to clear all data: {e}") + raise + finally: + cursor.close() + + async def close(self): + """Close database connection.""" + if self.conn: + self.conn.close() + self.conn = None + await super().close() diff --git a/reme_cli/component/file_watcher/__init__.py b/reme_cli/component/file_watcher/__init__.py new file mode 100644 index 00000000..020d0cfa --- /dev/null +++ b/reme_cli/component/file_watcher/__init__.py @@ -0,0 +1,19 @@ +"""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 .delta_file_watcher import DeltaFileWatcher +from .full_file_watcher import FullFileWatcher +from ..registry_factory import R + +__all__ = [ + "BaseFileWatcher", + "DeltaFileWatcher", + "FullFileWatcher", +] + +R.file_watchers.register("full")(FullFileWatcher) +R.file_watchers.register("delta")(DeltaFileWatcher) 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..b4644dd7 --- /dev/null +++ b/reme_cli/component/file_watcher/base_file_watcher.py @@ -0,0 +1,240 @@ +"""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 collections.abc import Coroutine +from pathlib import Path +from typing import Any, Callable + +from loguru import logger +from watchfiles import awatch, Change + +from ..enumeration import MemorySource +from ..file_store import BaseFileStore + + +class BaseFileWatcher: + """ + Minimal file watcher base class + + This base class provides basic 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: BaseFileStore | None = None, + callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None, + 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: File store instance + callback: Callback function for changes + 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. If > 300ms, force_polling will be enabled automatically. + **kwargs: Additional keyword arguments + """ + 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.file_store: BaseFileStore = file_store + self.callback = callback + self.rebuild_index_on_start: bool = rebuild_index_on_start + self.poll_delay_ms: int = poll_delay_ms + self.kwargs: dict = kwargs + + self._stop_event = asyncio.Event() + self._watch_task: asyncio.Task | None = None + self._running = False + + async def start(self): + """Start the file watcher""" + if self._running: + return + + self._running = True + + async def _initialize_and_watch(): + if self.rebuild_index_on_start: + await self.file_store.clear_all() + logger.info("Cleared all indexed data on start") + await self._scan_existing_files() + await self._watch_loop() + + self._watch_task = asyncio.create_task(_initialize_and_watch()) + logger.info(f"Started watching: {self.watch_paths}") + + async def close(self): + """Stop the file watcher""" + if not self._running: + return + + self._stop_event.set() + if self._watch_task: + await self._watch_task + self._running = False + logger.info("Stopped watching") + + def watch_filter(self, _change: Change, path: str) -> bool: + """Filter function for file watching.""" + # If no suffix filters are specified, watch all files + if not self.suffix_filters: + return True + + # Check if the file has one of the allowed suffixes + 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""" + 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(): + logger.warning(f"Watch path does not exist: {watch_path}") + continue + + if watch_path.is_file(): + # Single file + if self.watch_filter(Change.added, str(watch_path)): + existing_files.add((Change.added, str(watch_path))) + elif watch_path.is_dir(): + # Directory + if self.recursive: + # Recursive scan + 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: + # Non-recursive scan (only immediate children) + 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: + logger.info(f"[SCAN_ON_START] Found {len(existing_files)} existing files matching watch criteria") + await self.on_changes(existing_files) + logger.info(f"[SCAN_ON_START] Added {len(existing_files)} files to memory store") + else: + logger.info("[SCAN_ON_START] No existing files found matching watch criteria") + + if self.file_store is not None: + files: list[str] = await self.file_store.list_files(MemorySource.MEMORY) + for file_path in files: + chunks = await self.file_store.get_file_chunks(file_path, MemorySource.MEMORY) + logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks") + + async def _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 # Normal timeout, continue + + async def _watch_loop(self): + """Core monitoring loop with auto-restart on failure""" + if not self.watch_paths: + logger.warning("No watch paths specified") + return + + while not self._stop_event.is_set(): + # Filter out non-existent paths before each watch attempt + valid_paths = [p for p in self.watch_paths if Path(p).exists()] + + if not valid_paths: + 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: + logger.warning(f"Skipping non-existent paths: {invalid_paths}") + + try: + 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: + # Watch path was deleted during monitoring + 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: + # Log other exceptions and restart + 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]]): + """Callback method to handle file changes""" + + async def on_changes(self, changes: set[tuple[Change, str]]): + """Hook method to handle file changes""" + if self.callback: + result = self.callback(changes) + if asyncio.iscoroutine(result): + await result + else: + await self._on_changes(changes) + logger.info(f"[{self.__class__.__name__}] on_changes: {changes}") + + def is_running(self) -> bool: + """Check if the watcher is running""" + return self._running + + async def add_path(self, path: str): + """Dynamically add a path to monitor""" + if path not in self.watch_paths: + self.watch_paths.append(path) + if self._running: + await self.close() + await self.start() + + async def remove_path(self, path: str): + """Remove a monitored path""" + if path in self.watch_paths: + self.watch_paths.remove(path) + if self._running: + await self.close() + await self.start() diff --git a/reme_cli/component/file_watcher/delta_file_watcher.py b/reme_cli/component/file_watcher/delta_file_watcher.py new file mode 100644 index 00000000..6148bd07 --- /dev/null +++ b/reme_cli/component/file_watcher/delta_file_watcher.py @@ -0,0 +1,280 @@ +"""Delta file watcher for incremental file synchronization. + +This module provides a file watcher that detects append-only changes +and only processes newly added content, avoiding redundant operations. +""" + +import asyncio +import os + +from loguru import logger +from watchfiles import Change + +from .base_file_watcher import BaseFileWatcher +from ..enumeration import MemorySource +from ..schema import FileMetadata, MemoryChunk +from ..utils import chunk_markdown, hash_text + + +class DeltaFileWatcher(BaseFileWatcher): + """Delta file watcher implementation for incremental synchronization. + + This watcher detects append-only changes (e.g., log files) and only processes + the newly added content, avoiding redundant embedding requests for unchanged content. + + Strategy: + - Detect if file is append-only (new lines added at end) + - Find the safe cutoff point (considering chunk overlap) + - Only re-chunk and embed content from cutoff to end + - Delete affected old chunks and insert new chunks + """ + + def __init__(self, overlap_lines: int = 2, **kwargs): + """ + Initialize delta file watcher. + + Args: + chunk_tokens: Maximum tokens per chunk + chunk_overlap: Overlap tokens between chunks + """ + super().__init__(**kwargs) + self.overlap_lines = overlap_lines + self.dirty = False + + @staticmethod + async def _build_file_metadata(path: str) -> FileMetadata: + """Build file metadata from filesystem.""" + + def _read_file_sync(): + stat_t = os.stat(path) + with open(path, "r", encoding="utf-8") as f: + content_t = f.read() + return stat_t, content_t + + 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=path, + content=content, + ) + + def _find_cutoff_line( + self, + old_chunks: list[MemoryChunk], + old_file_meta: FileMetadata, + new_file_meta: FileMetadata, + ) -> int | None: + """Find the safe cutoff line for incremental update. + + Uses a heuristic approach: if file size increased and hash changed, + we verify by comparing content. For true append-only files (like logs), + the old content should be a prefix of new content. + + Args: + old_chunks: Existing chunks sorted by start_line + old_file_meta: Previous file metadata + new_file_meta: Current file metadata (with content) + + Returns: + Cutoff line number (1-indexed), or None if not append-only + """ + if not old_chunks: + return None + + # File shrunk - definitely not append-only + if new_file_meta.size < old_file_meta.size: + logger.debug("File shrunk, not append-only") + return None + + # File didn't grow much - might be a modification + size_growth = new_file_meta.size - old_file_meta.size + if size_growth < 10: # Less than 10 bytes growth + logger.debug("Minimal size growth, treating as modification") + return None + + # Verify append-only by checking if old content is prefix + # We need to read old file content from chunks + old_chunks_sorted = sorted(old_chunks, key=lambda c: c.start_line) + + # Simple heuristic: check if first few chunks' content matches + # This avoids reconstructing full old content + new_lines = new_file_meta.content.split("\n") + + # Sample check: verify first chunk still matches + first_chunk = old_chunks_sorted[0] + first_chunk_lines = first_chunk.text.split("\n") + new_first_lines = new_lines[first_chunk.start_line - 1 : first_chunk.end_line] + + # Compare (allowing for minor whitespace differences at boundaries) + if len(first_chunk_lines) > 0 and len(new_first_lines) > 0: + # Check if most of the lines match + matches = sum(1 for old, new in zip(first_chunk_lines, new_first_lines) if old == new) + if matches < len(first_chunk_lines) * 0.8: # Less than 80% match + logger.debug("First chunk content changed, not append-only") + return None + + # File appears to be append-only + # Find the last chunk and set cutoff considering overlap + last_chunk = max(old_chunks_sorted, key=lambda c: c.end_line) + cutoff_line = max(1, last_chunk.end_line - self.overlap_lines) + + logger.debug( + f"Append-only detected: size {old_file_meta.size} -> {new_file_meta.size}, " + f"cutoff at line {cutoff_line}", + ) + + return cutoff_line + + @staticmethod + def _extract_content_from_line(content: str, start_line: int) -> str: + """Extract content starting from a specific line number.""" + lines = content.split("\n") + if start_line <= 1: + return content + if start_line > len(lines): + return "" + # start_line is 1-indexed, array is 0-indexed + return "\n".join(lines[start_line - 1 :]) + + async def _on_changes(self, changes: set[tuple[Change, str]]): + """Handle file changes with incremental synchronization.""" + self.dirty = True + + for change_type, path in changes: + if change_type == Change.added: + # New file: process everything + file_meta = await self._build_file_metadata(path) + chunks = ( + chunk_markdown( + file_meta.content, + file_meta.path, + MemorySource.MEMORY, + 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.upsert_file(file_meta, MemorySource.MEMORY, chunks) + logger.info(f"File added: {path} ({len(chunks)} chunks)") + else: + logger.warning(f"No chunks generated for new file {path}") + + elif change_type == Change.modified: + # Get existing data + old_chunks = await self.file_store.get_file_chunks(path, MemorySource.MEMORY) + old_file_meta = await self.file_store.get_file_metadata(path, MemorySource.MEMORY) + + # Read new file + file_meta = await self._build_file_metadata(path) + + # If no old chunks, fallback to full update + if not old_chunks or not old_file_meta: + logger.debug(f"No existing chunks for {path}, doing full update") + chunks = ( + chunk_markdown( + file_meta.content, + file_meta.path, + MemorySource.MEMORY, + 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(path, MemorySource.MEMORY) + await self.file_store.upsert_file( + file_meta, + MemorySource.MEMORY, + chunks, + ) + logger.info(f"File modified (full): {path} ({len(chunks)} chunks)") + continue + + # Check if append-only and find cutoff line + old_chunks_sorted = sorted(old_chunks, key=lambda c: c.start_line) + cutoff_line = self._find_cutoff_line(old_chunks_sorted, old_file_meta, file_meta) + + if cutoff_line is None: + # Not append-only, do full update + logger.debug(f"File {path} has modifications, doing full update") + chunks = ( + chunk_markdown( + file_meta.content, + file_meta.path, + MemorySource.MEMORY, + 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(path, MemorySource.MEMORY) + await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) + logger.info(f"File modified (full): {path} ({len(chunks)} chunks)") + else: + # Append-only: incremental update + new_content_part = self._extract_content_from_line(file_meta.content, cutoff_line) + + new_chunks = ( + chunk_markdown( + new_content_part, + file_meta.path, + MemorySource.MEMORY, + self.chunk_tokens, + self.chunk_overlap, + ) + or [] + ) + + if not new_chunks: + logger.debug(f"No new chunks for {path}, skipping") + continue + + for idx, chunk in enumerate(new_chunks): + chunk.start_line += cutoff_line - 1 + chunk.end_line += cutoff_line - 1 + chunk.id = hash_text( + f"{chunk.source}:{chunk.path}:{chunk.start_line}:" f"{chunk.end_line}:{chunk.hash}:{idx}", + ) + + new_chunks = await self.file_store.get_chunk_embeddings(new_chunks) + + chunks_to_delete = [c.id for c in old_chunks_sorted if c.start_line >= cutoff_line] + + # Apply incremental updates + if chunks_to_delete: + await self.file_store.delete_file_chunks(path, chunks_to_delete) + + if new_chunks: + await self.file_store.upsert_chunks(new_chunks, MemorySource.MEMORY) + + # Update file metadata to reflect the changes + # Calculate new chunk count: old chunks - deleted + new chunks + new_chunk_count = len(old_chunks) - len(chunks_to_delete) + len(new_chunks) + file_meta.chunk_count = new_chunk_count + await self.file_store.update_file_metadata(file_meta, MemorySource.MEMORY) + + logger.info( + f"File modified (incremental): {path} " + f"(cutoff: line {cutoff_line}, " + f"+{len(new_chunks)} chunks, -{len(chunks_to_delete)} chunks)", + ) + + elif change_type == Change.deleted: + await self.file_store.delete_file(path, MemorySource.MEMORY) + logger.info(f"File deleted: {path}") + + else: + logger.warning(f"Unknown change type: {change_type}") + + self.dirty = False diff --git a/reme_cli/component/file_watcher/full_file_watcher.py b/reme_cli/component/file_watcher/full_file_watcher.py new file mode 100644 index 00000000..5b1852b9 --- /dev/null +++ b/reme_cli/component/file_watcher/full_file_watcher.py @@ -0,0 +1,79 @@ +"""Full file watcher for complete file synchronization. + +This module provides a file watcher that processes entire files +on any change, ensuring complete synchronization. +""" + +import asyncio +from pathlib import Path + +from loguru import logger +from watchfiles import Change + +from .base_file_watcher import BaseFileWatcher +from ..enumeration import MemorySource +from ..schema import FileMetadata +from ..utils import chunk_markdown, hash_text + + +class FullFileWatcher(BaseFileWatcher): + """Full file watcher implementation for full synchronization""" + + def __init__(self, **kwargs): + """ + Initialize full file watcher""" + super().__init__(**kwargs) + self.dirty = False + + @staticmethod + async def _build_file_metadata(path: str) -> FileMetadata: + file_path = Path(path) + + def _read_file_sync(): + return file_path.stat(), file_path.read_text(encoding="utf-8") + + 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, + ) + + async def _on_changes(self, changes: set[tuple[Change, str]]): + """Handle file changes with full synchronization""" + self.dirty = True + + for change_type, path in changes: + 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, + MemorySource.MEMORY, + 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, MemorySource.MEMORY) + logger.info(f"delete_file {file_meta.path}") + + await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks) + logger.info(f"Upserted {file_meta.chunk_count} chunks for {file_meta.path}") + + elif change_type == Change.deleted: + await self.file_store.delete_file(path, MemorySource.MEMORY) + logger.info(f"Deleted {path}") + + else: + logger.warning(f"Unknown change type: {change_type}") + + logger.info(f"File {change_type} changed: {path}") + self.dirty = False diff --git a/reme_cli/component/prompt_handler.py b/reme_cli/component/prompt_handler.py new file mode 100644 index 00000000..22e1b095 --- /dev/null +++ b/reme_cli/component/prompt_handler.py @@ -0,0 +1,146 @@ +"""Module for managing and formatting prompt templates from files or dictionaries.""" + +import json +from pathlib import Path +from string import Formatter +from typing import Any, Dict, Optional, Union + +import yaml +from loguru import logger + +from .base_dict import BaseDict + + +class PromptHandler(BaseDict): + """A context-aware handler for loading, retrieving, and formatting prompt templates.""" + + def __init__(self, language: str = "", **kwargs): + super().__init__(**kwargs) + # Use object.__setattr__ to avoid storing 'language' in the dict + object.__setattr__(self, "language", language.strip()) + + def load_prompt_by_file( + self, + prompt_file_path: Optional[Union[Path, str]] = None, + overwrite: bool = True, + ) -> "PromptHandler": + """Load prompt configurations from a YAML or JSON file.""" + if prompt_file_path is None: + return self + + if isinstance(prompt_file_path, str): + prompt_file_path = Path(prompt_file_path) + + if not prompt_file_path.exists(): + return self + + suffix = prompt_file_path.suffix.lower() + + with prompt_file_path.open(encoding="utf-8") as f: + if suffix in [".yaml", ".yml"]: + prompt_dict = yaml.safe_load(f) + elif suffix == ".json": + prompt_dict = json.load(f) + else: + raise ValueError(f"Unsupported file format: {suffix}") + + self.load_prompt_dict(prompt_dict, overwrite=overwrite) + return self + + def load_prompt_dict( + self, + prompt_dict: Optional[Dict[str, Any]] = None, + overwrite: bool = True, + ) -> "PromptHandler": + """Merge a dictionary of prompt strings into the current context.""" + if not prompt_dict: + return self + + for key, value in prompt_dict.items(): + if not isinstance(value, str): + continue + if key in self: + if overwrite: + logger.warning(f"Overwriting prompt '{key}'") + self[key] = value + else: + self[key] = value + + return self + + def get_prompt(self, prompt_name: str, fallback_to_base: bool = True) -> str: + """Retrieve a prompt by name with automatic language suffix handling.""" + if self.language and not prompt_name.endswith(f"_{self.language}"): + key_with_lang = f"{prompt_name}_{self.language}" + if key_with_lang in self: + return self[key_with_lang].strip() + + if prompt_name in self: + return self[prompt_name].strip() + + if fallback_to_base and self.language and prompt_name.endswith(f"_{self.language}"): + base_name = prompt_name[: -(len(self.language) + 1)] + if base_name in self: + return self[base_name].strip() + + raise KeyError(f"Prompt '{prompt_name}' not found. Available: {list(self.keys())[:10]}") + + def has_prompt(self, prompt_name: str) -> bool: + """Check if a prompt exists.""" + try: + self.get_prompt(prompt_name) + return True + except KeyError: + return False + + def list_prompts(self, language_filter: Optional[str] = None) -> list[str]: + """List all available prompt names.""" + if language_filter is None: + return list(self.keys()) + suffix = f"_{language_filter.strip()}" + return [key for key in self.keys() if key.endswith(suffix)] + + @staticmethod + def _extract_format_fields(template: str) -> set[str]: + """Extract all format field names from a template string.""" + return {field_name for _, field_name, _, _ in Formatter().parse(template) if field_name is not None} + + @staticmethod + def _filter_conditional_lines(prompt: str, flags: Dict[str, bool]) -> str: + """Filter lines based on boolean flags.""" + filtered_lines = [] + for line in prompt.split("\n"): + matched_flag = None + for flag_name in flags: + if line.startswith(f"[{flag_name}]"): + matched_flag = flag_name + break + if matched_flag is None: + filtered_lines.append(line) + elif flags[matched_flag]: + filtered_lines.append(line[len(f"[{matched_flag}]") :]) + return "\n".join(filtered_lines) + + def prompt_format(self, prompt_name: str, validate: bool = True, **kwargs) -> str: + """Format a prompt with conditional line filtering and variable substitution.""" + prompt = self.get_prompt(prompt_name) + + flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + format_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + + if flag_kwargs: + prompt = self._filter_conditional_lines(prompt, flag_kwargs) + + if validate: + required_fields = self._extract_format_fields(prompt) + missing_fields = required_fields - set(format_kwargs.keys()) + if missing_fields: + raise ValueError(f"Missing format variables for '{prompt_name}': {sorted(missing_fields)}") + + if format_kwargs: + prompt = prompt.format(**format_kwargs) + + return prompt.strip() + + def __repr__(self) -> str: + return f"PromptHandler(language='{self.language}', num_prompts={len(self)})" diff --git a/reme_cli/component/registry_factory.py b/reme_cli/component/registry_factory.py new file mode 100644 index 00000000..921f27a8 --- /dev/null +++ b/reme_cli/component/registry_factory.py @@ -0,0 +1,50 @@ +"""Module providing a registry class for managing class-to-name mappings via decorators.""" + +import inspect +from typing import Callable, TypeVar + +from .base_dict import BaseDict +from .utils import singleton + +T = TypeVar("T") + + +class Registry(BaseDict): + """A registry container that uses decorators to map and store class references.""" + + def register(self, name: str | type = "") -> Callable[[type[T]], type[T]] | type[T]: + """Return a decorator that registers a class under a specific name in the registry.""" + if inspect.isclass(name): + self[name.__name__] = name + return name + + else: + + def decorator(cls): + key: str = name if isinstance(name, str) and name else cls.__name__ + self[key] = cls + return cls + + return decorator + + +@singleton +class RegistryFactory: + """A factory class for creating registries.""" + + def __init__(self): + self.llms = Registry() + self.as_llms = Registry() + self.as_llm_formatters = Registry() + self.as_token_counters = Registry() + self.embedding_models = Registry() + self.vector_stores = Registry() + self.file_stores = Registry() + self.ops = Registry() + self.flows = Registry() + self.services = Registry() + self.token_counters = Registry() + self.file_watchers = Registry() + + +R = RegistryFactory() diff --git a/reme_cli/config/__init__.py b/reme_cli/config/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_cli/op/__init__.py b/reme_cli/op/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_cli/op/base_op.py b/reme_cli/op/base_op.py new file mode 100644 index 00000000..b016bb47 --- /dev/null +++ b/reme_cli/op/base_op.py @@ -0,0 +1,427 @@ +"""Base operator class for LLM workflow execution and composition.""" + +import asyncio +import copy +import inspect +from abc import ABCMeta +from pathlib import Path +from typing import Callable, Optional, Any + +from agentscope.formatter import FormatterBase +from agentscope.model import ChatModelBase +from agentscope.token import HuggingFaceTokenCounter +from loguru import logger +from tqdm import tqdm + +from ..embedding import BaseEmbeddingModel +from ..file_store import BaseFileStore +from ..llm import BaseLLM +from ..prompt_handler import PromptHandler +from ..runtime_context import RuntimeContext +from ..schema import Response, ServiceConfig +from ..schema.service_config import OpConfig +from ..service_context import ServiceContext +from ..token_counter import BaseTokenCounter +from ..utils import camel_to_snake, CacheHandler, timer +from ..vector_store import BaseVectorStore + + +class BaseOp(metaclass=ABCMeta): + """Base operator class for LLM workflow execution and composition.""" + + __alias_name__: str = "" + + def __new__(cls, *args, **kwargs): + """Capture initialization arguments 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 = "", + async_mode: bool = True, + language: str = "", + prompt_name: str = "", + prompt_path: str = "", + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | HuggingFaceTokenCounter = "default", + llm: str | BaseLLM = "default", + embedding_model: str | BaseEmbeddingModel = "default", + vector_store: str | BaseVectorStore = "default", + file_store: str | BaseFileStore = "default", + token_counter: str | BaseTokenCounter = "default", + enable_cache: bool = False, + cache_path: str = "cache/op", + cache_expire_hours: float | None = None, + sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"] = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + enable_parallel: bool = False, + max_retries: int = 1, + raise_exception: bool = False, + **kwargs, + ): + """Initialize operator configurations and internal state.""" + self.name = name or self.__alias_name__ or camel_to_snake(self.__class__.__name__) + self.async_mode = async_mode + self.language = language + self.prompt = self._get_prompt_handler(prompt_name, prompt_path) + + self._as_llm = as_llm + self._as_llm_formatter = as_llm_formatter + self._as_token_counter = as_token_counter + self._llm = llm + self._embedding_model = embedding_model + self._vector_store = vector_store + self._file_store = file_store + self._token_counter = token_counter + + self.enable_cache = enable_cache + self.cache_path = cache_path + self.cache_expire_hours = cache_expire_hours + + self.sub_ops: list["BaseOp"] = [] + self.add_sub_ops(sub_ops) + + self.input_mapping = input_mapping + self.output_mapping = output_mapping + self.enable_parallel = enable_parallel # Control whether to execute tasks in parallel + self.max_retries = max(1, max_retries) + self.raise_exception = raise_exception + self.op_params = kwargs + + self._pending_tasks: list = [] + self.context: RuntimeContext | None = None + self._cache: CacheHandler | None = None + + def _get_prompt_handler(self, prompt_name: str, prompt_path: str) -> PromptHandler: + """Load prompt configuration from the associated YAML file.""" + if prompt_path: + path = Path(prompt_path) + else: + path = Path(inspect.getfile(self.__class__)) + if prompt_name: + path = path.with_stem(prompt_name) + return PromptHandler(language=self.language).load_prompt_by_file(path.with_suffix(".yaml")) + + def _handle_failure(self, e: Exception, attempt: int) -> str | None: + """Log failures and handle final retry logic.""" + message = f"[{self.__class__.__name__}] failed (attempt {attempt + 1}): {e}" + if attempt == self.max_retries - 1: + logger.exception(message) + if self.raise_exception: + raise e + return f"[{self.__class__.__name__}] failed: {e}" + else: + logger.warning(message) + return None + + @property + def cache(self) -> CacheHandler: + """Access the operator-specific cache handler.""" + assert self.enable_cache, "Cache is disabled!" + if not self._cache: + self._cache = CacheHandler(f"{self.cache_path}/{self.name}") + return self._cache + + @property + def service_context(self) -> ServiceContext: + """Access the service context.""" + assert self.context, "Service context is not initialized!" + return self.context.service_context + + @property + def service_config(self) -> ServiceConfig: + """Access the service configuration.""" + return self.service_context.service_config + + @property + def as_llm(self) -> ChatModelBase: + """Get the AgentScope LLM instance from ServiceContext.""" + if isinstance(self._as_llm, str): + self._as_llm = self.service_context.as_llms[self._as_llm] + return self._as_llm + + @property + def as_llm_formatter(self) -> FormatterBase: + """Get the AgentScope LLM formatter instance from ServiceContext.""" + if isinstance(self._as_llm_formatter, str): + self._as_llm_formatter = self.service_context.as_llm_formatters[self._as_llm_formatter] + return self._as_llm_formatter + + @property + def as_token_counter(self) -> HuggingFaceTokenCounter: + """Get the token counter instance from ServiceContext.""" + if isinstance(self._as_token_counter, str): + self._as_token_counter = self.service_context.as_token_counters[self._as_token_counter] + return self._as_token_counter + + @property + def llm(self) -> BaseLLM: + """Get the LLM instance from ServiceContext.""" + if isinstance(self._llm, str): + self._llm = self.service_context.llms[self._llm] + return self._llm + + @property + def embedding_model(self) -> BaseEmbeddingModel: + """Get the embedding model instance from ServiceContext.""" + if isinstance(self._embedding_model, str): + self._embedding_model = self.service_context.embedding_models[self._embedding_model] + return self._embedding_model + + @property + def vector_store(self) -> BaseVectorStore: + """Lazily initialize and return the vector store instance.""" + if isinstance(self._vector_store, str): + self._vector_store = self.service_context.vector_stores[self._vector_store] + return self._vector_store + + @property + def file_store(self) -> BaseFileStore: + """Lazily initialize and return the file store instance.""" + if isinstance(self._file_store, str): + self._file_store = self.service_context.file_stores[self._file_store] + return self._file_store + + @property + def token_counter(self) -> BaseTokenCounter: + """Get the token counter instance from ServiceContext.""" + if isinstance(self._token_counter, str): + self._token_counter = self.service_context.token_counters[self._token_counter] + return self._token_counter + + @property + def service_metadata(self) -> dict: + """Get service configuration metadata.""" + return self.service_context.service_config.metadata + + @property + def response(self) -> Response: + """Access the response object.""" + return self.context.response + + def before_execute_sync(self): + """Prepare context and validate before sync execution. + + This method performs the following steps: + 1. Apply input mapping to transform context variables + 2. Load operator-specific configuration from service config if available + 3. Override operator parameters and prompts based on config + """ + self.context.apply_mapping(self.input_mapping) + + if self.context.service_context is None: + return + + service_config = self.service_context.service_config + if self.name not in service_config.ops: + return + + op_config: OpConfig = service_config.ops[self.name] + + # Override operator parameters from config + if op_config.params: + for k, v in op_config.params.items(): + if hasattr(self, k): + setattr(self, k, v) + logger.info(f"[{self.__class__.__name__}] Set attribute '{k}' = {v}") + else: + self.op_params[k] = v + logger.info(f"[{self.__class__.__name__}] Set op_param '{k}' = {v}") + + # Load custom prompt templates from config + if op_config.prompt_dict: + self.prompt.load_prompt_dict(op_config.prompt_dict) + logger.info(f"[{self.__class__.__name__}] Loaded prompt keys={list(op_config.prompt_dict.keys())}") + + async def before_execute(self): + """Prepare context and validate before async execution.""" + self.before_execute_sync() + + def execute_sync(self): + """Define core sync logic in subclasses.""" + + async def execute(self): + """Define core async logic in subclasses.""" + + def after_execute_sync(self, response: Any): + """Finalize context and mappings after sync execution.""" + self.context.apply_mapping(self.output_mapping) + if response is not None: + if isinstance(response, dict): + for k, v in response.items(): + if k == "answer": + self.response.answer = v + elif k == "success": + self.response.success = v if isinstance(v, bool) else v.lower() == "true" + else: + self.response.metadata[k] = v + else: + self.response.answer = response + return response + + async def after_execute(self, output: Any): + """Finalize context and mappings after async execution.""" + return self.after_execute_sync(output) + + @timer + def call_sync(self, context: RuntimeContext = None, **kwargs): + """Execute the operator synchronously with retry logic.""" + self.context = RuntimeContext.from_context(context, **kwargs) + response = None + for i in range(self.max_retries): + try: + self.before_execute_sync() + response = self.execute_sync() + response = self.after_execute_sync(response) + break + except Exception as e: + response = self._handle_failure(e, i) + + return response + + @timer + async def call(self, context: RuntimeContext = None, **kwargs): + """Execute the operator asynchronously with retry logic.""" + self.context = RuntimeContext.from_context(context, **kwargs) + response = None + for i in range(self.max_retries): + try: + await self.before_execute() + response = await self.execute() + response = await self.after_execute(response) + break + except Exception as e: + response = self._handle_failure(e, i) + return response + + def submit_sync_task(self, fn: Callable, *args, **kwargs) -> "BaseOp": + """Submit a task to the thread pool or local queue.""" + if self.enable_parallel and self.service_context.thread_pool is not None: + task = self.service_context.thread_pool.submit(fn, *args, **kwargs) + else: + task = (fn, args, kwargs) + self._pending_tasks.append(task) + return self + + def submit_async_task(self, coro_fn: Callable, *args, **kwargs) -> "BaseOp": + """Submit an async task to the pending tasks queue.""" + task = coro_fn(*args, **kwargs) + self._pending_tasks.append(task) + return self + + def join_sync_tasks(self, task_desc: str = None) -> list: + """Wait for all pending sync tasks and return flattened results.""" + results = [] + for task in tqdm(self._pending_tasks, desc=task_desc or self.name): + if self.enable_parallel: + result = task.result() + else: + result = task[0](*task[1], **task[2]) + if result: + if isinstance(result, list): + results.extend(result) + else: + results.append(result) + self._pending_tasks.clear() + return results + + async def join_async_tasks(self, return_exceptions: bool = True) -> list: + """Wait for all pending async tasks and aggregate results.""" + if self.enable_parallel: + raw_results = await asyncio.gather(*self._pending_tasks, return_exceptions=return_exceptions) + else: + raw_results = [] + for task in self._pending_tasks: + try: + result = await task + raw_results.append(result) + except Exception as e: + if return_exceptions: + raw_results.append(e) + else: + raise + + results = [] + for result in raw_results: + if isinstance(result, Exception): + logger.error(f"[{self.__class__.__name__}] Async task failed: {result}") + elif result: + if isinstance(result, list): + results.extend(result) + else: + results.append(result) + self._pending_tasks.clear() + return results + + def add_sub_ops(self, sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"]): + """Add child operators to this operator's sub_ops.""" + if not sub_ops: + return + + if isinstance(sub_ops, dict): + for name, op in sub_ops.items(): + assert self.async_mode == op.async_mode, "Async mode mismatch!" + op.name = name + if self.language: + op.language = self.language + self.sub_ops.append(op) + + elif isinstance(sub_ops, list): + for op in sub_ops: + assert self.async_mode == op.async_mode, "Async mode mismatch!" + if self.language: + op.language = self.language + self.sub_ops.append(op) + + else: + assert self.async_mode == sub_ops.async_mode, "Async mode mismatch!" + if self.language: + sub_ops.language = self.language + self.sub_ops.append(sub_ops) + + def add_sub_op(self, sub_op: "BaseOp"): + """Add a single child operator to this operator's sub_ops.""" + self.sub_ops.append(sub_op) + + def __lshift__(self, ops): + """Operator overload for adding sub-operators.""" + self.add_sub_ops(ops) + return self + + def __rshift__(self, op: "BaseOp"): + """Operator overload for sequential execution composition.""" + from .sequential_op import SequentialOp + + seq = SequentialOp(sub_ops=[self], async_mode=self.async_mode) + seq.add_sub_ops(op.sub_ops if isinstance(op, SequentialOp) else op) + return seq + + def __or__(self, op: "BaseOp"): + """Operator overload for parallel execution composition.""" + from .parallel_op import ParallelOp + + par = ParallelOp(sub_ops=[self], async_mode=self.async_mode) + par.add_sub_ops(op.sub_ops if isinstance(op, ParallelOp) else op) + return par + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Format a prompt template with provided keyword arguments.""" + 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): + """Create a copy of this operator with optional parameter overrides.""" + copy_op = self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs}) + if self.sub_ops: + copy_op.sub_ops.clear() + for op in self.sub_ops: + copy_op.add_sub_op(op.copy()) + return copy_op diff --git a/reme_cli/reme_cli.py b/reme_cli/reme_cli.py new file mode 100644 index 00000000..ccd2d0d6 --- /dev/null +++ b/reme_cli/reme_cli.py @@ -0,0 +1,78 @@ +from agentscope.message import Msg + +from reme_cli import Application + + +class ReMeCli(Application): + + async def init(self, **kwargs) -> None: + """Initialize the application.""" + ... + + async def read(self, file: str, **kwargs) -> None: + """Read a note file.""" + # obsidian read file="My Note" + ... + + async def create(self, name: str, content: str, template: str, silent: bool, **kwargs) -> None: + """Create a new note.""" + # obsidian create name="New Note" content="# Hello" template="Template" silent + ... + + async def append(self, file: str, content: str, **kwargs) -> None: + """Append content to a note.""" + # obsidian append file="My Note" content="New line" + ... + + async def search(self, query: str, limit: int, **kwargs) -> None: + """Search for notes.""" + # obsidian search query="search term" limit=10 + ... + + async def daily_read(self, **kwargs) -> None: + """Read daily note.""" + # obsidian daily:read + ... + + async def daily_append(self, content: str, **kwargs) -> None: + """Append content to daily note.""" + # obsidian daily:append content="- [ ] New task" + ... + + async def property_set(self, name: str, value: str, file: str, **kwargs) -> None: + """Set a property on a note.""" + # obsidian property:set name="status" value="done" file="My Note" + ... + + async def tasks(self, daily: bool, todo: bool, **kwargs) -> None: + """Manage tasks.""" + # obsidian tasks daily todo + ... + + async def tags(self, sort: str, counts: bool, **kwargs) -> None: + """Manage tags.""" + # obsidian tags sort=count counts + ... + + async def backlinks(self, file: str, **kwargs) -> None: + """Get backlinks for a note.""" + # obsidian backlinks file="My Note" + ... + + async def summary(self, messages: list[Msg], **kwargs): + ... + + async def dream(self) -> dict: + ... + + async def proactive(self, messages: list[Msg], **kwargs) -> dict: + ... + + +def main(): + """Main entry point for running ReMe from command line.""" + ReMeCli(*sys.argv[1:], config_path="service").run_service() + + +if __name__ == "__main__": + main() diff --git a/reme_cli/schema/__init__.py b/reme_cli/schema/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_cli/schema/service_config.py b/reme_cli/schema/service_config.py new file mode 100644 index 00000000..48b97e0a --- /dev/null +++ b/reme_cli/schema/service_config.py @@ -0,0 +1,144 @@ +"""Configuration schemas for service components using Pydantic models.""" + +import os + +from pydantic import BaseModel, Field, ConfigDict + +from .tool_call import ToolCall + + +class MCPConfig(BaseModel): + """Configuration for Model Context Protocol transport and network settings.""" + + model_config = ConfigDict(extra="allow") + + transport: str = Field(default="stdio") + host: str = Field(default="0.0.0.0") + port: int = Field(default=8001) + + +class HttpConfig(BaseModel): + """Configuration for the HTTP server interface and connection lifecycle.""" + + model_config = ConfigDict(extra="allow") + + host: str = Field(default="0.0.0.0") + port: int = Field(default=8001) + timeout_keep_alive: int = Field(default=3600) + limit_concurrency: int = Field(default=1000) + + +class CmdConfig(BaseModel): + """Configuration for command-line flow execution parameters.""" + + model_config = ConfigDict(extra="allow") + + flow: str = Field(default="") + + +class OpConfig(BaseModel): + """Configuration for op settings and parameters.""" + + model_config = ConfigDict(extra="allow") + + prompt_dict: dict[str, str] = Field(default_factory=dict) + params: dict = Field(default_factory=dict) + + +class FlowConfig(ToolCall): + """Configuration for workflow execution, caching, and error handling.""" + + model_config = ConfigDict(extra="allow") + + flow_content: str = Field(default="") + stream: bool = Field(default=False) + raise_exception: bool = Field(default=True) + enable_cache: bool = Field(default=False) + cache_path: str = Field(default="cache/flow") + cache_expire_hours: float = Field(default=0.1) + + +class BasicConfig(BaseModel): + """Configuration for basic service settings and parameters.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + + +class ModelConfig(BasicConfig): + """Configuration for model-based services with backend and model name.""" + + model_name: str = Field(default="") + + +class LLMConfig(ModelConfig): + """Configuration for Large Language Model backend and model identification.""" + + +class EmbeddingModelConfig(ModelConfig): + """Configuration for embedding model backends and identity.""" + + +class TokenCounterConfig(ModelConfig): + """Configuration for token counting services and model mapping.""" + + +class StoreConfig(BasicConfig): + """Configuration for storage services with embedding model support.""" + + embedding_model: str = Field(default="default") + + +class VectorStoreConfig(StoreConfig): + """Configuration for vector database storage and associated embeddings.""" + + collection_name: str = Field(default="reme") + + +class FileStoreConfig(StoreConfig): + """Configuration for file store database storage and associated embeddings.""" + + store_name: str = Field(default="reme") + + +class FileWatcherConfig(BasicConfig): + """Configuration for file watcher service.""" + + file_store: str = Field(default="") + watch_paths: list[str] = Field(default_factory=list) + + +class ServiceConfig(BasicConfig): + """Root configuration schema aggregating all service-level settings and components.""" + + app_name: str = Field(default=os.getenv("APP_NAME", "ReMe")) + working_dir: str = Field(default=".reme") + enable_logo: bool = Field(default=True) + language: str = Field(default="") + thread_pool_max_workers: int = Field( + default=16, + description="Number of thread pool workers. Set to -1 to disable thread pool.", + ) + ray_max_workers: int = Field(default=-1) + log_to_console: bool = Field(default=True) + disabled_flows: list[str] = Field(default_factory=list) + enabled_flows: list[str] = Field(default_factory=list) + + mcp_servers: dict[str, dict] = Field(default_factory=dict) + mcp: MCPConfig = Field(default_factory=MCPConfig) + http: HttpConfig = Field(default_factory=HttpConfig) + cmd: CmdConfig = Field(default_factory=CmdConfig) + ops: dict[str, OpConfig] = Field(default_factory=dict) + flows: dict[str, FlowConfig] = Field(default_factory=dict) + as_llms: dict[str, BasicConfig] = Field(default_factory=dict) + as_llm_formatters: dict[str, BasicConfig] = Field(default_factory=dict) + as_token_counters: dict[str, BasicConfig] = Field(default_factory=dict) + llms: dict[str, LLMConfig] = Field(default_factory=dict) + embedding_models: dict[str, EmbeddingModelConfig] = Field(default_factory=dict) + vector_stores: dict[str, VectorStoreConfig] = Field(default_factory=dict) + file_stores: dict[str, FileStoreConfig] = Field(default_factory=dict) + token_counters: dict[str, TokenCounterConfig] = Field(default_factory=dict) + file_watchers: dict[str, FileWatcherConfig] = Field(default_factory=dict) + + metadata: dict = Field(default_factory=dict) diff --git a/reme_cli/utils/__init__.py b/reme_cli/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_cli/utils/logger_utils.py b/reme_cli/utils/logger_utils.py new file mode 100644 index 00000000..f6928a4a --- /dev/null +++ b/reme_cli/utils/logger_utils.py @@ -0,0 +1,59 @@ +"""Logging configuration module for application-wide tracing.""" + +import os +import sys +from datetime import datetime + + +def init_logger(log_dir: str = "logs", level: str = "INFO", log_to_console: bool = True) -> None: + """Initialize the logger with both file and console handlers. + + Args: + log_dir: Directory path for log files + level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL) + log_to_console: Whether to print logs to console/screen + """ + from loguru import logger + + # Remove default handler to avoid duplicate logs + logger.remove() + + # Configure colorized standard output logging if enabled + if log_to_console: + logger.add( + sink=sys.stdout, + level=level, + format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}", + colorize=True, + ) + + # Try to configure file-based logging (skip if permission denied) + try: + # Ensure the logging directory exists + os.makedirs(log_dir, exist_ok=True) + + # Generate filename based on the current timestamp + # Use dashes instead of colons for Windows compatibility + current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_filename = f"{current_ts}.log" + log_filepath = os.path.join(log_dir, log_filename) + + # Configure file-based logging with rotation and compression + logger.add( + log_filepath, + level=level, + rotation="00:00", + retention="7 days", + compression="zip", + encoding="utf-8", + format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}", + ) + except Exception as e: + logger.error(f"Error configuring file logging: {e}") + + +def get_logger(): + """Get a configured logger instance using loguru.""" + from loguru import logger + + return logger diff --git a/reme_cli/utils/singleton.py b/reme_cli/utils/singleton.py new file mode 100644 index 00000000..8c2c071b --- /dev/null +++ b/reme_cli/utils/singleton.py @@ -0,0 +1,21 @@ +"""Module providing a decorator to implement the Singleton design pattern.""" + +import threading + + +def singleton(cls): + """A class decorator that ensures only one instance of a class exists.""" + + # Dictionary to cache the single instance of the class + _instance = {} + _lock = threading.Lock() + + def _singleton(*args, **kwargs): + """Return the existing instance or create a new one if it doesn't exist.""" + with _lock: + if cls not in _instance: + # Create and store the instance if it's the first call + _instance[cls] = cls(*args, **kwargs) + return _instance[cls] + + return _singleton