mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
feat(memory): add memory management and file system tools
This commit is contained in:
parent
0fb898ca3c
commit
857cd52a8e
31 changed files with 4055 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
11
reme/core/enumeration/memory_source.py
Normal file
11
reme/core/enumeration/memory_source.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
"""Memory source types."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MemorySource(str, Enum):
|
||||
"""Source of memory data."""
|
||||
|
||||
MEMORY = "memory"
|
||||
|
||||
SESSIONS = "sessions"
|
||||
0
reme/core/memory_manager/__init__.py
Normal file
0
reme/core/memory_manager/__init__.py
Normal file
0
reme/core/memory_manager/ingestion/__init__.py
Normal file
0
reme/core/memory_manager/ingestion/__init__.py
Normal file
126
reme/core/memory_manager/ingestion/chunking.py
Normal file
126
reme/core/memory_manager/ingestion/chunking.py
Normal 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
|
||||
890
reme/core/memory_manager/manager.py
Normal file
890
reme/core/memory_manager/manager.py
Normal 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
|
||||
0
reme/core/memory_manager/memory_storage/__init__.py
Normal file
0
reme/core/memory_manager/memory_storage/__init__.py
Normal file
172
reme/core/memory_manager/memory_storage/base_memory_store.py
Normal file
172
reme/core/memory_manager/memory_storage/base_memory_store.py
Normal 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."""
|
||||
582
reme/core/memory_manager/memory_storage/sqlite_memory_store.py
Normal file
582
reme/core/memory_manager/memory_storage/sqlite_memory_store.py
Normal 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
|
||||
0
reme/core/memory_manager/utils/__init__.py
Normal file
0
reme/core/memory_manager/utils/__init__.py
Normal file
15
reme/core/memory_manager/utils/hashing.py
Normal file
15
reme/core/memory_manager/utils/hashing.py
Normal 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()
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
20
reme/core/schema/file_metadata.py
Normal file
20
reme/core/schema/file_metadata.py
Normal 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")
|
||||
19
reme/core/schema/memory_chunk.py
Normal file
19
reme/core/schema/memory_chunk.py
Normal 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")
|
||||
14
reme/core/schema/memory_index_meta.py
Normal file
14
reme/core/schema/memory_index_meta.py
Normal 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")
|
||||
24
reme/core/schema/memory_search_result.py
Normal file
24
reme/core/schema/memory_search_result.py
Normal 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}"
|
||||
35
reme/core/schema/truncation_result.py
Normal file
35
reme/core/schema/truncation_result.py
Normal 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",
|
||||
)
|
||||
|
|
@ -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
24
reme/tool/fs/__init__.py
Normal 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
191
reme/tool/fs/bash_tool.py
Normal 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
164
reme/tool/fs/edit_diff.py
Normal 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
141
reme/tool/fs/edit_tool.py
Normal 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
183
reme/tool/fs/find_tool.py
Normal 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
276
reme/tool/fs/grep_tool.py
Normal 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
128
reme/tool/fs/ls_tool.py
Normal 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
219
reme/tool/fs/read_tool.py
Normal 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
209
reme/tool/fs/truncate.py
Normal 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,
|
||||
)
|
||||
81
reme/tool/fs/write_tool.py
Normal file
81
reme/tool/fs/write_tool.py
Normal 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}"
|
||||
445
tests/test_file_system_tool.py
Normal file
445
tests/test_file_system_tool.py
Normal 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())
|
||||
Loading…
Add table
Reference in a new issue