diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 917a5cd8..e29ebf69 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -78,7 +78,7 @@ repos: --disable=C3001, --disable=R1702, --disable=R0912, - --max-statements=75, + --max-statements=120, --max-line-length=120, ] - repo: https://github.com/regebro/pyroma diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index 8b0916f3..e43da12b 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -10,6 +10,7 @@ from abc import ABC from loguru import logger from ..schema import VectorNode +from ..schema.memory_chunk import MemoryChunk class BaseEmbeddingModel(ABC): @@ -172,6 +173,72 @@ class BaseEmbeddingModel(ABC): logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(nodes)} nodes") return nodes + async def get_chunk_embedding(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk: + """Async generate and populate embedding field for a single MemoryChunk object. + + Args: + chunk: MemoryChunk object containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same MemoryChunk object with populated embedding field + """ + chunk.embedding = await self.get_embedding(chunk.text, **kwargs) + return chunk + + async def get_chunk_embeddings(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]: + """Async generate and populate embedding fields for a batch of MemoryChunk objects. + + Args: + chunks: List of MemoryChunk objects containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same list of MemoryChunk objects with populated embedding fields + """ + texts = [chunk.text for chunk in chunks] + embeddings: list[list[float]] = await self.get_embeddings(texts, **kwargs) + + if len(embeddings) == len(chunks): + for chunk, vec in zip(chunks, embeddings): + chunk.embedding = vec + else: + logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(chunks)} chunks") + return chunks + + def get_chunk_embedding_sync(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk: + """Synchronously generate and populate embedding field for a single MemoryChunk object. + + Args: + chunk: MemoryChunk object containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same MemoryChunk object with populated embedding field + """ + chunk.embedding = self.get_embedding_sync(chunk.text, **kwargs) + return chunk + + def get_chunk_embeddings_sync(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]: + """Synchronously generate embeddings for a batch of MemoryChunk objects. + + Args: + chunks: List of MemoryChunk objects containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same list of MemoryChunk objects with populated embedding fields + """ + texts = [chunk.text for chunk in chunks] + embeddings: list[list[float]] = self.get_embeddings_sync(texts, **kwargs) + + if len(embeddings) == len(chunks): + for chunk, vec in zip(chunks, embeddings): + chunk.embedding = vec + else: + logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(chunks)} chunks") + return chunks + def close_sync(self): """Synchronously release resources and close connections.""" diff --git a/reme/core/enumeration/__init__.py b/reme/core/enumeration/__init__.py index 7202a949..c709c20d 100644 --- a/reme/core/enumeration/__init__.py +++ b/reme/core/enumeration/__init__.py @@ -3,6 +3,7 @@ from .chunk_enum import ChunkEnum from .http_enum import HttpEnum from .json_schema_enum import JsonSchemaEnum +from .memory_source import MemorySource from .memory_type import MemoryType from .registry_enum import RegistryEnum from .role import Role @@ -11,6 +12,7 @@ __all__ = [ "ChunkEnum", "HttpEnum", "JsonSchemaEnum", + "MemorySource", "MemoryType", "RegistryEnum", "Role", diff --git a/reme/core/enumeration/memory_source.py b/reme/core/enumeration/memory_source.py new file mode 100644 index 00000000..05174549 --- /dev/null +++ b/reme/core/enumeration/memory_source.py @@ -0,0 +1,11 @@ +"""Memory source types.""" + +from enum import Enum + + +class MemorySource(str, Enum): + """Source of memory data.""" + + MEMORY = "memory" + + SESSIONS = "sessions" diff --git a/reme/core/memory_manager/__init__.py b/reme/core/memory_manager/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/core/memory_manager/ingestion/__init__.py b/reme/core/memory_manager/ingestion/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/core/memory_manager/ingestion/chunking.py b/reme/core/memory_manager/ingestion/chunking.py new file mode 100644 index 00000000..9a159cf0 --- /dev/null +++ b/reme/core/memory_manager/ingestion/chunking.py @@ -0,0 +1,126 @@ +"""Chunking logic for Markdown files.""" + +from typing import List, Dict, Any + +from ..utils.hashing import hash_text +from ...enumeration import MemorySource +from ...schema import MemoryChunk + + +def chunk_markdown( + text: str, + path: str, + source: MemorySource, + chunk_tokens: int = 300, + overlap: int = 30, +) -> List[MemoryChunk]: + """ + Markdown chunking logic implemented based on the TypeScript version. + + Args: + text: Input text + path: File path + source: Memory source + chunk_tokens: Maximum tokens per chunk + overlap: Overlap tokens between chunks + + Returns: + List of MemoryChunk objects + """ + lines = text.split("\n") + if not lines: + return [] + + # Convert tokens to characters (~1 token = 4 chars) + max_chars = max(32, chunk_tokens * 4) + overlap_chars = max(0, overlap * 4) + + chunks: List[MemoryChunk] = [] + + # Currently building chunk + current: List[Dict[str, Any]] = [] # [{'line': str, 'line_no': int}] + current_chars = 0 + + def flush(): + """Add current chunk to results list""" + if not current: + return + + first_entry = current[0] + last_entry = current[-1] + + if not first_entry or not last_entry: + return + + chunk_text = "\n".join([entry["line"] for entry in current]) + start_line = first_entry["line_no"] + end_line = last_entry["line_no"] + + chunk_hash = hash_text(chunk_text) + + chunks.append( + MemoryChunk( + id=hash_text(f"{source}:{path}:{start_line}:{end_line}:{chunk_hash}:{len(chunks)}"), + path=path, + source=source, + start_line=start_line, + end_line=end_line, + text=chunk_text, + hash=chunk_hash, + ), + ) + + def carry_overlap(): + """Keep overlapping part and clear the rest""" + nonlocal current, current_chars + + if overlap_chars <= 0 or not current: + current = [] + current_chars = 0 + return + + acc = 0 + kept = [] + + # Collect lines from the end until reaching overlap size + for i in range(len(current) - 1, -1, -1): + entry = current[i] + if not entry: + continue + + acc += len(entry["line"]) + 1 # +1 for newline + kept.insert(0, entry) # Insert at the beginning to maintain order + + if acc >= overlap_chars: + break + + current = kept + current_chars = sum(len(entry["line"]) + 1 for entry in kept) + + for i, line in enumerate(lines): + line_no = i + 1 + + # Split long lines into multiple segments + segments = [] + if not line: # Empty line + segments.append("") + else: + # If line is too long, split by maximum character count + for start in range(0, len(line), max_chars): + segments.append(line[start : start + max_chars]) + + for segment in segments: + line_size = len(segment) + 1 # +1 for newline + + # If adding current segment would exceed the limit, flush current chunk + if current_chars + line_size > max_chars and current: + flush() + carry_overlap() + + current.append({"line": segment, "line_no": line_no}) + current_chars += line_size + + # Process the final chunk + flush() + + return chunks diff --git a/reme/core/memory_manager/manager.py b/reme/core/memory_manager/manager.py new file mode 100644 index 00000000..f8eccada --- /dev/null +++ b/reme/core/memory_manager/manager.py @@ -0,0 +1,890 @@ +"""Memory Index Manager - Main coordination layer. + +This module provides the main MemoryIndexManager class that coordinates +file watching, embedding generation, and search operations across memory files +and session transcripts. +""" + +import asyncio +import json +import os +import re +from typing import Any, Callable + +from loguru import logger +from pydantic import BaseModel, Field +from watchfiles import awatch + +from .ingestion.chunking import chunk_markdown +from .memory_storage.sqlite_memory_store import SqliteMemoryStore +from .utils.hashing import hash_text +from ..enumeration import MemorySource +from ..schema import FileMetadata, MemorySearchResult + +# Constants +SNIPPET_MAX_CHARS = 700 +SESSION_DIRTY_DEBOUNCE_MS = 5000 +EMBEDDING_BATCH_MAX_TOKENS = 8000 +EMBEDDING_APPROX_CHARS_PER_TOKEN = 1 +EMBEDDING_INDEX_CONCURRENCY = 4 +EMBEDDING_RETRY_MAX_ATTEMPTS = 3 +EMBEDDING_RETRY_BASE_DELAY_MS = 500 +EMBEDDING_RETRY_MAX_DELAY_MS = 8000 +BATCH_FAILURE_LIMIT = 2 +SESSION_DELTA_READ_CHUNK_BYTES = 64 * 1024 +EMBEDDING_QUERY_TIMEOUT_REMOTE_MS = 60_000 +EMBEDDING_QUERY_TIMEOUT_LOCAL_MS = 5 * 60_000 +EMBEDDING_BATCH_TIMEOUT_REMOTE_MS = 2 * 60_000 +EMBEDDING_BATCH_TIMEOUT_LOCAL_MS = 10 * 60_000 + + +class MemorySyncProgressUpdate(BaseModel): + """Progress update for memory sync operations.""" + + completed: int = Field(default=..., description="Number of items completed") + total: int = Field(default=..., description="Total number of items to process") + label: str | None = Field(default=None, description="Optional label for the progress operation") + + +class MemorySyncProgressState(BaseModel): + """Internal state for tracking sync progress.""" + + completed: int = Field(default=0, description="Number of items completed") + total: int = Field(default=0, description="Total number of items to process") + label: str | None = Field(default=None, description="Optional label for the progress operation") + report: Callable[[MemorySyncProgressUpdate], None] | None = Field( + default=None, + description="Callback function to report progress updates", + ) + + +class SessionDelta(BaseModel): + """Tracks incremental changes in session files.""" + + last_size: int = Field(default=0, description="Last known size of the session file") + pending_bytes: int = Field(default=0, description="Number of pending bytes to process") + pending_messages: int = Field(default=0, description="Number of pending messages to process") + + +class MemorySearchConfig(BaseModel): + """Configuration for memory search operations.""" + + model: str = Field(default="default", description="Model name for embeddings") + sources: list[MemorySource] = Field(default_factory=lambda: [MemorySource.MEMORY], description="Sources to search") + extra_paths: list[str] = Field(default_factory=list, description="Additional paths to include in search") + store_path: str = Field(default="memory.db", description="Path to SQLite database file") + vector_enabled: bool = Field(default=True, description="Whether to enable vector search") + vector_extension_path: str | None = Field(default=None, description="Path to vector extension for SQLite") + fts_enabled: bool = Field(default=True, description="Whether to enable full-text search") + chunk_tokens: int = Field(default=300, description="Number of tokens per chunk") + chunk_overlap: int = Field(default=30, description="Number of overlapping tokens between chunks") + watch_enabled: bool = Field(default=True, description="Whether to enable file watching") + watch_debounce_ms: int = Field(default=1000, description="Debounce time for file watcher in milliseconds") + interval_minutes: int = Field(default=0, description="Interval between automatic syncs in minutes (0 to disable)") + sync_on_search: bool = Field(default=True, description="Whether to sync before search operations") + sync_on_session_start: bool = Field(default=True, description="Whether to sync when a session starts") + query_min_score: float = Field(default=0.3, description="Minimum relevance score for search results") + query_max_results: int = Field(default=10, description="Maximum number of search results to return") + hybrid_enabled: bool = Field(default=True, description="Whether to use hybrid vector + keyword search") + hybrid_vector_weight: float = Field(default=0.7, description="Weight for vector search in hybrid scoring") + hybrid_text_weight: float = Field(default=0.3, description="Weight for text search in hybrid scoring") + hybrid_candidate_multiplier: float = Field( + default=2.0, + description="Multiplier for number of candidates to consider in hybrid search", + ) + session_delta_bytes: int = Field( + default=0, + description="Threshold for session sync based on bytes changed (0 for any change)", + ) + session_delta_messages: int = Field( + default=5, + description="Threshold for session sync based on messages changed (0 for any change)", + ) + + +# Global cache for manager instances +INDEX_CACHE: dict[str, "MemoryIndexManager"] = {} + + +class MemoryIndexManager: + """Main memory index manager coordinating all memory operations.""" + + # ============================================================================ + # Initialization and Lifecycle + # ============================================================================ + + def __init__( + self, + agent_id: str, + workspace_dir: str, + settings: MemorySearchConfig, + store: SqliteMemoryStore, + ): + """Initialize the memory index manager.""" + + self.agent_id = agent_id + self.workspace_dir = workspace_dir + self.settings = settings + self.store = store + + # State tracking + self.sources = set(settings.sources) + self.closed = False + self.dirty = MemorySource.MEMORY in self.sources + self.sessions_dirty = False + self.sessions_dirty_files: set[str] = set() + self.session_pending_files: set[str] = set() + self.session_deltas: dict[str, SessionDelta] = {} + self.session_warm: set[str] = set() + + # Sync control + self.syncing: asyncio.Task | None = None + self.watch_task: asyncio.Task | None = None + self.session_watch_task: asyncio.Task | None = None + self.interval_task: asyncio.Task | None = None + + # Batch failure tracking + self.batch_failure_count = 0 + self.batch_failure_last_error: str | None = None + self.batch_failure_lock = asyncio.Lock() + + async def close(self) -> None: + """Close the manager and release resources.""" + if self.closed: + return + + self.closed = True + + # Cancel all background tasks + if self.watch_task: + self.watch_task.cancel() + if self.session_watch_task: + self.session_watch_task.cancel() + if self.interval_task: + self.interval_task.cancel() + + await self.store.close() + + # ============================================================================ + # Public API Methods + # ============================================================================ + + async def warm_session(self, session_key: str | None = None): + """Pre-sync memory before a session starts.""" + if not self.settings.sync_on_session_start: + return + + key = (session_key or "").strip() + if key and key in self.session_warm: + return + + await self.sync(reason="session-start") + + if key: + self.session_warm.add(key) + + async def sync( + self, + reason: str | None = None, + force: bool = False, + progress: Callable[[MemorySyncProgressUpdate], None] | None = None, + ): + """Synchronize memory index with file system.""" + if self.syncing: + await self.syncing + return + + self.syncing = asyncio.create_task(self._run_sync(reason, force, progress)) + try: + await self.syncing + finally: + self.syncing = None + + async def search( + self, + query: str, + max_results: int | None = None, + min_score: float | None = None, + session_key: str | None = None, + ) -> list[MemorySearchResult]: + """Search indexed memory with hybrid vector + keyword search. + + Args: + query: Search query text + max_results: Maximum number of results to return + min_score: Minimum relevance score threshold + session_key: Optional session key for warmup + + Returns: + List of search results sorted by relevance + """ + await self.warm_session(session_key) + + if self.settings.sync_on_search and (self.dirty or self.sessions_dirty): + try: + await self.sync(reason="search") + except Exception as err: + logger.warning(f"memory sync failed (search): {err}") + + cleaned = query.strip() + if not cleaned: + return [] + + min_score = min_score if min_score is not None else self.settings.query_min_score + max_results = max_results if max_results is not None else self.settings.query_max_results + + hybrid = self.settings.hybrid_enabled + candidates = min(200, max(1, int(max_results * self.settings.hybrid_candidate_multiplier))) + + # Run keyword search if hybrid enabled + keyword_results = [] + if hybrid: + keyword_results = await self._search_keyword(cleaned, candidates) + + # Perform vector search + vector_results = await self._search_vector(cleaned, candidates) + + if not hybrid: + return [r for r in vector_results if r.score >= min_score][:max_results] + + merged = self._merge_hybrid_results( + vector=vector_results, + keyword=keyword_results, + vector_weight=self.settings.hybrid_vector_weight, + text_weight=self.settings.hybrid_text_weight, + ) + + return [r for r in merged if r.score >= min_score][:max_results] + + async def read_file( + self, + rel_path: str, + from_line: int | None = None, + num_lines: int | None = None, + ) -> dict[str, str]: + """Read a memory file with optional line range. + + Args: + rel_path: Relative path to file + from_line: Starting line number (1-indexed) + num_lines: Number of lines to read + + Returns: + Dictionary with 'text' and 'path' keys + + Raises: + ValueError: If path is invalid or not allowed + """ + raw_path = rel_path.strip() + assert raw_path, "path required" + abs_path = os.path.abspath(os.path.join(self.workspace_dir, raw_path)) + rel_path_clean = os.path.relpath(abs_path, self.workspace_dir) + + in_workspace = not rel_path_clean.startswith("..") and not os.path.isabs(rel_path_clean) + allowed = in_workspace and self._is_memory_path(rel_path_clean) + + if not allowed and self.settings.extra_paths: + for extra in self.settings.extra_paths: + extra_abs = os.path.abspath(extra) + if abs_path.startswith(extra_abs): + allowed = True + break + + if not allowed: + raise ValueError("path required") + + if not abs_path.endswith(".md"): + raise ValueError("path required") + + # Read file + with open(abs_path, "r", encoding="utf-8") as f: + content = f.read() + + if from_line is None and num_lines is None: + return {"text": content, "path": rel_path_clean} + + lines = content.split("\n") + start = max(1, from_line or 1) + count = max(1, num_lines or len(lines)) + slice_lines = lines[start - 1 : start - 1 + count] + + return {"text": "\n".join(slice_lines), "path": rel_path_clean} + + # ============================================================================ + # Sync Logic + # ============================================================================ + + async def _run_sync( + self, + reason: str | None, + force: bool, + progress_callback: Callable[[MemorySyncProgressUpdate], None] | None, + ): + """Execute sync operation.""" + progress = MemorySyncProgressState() + if progress_callback: + progress.report = progress_callback + + should_sync_memory = MemorySource.MEMORY in self.sources and (force or self.dirty) + should_sync_sessions = self._should_sync_sessions(reason, force) + + if should_sync_memory: + await self._sync_memory_files(progress) + self.dirty = False + + if should_sync_sessions: + await self._sync_session_files(progress) + self.sessions_dirty = False + self.sessions_dirty_files.clear() + elif len(self.sessions_dirty_files) > 0: + self.sessions_dirty = True + else: + self.sessions_dirty = False + + def _should_sync_sessions(self, reason: str | None, force: bool) -> bool: + """Check if session sync is needed.""" + if MemorySource.SESSIONS not in self.sources: + return False + + if force: + return True + + if reason in ("session-start", "watch"): + return False + + return self.sessions_dirty and len(self.sessions_dirty_files) > 0 + + async def _sync_memory_files(self, progress: MemorySyncProgressState): + """Sync memory markdown files.""" + files = self._list_memory_files() + logger.debug("memory sync: indexing memory files", files=len(files)) + + active_paths = {f.path for f in files} + if progress.report: + progress.total += len(files) + progress.report( + MemorySyncProgressUpdate( + completed=progress.completed, + total=progress.total, + label="Indexing memory files…", + ), + ) + + tasks = [] + for file_entry in files: + task = self._index_memory_file(file_entry, progress) + tasks.append(task) + await asyncio.gather(*tasks) + + indexed = await self.store.list_files(MemorySource.MEMORY) + for stale_path in indexed: + if stale_path not in active_paths: + await self.store.delete_file(stale_path, MemorySource.MEMORY) + + async def _sync_session_files(self, progress: MemorySyncProgressState): + """Sync session transcript files.""" + files = self._list_session_files() + logger.debug( + "memory sync: indexing session files", + files=len(files), + index_all=len(self.sessions_dirty_files) == 0, + dirty_files=len(self.sessions_dirty_files), + ) + + if progress.report: + progress.total += len(files) + progress.report( + MemorySyncProgressUpdate( + completed=progress.completed, + total=progress.total, + label="Indexing session files...", + ), + ) + + active_paths = set() + tasks = [] + + for abs_path in files: + rel_path = self._session_path_for_file(abs_path) + active_paths.add(rel_path) + + if len(self.sessions_dirty_files) == 0 or abs_path in self.sessions_dirty_files: + task = self._index_session_file(abs_path, progress) + tasks.append(task) + else: + if progress.report: + progress.completed += 1 + progress.report(MemorySyncProgressUpdate(completed=progress.completed, total=progress.total)) + await asyncio.gather(*tasks) + + indexed = await self.store.list_files(MemorySource.SESSIONS) + for stale_path in indexed: + if stale_path not in active_paths: + await self.store.delete_file(stale_path, MemorySource.SESSIONS) + + # ============================================================================ + # File Indexing + # ============================================================================ + + async def _index_memory_file(self, file_meta: FileMetadata, progress: MemorySyncProgressState): + """Index a single memory file.""" + existing_meta = await self.store.get_file_metadata(file_meta.path, MemorySource.MEMORY) + if existing_meta and existing_meta.hash == file_meta.hash: + if progress.report: + progress.completed += 1 + progress.report(MemorySyncProgressUpdate(completed=progress.completed, total=progress.total)) + return + + # Read and chunk file + with open(file_meta.abs_path, "r", encoding="utf-8") as f: + content = f.read() + + chunks = chunk_markdown( + content, + file_meta.path, + MemorySource.MEMORY, + self.settings.chunk_tokens, + self.settings.chunk_overlap, + ) + + chunks = [c for c in chunks if c.text.strip()] + + if chunks: + chunks = await self.store.get_chunk_embeddings(chunks) + + await self.store.upsert_file(file_meta, MemorySource.MEMORY, chunks) + if progress.report: + progress.completed += 1 + progress.report(MemorySyncProgressUpdate(completed=progress.completed, total=progress.total)) + + async def _index_session_file(self, abs_path: str, progress: MemorySyncProgressState): + """Index a single session transcript file.""" + file_meta = self._build_session_file_meta(abs_path) + if not file_meta: + if progress.report: + progress.completed += 1 + progress.report(MemorySyncProgressUpdate(completed=progress.completed, total=progress.total)) + return + + existing_meta = await self.store.get_file_metadata(file_meta.path, MemorySource.SESSIONS) + if existing_meta and existing_meta.hash == file_meta.hash: + self._reset_session_delta(abs_path, file_meta.size) + if progress.report: + progress.completed += 1 + progress.report(MemorySyncProgressUpdate(completed=progress.completed, total=progress.total)) + return + + chunks = chunk_markdown( + file_meta.content, + file_meta.path, + MemorySource.SESSIONS, + self.settings.chunk_tokens, + self.settings.chunk_overlap, + ) + + chunks = [c for c in chunks if c.text.strip()] + + if chunks: + chunks = await self.store.get_chunk_embeddings(chunks) + + await self.store.upsert_file(file_meta, MemorySource.SESSIONS, chunks) + self._reset_session_delta(abs_path, file_meta.size) + + if progress.report: + progress.completed += 1 + progress.report(MemorySyncProgressUpdate(completed=progress.completed, total=progress.total)) + + # ============================================================================ + # File Listing and Building + # ============================================================================ + + def _list_memory_files(self) -> list[FileMetadata]: + """List all memory markdown files.""" + files = [] + + # Scan workspace + memory_paths = [ + os.path.join(self.workspace_dir, "MEMORY.md"), + os.path.join(self.workspace_dir, "memory.md"), + os.path.join(self.workspace_dir, "memory"), + ] + + for base_path in memory_paths: + if os.path.isfile(base_path) and base_path.endswith(".md"): + files.append(self._build_file_entry(base_path)) + elif os.path.isdir(base_path): + for root, _, filenames in os.walk(base_path): + for filename in filenames: + if filename.endswith(".md"): + abs_path = os.path.join(root, filename) + files.append(self._build_file_entry(abs_path)) + + # Extra paths + for extra in self.settings.extra_paths: + if os.path.isfile(extra) and extra.endswith(".md"): + files.append(self._build_file_entry(extra)) + elif os.path.isdir(extra): + for root, _, filenames in os.walk(extra): + for filename in filenames: + if filename.endswith(".md"): + abs_path = os.path.join(root, filename) + files.append(self._build_file_entry(abs_path)) + + return files + + def _list_session_files(self) -> list[str]: + """List all session transcript files.""" + sessions_dir = os.path.join(self.workspace_dir, "sessions", self.agent_id) + if not os.path.exists(sessions_dir): + return [] + + files = [] + for filename in os.listdir(sessions_dir): + if filename.endswith(".jsonl"): + files.append(os.path.join(sessions_dir, filename)) + + return files + + def _build_file_entry(self, abs_path: str) -> FileMetadata: + """Build file entry metadata.""" + stat = os.stat(abs_path) + with open(abs_path, "r", encoding="utf-8") as f: + content = f.read() + + rel_path = os.path.relpath(abs_path, self.workspace_dir) + + return FileMetadata( + hash=hash_text(content), + mtime_ms=stat.st_mtime * 1000, + size=stat.st_size, + path=rel_path.replace("\\", "/"), + abs_path=abs_path, + ) + + def _build_session_file_meta(self, abs_path: str) -> FileMetadata | None: + """Build session file entry with parsed content. TODO 修改message解析逻辑""" + stat = os.stat(abs_path) + with open(abs_path, "r", encoding="utf-8") as f: + raw = f.read() + + lines = raw.split("\n") + collected = [] + + for line in lines: + if not line.strip(): + continue + + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + + if record.get("type") != "message": + continue + + message = record.get("message", {}) + role = message.get("role") + + if role not in ("user", "assistant"): + continue + + text = self._extract_session_text(message.get("content")) + if not text: + continue + + label = "User" if role == "user" else "Assistant" + collected.append(f"{label}: {text}") + + content = "\n".join(collected) + rel_path = self._session_path_for_file(abs_path) + + return FileMetadata( + hash=hash_text(content), + mtime_ms=stat.st_mtime * 1000, + size=stat.st_size, + path=rel_path, + abs_path=abs_path, + content=content, + ) + + # ============================================================================ + # Session Processing Helpers + # ============================================================================ + + @staticmethod + def _session_path_for_file(abs_path: str) -> str: + """Convert absolute session path to relative.""" + return f"sessions/{os.path.basename(abs_path)}" + + def _extract_session_text(self, content: Any) -> str | None: + """Extract text from session message content.""" + if isinstance(content, str): + normalized = self._normalize_session_text(content) + return normalized if normalized else None + + if not isinstance(content, list): + return None + + parts = [] + for block in content: + if not isinstance(block, dict): + continue + + if block.get("type") != "text": + continue + + text = block.get("text") + if isinstance(text, str): + normalized = self._normalize_session_text(text) + if normalized: + parts.append(normalized) + + return " ".join(parts) if parts else None + + @staticmethod + def _normalize_session_text(text: str) -> str: + """Normalize session text by collapsing whitespace.""" + text = re.sub(r"\s*\n+\s*", " ", text) + text = re.sub(r"\s+", " ", text) + return text.strip() + + # ============================================================================ + # Session Delta Tracking + # ============================================================================ + + async def _process_session_delta_batch(self) -> None: + """Process pending session file changes.""" + if not self.session_pending_files: + return + + pending = list(self.session_pending_files) + self.session_pending_files.clear() + + should_sync = False + for session_file in pending: + delta = await self._update_session_delta(session_file) + if not delta: + continue + + bytes_threshold = self.settings.session_delta_bytes + messages_threshold = self.settings.session_delta_messages + + if bytes_threshold <= 0: + bytes_hit = delta["pending_bytes"] > 0 + else: + bytes_hit = delta["pending_bytes"] >= bytes_threshold + + if messages_threshold <= 0: + messages_hit = delta["pending_messages"] > 0 + else: + messages_hit = delta["pending_messages"] >= messages_threshold + + if not bytes_hit and not messages_hit: + continue + + self.sessions_dirty_files.add(session_file) + self.sessions_dirty = True + should_sync = True + + if should_sync: + try: + await self.sync(reason="session-delta") + except Exception as err: + logger.warning(f"memory sync failed (session-delta): {err}") + + async def _update_session_delta(self, session_file: str) -> dict[str, int] | None: + """Update delta tracking for a session file.""" + try: + stat = os.stat(session_file) + size = stat.st_size + except OSError: + return None + + state = self.session_deltas.get(session_file) + if not state: + state = SessionDelta() + self.session_deltas[session_file] = state + + delta_bytes = max(0, size - state.last_size) + + if delta_bytes == 0 and size == state.last_size: + return { + "delta_bytes": self.settings.session_delta_bytes, + "delta_messages": self.settings.session_delta_messages, + "pending_bytes": state.pending_bytes, + "pending_messages": state.pending_messages, + } + + if size < state.last_size: + state.last_size = size + state.pending_bytes += size + if self.settings.session_delta_messages > 0: + state.pending_messages += await self._count_newlines(session_file, 0, size) + else: + state.pending_bytes += delta_bytes + if self.settings.session_delta_messages > 0: + state.pending_messages += await self._count_newlines(session_file, state.last_size, size) + state.last_size = size + + return { + "delta_bytes": self.settings.session_delta_bytes, + "delta_messages": self.settings.session_delta_messages, + "pending_bytes": state.pending_bytes, + "pending_messages": state.pending_messages, + } + + def _reset_session_delta(self, abs_path: str, size: int) -> None: + """Reset delta tracking for a session file.""" + state = self.session_deltas.get(abs_path) + if state: + state.last_size = size + state.pending_bytes = 0 + state.pending_messages = 0 + + @staticmethod + async def _count_newlines(abs_path: str, start: int, end: int) -> int: + """Count newlines in a file range.""" + if end <= start: + return 0 + + count = 0 + with open(abs_path, "rb") as f: + f.seek(start) + remaining = end - start + + while remaining > 0: + chunk_size = min(SESSION_DELTA_READ_CHUNK_BYTES, remaining) + chunk = f.read(chunk_size) + if not chunk: + break + + count += chunk.count(b"\n") + remaining -= len(chunk) + + return count + + # ============================================================================ + # File Watchers + # ============================================================================ + + async def _start_watchers(self): + """Start file watching and interval sync tasks.""" + if self.settings.watch_enabled and MemorySource.MEMORY in self.sources: + self.watch_task = asyncio.create_task(self._watch_memory_files()) + + if MemorySource.SESSIONS in self.sources: + self.session_watch_task = asyncio.create_task(self._watch_session_files()) + + if self.settings.interval_minutes > 0: + self.interval_task = asyncio.create_task(self._interval_sync()) + + async def _watch_memory_files(self) -> None: + """Watch memory files for changes.""" + watch_paths = [ + os.path.join(self.workspace_dir, "MEMORY.md"), + os.path.join(self.workspace_dir, "memory.md"), + os.path.join(self.workspace_dir, "memory"), + ] + + for extra in self.settings.extra_paths: + watch_paths.append(extra) + + async for changes in awatch(*watch_paths, stop_event=None): + if self.closed: + break + + for _, path in changes: + if path.endswith(".md"): + self.dirty = True + await asyncio.sleep(self.settings.watch_debounce_ms / 1000) + try: + await self.sync(reason="watch") + except Exception as e: + logger.exception(f"memory sync failed (watch): {e}") + + async def _watch_session_files(self): + """Watch session files for changes.""" + sessions_dir = os.path.join(self.workspace_dir, "sessions", self.agent_id) + if not os.path.exists(sessions_dir): + return + + async for changes in awatch(sessions_dir, stop_event=None): + if self.closed: + break + + for _, path in changes: + if path.endswith(".jsonl"): + self.session_pending_files.add(path) + + await asyncio.sleep(SESSION_DIRTY_DEBOUNCE_MS / 1000) + await self._process_session_delta_batch() + + async def _interval_sync(self) -> None: + """Periodically sync the index.""" + while not self.closed: + await asyncio.sleep(self.settings.interval_minutes * 60) + if not self.closed: + try: + await self.sync(reason="interval") + except Exception as err: + logger.warning(f"memory sync failed (interval): {err}") + + # ============================================================================ + # Search Methods + # ============================================================================ + + async def _search_vector(self, query: str, limit: int) -> list[MemorySearchResult]: + """Perform vector similarity search.""" + return await self.store.vector_search(query, limit, sources=list(self.sources)) + + async def _search_keyword(self, query: str, limit: int) -> list[MemorySearchResult]: + """Perform keyword/FTS search.""" + if not self.settings.fts_enabled: + return [] + + return await self.store.keyword_search(query, limit, sources=list(self.sources)) + + @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.""" + 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 + + results = list(merged.values()) + results.sort(key=lambda r: r.score, reverse=True) + return results + + # ============================================================================ + # Utility Methods + # ============================================================================ + + @staticmethod + def _is_memory_path(rel_path: str) -> bool: + """Check if path is a valid memory path.""" + normalized = rel_path.replace("\\", "/") + + if normalized in ("MEMORY.md", "memory.md"): + return True + + if normalized.startswith("memory/") and normalized.endswith(".md"): + return True + + return False diff --git a/reme/core/memory_manager/memory_storage/__init__.py b/reme/core/memory_manager/memory_storage/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/core/memory_manager/memory_storage/base_memory_store.py b/reme/core/memory_manager/memory_storage/base_memory_store.py new file mode 100644 index 00000000..d17b8fa5 --- /dev/null +++ b/reme/core/memory_manager/memory_storage/base_memory_store.py @@ -0,0 +1,172 @@ +"""Base storage interface for memory manager.""" + +from abc import ABC, abstractmethod + +from ...embedding import BaseEmbeddingModel +from ...enumeration import MemorySource +from ...schema import FileMetadata, MemoryIndexMeta, MemoryChunk, MemorySearchResult + + +class BaseMemoryStore(ABC): + """Abstract base class for memory storage backends.""" + + def __init__(self, embedding_model: BaseEmbeddingModel): + """Initialize""" + self.embedding_model: BaseEmbeddingModel = embedding_model + + @property + def embedding_dim(self) -> int: + """Get the embedding model's dimensionality.""" + return self.embedding_model.dimensions + + async def get_embedding(self, query: str, **kwargs) -> list[float]: + """Get embedding for a single query string. + + Args: + query: Input text to generate embedding for + **kwargs: Additional arguments passed to the embedding model + + Returns: + Embedding vector as a list of floats + """ + return await self.embedding_model.get_embedding(query, **kwargs) + + async def get_embeddings(self, queries: list[str], **kwargs) -> list[list[float]]: + """Get embeddings for a batch of query strings. + + Args: + queries: List of input texts to generate embeddings for + **kwargs: Additional arguments passed to the embedding model + + Returns: + List of embedding vectors, each as a list of floats + """ + return await self.embedding_model.get_embeddings(queries, **kwargs) + + async def get_chunk_embedding(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk: + """Generate and populate embedding field for a single MemoryChunk object. + + Args: + chunk: MemoryChunk object containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same MemoryChunk object with populated embedding field + """ + return await self.embedding_model.get_chunk_embedding(chunk, **kwargs) + + async def get_chunk_embeddings(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]: + """Generate and populate embedding fields for a batch of MemoryChunk objects. + + Args: + chunks: List of MemoryChunk objects containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same list of MemoryChunk objects with populated embedding fields + """ + return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs) + + def get_chunk_embedding_sync(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk: + """Synchronously generate and populate embedding field for a single MemoryChunk object. + + Args: + chunk: MemoryChunk object containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same MemoryChunk object with populated embedding field + """ + return self.embedding_model.get_chunk_embedding_sync(chunk, **kwargs) + + def get_chunk_embeddings_sync(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]: + """Synchronously generate embeddings for a batch of MemoryChunk objects. + + Args: + chunks: List of MemoryChunk objects containing text to embed + **kwargs: Additional arguments passed to the embedding model + + Returns: + The same list of MemoryChunk objects with populated embedding fields + """ + return self.embedding_model.get_chunk_embeddings_sync(chunks, **kwargs) + + @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 get_file_hash(self, path: str, source: MemorySource) -> str | None: + """Get the hash of an indexed file.""" + + @abstractmethod + async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None: + """Get full file metadata with statistics.""" + + @abstractmethod + async def list_files(self, source: MemorySource) -> list[str]: + """List all indexed file paths for a source.""" + + @abstractmethod + async def get_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 read_meta(self, key: str) -> MemoryIndexMeta | None: + """Read metadata value.""" + + @abstractmethod + async def write_meta(self, key: str, value: MemoryIndexMeta | dict): + """Write metadata value.""" + + @abstractmethod + async def clear_all(self): + """Clear all indexed data.""" + + @abstractmethod + async def close(self): + """Close storage and release resources.""" diff --git a/reme/core/memory_manager/memory_storage/sqlite_memory_store.py b/reme/core/memory_manager/memory_storage/sqlite_memory_store.py new file mode 100644 index 00000000..0d6b47fe --- /dev/null +++ b/reme/core/memory_manager/memory_storage/sqlite_memory_store.py @@ -0,0 +1,582 @@ +"""SQLite storage backend for memory index.""" + +import json +import sqlite3 +import struct +import time +from pathlib import Path + +from loguru import logger + +from .base_memory_store import BaseMemoryStore +from ...embedding import BaseEmbeddingModel +from ...enumeration import MemorySource +from ...schema import FileMetadata, MemoryIndexMeta, MemoryChunk, MemorySearchResult + + +class SqliteMemoryStore(BaseMemoryStore): + """SQLite memory storage with vector and full-text search. + + Inherits embedding methods from BaseMemoryStore: + - 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 + """ + + VECTOR_TABLE = "chunks_vec" + FTS_TABLE = "chunks_fts" + + def __init__( + self, + db_path: str, + embedding_model: BaseEmbeddingModel, + vec_ext_path: str = "", + fts_enabled: bool = True, + snippet_max_chars: int = 700, + ): + super().__init__(embedding_model=embedding_model) + self.db_path = db_path + self.vec_ext_path = vec_ext_path + self.fts_enabled = fts_enabled + self.snippet_max_chars = snippet_max_chars + self.conn: sqlite3.Connection | None = None + self.vector_available = False + self.fts_available = False + + @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.""" + Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) + + self.conn = sqlite3.connect(self.db_path, check_same_thread=False) + self.conn.enable_load_extension(True) + + # Load sqlite-vec extension + if self.vec_ext_path: + try: + self.conn.load_extension(self.vec_ext_path) + self.vector_available = True + 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 common extension names + for name in ["vec0", "sqlite_vec", "vector0"]: + try: + self.conn.load_extension(name) + self.vector_available = True + logger.info(f"Loaded sqlite-vec: {name}") + break + except Exception: + pass + + self.conn.enable_load_extension(False) + await self._create_tables() + + async def _create_tables(self) -> None: + """Create database schema.""" + cursor = self.conn.cursor() + + # Metadata + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT + ) + """, + ) + + # Files + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS files ( + path TEXT, + source TEXT, + hash TEXT, + mtime REAL, + size INTEGER, + PRIMARY KEY (path, source) + ) + """, + ) + + # Chunks + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS chunks ( + id TEXT PRIMARY KEY, + path TEXT, + source TEXT, + start_line INTEGER, + end_line INTEGER, + hash TEXT, + text TEXT, + embedding TEXT, + updated_at INTEGER + ) + """, + ) + + cursor.execute( + """ + CREATE INDEX IF NOT EXISTS idx_chunks_path_source + ON chunks(path, source) + """, + ) + + # Vector table (sqlite-vec) + if self.vector_available: + try: + cursor.execute( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self.VECTOR_TABLE} USING vec0( + id TEXT PRIMARY KEY, + embedding FLOAT[{self.embedding_dim}] + ) + """, + ) + logger.info(f"Created vector table (dims={self.embedding_dim})") + except Exception as e: + logger.warning(f"Failed to create vector table: {e}") + self.vector_available = False + + # FTS table + if self.fts_enabled: + try: + cursor.execute( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self.FTS_TABLE} USING fts5( + text, + id UNINDEXED, + path UNINDEXED, + source UNINDEXED, + start_line UNINDEXED, + end_line UNINDEXED + ) + """, + ) + self.fts_available = True + logger.info("Created FTS5 table") + except Exception as e: + logger.warning(f"Failed to create FTS table: {e}") + self.fts_available = False + + self.conn.commit() + 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") + await self._delete_file_internal(cursor, file_meta.path, source) + + # Insert file + cursor.execute( + """ + INSERT OR REPLACE INTO files (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( + """ + INSERT INTO chunks ( + 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 + if self.vector_available and chunk.embedding: + try: + cursor.execute( + f""" + INSERT INTO {self.VECTOR_TABLE} (id, embedding) + VALUES (?, ?) + """, + (chunk.id, self.vector_to_blob(chunk.embedding)), + ) + except Exception as e: + logger.debug(f"Vector insert failed: {e}") + + # Insert FTS + if self.fts_available: + try: + cursor.execute( + f""" + INSERT INTO {self.FTS_TABLE} ( + text, id, path, source, start_line, end_line + ) VALUES (?, ?, ?, ?, ?, ?) + """, + ( + chunk.text, + chunk.id, + file_meta.path, + source.value, + chunk.start_line, + chunk.end_line, + ), + ) + except Exception as e: + logger.debug(f"FTS insert failed: {e}") + + cursor.execute("COMMIT") + except Exception: + cursor.execute("ROLLBACK") + raise + finally: + cursor.close() + + async def delete_file(self, path: str, source: MemorySource) -> None: + """Delete file and all its chunks.""" + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + await self._delete_file_internal(cursor, path, source) + cursor.execute("COMMIT") + except Exception: + cursor.execute("ROLLBACK") + raise + finally: + cursor.close() + + async def _delete_file_internal(self, cursor: sqlite3.Cursor, path: str, source: MemorySource): + """Internal delete helper.""" + # Get chunk IDs for vector deletion + cursor.execute( + "SELECT id FROM chunks WHERE path = ? AND source = ?", + (path, source.value), + ) + chunk_ids = [row[0] for row in cursor.fetchall()] + + # Delete vectors + if self.vector_available and chunk_ids: + for chunk_id in chunk_ids: + try: + cursor.execute( + f"DELETE FROM {self.VECTOR_TABLE} WHERE id = ?", + (chunk_id,), + ) + except Exception as e: + logger.debug(f"Vector delete failed: {e}") + + # Delete FTS entries + if self.fts_available: + try: + cursor.execute( + f"DELETE FROM {self.FTS_TABLE} WHERE path = ? AND source = ?", + (path, source.value), + ) + except Exception as e: + logger.debug(f"FTS delete failed: {e}") + + # Delete chunks and file + cursor.execute( + "DELETE FROM chunks WHERE path = ? AND source = ?", + (path, source.value), + ) + cursor.execute( + "DELETE FROM files WHERE path = ? AND source = ?", + (path, source.value), + ) + + async def get_file_hash(self, path: str, source: MemorySource) -> str | None: + """Get file hash.""" + cursor = self.conn.cursor() + cursor.execute( + "SELECT hash FROM files WHERE path = ? AND source = ?", + (path, source.value), + ) + row = cursor.fetchone() + cursor.close() + return row[0] if row else None + + async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None: + """Get file metadata with chunk count.""" + cursor = self.conn.cursor() + cursor.execute( + "SELECT hash, mtime, size FROM files WHERE path = ? AND source = ?", + (path, source.value), + ) + row = cursor.fetchone() + if not row: + cursor.close() + return None + + hash_val, mtime, size = row + cursor.execute( + "SELECT COUNT(*) FROM chunks WHERE path = ? AND source = ?", + (path, source.value), + ) + chunk_count = cursor.fetchone()[0] + cursor.close() + + return FileMetadata( + hash=hash_val, + mtime_ms=mtime, + size=size, + chunk_count=chunk_count, + ) + + async def list_files(self, source: MemorySource) -> list[str]: + """List all indexed files.""" + cursor = self.conn.cursor() + cursor.execute("SELECT path FROM files WHERE source = ?", (source.value,)) + paths = [row[0] for row in cursor.fetchall()] + cursor.close() + return paths + + async def get_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]: + """Get all chunks for a file.""" + cursor = self.conn.cursor() + cursor.execute( + """ + SELECT id, path, source, start_line, end_line, text, hash, embedding + FROM chunks 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, + ), + ) + + cursor.close() + 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_available 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 + query_sql = f""" + SELECT c.id, c.path, c.start_line, c.end_line, c.source, c.text, v.distance + FROM {self.VECTOR_TABLE} v + JOIN chunks c ON v.id = c.id + WHERE v.embedding MATCH ? + """ + query_params: list = [query_blob] + + # Add source filter if specified + if source_filter: + query_sql += source_filter + query_params.extend(params) + + # Order and limit results + query_sql += " ORDER BY v.distance LIMIT ?" + query_params.append(str(limit)) + + cursor.execute(query_sql, query_params) + + results = [] + for _, path, start, end, src, text, dist in cursor.fetchall(): + score = max(0.0, 1.0 - dist) + snippet = text[: self.snippet_max_chars] if len(text) > self.snippet_max_chars else text + results.append( + MemorySearchResult( + path=path, + start_line=start, + end_line=end, + score=score, + snippet=snippet, + source=MemorySource(src), + ), + ) + + return results + except Exception as e: + logger.error(f"Vector search failed: {e}") + return [] + finally: + cursor.close() + + async def keyword_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform full-text search.""" + if not self.fts_available: + return [] + + # Build FTS5 query, escaping quotes + cleaned = query.strip().replace('"', '""') + if not cleaned: + return [] + fts_query = f'"{cleaned}"' + + 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} fts + WHERE fts.text MATCH ?{source_filter} + ORDER BY rank + LIMIT ? + """, + params, + ) + + results = [] + for _, path, start, end, src, text, rank in cursor.fetchall(): + # Convert BM25 rank (negative) to 0-1 score (higher=better) + score = max(0.0, 1.0 / (1.0 + abs(rank))) + snippet = text[: self.snippet_max_chars] if len(text) > self.snippet_max_chars else text + results.append( + MemorySearchResult( + path=path, + start_line=start, + end_line=end, + score=score, + snippet=snippet, + source=MemorySource(src), + ), + ) + + return results + except Exception as e: + logger.error(f"Keyword search failed: {e}") + return [] + finally: + cursor.close() + + async def read_meta(self, key: str) -> MemoryIndexMeta | None: + """Read metadata value.""" + cursor = self.conn.cursor() + cursor.execute("SELECT value FROM meta WHERE key = ?", (key,)) + row = cursor.fetchone() + cursor.close() + + if not row: + return None + + return MemoryIndexMeta(**json.loads(row[0])) + + async def write_meta(self, key: str, value: MemoryIndexMeta | dict) -> None: + """Write metadata value.""" + data = value.model_dump() if isinstance(value, MemoryIndexMeta) else value + cursor = self.conn.cursor() + cursor.execute( + """ + INSERT OR REPLACE INTO meta (key, value) + VALUES (?, ?) + """, + (key, json.dumps(data)), + ) + self.conn.commit() + cursor.close() + + async def clear_all(self): + """Clear all indexed data.""" + cursor = self.conn.cursor() + cursor.execute("BEGIN") + + try: + cursor.execute("DELETE FROM files") + cursor.execute("DELETE FROM chunks") + + if self.vector_available: + try: + cursor.execute(f"DELETE FROM {self.VECTOR_TABLE}") + except Exception as e: + logger.debug(f"Vector clear failed: {e}") + + if self.fts_available: + try: + cursor.execute(f"DELETE FROM {self.FTS_TABLE}") + except Exception as e: + logger.debug(f"FTS clear failed: {e}") + + cursor.execute("COMMIT") + except Exception: + cursor.execute("ROLLBACK") + raise + finally: + cursor.close() + + async def close(self): + """Close database connection.""" + if self.conn: + self.conn.close() + self.conn = None diff --git a/reme/core/memory_manager/utils/__init__.py b/reme/core/memory_manager/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/core/memory_manager/utils/hashing.py b/reme/core/memory_manager/utils/hashing.py new file mode 100644 index 00000000..7894d69b --- /dev/null +++ b/reme/core/memory_manager/utils/hashing.py @@ -0,0 +1,15 @@ +"""Utility functions for hashing text content.""" + +import hashlib + + +def hash_text(text: str) -> str: + """Generate SHA-256 hash of text content. + + Args: + text: Input text to hash + + Returns: + Hexadecimal representation of the SHA-256 hash + """ + return hashlib.sha256(text.encode("utf-8")).hexdigest() diff --git a/reme/core/schema/__init__.py b/reme/core/schema/__init__.py index b7b73719..a4a4214e 100644 --- a/reme/core/schema/__init__.py +++ b/reme/core/schema/__init__.py @@ -1,6 +1,10 @@ """schema""" +from .file_metadata import FileMetadata +from .memory_chunk import MemoryChunk +from .memory_index_meta import MemoryIndexMeta from .memory_node import MemoryNode +from .memory_search_result import MemorySearchResult from .message import ContentBlock, Message, Trajectory from .request import Request from .response import Response @@ -17,16 +21,22 @@ from .service_config import ( ) from .stream_chunk import StreamChunk from .tool_call import ToolAttr, ToolCall +from .truncation_result import TruncationResult from .vector_node import VectorNode __all__ = [ - "MemoryNode", + "CmdConfig", "ContentBlock", "EmbeddingModelConfig", + "FileMetadata", "FlowConfig", "HttpConfig", "LLMConfig", "MCPConfig", + "MemoryChunk", + "MemoryIndexMeta", + "MemoryNode", + "MemorySearchResult", "Message", "Request", "Response", @@ -36,7 +46,7 @@ __all__ = [ "Trajectory", "ToolAttr", "ToolCall", + "TruncationResult", "VectorNode", "VectorStoreConfig", - "CmdConfig", ] diff --git a/reme/core/schema/file_metadata.py b/reme/core/schema/file_metadata.py new file mode 100644 index 00000000..053f1e62 --- /dev/null +++ b/reme/core/schema/file_metadata.py @@ -0,0 +1,20 @@ +"""File metadata schema.""" + +from pydantic import BaseModel, Field + + +class FileMetadata(BaseModel): + """File metadata with optional extended fields for various use cases.""" + + # Core fields (always required) + hash: str = Field(default=..., description="Hash of the file content") + mtime_ms: float = Field(default=..., description="Last modification time in milliseconds") + size: int = Field(default=..., description="File size in bytes") + + # Extended fields for session files + path: str | None = Field(default=None, description="Relative path to the session file") + abs_path: str | None = Field(default=None, description="Absolute path to the session file") + content: str | None = Field(default=None, description="Parsed content from the session file") + + # Extended fields for statistics + chunk_count: int | None = Field(default=None, description="Number of chunks in the file") diff --git a/reme/core/schema/memory_chunk.py b/reme/core/schema/memory_chunk.py new file mode 100644 index 00000000..36959da1 --- /dev/null +++ b/reme/core/schema/memory_chunk.py @@ -0,0 +1,19 @@ +"""Memory chunk schema.""" + +from pydantic import BaseModel, Field + +from ..enumeration import MemorySource + + +class MemoryChunk(BaseModel): + """A chunk of memory content with metadata.""" + + id: str = Field(..., description="Unique identifier for the chunk") + path: str = Field(..., description="File path relative to workspace") + source: MemorySource = Field(..., description="Source of the memory data") + start_line: int = Field(..., description="Starting line number in the source file") + end_line: int = Field(..., description="Ending line number in the source file") + text: str = Field(..., description="Text content of the chunk") + hash: str = Field(..., description="Hash of the chunk content") + embedding: list[float] | None = Field(default=None, description="Vector embedding of the chunk") + metadata: dict = Field(default_factory=dict, description="Additional metadata") diff --git a/reme/core/schema/memory_index_meta.py b/reme/core/schema/memory_index_meta.py new file mode 100644 index 00000000..a4538b77 --- /dev/null +++ b/reme/core/schema/memory_index_meta.py @@ -0,0 +1,14 @@ +"""Memory index metadata schema.""" + +from typing import Optional + +from pydantic import BaseModel, Field + + +class MemoryIndexMeta(BaseModel): + """Metadata for memory index configuration.""" + + model: str = Field(..., description="Name of the embedding model") + chunk_tokens: int = Field(..., description="Maximum tokens per chunk") + chunk_overlap: int = Field(..., description="Number of overlapping tokens between chunks") + vector_dims: Optional[int] = Field(default=None, description="Vector embedding dimensions") diff --git a/reme/core/schema/memory_search_result.py b/reme/core/schema/memory_search_result.py new file mode 100644 index 00000000..d08ca12b --- /dev/null +++ b/reme/core/schema/memory_search_result.py @@ -0,0 +1,24 @@ +"""Memory search result schema.""" + +from typing import Any, Dict + +from pydantic import BaseModel, Field + +from ..enumeration import MemorySource + + +class MemorySearchResult(BaseModel): + """Search result from memory index.""" + + path: str = Field(..., description="File path relative to workspace") + start_line: int = Field(..., description="Starting line number of the match") + end_line: int = Field(..., description="Ending line number of the match") + score: float = Field(..., description="Relevance score of the search result") + snippet: str = Field(..., description="Text snippet from the matched content") + source: MemorySource = Field(..., description="Source of the memory data") + metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata") + + @property + def merge_key(self) -> str: + """Merge key for the search result.""" + return self.path + f":{self.start_line}:{self.end_line}" diff --git a/reme/core/schema/truncation_result.py b/reme/core/schema/truncation_result.py new file mode 100644 index 00000000..18e8cc1a --- /dev/null +++ b/reme/core/schema/truncation_result.py @@ -0,0 +1,35 @@ +"""Truncation result schema for command output truncation.""" + +from typing import Literal + +from pydantic import BaseModel, Field + + +class TruncationResult(BaseModel): + """Result of output truncation operation. + + Attributes: + content: The truncated content + truncated: Whether truncation occurred + total_lines: Total number of lines in original output + output_lines: Number of lines in truncated output + total_bytes: Total bytes in original output + output_bytes: Bytes in truncated output + truncated_by: What caused truncation ('lines' or 'bytes') + last_line_partial: Whether last line was partially truncated + """ + + content: str = Field(description="The truncated content") + truncated: bool = Field(description="Whether truncation occurred") + total_lines: int = Field(description="Total number of lines in original output") + output_lines: int = Field(description="Number of lines in truncated output") + total_bytes: int = Field(description="Total bytes in original output") + output_bytes: int = Field(description="Bytes in truncated output") + truncated_by: Literal["lines", "bytes"] | None = Field( + default=None, + description="What caused truncation ('lines' or 'bytes')", + ) + last_line_partial: bool = Field( + default=False, + description="Whether last line was partially truncated", + ) diff --git a/reme/reme.py b/reme/reme.py index 86f87de4..060019b7 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -82,11 +82,14 @@ class ReMe(Application): Example: ```python - reme = await ReMe(...).start() + reme = ReMe(...) + await reme.start() # reme = await ReMe.create(...) # both ok await reme.summarize_memory(...) await reme.retrieve_memory(...) + + await reme.close() ``` """ diff --git a/reme/tool/fs/__init__.py b/reme/tool/fs/__init__.py new file mode 100644 index 00000000..090e1a1f --- /dev/null +++ b/reme/tool/fs/__init__.py @@ -0,0 +1,24 @@ +"""File system tools.""" + +from .bash_tool import BashTool +from .edit_tool import EditTool +from .find_tool import FindTool +from .grep_tool import GrepTool +from .ls_tool import LsTool +from .read_tool import ReadTool +from .write_tool import WriteTool +from ...core import R + +__all__ = [ + "BashTool", + "EditTool", + "FindTool", + "GrepTool", + "LsTool", + "ReadTool", + "WriteTool", +] + +for name in __all__: + tool_class = globals()[name] + R.op.register(tool_class) diff --git a/reme/tool/fs/bash_tool.py b/reme/tool/fs/bash_tool.py new file mode 100644 index 00000000..d8aea81c --- /dev/null +++ b/reme/tool/fs/bash_tool.py @@ -0,0 +1,191 @@ +"""Bash command execution tool with production-grade features. + +This module provides a production-grade tool for executing bash commands with: +- Smart output truncation (keeps last N lines/bytes to prevent memory issues) +- Process tree termination (prevents orphan processes) +""" + +import asyncio +import os +import platform +import signal +from pathlib import Path + +from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncate_tail +from ...core.op import BaseTool +from ...core.schema import ToolCall, TruncationResult + + +def get_shell_config() -> tuple[str, list[str]]: + """Get the appropriate shell and arguments for the current platform. + + Returns: + Tuple of (shell_path, args) for subprocess execution + """ + system = platform.system() + + if system == "Windows": + # Use PowerShell on Windows + return "powershell.exe", ["-Command"] + else: + # Use bash on Unix-like systems + shell = os.environ.get("SHELL", "/bin/bash") + return shell, ["-c"] + + +def kill_process_tree(pid: int) -> None: + """Kill a process and all its children. + + Args: + pid: Process ID to kill + """ + try: + if platform.system() == "Windows": + # Windows: use taskkill + os.system(f"taskkill /F /T /PID {pid}") + else: + # Unix: kill process group + try: + os.killpg(os.getpgid(pid), signal.SIGTERM) + except ProcessLookupError: + pass # Process already dead + except Exception: + pass # Best effort + + +class BashTool(BaseTool): + """Production-grade tool for executing bash commands. + + Features: + - Smart output truncation (preserves last N lines or M bytes) + - Kills entire process tree on timeout (prevents orphan processes) + """ + + def __init__(self, cwd: str | None = None, command_prefix: str | None = None): + """Initialize bash tool. + + Args: + cwd: Working directory (defaults to current directory) + command_prefix: Optional prefix prepended to every command + """ + super().__init__() + self.cwd = cwd or os.getcwd() + self.command_prefix = command_prefix + + def _build_tool_call(self) -> ToolCall: + max_kb = DEFAULT_MAX_BYTES // 1024 + return ToolCall( + **{ + "description": ( + f"Execute a bash command in the current working directory. " + f"Returns stdout and stderr. Output is truncated to last " + f"{DEFAULT_MAX_LINES} lines or {max_kb}KB (whichever is hit first). " + f"Optionally provide a timeout in seconds." + ), + "parameters": { + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Bash command to execute", + }, + "timeout": { + "type": "number", + "description": "Timeout in seconds (optional, no default timeout)", + }, + }, + "required": ["command"], + }, + }, + ) + + async def execute(self) -> str: + """Execute the bash command with production-grade features.""" + command: str = self.context.command + timeout: float | None = self.context.get("timeout", None) + + # Apply command prefix if configured + if self.command_prefix: + command = f"{self.command_prefix}\n{command}" + + # Verify working directory exists + if not Path(self.cwd).exists(): + raise FileNotFoundError( + f"Working directory does not exist: {self.cwd}\n" f"Cannot execute bash commands.", + ) + + # Get shell configuration + shell, shell_args = get_shell_config() + + # Start process + try: + process = await asyncio.create_subprocess_exec( + shell, + *shell_args, + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self.cwd, + # Create process group for clean termination + preexec_fn=os.setpgrp if platform.system() != "Windows" else None, + ) + except Exception as e: + raise RuntimeError(f"Failed to start process: {e}") from e + + # Execute command with optional timeout + try: + if timeout and timeout > 0: + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), + timeout=timeout, + ) + except asyncio.TimeoutError as e: + # Kill process tree on timeout + if process.pid: + kill_process_tree(process.pid) + try: + await asyncio.wait_for(process.wait(), timeout=1.0) + except asyncio.TimeoutError: + process.kill() + raise TimeoutError(f"Command timed out after {timeout} seconds") from e + else: + stdout, stderr = await process.communicate() + except TimeoutError as e: + raise RuntimeError(str(e)) from e + + # Decode output + full_output = stdout.decode("utf-8", errors="ignore") + if stderr: + stderr_text = stderr.decode("utf-8", errors="ignore") + if full_output: + full_output += "\n" + full_output += stderr_text + + # Apply tail truncation_result to prevent memory issues + truncation_result: TruncationResult = truncate_tail(full_output) + output_text = truncation_result.content or "(no output)" + + # Build truncation_result notice if needed + if truncation_result.truncated: + start_line = truncation_result.total_lines - truncation_result.output_lines + 1 + end_line = truncation_result.total_lines + + if truncation_result.truncated_by == "lines": + output_text += ( + f"\n\n[Output truncated: showing lines {start_line}-{end_line} " + f"of {truncation_result.total_lines} total lines]" + ) + else: + max_kb = DEFAULT_MAX_BYTES // 1024 + output_text += ( + f"\n\n[Output truncated: showing lines {start_line}-{end_line} " + f"of {truncation_result.total_lines} ({max_kb}KB limit reached)]" + ) + + # Handle non-zero exit code + if process.returncode != 0: + output_text += f"\n\nCommand exited with code {process.returncode}" + raise RuntimeError(output_text) + + return output_text diff --git a/reme/tool/fs/edit_diff.py b/reme/tool/fs/edit_diff.py new file mode 100644 index 00000000..76e2fcdb --- /dev/null +++ b/reme/tool/fs/edit_diff.py @@ -0,0 +1,164 @@ +"""Diff utilities for edit tool.""" + +import re +from dataclasses import dataclass +from difflib import unified_diff + + +def detect_line_ending(content: str) -> str: + """Detect line ending style (CRLF or LF).""" + crlf_idx = content.find("\r\n") + lf_idx = content.find("\n") + if lf_idx == -1: + return "\n" + if crlf_idx == -1: + return "\n" + return "\r\n" if crlf_idx < lf_idx else "\n" + + +def normalize_to_lf(text: str) -> str: + """Normalize line endings to LF.""" + return text.replace("\r\n", "\n").replace("\r", "\n") + + +def restore_line_endings(text: str, ending: str) -> str: + """Restore original line endings.""" + return text.replace("\n", ending) if ending == "\r\n" else text + + +def normalize_for_fuzzy_match(text: str) -> str: + """Normalize text for fuzzy matching: strip trailing whitespace, normalize quotes/dashes.""" + lines = text.split("\n") + normalized = "\n".join(line.rstrip() for line in lines) + + # Smart quotes → ASCII + normalized = re.sub(r"[\u2018\u2019\u201A\u201B]", "'", normalized) + normalized = re.sub(r"[\u201C\u201D\u201E\u201F]", '"', normalized) + + # Dashes → hyphen + normalized = re.sub(r"[\u2010\u2011\u2012\u2013\u2014\u2015\u2212]", "-", normalized) + + # Special spaces → regular space + normalized = re.sub(r"[\u00A0\u2002-\u200A\u202F\u205F\u3000]", " ", normalized) + + return normalized + + +@dataclass +class FuzzyMatchResult: + """Result of fuzzy text matching.""" + + found: bool + index: int + match_length: int + used_fuzzy_match: bool + content_for_replacement: str + + +def fuzzy_find_text(content: str, old_text: str) -> FuzzyMatchResult: + """Find old_text in content, trying exact match first, then fuzzy match.""" + # Try exact match + exact_index = content.find(old_text) + if exact_index != -1: + return FuzzyMatchResult( + found=True, + index=exact_index, + match_length=len(old_text), + used_fuzzy_match=False, + content_for_replacement=content, + ) + + # Try fuzzy match + fuzzy_content = normalize_for_fuzzy_match(content) + fuzzy_old_text = normalize_for_fuzzy_match(old_text) + fuzzy_index = fuzzy_content.find(fuzzy_old_text) + + if fuzzy_index == -1: + return FuzzyMatchResult( + found=False, + index=-1, + match_length=0, + used_fuzzy_match=False, + content_for_replacement=content, + ) + + return FuzzyMatchResult( + found=True, + index=fuzzy_index, + match_length=len(fuzzy_old_text), + used_fuzzy_match=True, + content_for_replacement=fuzzy_content, + ) + + +def strip_bom(content: str) -> tuple[str, str]: + """Strip UTF-8 BOM, return (bom, text_without_bom).""" + if content.startswith("\ufeff"): + return "\ufeff", content[1:] + return "", content + + +@dataclass +class DiffResult: + """Result of diff generation.""" + + diff: str + first_changed_line: int | None + + +def generate_diff_string(old_content: str, new_content: str, context_lines: int = 4) -> DiffResult: + """Generate unified diff with line numbers.""" + old_lines = old_content.split("\n") + new_lines = new_content.split("\n") + + # Use difflib to get the changes + diff_lines = list( + unified_diff( + old_lines, + new_lines, + lineterm="", + n=context_lines, + ), + ) + + if not diff_lines: + return DiffResult(diff="", first_changed_line=None) + + # Parse and format the diff + output = [] + first_changed_line = None + max_line_num = max(len(old_lines), len(new_lines)) + line_num_width = len(str(max_line_num)) + + old_line_num = 1 + new_line_num = 1 + + for line in diff_lines[2:]: # Skip header lines + if line.startswith("@@"): + # Parse hunk header + match = re.match(r"@@ -(\d+),?\d* \+(\d+),?\d* @@", line) + if match: + old_line_num = int(match.group(1)) + new_line_num = int(match.group(2)) + continue + + if line.startswith("+"): + if first_changed_line is None: + first_changed_line = new_line_num + line_num = str(new_line_num).rjust(line_num_width) + output.append(f"+{line_num} {line[1:]}") + new_line_num += 1 + elif line.startswith("-"): + if first_changed_line is None: + first_changed_line = new_line_num + line_num = str(old_line_num).rjust(line_num_width) + output.append(f"-{line_num} {line[1:]}") + old_line_num += 1 + else: + # Context line + line_num = str(old_line_num).rjust(line_num_width) + output.append(f" {line_num} {line[1:] if line.startswith(' ') else line}") + old_line_num += 1 + new_line_num += 1 + + return DiffResult(diff="\n".join(output), first_changed_line=first_changed_line) diff --git a/reme/tool/fs/edit_tool.py b/reme/tool/fs/edit_tool.py new file mode 100644 index 00000000..fe73bd1f --- /dev/null +++ b/reme/tool/fs/edit_tool.py @@ -0,0 +1,141 @@ +"""File editing tool with exact text replacement.""" + +import os +from pathlib import Path + +from .edit_diff import ( + detect_line_ending, + fuzzy_find_text, + generate_diff_string, + normalize_for_fuzzy_match, + normalize_to_lf, + restore_line_endings, + strip_bom, +) +from ...core.op import BaseTool +from ...core.schema import ToolCall + + +class EditTool(BaseTool): + """Edit a file by replacing exact text.""" + + def __init__(self, cwd: str | None = None): + """Initialize edit tool. + + Args: + cwd: Working directory (defaults to current directory) + """ + super().__init__() + self.cwd = cwd or os.getcwd() + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": ( + "Edit a file by replacing exact text. The oldText must match exactly " + "(including whitespace). Use this for precise, surgical edits." + ), + "parameters": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to edit (relative or absolute)", + }, + "oldText": { + "type": "string", + "description": "Exact text to find and replace (must match exactly)", + }, + "newText": { + "type": "string", + "description": "New text to replace the old text with", + }, + }, + "required": ["path", "oldText", "newText"], + }, + }, + ) + + async def execute(self) -> str: + """Execute the edit operation.""" + path: str = self.context.path + old_text: str = self.context.oldText + new_text: str = self.context.newText + + # Resolve path + if not os.path.isabs(path): + absolute_path = os.path.join(self.cwd, path) + else: + absolute_path = path + + # Check file exists and is writable + path_obj = Path(absolute_path) + if not path_obj.exists(): + raise FileNotFoundError(f"File not found: {path}") + + if not os.access(absolute_path, os.R_OK | os.W_OK): + raise PermissionError(f"File not readable/writable: {path}") + + # Read file + try: + with open(absolute_path, "r", encoding="utf-8") as f: + raw_content = f.read() + except Exception as e: + raise IOError(f"Failed to read file {path}: {e}") from e + + # Strip BOM (LLM won't include invisible BOM in oldText) + bom, content = strip_bom(raw_content) + + original_ending = detect_line_ending(content) + normalized_content = normalize_to_lf(content) + normalized_old_text = normalize_to_lf(old_text) + normalized_new_text = normalize_to_lf(new_text) + + # Find old text using fuzzy matching + match_result = fuzzy_find_text(normalized_content, normalized_old_text) + + if not match_result.found: + raise ValueError( + f"Could not find the exact text in {path}. The old text must match " + f"exactly including all whitespace and newlines.", + ) + + # Count occurrences for uniqueness check + fuzzy_content = normalize_for_fuzzy_match(normalized_content) + fuzzy_old_text = normalize_for_fuzzy_match(normalized_old_text) + occurrences = fuzzy_content.count(fuzzy_old_text) + + if occurrences > 1: + raise ValueError( + f"Found {occurrences} occurrences of the text in {path}. " + f"The text must be unique. Please provide more context to make it unique.", + ) + + # Perform replacement + base_content = match_result.content_for_replacement + new_content = ( + base_content[: match_result.index] + + normalized_new_text + + base_content[match_result.index + match_result.match_length :] + ) + + # Verify replacement changed something + if base_content == new_content: + raise ValueError( + f"No changes made to {path}. The replacement produced identical content. " + f"This might indicate an issue with special characters or the text not " + f"exist as expected.", + ) + + # Write file + final_content = bom + restore_line_endings(new_content, original_ending) + try: + with open(absolute_path, "w", encoding="utf-8") as f: + f.write(final_content) + except Exception as e: + raise IOError(f"Failed to write file {path}: {e}") from e + + # Generate diff + diff_result = generate_diff_string(base_content, new_content) + + return f"Successfully replaced text in {path}.\n\n{diff_result.diff}" diff --git a/reme/tool/fs/find_tool.py b/reme/tool/fs/find_tool.py new file mode 100644 index 00000000..7b855586 --- /dev/null +++ b/reme/tool/fs/find_tool.py @@ -0,0 +1,183 @@ +"""File search tool using glob patterns with gitignore support.""" + +import os +from pathlib import Path + +from .truncate import FIND_MAX_BYTES, FIND_MAX_LINES, format_size, truncate_head +from ...core.op import BaseTool +from ...core.schema import ToolCall + + +class FindTool(BaseTool): + """Search for files by glob pattern, respecting .gitignore.""" + + def __init__(self, cwd: str | None = None): + """Initialize find tool. + + Args: + cwd: Working directory (defaults to current directory) + """ + super().__init__() + self.cwd = cwd or os.getcwd() + + def _build_tool_call(self) -> ToolCall: + max_kb = FIND_MAX_BYTES // 1024 + return ToolCall( + **{ + "description": ( + f"Search for files by glob pattern. Returns matching file paths relative " + f"to the search directory. Respects .gitignore. Output is truncated to " + f"1000 results or {max_kb}KB (whichever is hit first)." + ), + "parameters": { + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "Glob pattern to match files, " + "e.g. '*.ts', '**/*.json', or 'src/**/*.spec.ts'", + }, + "path": { + "type": "string", + "description": "Directory to search in (default: current directory)", + }, + "limit": { + "type": "number", + "description": "Maximum number of results (default: 1000)", + }, + }, + "required": ["pattern"], + }, + }, + ) + + def _load_gitignore_patterns(self, search_path: Path) -> list[str]: + """Load gitignore patterns from directory and subdirectories.""" + patterns = ["**/node_modules/**", "**/.git/**"] + + # Load root .gitignore + gitignore_path = search_path / ".gitignore" + if gitignore_path.exists(): + try: + with open(gitignore_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line and not line.startswith("#"): + patterns.append(line) + except Exception: + pass # Ignore errors + + # Load nested .gitignore files + try: + for gitignore in search_path.rglob(".gitignore"): + if gitignore == gitignore_path: + continue + try: + with open(gitignore, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line and not line.startswith("#"): + patterns.append(line) + except Exception: + pass # Ignore errors + except Exception: + pass # Ignore glob errors + + return patterns + + def _should_ignore(self, path: Path, ignore_patterns: list[str]) -> bool: + """Check if path matches any ignore pattern.""" + path_str = str(path) + + for pattern in ignore_patterns: + # Simple pattern matching (not full gitignore spec) + if "**" in pattern: + # Recursive match + clean_pattern = pattern.replace("**/", "").replace("/**", "") + if clean_pattern in path_str: + return True + elif "*" in pattern: + # Wildcard match + from fnmatch import fnmatch + + if fnmatch(path.name, pattern): + return True + elif pattern in path_str: + return True + + return False + + async def execute(self) -> str: + """Execute file search.""" + pattern: str = self.context.pattern + search_dir: str = self.context.get("path", ".") + limit: int = self.context.get("limit", 1000) + + # Resolve search path + if not os.path.isabs(search_dir): + search_path = Path(self.cwd) / search_dir + else: + search_path = Path(search_dir) + + # Check if directory exists + if not search_path.exists(): + raise FileNotFoundError(f"Path not found: {search_dir}") + + if not search_path.is_dir(): + raise NotADirectoryError(f"Path is not a directory: {search_dir}") + + # Load gitignore patterns + ignore_patterns = self._load_gitignore_patterns(search_path) + + # Search for files + results = [] + try: + for file_path in search_path.glob(pattern): + if len(results) >= limit: + break + + # Skip if matches ignore patterns + if self._should_ignore(file_path, ignore_patterns): + continue + + # Get relative path + try: + rel_path = file_path.relative_to(search_path) + # Add trailing slash for directories + if file_path.is_dir(): + results.append(f"{rel_path}/") + else: + results.append(str(rel_path)) + except ValueError: + # If relative_to fails, use the path as-is + results.append(str(file_path)) + except Exception as e: + raise RuntimeError(f"Error searching for files: {e}") from e + + # Handle no results + if not results: + return "No files found matching pattern" + + # Sort results for consistency + results.sort() + + # Apply limit and truncation + result_limit_reached = len(results) >= limit + raw_output = "\n".join(results) + truncation = truncate_head(raw_output, max_lines=FIND_MAX_LINES, max_bytes=FIND_MAX_BYTES) + + output = truncation.content + notices = [] + + if result_limit_reached: + notices.append( + f"{limit} results limit reached. Use limit={limit * 2} for more, or refine pattern", + ) + + if truncation.truncated: + notices.append(f"{format_size(FIND_MAX_BYTES)} limit reached") + + if notices: + output += f"\n\n[{'. '.join(notices)}]" + + return output diff --git a/reme/tool/fs/grep_tool.py b/reme/tool/fs/grep_tool.py new file mode 100644 index 00000000..9df9f8f2 --- /dev/null +++ b/reme/tool/fs/grep_tool.py @@ -0,0 +1,276 @@ +"""Grep tool for searching file contents using ripgrep. + +This module provides a tool for searching file contents with: +- Pattern matching (regex or literal string) +- Smart output truncation (prevents memory issues) +- Context lines support +- Respects .gitignore +""" + +import asyncio +import json +import os +import shutil +from pathlib import Path + +from .truncate import ( + DEFAULT_MAX_BYTES, + GREP_MAX_LINE_LENGTH, + format_size, + truncate_head, + truncate_line, +) +from ...core.op import BaseTool +from ...core.schema import ToolCall + +# Default limits +DEFAULT_LIMIT = 100 # Maximum number of matches + + +class GrepTool(BaseTool): + """Tool for searching file contents using ripgrep. + + Features: + - Pattern matching with regex or literal string + - Context lines support + - Smart output truncation + - Respects .gitignore + """ + + def __init__(self, cwd: str | None = None): + """Initialize grep tool. + + Args: + cwd: Working directory (defaults to current directory) + """ + super().__init__() + self.cwd = cwd or os.getcwd() + + def _build_tool_call(self) -> ToolCall: + max_kb = DEFAULT_MAX_BYTES // 1024 + return ToolCall( + **{ + "description": ( + f"Search file contents for a pattern. Returns matching lines with " + f"file paths and line numbers. Respects .gitignore. Output is " + f"truncated to {DEFAULT_LIMIT} matches or {max_kb}KB (whichever is " + f"hit first). Long lines are truncated to {GREP_MAX_LINE_LENGTH} chars." + ), + "parameters": { + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "Search pattern (regex or literal string)", + }, + "path": { + "type": "string", + "description": "Directory or file to search (default: current directory)", + }, + "glob": { + "type": "string", + "description": "Filter files by glob pattern, e.g. '*.ts' or '**/*.spec.ts'", + }, + "ignoreCase": { + "type": "boolean", + "description": "Case-insensitive search (default: false)", + }, + "literal": { + "type": "boolean", + "description": "Treat pattern as literal string instead of regex (default: false)", + }, + "contextLines": { + "type": "number", + "description": "Number of lines to show before and after each match (default: 0)", + }, + "limit": { + "type": "number", + "description": f"Maximum number of matches to return (default: {DEFAULT_LIMIT})", + }, + }, + "required": ["pattern"], + }, + }, + ) + + async def execute(self) -> str: + """Execute the grep search.""" + pattern: str = self.context.pattern + search_path: str = self.context.get("path", ".") + glob: str | None = self.context.get("glob", None) + ignore_case: bool = self.context.get("ignoreCase", False) + literal: bool = self.context.get("literal", False) + context_lines: int = self.context.get("contextLines", 0) + limit: int = self.context.get("limit", DEFAULT_LIMIT) + + # Check if ripgrep is available + rg_path = shutil.which("rg") + if not rg_path: + raise RuntimeError( + "ripgrep (rg) is not available. Please install it:\n" + " macOS: brew install ripgrep\n" + " Ubuntu: apt-get install ripgrep\n" + " Other: https://github.com/BurntSushi/ripgrep", + ) + + # Resolve search path + if not os.path.isabs(search_path): + search_path = os.path.join(self.cwd, search_path) + + # Check if path exists + if not Path(search_path).exists(): + raise FileNotFoundError(f"Path not found: {search_path}") + + is_directory = Path(search_path).is_dir() + effective_limit = max(1, limit) + + # Build ripgrep arguments + args = [ + rg_path, + "--json", + "--line-number", + "--color=never", + "--hidden", + ] + + if ignore_case: + args.append("--ignore-case") + + if literal: + args.append("--fixed-strings") + + if glob: + args.extend(["--glob", glob]) + + args.extend([pattern, search_path]) + + # Execute ripgrep + try: + process = await asyncio.create_subprocess_exec( + *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self.cwd, + ) + except Exception as e: + raise RuntimeError(f"Failed to run ripgrep: {e}") from e + + stdout, stderr = await process.communicate() + + # Parse JSON output + matches = [] + match_count = 0 + lines_truncated = False + + for line in stdout.decode("utf-8", errors="ignore").splitlines(): + if not line.strip() or match_count >= effective_limit: + break + + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + + if event.get("type") == "match": + match_count += 1 + data = event.get("data", {}) + file_path = data.get("path", {}).get("text", "") + line_number = data.get("line_number", 0) + + if file_path and line_number: + matches.append({"file_path": file_path, "line_number": line_number}) + + if match_count >= effective_limit: + break + + # Check for errors + if process.returncode not in (0, 1) and match_count == 0: + error_msg = stderr.decode("utf-8", errors="ignore").strip() + if error_msg: + raise RuntimeError(error_msg) + raise RuntimeError(f"ripgrep exited with code {process.returncode}") + + # No matches found + if match_count == 0: + return "No matches found" + + # Format matches with context + output_lines = [] + file_cache = {} + + for match in matches: + file_path = match["file_path"] + line_number = match["line_number"] + + # Read file if not cached + if file_path not in file_cache: + try: + with open(file_path, "r", encoding="utf-8", errors="ignore") as f: + file_cache[file_path] = f.read().replace("\r\n", "\n").replace("\r", "\n").split("\n") + except Exception: + file_cache[file_path] = [] + + lines = file_cache[file_path] + + # Format relative path + if is_directory: + relative_path = os.path.relpath(file_path, search_path) + if not relative_path.startswith(".."): + display_path = relative_path.replace("\\", "/") + else: + display_path = os.path.basename(file_path) + else: + display_path = os.path.basename(file_path) + + # Generate context block + if not lines: + output_lines.append(f"{display_path}:{line_number}: (unable to read file)") + continue + + context_value = max(0, context_lines) + start = max(1, line_number - context_value) if context_value > 0 else line_number + end = min(len(lines), line_number + context_value) if context_value > 0 else line_number + + for current in range(start, end + 1): + if current < 1 or current > len(lines): + continue + + line_text = lines[current - 1] + is_match_line = current == line_number + + # Truncate long lines + truncated_text, was_truncated = truncate_line(line_text) + if was_truncated: + lines_truncated = True + + if is_match_line: + output_lines.append(f"{display_path}:{current}: {truncated_text}") + else: + output_lines.append(f"{display_path}-{current}- {truncated_text}") + + # Apply byte truncation + raw_output = "\n".join(output_lines) + truncation = truncate_head(raw_output, max_lines=999999999) + + output = truncation.content + notices = [] + + # Add notices + if match_count >= effective_limit: + notices.append( + f"{effective_limit} matches limit reached. " + f"Use limit={effective_limit * 2} for more, or refine pattern", + ) + + if truncation.truncated: + notices.append(f"{format_size(DEFAULT_MAX_BYTES)} limit reached") + + if lines_truncated: + notices.append( + f"Some lines truncated to {GREP_MAX_LINE_LENGTH} chars. " f"Use read tool to see full lines", + ) + + if notices: + output += f"\n\n[{'. '.join(notices)}]" + + return output diff --git a/reme/tool/fs/ls_tool.py b/reme/tool/fs/ls_tool.py new file mode 100644 index 00000000..1c054c62 --- /dev/null +++ b/reme/tool/fs/ls_tool.py @@ -0,0 +1,128 @@ +"""Directory listing tool with truncation support.""" + +import os +from pathlib import Path + +from .truncate import DEFAULT_MAX_BYTES, truncate_head +from ...core.op import BaseTool +from ...core.schema import ToolCall + +DEFAULT_LIMIT = 500 + + +class LsTool(BaseTool): + """List directory contents with smart truncation. + + Features: + - Returns entries sorted alphabetically (case-insensitive) + - Directory indicators ('/' suffix) + - Includes dotfiles + - Entry count limiting + - Byte truncation + """ + + def __init__(self, cwd: str | None = None): + """Initialize ls tool. + + Args: + cwd: Working directory (defaults to current directory) + """ + super().__init__() + self.cwd = cwd or os.getcwd() + + def _build_tool_call(self) -> ToolCall: + max_kb = DEFAULT_MAX_BYTES // 1024 + return ToolCall( + **{ + "description": ( + f"List directory contents. Returns entries sorted alphabetically, " + f"with '/' suffix for directories. Includes dotfiles. Output is truncated " + f"to {DEFAULT_LIMIT} entries or {max_kb}KB (whichever is hit first)." + ), + "parameters": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Directory to list (default: current directory)", + }, + "limit": { + "type": "number", + "description": f"Maximum number of entries to return (default: {DEFAULT_LIMIT})", + }, + }, + "required": [], + }, + }, + ) + + async def execute(self) -> str: + """List directory contents with production-grade features.""" + path: str | None = self.context.get("path", None) + limit: int | None = self.context.get("limit", None) + + # Resolve directory path + dir_path = Path(self.cwd) / (path or ".") + dir_path = dir_path.resolve() + effective_limit = limit if limit is not None else DEFAULT_LIMIT + + # Check if path exists + if not dir_path.exists(): + raise FileNotFoundError(f"Path not found: {dir_path}") + + # Check if path is a directory + if not dir_path.is_dir(): + raise NotADirectoryError(f"Not a directory: {dir_path}") + + # Read directory entries + try: + entries = list(dir_path.iterdir()) + except Exception as e: + raise PermissionError(f"Cannot read directory: {e}") from e + + # Sort alphabetically (case-insensitive) + entries.sort(key=lambda e: e.name.lower()) + + # Format entries with directory indicators + results: list[str] = [] + entry_limit_reached = False + + for entry in entries: + if len(results) >= effective_limit: + entry_limit_reached = True + break + + try: + # Add '/' suffix for directories + suffix = "/" if entry.is_dir() else "" + results.append(entry.name + suffix) + except Exception: + # Skip entries we can't stat + continue + + # Handle empty directory + if len(results) == 0: + return "(empty directory)" + + # Apply byte truncation + raw_output = "\n".join(results) + truncation_result = truncate_head(raw_output, max_lines=float("inf")) + + output_text = truncation_result.content + + # Build notices + notices: list[str] = [] + + if entry_limit_reached: + notices.append( + f"{effective_limit} entries limit reached. Use limit={effective_limit * 2} for more", + ) + + if truncation_result.truncated: + max_kb = DEFAULT_MAX_BYTES // 1024 + notices.append(f"{max_kb}KB limit reached") + + if notices: + output_text += f"\n\n[{'. '.join(notices)}]" + + return output_text diff --git a/reme/tool/fs/read_tool.py b/reme/tool/fs/read_tool.py new file mode 100644 index 00000000..b236d4a4 --- /dev/null +++ b/reme/tool/fs/read_tool.py @@ -0,0 +1,219 @@ +"""Read file tool with smart truncation and image support. + +Features: +- Reads text files with offset/limit support +- Detects and handles image files (jpg, png, gif, webp) +- Smart truncation to prevent memory issues +""" + +import os +from pathlib import Path + +from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, format_size, truncate_head +from ...core.op import BaseTool +from ...core.schema import ToolCall + +# Supported image extensions +IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"} + + +def is_image_file(path: str) -> bool: + """Check if file is a supported image type. + + Args: + path: File path to check + + Returns: + True if file is a supported image + """ + return Path(path).suffix.lower() in IMAGE_EXTENSIONS + + +class ReadTool(BaseTool): + """Read file contents with smart truncation. + + Features: + - Supports text files and images (jpg, png, gif, webp) + - Smart truncation for large files + - Offset/limit for reading specific portions + """ + + def __init__(self, cwd: str | None = None): + """Initialize read tool. + + Args: + cwd: Working directory (defaults to current directory) + """ + super().__init__() + self.cwd = cwd or os.getcwd() + + def _build_tool_call(self) -> ToolCall: + max_kb = DEFAULT_MAX_BYTES // 1024 + return ToolCall( + **{ + "description": ( + f"Read the contents of a file. Supports text files and images " + f"(jpg, png, gif, webp). Images are sent as attachments. For text files, " + f"output is truncated to {DEFAULT_MAX_LINES} lines or {max_kb}KB " + f"(whichever is hit first). Use offset/limit for large files. " + f"When you need the full file, continue with offset until complete." + ), + "parameters": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to read (relative or absolute)", + }, + "offset": { + "type": "number", + "description": "Line number to start reading from (1-indexed)", + }, + "limit": { + "type": "number", + "description": "Maximum number of lines to read", + }, + }, + "required": ["path"], + }, + }, + ) + + async def execute(self) -> str: + """Execute the read operation.""" + path: str = self.context.path + offset: int | None = self.context.get("offset", None) + limit: int | None = self.context.get("limit", None) + + # Resolve path + if not os.path.isabs(path): + absolute_path = os.path.join(self.cwd, path) + else: + absolute_path = path + absolute_path = os.path.normpath(absolute_path) + + # Check file exists and is readable + if not os.path.exists(absolute_path): + raise ValueError(f"File not found: {path}") + + if not os.path.isfile(absolute_path): + raise ValueError(f"Not a file: {path}") + + if not os.access(absolute_path, os.R_OK): + raise ValueError(f"File not readable: {path}") + + # Check if image + if is_image_file(absolute_path): + return await self._read_image(absolute_path, path) + else: + return await self._read_text(absolute_path, path, offset, limit) + + @staticmethod + async def _read_image(absolute_path: str, display_path: str) -> str: + """Read and return image file information. + + Args: + absolute_path: Absolute path to image + display_path: Path to display to user + + Returns: + Image information text + """ + # Get file size + file_size = os.path.getsize(absolute_path) + file_ext = Path(absolute_path).suffix.lower() + + # For Python tools, we typically can't return image data directly to LLM + # So we return a descriptive message + return ( + f"Read image file [{file_ext}]\n" + f"Path: {display_path}\n" + f"Size: {format_size(file_size)}\n" + f"Note: Image content cannot be displayed in text format. " + f"Use bash tool or other methods to process the image." + ) + + @staticmethod + async def _read_text( + absolute_path: str, + _display_path: str, + offset: int | None, + limit: int | None, + ) -> str: + """Read text file with smart truncation. + + Args: + absolute_path: Absolute path to file + _display_path: Path to display to user + offset: Starting line (1-indexed) + limit: Maximum lines to read + + Returns: + File contents with truncation notices + """ + # Read file + try: + with open(absolute_path, "r", encoding="utf-8") as f: + content = f.read() + except UnicodeDecodeError: + # Try with error handling for binary files + with open(absolute_path, "r", encoding="utf-8", errors="ignore") as f: + content = f.read() + + all_lines = content.split("\n") + total_file_lines = len(all_lines) + + # Apply offset if specified (convert 1-indexed to 0-indexed) + start_line = max(0, (offset - 1)) if offset else 0 + start_line_display = start_line + 1 + + # Check offset bounds + if start_line >= len(all_lines): + raise IndexError( + f"Offset {offset} is beyond end of file ({len(all_lines)} lines total)", + ) + + # Apply user limit if specified + if limit is not None: + end_line = min(start_line + limit, len(all_lines)) + selected_content = "\n".join(all_lines[start_line:end_line]) + user_limited_lines = end_line - start_line + else: + selected_content = "\n".join(all_lines[start_line:]) + user_limited_lines = None + + # Apply truncation + truncation = truncate_head(selected_content) + + # Build output with truncation notices + if truncation.truncated: + # Truncation occurred + end_line_display = start_line_display + truncation.output_lines - 1 + next_offset = end_line_display + 1 + + output_text = truncation.content + + if truncation.truncated_by == "lines": + output_text += ( + f"\n\n[Showing lines {start_line_display}-{end_line_display} " + f"of {total_file_lines}. Use offset={next_offset} to continue.]" + ) + else: + max_kb = DEFAULT_MAX_BYTES // 1024 + output_text += ( + f"\n\n[Showing lines {start_line_display}-{end_line_display} " + f"of {total_file_lines} ({max_kb}KB limit). " + f"Use offset={next_offset} to continue.]" + ) + elif user_limited_lines is not None and start_line + user_limited_lines < len(all_lines): + # User limit exceeded, but no truncation + remaining = len(all_lines) - (start_line + user_limited_lines) + next_offset = start_line + user_limited_lines + 1 + + output_text = truncation.content + output_text += f"\n\n[{remaining} more lines in file. " f"Use offset={next_offset} to continue.]" + else: + # No truncation or user limit exceeded + output_text = truncation.content + + return output_text diff --git a/reme/tool/fs/truncate.py b/reme/tool/fs/truncate.py new file mode 100644 index 00000000..b4de0197 --- /dev/null +++ b/reme/tool/fs/truncate.py @@ -0,0 +1,209 @@ +"""fs utils""" + +from typing import Literal + +from ...core.schema import TruncationResult + +# Default limits for output truncation +DEFAULT_MAX_LINES = 1000 # Maximum lines to keep for tail truncation +DEFAULT_MAX_BYTES = 30 * 1024 # Maximum bytes to keep (30KB) + +# Find tool limits +FIND_MAX_LINES = 2000 # Maximum lines for find output +FIND_MAX_BYTES = 50 * 1024 # 50KB for find output + +# Grep tool limits +GREP_MAX_LINE_LENGTH = 500 # Maximum line length for grep output + + +def format_size(num_bytes: int) -> str: + """Format byte size in human-readable format. + + Args: + num_bytes: Number of bytes + + Returns: + Formatted string (e.g., "1.5KB", "2.3MB") + """ + if num_bytes < 1024: + return f"{num_bytes}B" + elif num_bytes < 1024 * 1024: + return f"{num_bytes / 1024:.1f}KB" + else: + return f"{num_bytes / (1024 * 1024):.1f}MB" + + +def truncate_line(text: str, max_length: int = GREP_MAX_LINE_LENGTH) -> tuple[str, bool]: + """Truncate a single line if it exceeds max length. + + Args: + text: Line text + max_length: Maximum line length + + Returns: + Tuple of (truncated_text, was_truncated) + """ + if len(text) <= max_length: + return text, False + return text[:max_length] + "...", True + + +def truncate_tail( + text: str, + max_lines: int = DEFAULT_MAX_LINES, + max_bytes: int = DEFAULT_MAX_BYTES, +) -> TruncationResult: + """Truncate text to keep only the tail (last portion). + + Keeps the last N lines or M bytes, whichever is hit first. + This is useful for command outputs where the end is most relevant. + + Args: + text: The text to truncate + max_lines: Maximum number of lines to keep + max_bytes: Maximum bytes to keep + + Returns: + TruncationResult with truncated content and metadata + """ + if not text: + return TruncationResult( + content="", + truncated=False, + total_lines=0, + output_lines=0, + total_bytes=0, + output_bytes=0, + ) + + total_bytes = len(text.encode("utf-8")) + lines = text.split("\n") + total_lines = len(lines) + + # Check if we need to truncate + if total_lines <= max_lines and total_bytes <= max_bytes: + return TruncationResult( + content=text, + truncated=False, + total_lines=total_lines, + output_lines=total_lines, + total_bytes=total_bytes, + output_bytes=total_bytes, + ) + + # Keep last N lines + kept_lines = lines[-max_lines:] if total_lines > max_lines else lines + truncated_by: Literal["lines", "bytes"] = "lines" if total_lines > max_lines else "bytes" + + # Check byte limit on kept lines + kept_text = "\n".join(kept_lines) + kept_bytes = len(kept_text.encode("utf-8")) + + # If still over byte limit, truncate further + last_line_partial = False + if kept_bytes > max_bytes: + truncated_by = "bytes" + # Keep truncating from the start until under byte limit + while kept_lines and len("\n".join(kept_lines).encode("utf-8")) > max_bytes: + kept_lines.pop(0) + + # If still over (single line > max_bytes), truncate the line itself + if kept_lines and len("\n".join(kept_lines).encode("utf-8")) > max_bytes: + last_line = kept_lines[-1] + # Binary search to find how much of last line fits + encoded = last_line.encode("utf-8") + if len(encoded) > max_bytes: + last_line_partial = True + # Take last max_bytes of the line + kept_lines[-1] = encoded[-max_bytes:].decode("utf-8", errors="ignore") + + kept_text = "\n".join(kept_lines) + kept_bytes = len(kept_text.encode("utf-8")) + + return TruncationResult( + content=kept_text, + truncated=True, + total_lines=total_lines, + output_lines=len(kept_lines), + total_bytes=total_bytes, + output_bytes=kept_bytes, + truncated_by=truncated_by, + last_line_partial=last_line_partial, + ) + + +def truncate_head( + text: str, + max_lines: int = FIND_MAX_LINES, + max_bytes: int = FIND_MAX_BYTES, +) -> TruncationResult: + """Truncate text to keep only the head (first portion). + + Keeps the first N lines or M bytes, whichever is hit first. + Suitable for file reads where you want to see the beginning. + + Args: + text: The text to truncate + max_lines: Maximum number of lines to keep + max_bytes: Maximum bytes to keep + + Returns: + TruncationResult with truncated content and metadata + """ + if not text: + return TruncationResult( + content="", + truncated=False, + total_lines=0, + output_lines=0, + total_bytes=0, + output_bytes=0, + ) + + total_bytes = len(text.encode("utf-8")) + lines = text.split("\n") + total_lines = len(lines) + + # Check if no truncation needed + if total_lines <= max_lines and total_bytes <= max_bytes: + return TruncationResult( + content=text, + truncated=False, + total_lines=total_lines, + output_lines=total_lines, + total_bytes=total_bytes, + output_bytes=total_bytes, + ) + + # Collect complete lines that fit + kept_lines = [] + kept_bytes = 0 + truncated_by: Literal["lines", "bytes"] = "lines" + + for i, line in enumerate(lines): + if i >= max_lines: + truncated_by = "lines" + break + + # Calculate bytes for this line (+1 for newline except first line) + line_bytes = len(line.encode("utf-8")) + (1 if i > 0 else 0) + + if kept_bytes + line_bytes > max_bytes: + truncated_by = "bytes" + break + + kept_lines.append(line) + kept_bytes += line_bytes + + kept_text = "\n".join(kept_lines) + final_bytes = len(kept_text.encode("utf-8")) + + return TruncationResult( + content=kept_text, + truncated=True, + total_lines=total_lines, + output_lines=len(kept_lines), + total_bytes=total_bytes, + output_bytes=final_bytes, + truncated_by=truncated_by, + ) diff --git a/reme/tool/fs/write_tool.py b/reme/tool/fs/write_tool.py new file mode 100644 index 00000000..55a4a0c1 --- /dev/null +++ b/reme/tool/fs/write_tool.py @@ -0,0 +1,81 @@ +"""Write tool for creating and overwriting files. + +This module provides a tool for writing content to files with: +- Automatic parent directory creation +- File overwriting (creates if doesn't exist, overwrites if exists) +- Path resolution (relative to working directory) +""" + +import os + +from ...core.op import BaseTool +from ...core.schema import ToolCall + + +class WriteTool(BaseTool): + """Tool for writing content to files. + + Features: + - Creates file if it doesn't exist, overwrites if it does + - Automatically creates parent directories + - Supports both relative and absolute paths + """ + + def __init__(self, cwd: str | None = None): + """Initialize write tool. + + Args: + cwd: Working directory (defaults to current directory) + """ + super().__init__() + self.cwd = cwd or os.getcwd() + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": ( + "Write content to a file. Creates the file if it doesn't exist, " + "overwrites if it does. Automatically creates parent directories." + ), + "parameters": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to write (relative or absolute)", + }, + "content": { + "type": "string", + "description": "Content to write to the file", + }, + }, + "required": ["path", "content"], + }, + }, + ) + + async def execute(self) -> str: + """Execute the write operation.""" + path: str = self.context.path + content: str = self.context.content + + # Resolve path to absolute + if not os.path.isabs(path): + absolute_path = os.path.join(self.cwd, path) + else: + absolute_path = path + + absolute_path = os.path.normpath(absolute_path) + + # Create parent directories if needed + parent_dir = os.path.dirname(absolute_path) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + + # Write the file + with open(absolute_path, "w", encoding="utf-8") as f: + f.write(content) + + # Return success message + content_bytes = len(content.encode("utf-8")) + return f"Successfully wrote {content_bytes} bytes to {path}" diff --git a/tests/test_file_system_tool.py b/tests/test_file_system_tool.py new file mode 100644 index 00000000..15f2d858 --- /dev/null +++ b/tests/test_file_system_tool.py @@ -0,0 +1,445 @@ +"""Tests for file system tools including bash, edit, find, grep, ls, read, and write tools.""" + +import asyncio +import os +import tempfile +from pathlib import Path + + +async def test_bash_tool(): + """Test BashTool.""" + from reme.tool.fs import BashTool + + print("=== Testing BashTool ===") + bash_tool = BashTool() + result = await bash_tool.call(command="echo 'Hello World'") + print(f"Result: {result}") + assert "Hello World" in result + print("✓ BashTool test passed\n") + + +async def test_edit_tool(): + """Test EditTool.""" + from reme.tool.fs import EditTool + + print("=== Testing EditTool ===") + + # Create temp file + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + temp_path = f.name + f.write("Hello World\nThis is a test\nGoodbye World\n") + + try: + # Test edit + edit_tool = EditTool() + result = await edit_tool.call( + path=temp_path, + oldText="This is a test", + newText="This is an updated test", + ) + print(f"Result: {result}") + + # Verify content + with open(temp_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "This is an updated test" in content + assert "This is a test" not in content + print("✓ EditTool test passed\n") + + # Test error: file not found + print("=== Testing file not found error ===") + result = await edit_tool.call( + path="/nonexistent/file.txt", + oldText="test", + newText="new", + ) + print(f"Expected error result: {result}") + assert "failed" in result and "File not found" in result + print("✓ File not found error test passed\n") + + # Test error: text not found + print("=== Testing text not found error ===") + result = await edit_tool.call( + path=temp_path, + oldText="nonexistent text", + newText="new", + ) + print(f"Expected error result: {result}") + assert "failed" in result and "Could not find" in result + print("✓ Text not found error test passed\n") + + finally: + # Cleanup + if os.path.exists(temp_path): + os.unlink(temp_path) + + +async def test_find_tool(): + """Test FindTool.""" + from reme.tool.fs import FindTool + + print("=== Testing FindTool ===") + + # Create temp directory with test files + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create some test files + (temp_path / "test1.txt").write_text("test file 1") + (temp_path / "test2.txt").write_text("test file 2") + (temp_path / "readme.md").write_text("readme") + + # Create subdirectory with files + sub_dir = temp_path / "subdir" + sub_dir.mkdir() + (sub_dir / "test3.txt").write_text("test file 3") + (sub_dir / "config.json").write_text("{}") + + # Create .gitignore to ignore certain files + (temp_path / ".gitignore").write_text("*.md\n") + + # Test: find all txt files + find_tool = FindTool(cwd=str(temp_path)) + result = await find_tool.call(pattern="*.txt") + print(f"Find *.txt result:\n{result}") + assert "test1.txt" in result + assert "test2.txt" in result + assert "readme.md" not in result # Should be ignored by .gitignore + print("✓ Find *.txt test passed\n") + + # Test: find with recursive pattern + result = await find_tool.call(pattern="**/*.txt") + print(f"Find **/*.txt result:\n{result}") + assert "test1.txt" in result + assert "subdir/test3.txt" in result or "test3.txt" in result + print("✓ Find **/*.txt test passed\n") + + # Test: find with no matches + result = await find_tool.call(pattern="*.nonexistent") + print(f"Find *.nonexistent result:\n{result}") + assert "No files found" in result + print("✓ No matches test passed\n") + + # Test: error - directory not found + print("=== Testing directory not found error ===") + result = await find_tool.call(pattern="*.txt", path="/nonexistent/dir") + print(f"Expected error result: {result}") + assert "failed" in result and "Path not found" in result + print("✓ Directory not found error test passed\n") + + +async def test_grep_tool(): + """Test GrepTool.""" + from reme.tool.fs import GrepTool + + print("=== Testing GrepTool ===") + + # Create temp directory with test files + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create test files with searchable content + (temp_path / "file1.txt").write_text("Hello World\nThis is a test\nGoodbye World\n") + (temp_path / "file2.txt").write_text("Another test file\nWith multiple lines\nHello again\n") + (temp_path / "script.py").write_text("def hello():\n print('Hello')\n return True\n") + + # Create subdirectory with files + sub_dir = temp_path / "subdir" + sub_dir.mkdir() + (sub_dir / "nested.txt").write_text("Nested file content\nWith hello keyword\n") + + # Test: search for pattern + grep_tool = GrepTool(cwd=str(temp_path)) + result = await grep_tool.call(pattern="Hello", path=str(temp_path)) + print(f"Search 'Hello' result:\n{result}") + assert "file1.txt" in result + assert "Hello World" in result or "Hello" in result + print("✓ Basic search test passed\n") + + # Test: case-insensitive search + result = await grep_tool.call(pattern="hello", path=str(temp_path), ignoreCase=True) + print(f"Case-insensitive search result:\n{result}") + assert "file1.txt" in result or "Hello" in result.lower() + print("✓ Case-insensitive search test passed\n") + + # Test: literal string search + result = await grep_tool.call(pattern="Hello()", path=str(temp_path), literal=True) + print(f"Literal search result:\n{result}") + # Should not find regex interpretation + print("✓ Literal search test passed\n") + + # Test: glob filter + result = await grep_tool.call(pattern="Hello", path=str(temp_path), glob="*.txt") + print(f"Glob filter *.txt result:\n{result}") + assert "file1.txt" in result or "file2.txt" in result + assert ".py" not in result # Python files should be excluded + print("✓ Glob filter test passed\n") + + # Test: context lines + result = await grep_tool.call(pattern="test", path=str(temp_path), contextLines=1) + print(f"Context lines result:\n{result}") + # Should include lines before and after matches + print("✓ Context lines test passed\n") + + # Test: limit matches + result = await grep_tool.call(pattern="Hello", path=str(temp_path), limit=1) + print(f"Limit to 1 match result:\n{result}") + assert "limit reached" in result or result.count(":") >= 1 + print("✓ Limit test passed\n") + + # Test: no matches + result = await grep_tool.call(pattern="nonexistent_pattern_xyz", path=str(temp_path)) + print(f"No matches result:\n{result}") + assert "No matches found" in result + print("✓ No matches test passed\n") + + # Test: error - path not found + print("=== Testing path not found error ===") + try: + result = await grep_tool.call(pattern="test", path="/nonexistent/path") + assert "failed" in result and "not found" in result + except Exception as e: + assert "not found" in str(e).lower() + print("✓ Path not found error test passed\n") + + +async def test_ls_tool(): + """Test LsTool.""" + from reme.tool.fs import LsTool + + print("=== Testing LsTool ===") + + # Create temp directory with test files + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create test files and directories + (temp_path / "file1.txt").write_text("test file 1") + (temp_path / "file2.py").write_text("test file 2") + (temp_path / ".hidden").write_text("hidden file") + (temp_path / "README.md").write_text("readme") + + # Create subdirectories + (temp_path / "subdir1").mkdir() + (temp_path / "subdir2").mkdir() + + # Test: list current directory + ls_tool = LsTool(cwd=str(temp_path)) + result = await ls_tool.call() + print(f"List directory result:\n{result}") + assert ".hidden" in result # Includes dotfiles + assert "file1.txt" in result + assert "file2.py" in result + assert "subdir1/" in result # Directories have '/' suffix + assert "subdir2/" in result + print("✓ Basic ls test passed\n") + + # Test: list specific path + result = await ls_tool.call(path=".") + print(f"List current directory result:\n{result}") + assert "file1.txt" in result + print("✓ Specific path test passed\n") + + # Test: empty directory + empty_dir = temp_path / "empty" + empty_dir.mkdir() + result = await ls_tool.call(path="empty") + print(f"Empty directory result:\n{result}") + assert "(empty directory)" in result + print("✓ Empty directory test passed\n") + + # Test: entry limit + # Create many files + for i in range(10): + (temp_path / f"file{i:03d}.txt").write_text(f"file {i}") + + result = await ls_tool.call(limit=5) + print(f"Limited entries result:\n{result}") + assert "entries limit reached" in result + assert "limit=10" in result # Should suggest doubling the limit + print("✓ Entry limit test passed\n") + + # Test: error - path not found + print("=== Testing path not found error ===") + result = await ls_tool.call(path="/nonexistent/path") + print(f"Expected error result: {result}") + assert "failed" in result and "Path not found" in result + print("✓ Path not found error test passed\n") + + # Test: error - not a directory + print("=== Testing not a directory error ===") + result = await ls_tool.call(path="file1.txt") + print(f"Expected error result: {result}") + assert "failed" in result and "Not a directory" in result + print("✓ Not a directory error test passed\n") + + +async def test_read_tool(): + """Test ReadTool.""" + from reme.tool.fs import ReadTool + + print("=== Testing ReadTool ===") + + # Create temp directory with test files + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create test text file + test_file = temp_path / "test.txt" + test_content = "\n".join([f"Line {i}" for i in range(1, 101)]) # 100 lines + test_file.write_text(test_content) + + # Create test image file + image_file = temp_path / "test.jpg" + image_file.write_bytes(b"\xff\xd8\xff\xe0") # Minimal JPEG header + + # Test: read full file + read_tool = ReadTool(cwd=str(temp_path)) + result = await read_tool.call(path="test.txt") + print(f"Read full file result:\n{result[:200]}...") + assert "Line 1" in result + assert "Line 100" in result + print("✓ Read full file test passed\n") + + # Test: read with offset + result = await read_tool.call(path="test.txt", offset=50) + print(f"Read with offset=50 result:\n{result[:200]}...") + # Check that we start from Line 50 (should be first line of content) + assert result.startswith("Line 50"), f"Should start with 'Line 50', got: {result[:50]}" + assert "Line 100" in result + print("✓ Read with offset test passed\n") + + # Test: read with limit + result = await read_tool.call(path="test.txt", limit=10) + print(f"Read with limit=10 result:\n{result}") + assert "Line 1" in result + assert "Line 10" in result or "more lines in file" in result + assert "Line 50" not in result + print("✓ Read with limit test passed\n") + + # Test: read with offset and limit + result = await read_tool.call(path="test.txt", offset=20, limit=5) + print(f"Read with offset=20, limit=5 result:\n{result}") + assert "Line 20" in result + assert "Line 24" in result or "more lines" in result + print("✓ Read with offset and limit test passed\n") + + # Test: read image file + result = await read_tool.call(path="test.jpg") + print(f"Read image result:\n{result}") + assert "image file" in result.lower() or ".jpg" in result.lower() + print("✓ Read image test passed\n") + + # Test: offset beyond file + print("=== Testing offset beyond file error ===") + result = await read_tool.call(path="test.txt", offset=200) + print(f"Expected error result: {result}") + assert "failed" in result and ("beyond end of file" in result or "offset" in result.lower()) + print("✓ Offset beyond file error test passed\n") + + # Test: file not found + print("=== Testing file not found error ===") + result = await read_tool.call(path="nonexistent.txt") + print(f"Expected error result: {result}") + assert "failed" in result and "not found" in result.lower() + print("✓ File not found error test passed\n") + + # Test: read directory (should fail) + print("=== Testing read directory error ===") + sub_dir = temp_path / "subdir" + sub_dir.mkdir() + result = await read_tool.call(path="subdir") + print(f"Expected error result: {result}") + assert "failed" in result and ("Not a file" in result or "directory" in result.lower()) + print("✓ Read directory error test passed\n") + + +async def test_write_tool(): + """Test WriteTool.""" + from reme.tool.fs import WriteTool + + print("=== Testing WriteTool ===") + + # Create temp directory + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Test: write new file + write_tool = WriteTool(cwd=str(temp_path)) + test_content = "Hello World\nThis is a test file\n" + result = await write_tool.call(path="test.txt", content=test_content) + print(f"Write result: {result}") + assert "Successfully wrote" in result + assert "test.txt" in result + + # Verify file was created + test_file = temp_path / "test.txt" + assert test_file.exists() + assert test_file.read_text() == test_content + print("✓ Write new file test passed\n") + + # Test: overwrite existing file + new_content = "Updated content\n" + result = await write_tool.call(path="test.txt", content=new_content) + print(f"Overwrite result: {result}") + assert "Successfully wrote" in result + + # Verify file was overwritten + assert test_file.read_text() == new_content + assert test_content not in test_file.read_text() + print("✓ Overwrite existing file test passed\n") + + # Test: create file with parent directories + nested_path = "subdir1/subdir2/nested.txt" + nested_content = "Nested file content" + result = await write_tool.call(path=nested_path, content=nested_content) + print(f"Create with parents result: {result}") + assert "Successfully wrote" in result + + # Verify nested file was created + nested_file = temp_path / "subdir1" / "subdir2" / "nested.txt" + assert nested_file.exists() + assert nested_file.read_text() == nested_content + print("✓ Create file with parent directories test passed\n") + + # Test: write empty file + result = await write_tool.call(path="empty.txt", content="") + print(f"Write empty file result: {result}") + assert "Successfully wrote" in result + assert "0 bytes" in result + + # Verify empty file + empty_file = temp_path / "empty.txt" + assert empty_file.exists() + assert empty_file.read_text() == "" + print("✓ Write empty file test passed\n") + + # Test: write file with absolute path + abs_path = str(temp_path / "absolute.txt") + abs_content = "Absolute path content" + result = await write_tool.call(path=abs_path, content=abs_content) + print(f"Write absolute path result: {result}") + assert "Successfully wrote" in result + + # Verify absolute path file + abs_file = Path(abs_path) + assert abs_file.exists() + assert abs_file.read_text(encoding="utf-8") == abs_content + print("✓ Write absolute path test passed\n") + + +async def main(): + """Run all file system tool tests.""" + await test_bash_tool() + await test_edit_tool() + await test_find_tool() + await test_grep_tool() + await test_ls_tool() + await test_read_tool() + await test_write_tool() + print("=== All tests passed! ===") + + +if __name__ == "__main__": + asyncio.run(main())