feat(memory): add memory management and file system tools

This commit is contained in:
jinli.yl 2026-02-03 22:29:45 +08:00
parent 0fb898ca3c
commit 857cd52a8e
31 changed files with 4055 additions and 4 deletions

View file

@ -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

View file

@ -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."""

View file

@ -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",

View file

@ -0,0 +1,11 @@
"""Memory source types."""
from enum import Enum
class MemorySource(str, Enum):
"""Source of memory data."""
MEMORY = "memory"
SESSIONS = "sessions"

View file

View file

@ -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

View file

@ -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

View file

@ -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."""

View file

@ -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

View file

@ -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()

View file

@ -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",
]

View file

@ -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")

View file

@ -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")

View file

@ -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")

View file

@ -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}"

View file

@ -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",
)

View file

@ -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()
```
"""

24
reme/tool/fs/__init__.py Normal file
View file

@ -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)

191
reme/tool/fs/bash_tool.py Normal file
View file

@ -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

164
reme/tool/fs/edit_diff.py Normal file
View file

@ -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)

141
reme/tool/fs/edit_tool.py Normal file
View file

@ -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}"

183
reme/tool/fs/find_tool.py Normal file
View file

@ -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

276
reme/tool/fs/grep_tool.py Normal file
View file

@ -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

128
reme/tool/fs/ls_tool.py Normal file
View file

@ -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

219
reme/tool/fs/read_tool.py Normal file
View file

@ -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

209
reme/tool/fs/truncate.py Normal file
View file

@ -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,
)

View file

@ -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}"

View file

@ -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())