This commit is contained in:
jinli.yl 2026-04-08 16:19:27 +08:00
parent 935e886af3
commit 16d139583f
28 changed files with 4195 additions and 0 deletions

6
reme_cli/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""ReMe CLI package."""
from reme_cli.component import BaseComponent
from reme_cli.application import Application
__all__ = ["BaseComponent", "Application"]

20
reme_cli/application.py Normal file
View file

@ -0,0 +1,20 @@
from reme_cli.component import BaseComponent
class Application(BaseComponent):
"""Application component for managing the main application."""
def __init__(self) -> None:
super().__init__()
...
async def start(self) -> None:
"""Start the application."""
# 初始化llm formater
#
pass
async def close(self) -> None:
"""Close the application."""
pass

View file

@ -0,0 +1,5 @@
from .base_component import BaseComponent
__all__ = [
"BaseComponent",
]

View file

@ -0,0 +1,9 @@
"""Module for registering AgentScope LLM models."""
from agentscope.model import DashScopeChatModel
from agentscope.model import OpenAIChatModel
from ..registry_factory import R
R.as_llms.register("openai")(OpenAIChatModel)
R.as_llms.register("dashscope")(DashScopeChatModel)

View file

@ -0,0 +1,9 @@
"""Module for registering AgentScope LLM formatters."""
from agentscope.formatter import DashScopeChatFormatter
from .reme_openai_chat_formatter import ReMeOpenAIChatFormatter
from ..registry_factory import R
R.as_llm_formatters.register("openai")(ReMeOpenAIChatFormatter)
R.as_llm_formatters.register("dashscope")(DashScopeChatFormatter)

View file

@ -0,0 +1,215 @@
"""ReMeOpenAIChatFormatter"""
import json
from typing import Any
from agentscope.formatter import OpenAIChatFormatter
from agentscope.formatter._openai_formatter import (
_format_openai_image_block,
_to_openai_audio_data,
)
from agentscope.message import Msg, TextBlock, ImageBlock, URLSource
from loguru import logger
def _format_openai_video_block(video_block: dict) -> dict[str, Any]:
"""Format a video block for OpenAI API.
Args:
video_block: The video block to format.
Returns:
A dictionary with video content in OpenAI format.
"""
source = video_block["source"]
if source["type"] == "url":
url = source["url"]
elif source["type"] == "base64":
data = source["data"]
media_type = source["media_type"]
url = f"data:{media_type};base64,{data}"
else:
raise ValueError(f"Unsupported video source type: {source['type']}")
return {
"type": "video_url",
"video_url": {
"url": url,
},
}
class ReMeOpenAIChatFormatter(OpenAIChatFormatter):
"""ReMeOpenAIChatFormatter"""
async def _format(
self,
msgs: list[Msg],
) -> list[dict[str, Any]]:
"""Format message objects into OpenAI API required format.
Args:
msgs (`list[Msg]`):
The list of Msg objects to format.
Returns:
`list[dict[str, Any]]`:
A list of dictionaries, where each dictionary has "name",
"role", and "content" keys.
"""
self.assert_list_of_msgs(msgs)
messages: list[dict] = []
i = 0
while i < len(msgs):
msg = msgs[i]
content_blocks = []
tool_calls = []
reasoning_content_blocks = []
for block in msg.get_content_blocks():
typ = block.get("type")
if typ == "text":
content_blocks.append({**block})
elif typ == "thinking":
# Collect thinking blocks for reasoning_content field
# This is compatible with models like DeepSeek that support
# extended thinking via reasoning_content field
reasoning_content_blocks.append({**block})
elif typ == "tool_use":
tool_calls.append(
{
"id": block.get("id"),
"type": "function",
"function": {
"name": block.get("name"),
"arguments": json.dumps(
block.get("input", {}),
ensure_ascii=False,
),
},
},
)
elif typ == "tool_result":
(
textual_output,
multimodal_data,
) = self.convert_tool_result_to_string(block["output"])
messages.append(
{
"role": "tool",
"tool_call_id": block.get("id"),
"content": (textual_output), # type: ignore[arg-type]
"name": block.get("name"),
},
)
# Then, handle the multimodal data if any
promoted_blocks: list = []
for url, multimodal_block in multimodal_data:
if multimodal_block["type"] == "image" and self.promote_tool_result_images:
promoted_blocks.extend(
[
TextBlock(
type="text",
text=f"\n- The image from '{url}': ",
),
ImageBlock(
type="image",
source=URLSource(
type="url",
url=url,
),
),
],
)
if promoted_blocks:
# Insert promoted blocks as new user message(s)
promoted_blocks = [
TextBlock(
type="text",
text="<system-info>The following are "
"the image contents from the tool "
f"result of '{block['name']}':",
),
*promoted_blocks,
TextBlock(
type="text",
text="</system-info>",
),
]
msgs.insert(
i + 1,
Msg(
name="user",
content=promoted_blocks,
role="user",
),
)
elif typ == "image":
content_blocks.append(
_format_openai_image_block(
block, # type: ignore[arg-type]
),
)
elif typ == "audio":
# Filter out audio content when the multimodal model
# outputs both text and audio, to prevent errors in
# subsequent model calls
if msg.role == "assistant":
continue
input_audio = _to_openai_audio_data(block["source"])
content_blocks.append(
{
"type": "input_audio",
"input_audio": input_audio,
},
)
elif typ == "video":
# Filter out video content when the multimodal model
# outputs both text and video, to prevent errors in
# subsequent model calls
if msg.role == "assistant":
continue
content_blocks.append(
_format_openai_video_block(block),
)
else:
logger.warning(
"Unsupported block type %s in the message, skipped.",
typ,
)
msg_openai = {
"role": msg.role,
"name": msg.name,
"content": content_blocks or None,
}
if tool_calls:
msg_openai["tool_calls"] = tool_calls
# Add reasoning_content for thinking blocks (compatible with DeepSeek, etc.)
if reasoning_content_blocks:
reasoning_msg = "\n".join(reasoning.get("thinking", "") for reasoning in reasoning_content_blocks)
if reasoning_msg:
msg_openai["reasoning_content"] = reasoning_msg
# When both content and tool_calls are None, skipped
if msg_openai["content"] or msg_openai.get("tool_calls"):
messages.append(msg_openai)
# Move to next message
i += 1
return messages

View file

@ -0,0 +1,25 @@
"""Base class for components."""
from abc import ABC, abstractmethod
class BaseComponent(ABC):
"""Base class supporting async start/close and async context management."""
@abstractmethod
async def start(self) -> None:
"""Start the component asynchronously."""
@abstractmethod
async def close(self) -> None:
"""Close the component asynchronously."""
...
async def __aenter__(self) -> "BaseComponent":
"""Enter async context manager."""
await self.start()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
"""Exit async context manager."""
await self.close()

View file

@ -0,0 +1,41 @@
"""Module providing a dictionary subclass with attribute-style access and pickling support."""
from typing import Generic, TypeVar
_KT = TypeVar("_KT")
_VT = TypeVar("_VT")
class BaseDict(dict, Generic[_KT, _VT]):
"""A dictionary subclass that enables accessing and modifying keys as attributes."""
def __getattr__(self, name: str) -> _VT:
"""Retrieve a dictionary item as an attribute."""
try:
return self[name]
except KeyError as e:
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e
def __setattr__(self, name: str, value: _VT) -> None:
"""Assign a value to a dictionary item using attribute syntax."""
self[name] = value
def __delattr__(self, name: str) -> None:
"""Remove a dictionary item using attribute syntax."""
try:
# Delete item from dict via key
del self[name]
except KeyError as e:
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e
def __getstate__(self) -> dict:
"""Return the dictionary representation for pickling."""
return dict(self)
def __setstate__(self, state: dict) -> None:
"""Restore the dictionary state from a pickled object."""
self.update(state)
def __reduce__(self):
"""Define the reconstruction logic for pickling processes."""
return self.__class__, (), self.__getstate__()

View file

@ -0,0 +1,23 @@
"""File store module for persistent memory management.
This module provides storage backends for memory chunks and file metadata,
including SQLite-based, ChromaDB-based, and pure-Python local implementations
with vector and full-text search.
"""
from .base_file_store import BaseFileStore
from .chroma_file_store import ChromaFileStore
from .local_file_store import LocalFileStore
from .sqlite_file_store import SqliteFileStore
from ..registry_factory import R
__all__ = [
"BaseFileStore",
"ChromaFileStore",
"LocalFileStore",
"SqliteFileStore",
]
R.file_stores.register("sqlite")(SqliteFileStore)
R.file_stores.register("chroma")(ChromaFileStore)
R.file_stores.register("local")(LocalFileStore)

View file

@ -0,0 +1,227 @@
"""Base storage interface for file store."""
import re
from abc import ABC, abstractmethod
from pathlib import Path
from ..embedding import BaseEmbeddingModel
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
from ..utils import get_logger
logger = get_logger()
class BaseFileStore(ABC):
"""Abstract base class for file storage backends."""
def __init__(
self,
store_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel | None = None,
vector_enabled: bool = False,
fts_enabled: bool = True,
**kwargs,
):
"""Initialize"""
# Validate store_name to prevent SQL injection
# Only allow alphanumeric characters and underscores
if not re.match(r"^[a-zA-Z0-9_]+$", store_name):
raise ValueError(f"Invalid '{store_name}'. Only alphanumeric characters and underscores are allowed.")
# Ensure at least one search method is enabled
if not vector_enabled and not fts_enabled:
raise ValueError("At least one of vector_enabled or fts_enabled must be True.")
# Ensure embedding_model is provided when vector search is enabled
if vector_enabled and embedding_model is None:
raise ValueError("embedding_model is required when vector_enabled is True.")
self.store_name: str = store_name
self.db_path: Path = Path(db_path)
self.db_path.mkdir(parents=True, exist_ok=True)
self.embedding_model: BaseEmbeddingModel | None = embedding_model
self.vector_enabled: bool = vector_enabled
self.fts_enabled: bool = fts_enabled
self.kwargs: dict = kwargs
@property
def embedding_dim(self) -> int:
"""Get the embedding model's dimensionality."""
if self.embedding_model is None:
return 1024
return self.embedding_model.dimensions
def _get_mock_embedding(self) -> list[float]:
"""Generate a zero vector based on embedding model dimensions."""
return [0.0] * self.embedding_dim
def _disable_vector_search(self, reason: str = "embedding API error") -> None:
"""Disable vector search and log a warning."""
if self.vector_enabled:
logger.warning(
f"[{self.store_name}] Disabling vector search due to {reason}. "
"Falling back to full-text search only.",
)
self.vector_enabled = False
async def get_embedding(self, query: str, **kwargs) -> list[float]:
"""Get embedding for a single query string."""
if not self.vector_enabled:
return self._get_mock_embedding()
try:
return await self.embedding_model.get_embedding(query, **kwargs)
except Exception as e:
self._disable_vector_search(str(e))
return self._get_mock_embedding()
async def get_embeddings(self, queries: list[str], **kwargs) -> list[list[float]]:
"""Get embeddings for a batch of query strings."""
if not self.vector_enabled:
return [self._get_mock_embedding() for _ in queries]
try:
return await self.embedding_model.get_embeddings(queries, **kwargs)
except Exception as e:
self._disable_vector_search(str(e))
return [self._get_mock_embedding() for _ in queries]
async def get_chunk_embedding(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk:
"""Generate and populate embedding field for a single MemoryChunk object."""
if not self.vector_enabled:
chunk.embedding = self._get_mock_embedding()
return chunk
try:
return await self.embedding_model.get_chunk_embedding(chunk, **kwargs)
except Exception as e:
self._disable_vector_search(str(e))
chunk.embedding = self._get_mock_embedding()
return chunk
async def get_chunk_embeddings(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]:
"""Generate and populate embedding fields for a batch of MemoryChunk objects."""
if not self.vector_enabled:
mock_embedding = self._get_mock_embedding()
for chunk in chunks:
chunk.embedding = mock_embedding.copy()
return chunks
try:
return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs)
except Exception as e:
self._disable_vector_search(str(e))
mock_embedding = self._get_mock_embedding()
for chunk in chunks:
chunk.embedding = mock_embedding.copy()
return chunks
@abstractmethod
async def start(self):
"""Initialize the storage backend."""
@abstractmethod
async def upsert_file(self, file_meta: FileMetadata, source: MemorySource, chunks: list[MemoryChunk]):
"""Insert or update a file and its chunks."""
@abstractmethod
async def delete_file(self, path: str, source: MemorySource):
"""Delete a file and all its chunks."""
@abstractmethod
async def delete_file_chunks(self, path: str, chunk_ids: list[str]):
"""Delete chunks for a file."""
@abstractmethod
async def upsert_chunks(self, chunks: list[MemoryChunk], source: MemorySource):
"""Insert or update specific chunks without affecting other chunks."""
@abstractmethod
async def list_files(self, source: MemorySource) -> list[str]:
"""List all indexed file paths for a source."""
@abstractmethod
async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None:
"""Get full file metadata with statistics."""
@abstractmethod
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks.
This is useful for incremental updates where only metadata needs to be updated
(e.g., after adding/removing chunks in delta file watcher).
Args:
file_meta: Updated file metadata (hash, mtime_ms, size, chunk_count)
source: Memory source
"""
@abstractmethod
async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]:
"""Get all chunks for a file."""
@abstractmethod
async def vector_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform vector similarity search.
Args:
query: Query embedding vector
limit: Maximum number of results
sources: Optional list of sources to filter
Returns:
List of search results sorted by similarity
"""
@abstractmethod
async def keyword_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform keyword/full-text search.
Args:
query: Search query text
limit: Maximum number of results
sources: Optional list of sources to filter
Returns:
List of search results sorted by relevance
"""
@abstractmethod
async def hybrid_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
vector_weight: float = 0.7,
candidate_multiplier: float = 3.0,
) -> list[MemorySearchResult]:
"""Perform hybrid search combining vector and keyword search.
Args:
query: Search query text
limit: Maximum number of results
sources: Optional list of sources to filter
vector_weight: Weight for vector search results (0.0-1.0).
Keyword weight = 1.0 - vector_weight.
candidate_multiplier: Multiplier for candidate pool size.
candidates = limit * candidate_multiplier
Returns:
List of search results sorted by combined relevance score
"""
@abstractmethod
async def clear_all(self):
"""Clear all indexed data."""
@abstractmethod
async def close(self):
"""Close storage and release resources."""

View file

@ -0,0 +1,633 @@
"""ChromaDB storage backend for file store."""
import json
import random
import time
from pathlib import Path
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
from ..utils import get_logger
logger = get_logger()
try:
import chromadb
from chromadb.config import Settings
_CHROMADB_IMPORT_ERROR: Exception | None = None
except Exception as e:
_CHROMADB_IMPORT_ERROR = e
chromadb = None
Settings = None
class ChromaFileStore(BaseFileStore):
"""ChromaDB file storage with vector and full-text search.
Inherits embedding methods from BaseFileStore:
- get_chunk_embedding / get_chunk_embeddings (async)
- get_embedding / get_embeddings (async)
Provides ChromaDB-backed persistent storage with:
- Vector similarity search (native ChromaDB)
- Full-text search (via ChromaDB where_document filter)
- Efficient chunk and file metadata management
"""
def __init__(
self,
**kwargs,
):
if _CHROMADB_IMPORT_ERROR is not None:
raise _CHROMADB_IMPORT_ERROR
super().__init__(**kwargs)
self.client: "chromadb.ClientAPI | None" = None
self.chunks_collection: "chromadb.Collection | None" = None
# Initialize metadata file path (db_path and store_name are set by base class)
self._metadata_file: Path = self.db_path.parent / f"{self.store_name}_file_metadata.json"
self._metadata_cache: dict[str, dict[str, FileMetadata]] = {}
@property
def collection_name(self) -> str:
"""Get the name of the ChromaDB collection for this store."""
return f"chunks_{self.store_name}"
async def _load_metadata(self) -> dict[str, dict[str, FileMetadata]]:
"""Load file metadata from disk.
Returns:
Dictionary mapping source -> path -> FileMetadata
"""
if not self._metadata_file.exists():
return {}
try:
data = self._metadata_file.read_text(encoding="utf-8")
metadata_dict = json.loads(data)
# Convert dict to FileMetadata objects
result = {}
for source, files in metadata_dict.items():
result[source] = {}
for path, meta in files.items():
result[source][path] = FileMetadata(**meta)
logger.debug(f"Loaded file metadata from {self._metadata_file}")
return result
except Exception as e:
logger.warning(f"Failed to load file metadata from {self._metadata_file}: {e}")
return {}
async def _save_metadata(self, metadata: dict[str, dict[str, FileMetadata]]) -> None:
"""Save file metadata to disk.
Args:
metadata: Dictionary mapping source -> path -> FileMetadata
"""
try:
# Convert FileMetadata objects to dict for JSON serialization
metadata_dict = {}
for source, files in metadata.items():
metadata_dict[source] = {}
for path, meta in files.items():
metadata_dict[source][path] = {
"path": meta.path,
"hash": meta.hash,
"mtime_ms": meta.mtime_ms,
"size": meta.size,
"chunk_count": meta.chunk_count,
}
data = json.dumps(metadata_dict, indent=2, ensure_ascii=False)
self._metadata_file.write_text(data, encoding="utf-8")
logger.debug(f"Saved file metadata to {self._metadata_file}")
except Exception as e:
logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}")
async def start(self) -> None:
"""Initialize ChromaDB client and collection."""
if self.client is not None:
return
# Initialize persistent ChromaDB client
self.client = chromadb.PersistentClient(
path=str(self.db_path),
settings=Settings(
anonymized_telemetry=False,
allow_reset=True,
),
)
# Get or create the chunks collection
# ChromaDB uses cosine distance by default for similarity
self.chunks_collection = self.client.get_or_create_collection(
name=self.collection_name,
metadata={"hnsw:space": "cosine"},
)
# Load metadata into cache
self._metadata_cache = await self._load_metadata()
logger.info(f"ChromaDB initialized with collection: {self.collection_name}")
logger.info(f"File metadata will be persisted to: {self._metadata_file}")
async def upsert_file(
self,
file_meta: FileMetadata,
source: MemorySource,
chunks: list[MemoryChunk],
) -> None:
"""Insert or update file and its chunks."""
if not chunks:
return
# Delete existing chunks for this file first
await self.delete_file(file_meta.path, source)
# Batch generate embeddings for all chunks
# (base class returns mock embeddings when vector_enabled=False)
chunks = await self.get_chunk_embeddings(chunks)
# Prepare data for ChromaDB batch upsert
ids = []
documents = []
embeddings = []
metadatas = []
now = int(time.time() * 1000)
for chunk in chunks:
ids.append(chunk.id)
documents.append(chunk.text)
embeddings.append(chunk.embedding)
metadatas.append(
{
"path": file_meta.path,
"source": source.value,
"start_line": chunk.start_line,
"end_line": chunk.end_line,
"hash": chunk.hash,
"updated_at": now,
},
)
# Batch upsert to ChromaDB (always pass embeddings to prevent default embedding function)
self.chunks_collection.upsert(
ids=ids,
documents=documents,
embeddings=embeddings,
metadatas=metadatas,
)
# Update file metadata in cache
if source.value not in self._metadata_cache:
self._metadata_cache[source.value] = {}
self._metadata_cache[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=len(chunks),
)
async def delete_file(self, path: str, source: MemorySource) -> None:
"""Delete file and all its chunks."""
# Query for all chunks with this path and source
results = self.chunks_collection.get(
where={"$and": [{"path": path}, {"source": source.value}]},
include=[],
)
if results["ids"]:
self.chunks_collection.delete(
ids=results["ids"],
)
# Remove from file metadata cache
if source.value in self._metadata_cache:
self._metadata_cache[source.value].pop(path, None)
async def delete_file_chunks(self, path: str, chunk_ids: list[str]) -> None:
"""Delete specific chunks for a file."""
if not chunk_ids:
return
self.chunks_collection.delete(
ids=chunk_ids,
)
# Update chunk count in file metadata cache
for source_meta in self._metadata_cache.values():
if path in source_meta:
# Recalculate chunk count
results = self.chunks_collection.get(
where={"path": path},
include=[],
)
source_meta[path].chunk_count = len(results["ids"])
break
async def upsert_chunks(
self,
chunks: list[MemoryChunk],
source: MemorySource,
) -> None:
"""Insert or update specific chunks without affecting other chunks."""
if not chunks:
return
# Batch generate embeddings for all chunks
# (base class returns mock embeddings when vector_enabled=False)
chunks = await self.get_chunk_embeddings(chunks)
ids = []
documents = []
embeddings = []
metadatas = []
now = int(time.time() * 1000)
for chunk in chunks:
ids.append(chunk.id)
documents.append(chunk.text)
embeddings.append(chunk.embedding)
metadatas.append(
{
"path": chunk.path,
"source": source.value,
"start_line": chunk.start_line,
"end_line": chunk.end_line,
"hash": chunk.hash,
"updated_at": now,
},
)
# Always pass embeddings to prevent default embedding function
self.chunks_collection.upsert(
ids=ids,
documents=documents,
embeddings=embeddings,
metadatas=metadatas,
)
async def list_files(self, source: MemorySource) -> list[str]:
"""List all indexed files for a source."""
if source.value not in self._metadata_cache:
return []
return list(self._metadata_cache[source.value].keys())
async def get_file_metadata(
self,
path: str,
source: MemorySource,
) -> FileMetadata | None:
"""Get file metadata with chunk count."""
if source.value not in self._metadata_cache:
return None
return self._metadata_cache[source.value].get(path)
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
if source.value not in self._metadata_cache:
self._metadata_cache[source.value] = {}
self._metadata_cache[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,
)
async def get_file_chunks(
self,
path: str,
source: MemorySource,
) -> list[MemoryChunk]:
"""Get all chunks for a file."""
results = self.chunks_collection.get(
where={"$and": [{"path": path}, {"source": source.value}]},
include=["documents", "embeddings", "metadatas"],
)
chunks = []
for i, chunk_id in enumerate(results["ids"]):
metadata = results["metadatas"][i]
chunks.append(
MemoryChunk(
id=chunk_id,
path=metadata["path"],
source=MemorySource(metadata["source"]),
start_line=metadata["start_line"],
end_line=metadata["end_line"],
text=results["documents"][i],
hash=metadata["hash"],
embedding=results["embeddings"][i] if results["embeddings"] is not None else None,
),
)
# Sort by start_line
chunks.sort(key=lambda c: c.start_line)
return chunks
async def vector_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform vector similarity search."""
if not self.vector_enabled or not query:
return []
# Get query embedding
query_embedding = await self.get_embedding(query)
if not query_embedding:
return []
# Build where filter for sources
where_filter = None
if sources:
if len(sources) == 1:
where_filter = {"source": sources[0].value}
else:
where_filter = {"source": {"$in": [s.value for s in sources]}}
# Perform vector search
try:
results = self.chunks_collection.query(
query_embeddings=[query_embedding],
n_results=limit,
where=where_filter,
include=["documents", "metadatas", "distances"],
)
except Exception as e:
logger.error(f"Vector search failed: {e}, falling back to random results")
# Fallback: get some documents without vector search and assign random scores
try:
fallback_results = self.chunks_collection.get(
where=where_filter,
limit=limit,
include=["documents", "metadatas"],
)
search_results = []
if fallback_results["ids"]:
for i, _ in enumerate(fallback_results["ids"]):
metadata = fallback_results["metadatas"][i]
search_results.append(
MemorySearchResult(
path=metadata["path"],
start_line=metadata["start_line"],
end_line=metadata["end_line"],
score=random.uniform(0.3, 0.7), # Random score in middle range
snippet=fallback_results["documents"][i],
source=MemorySource(metadata["source"]),
raw_metric=None,
),
)
return search_results
except Exception as fallback_e:
logger.error(f"Fallback search also failed: {fallback_e}")
return []
search_results = []
if results["ids"] and results["ids"][0]:
for i, _ in enumerate(results["ids"][0]):
metadata = results["metadatas"][0][i]
distance = results["distances"][0][i]
# Convert cosine distance to similarity score
# Cosine distance range is [0, 2], convert to [1, 0] score
score = max(0.0, 1.0 - distance / 2.0)
search_results.append(
MemorySearchResult(
path=metadata["path"],
start_line=metadata["start_line"],
end_line=metadata["end_line"],
score=score,
snippet=results["documents"][0][i],
source=MemorySource(metadata["source"]),
raw_metric=distance,
),
)
# Sort by score descending
search_results.sort(key=lambda r: r.score, reverse=True)
return search_results
async def keyword_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform keyword/full-text search.
ChromaDB supports where_document filter for text matching.
Note: ChromaDB's $contains is case-sensitive, so we generate multiple
case variants (original, lowercase, capitalized) for each word to
improve recall while maintaining case-insensitive scoring.
"""
if not self.fts_enabled or not query:
return []
# Normalize whitespace and split into words
words = query.split()
if not words:
return []
# Generate case variants for each word to handle case-sensitive $contains
# Include: original, lowercase, and capitalized forms
word_variants = set()
for word in words:
word_variants.add(word) # original
word_variants.add(word.lower()) # lowercase
word_variants.add(word.capitalize()) # Capitalized
word_variants.add(word.upper()) # UPPERCASE
word_variants_list = list(word_variants)
# Build where filter for sources
where_filter = None
if sources:
if len(sources) == 1:
where_filter = {"source": sources[0].value}
else:
where_filter = {"source": {"$in": [s.value for s in sources]}}
# ChromaDB where_document uses $contains for substring matching (case-sensitive)
# Use multiple case variants to improve recall
if len(word_variants_list) == 1:
where_document: dict = {"$contains": word_variants_list[0]}
else:
where_document = {"$or": [{"$contains": w} for w in word_variants_list]}
# Get all matching documents
results = self.chunks_collection.get(
where=where_filter,
where_document=where_document,
include=["documents", "metadatas"],
)
search_results = []
query_lower = query.lower()
words_lower = [w.lower() for w in words] # lowercase words for scoring
n_words = len(words)
for i, _ in enumerate(results["ids"]):
metadata = results["metadatas"][i]
text = results["documents"][i]
text_lower = text.lower()
# Calculate relevance score based on word matches
match_count = sum(1 for w in words_lower if w in text_lower)
base_score = match_count / n_words
# Bonus for full phrase match (only applies to multi-word queries)
phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0
# Scale base_score and add phrase bonus, max score is 1.0
score = min(1.0, base_score + phrase_bonus)
search_results.append(
MemorySearchResult(
path=metadata["path"],
start_line=metadata["start_line"],
end_line=metadata["end_line"],
score=score,
snippet=text,
source=MemorySource(metadata["source"]),
),
)
# Sort by score descending and limit results
search_results.sort(key=lambda r: r.score, reverse=True)
return search_results[:limit]
async def hybrid_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
vector_weight: float = 0.7,
candidate_multiplier: float = 3.0,
) -> list[MemorySearchResult]:
"""Perform hybrid search combining vector and keyword search.
Args:
query: Search query text
limit: Maximum number of results
sources: Optional list of sources to filter
vector_weight: Weight for vector search results (0.0-1.0).
Keyword weight = 1.0 - vector_weight.
candidate_multiplier: Multiplier for candidate pool size.
Returns:
List of search results sorted by combined relevance score
"""
assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}"
candidates = min(200, max(1, int(limit * candidate_multiplier)))
text_weight = 1.0 - vector_weight
# Perform search based on enabled backends
if self.vector_enabled and self.fts_enabled:
keyword_results = await self.keyword_search(query, candidates, sources)
vector_results = await self.vector_search(query, candidates, sources)
# Log original vector results
logger.info("\n=== Vector Search Results ===")
for i, r in enumerate(vector_results[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
# Log original keyword results
logger.info("\n=== Keyword Search Results ===")
for i, r in enumerate(keyword_results[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
if not keyword_results:
return vector_results[:limit]
elif not vector_results:
return keyword_results[:limit]
else:
merged = self._merge_hybrid_results(
vector=vector_results,
keyword=keyword_results,
vector_weight=vector_weight,
text_weight=text_weight,
)
# Log merged results
logger.info("\n=== Merged Hybrid Results ===")
for i, r in enumerate(merged[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
return merged[:limit]
elif self.vector_enabled:
vector_results = await self.vector_search(query, limit, sources)
return vector_results
elif self.fts_enabled:
keyword_results = await self.keyword_search(query, limit, sources)
return keyword_results
else:
return []
@staticmethod
def _merge_hybrid_results(
vector: list[MemorySearchResult],
keyword: list[MemorySearchResult],
vector_weight: float,
text_weight: float,
) -> list[MemorySearchResult]:
"""Merge vector and keyword search results with weighted scoring."""
merged: dict[str, MemorySearchResult] = {}
# Process vector results
for result in vector:
result.score = result.score * vector_weight
merged[result.merge_key] = result
# Process keyword results
for result in keyword:
key = result.merge_key
if key in merged:
merged[key].score += result.score * text_weight
else:
result.score = result.score * text_weight
merged[key] = result
# Sort by score and return
results = list(merged.values())
results.sort(key=lambda r: r.score, reverse=True)
return results
async def clear_all(self) -> None:
"""Clear all indexed data."""
# Delete and recreate the collection
self.client.delete_collection(
name=self.collection_name,
)
self.chunks_collection = self.client.get_or_create_collection(
name=self.collection_name,
metadata={"hnsw:space": "cosine"},
)
# Clear file metadata cache and disk
self._metadata_cache = {}
await self._save_metadata({})
logger.info(f"Cleared all data from ChromaDB collection: {self.collection_name}")
async def close(self) -> None:
"""Close ChromaDB client and release resources."""
# Persist metadata cache to disk before closing
if self._metadata_cache:
await self._save_metadata(self._metadata_cache)
# ChromaDB PersistentClient handles persistence automatically
self.client = None
self.chunks_collection = None
await super().close()

View file

@ -0,0 +1,461 @@
"""Pure-Python in-memory storage backend for file store, with JSON file persistence."""
import json
from pathlib import Path
import numpy as np
from loguru import logger
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
from ..utils.common_utils import batch_cosine_similarity
class LocalFileStore(BaseFileStore):
"""Pure-Python in-memory file storage with JSONL file persistence.
No external dependencies required. All data lives in Python dicts;
writes are persisted to JSONL files on disk so state survives restarts.
Inherits embedding methods from BaseFileStore:
- get_chunk_embedding / get_chunk_embeddings (async)
- get_embedding / get_embeddings (async)
Provides:
- Vector similarity search (cosine similarity, pure Python)
- Full-text / keyword search (Python substring matching)
- Efficient chunk and file metadata management
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._started: bool = False
# In-memory indexes
self._chunks: dict[str, MemoryChunk] = {}
self._files: dict[str, dict[str, FileMetadata]] = {} # source -> path -> meta
# Persistence paths (mirror ChromaFileStore convention)
self._chunks_file: Path = self.db_path / f"{self.store_name}_chunks.jsonl"
self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json"
# ------------------------------------------------------------------
# Persistence helpers
# ------------------------------------------------------------------
async def _load_chunks(self) -> None:
"""Load chunks from JSONL file into memory."""
if not self._chunks_file.exists():
return
try:
data = self._chunks_file.read_text(encoding="utf-8")
self._chunks = {}
for line in data.strip().split("\n"):
if not line:
continue
rec = json.loads(line)
chunk = MemoryChunk.model_validate(rec)
self._chunks[chunk.id] = chunk
logger.debug(f"Loaded {len(self._chunks)} chunks from {self._chunks_file}")
except Exception as e:
logger.warning(f"Failed to load chunks from {self._chunks_file}: {e}")
async def _save_chunks(self) -> None:
"""Persist chunks to JSONL file."""
try:
lines = []
for chunk in self._chunks.values():
chunk_dict = chunk.model_dump(mode="json")
lines.append(json.dumps(chunk_dict, ensure_ascii=False))
data = "\n".join(lines)
self._chunks_file.write_text(data, encoding="utf-8")
logger.debug(f"Saved {len(self._chunks)} chunks to {self._chunks_file}")
except Exception as e:
logger.error(f"Failed to save chunks to {self._chunks_file}: {e}")
async def _load_metadata(self) -> None:
"""Load file metadata from JSON file into memory."""
if not self._metadata_file.exists():
return
try:
data = self._metadata_file.read_text(encoding="utf-8")
raw: dict = json.loads(data)
self._files = {
source: {path: FileMetadata(**meta) for path, meta in files.items()} for source, files in raw.items()
}
logger.debug(f"Loaded file metadata from {self._metadata_file}")
except Exception as e:
logger.warning(f"Failed to load file metadata from {self._metadata_file}: {e}")
async def _save_metadata(self) -> None:
"""Persist file metadata to JSON file."""
try:
raw: dict = {}
for source, files in self._files.items():
raw[source] = {
path: {
"path": meta.path,
"hash": meta.hash,
"mtime_ms": meta.mtime_ms,
"size": meta.size,
"chunk_count": meta.chunk_count,
}
for path, meta in files.items()
}
data = json.dumps(raw, indent=2, ensure_ascii=False)
self._metadata_file.write_text(data, encoding="utf-8")
logger.debug(f"Saved file metadata to {self._metadata_file}")
except Exception as e:
logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}")
# ------------------------------------------------------------------
# Lifecycle
# ------------------------------------------------------------------
async def start(self) -> None:
"""Load persisted data into memory."""
if self._started:
return
self._started = True
await self._load_metadata()
await self._load_chunks()
logger.info(
f"LocalFileStore '{self.store_name}' ready: "
f"{len(self._chunks)} chunks, metadata at {self._metadata_file}",
)
async def close(self) -> None:
"""Flush state to disk and release memory."""
await self._save_metadata()
await self._save_chunks()
self._chunks.clear()
self._files.clear()
self._started = False
# ------------------------------------------------------------------
# Write operations
# ------------------------------------------------------------------
async def upsert_file(
self,
file_meta: FileMetadata,
source: MemorySource,
chunks: list[MemoryChunk],
) -> None:
"""Insert or update file and its chunks."""
if not chunks:
return
# Remove existing chunks for this file/source first
await self.delete_file(file_meta.path, source)
# Batch generate embeddings (base class returns mock embeddings when vector_enabled=False)
chunks = await self.get_chunk_embeddings(chunks)
for chunk in chunks:
self._chunks[chunk.id] = chunk
if source.value not in self._files:
self._files[source.value] = {}
self._files[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=len(chunks),
)
async def delete_file(self, path: str, source: MemorySource) -> None:
"""Delete file and all its chunks."""
to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path and chunk.source == source]
for cid in to_delete:
del self._chunks[cid]
if source.value in self._files:
self._files[source.value].pop(path, None)
async def delete_file_chunks(self, path: str, chunk_ids: list[str]) -> None:
"""Delete specific chunks for a file."""
if not chunk_ids:
return
for cid in chunk_ids:
self._chunks.pop(cid, None)
# Recalculate chunk_count in file metadata (per source)
for source_key, source_meta in self._files.items():
if path in source_meta:
source_meta[path].chunk_count = sum(
1 for chunk in self._chunks.values() if chunk.path == path and chunk.source.value == source_key
)
async def upsert_chunks(
self,
chunks: list[MemoryChunk],
source: MemorySource,
) -> None:
"""Insert or update specific chunks without affecting other chunks."""
if not chunks:
return
chunks = await self.get_chunk_embeddings(chunks)
for chunk in chunks:
self._chunks[chunk.id] = chunk
# ------------------------------------------------------------------
# Read operations
# ------------------------------------------------------------------
async def list_files(self, source: MemorySource) -> list[str]:
"""List all indexed files for a source."""
return list(self._files.get(source.value, {}).keys())
async def get_file_metadata(
self,
path: str,
source: MemorySource,
) -> FileMetadata | None:
"""Get file metadata."""
return self._files.get(source.value, {}).get(path)
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
if source.value not in self._files:
self._files[source.value] = {}
self._files[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,
)
async def get_file_chunks(
self,
path: str,
source: MemorySource,
) -> list[MemoryChunk]:
"""Get all chunks for a file, sorted by start_line."""
chunks = [chunk for chunk in self._chunks.values() if chunk.path == path and chunk.source == source]
chunks.sort(key=lambda c: c.start_line)
return chunks
# ------------------------------------------------------------------
# Search
# ------------------------------------------------------------------
async def vector_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform cosine-similarity vector search over in-memory embeddings."""
if not self.vector_enabled or not query:
return []
query_embedding = await self.get_embedding(query)
if not query_embedding:
return []
expected_dim = self.embedding_dim
# Collect candidate chunks with embeddings
candidates = [
chunk for chunk in self._chunks.values() if (not sources or chunk.source in sources) and chunk.embedding
]
if not candidates:
return []
# Validate and fix chunk embedding dimensions
valid_embeddings = []
for chunk in candidates:
emb = chunk.embedding
emb_len = len(emb)
if emb_len != expected_dim:
if emb_len < expected_dim:
emb = emb + [0.0] * (expected_dim - emb_len)
logger.warning(
f"Chunk embedding dimension {emb_len} < expected {expected_dim}, "
f"padded with zeros (chunk_id={chunk.id})",
)
else:
emb = emb[:expected_dim]
logger.warning(
f"Chunk embedding dimension {emb_len} > expected {expected_dim}, "
f"truncated to {expected_dim} (chunk_id={chunk.id})",
)
valid_embeddings.append(emb)
# Build embedding matrix and compute similarities in batch
query_array = np.array([query_embedding]) # Shape: (1, emb_size)
chunk_embeddings = np.array(valid_embeddings) # Shape: (n, emb_size)
similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] # Shape: (n,)
# Build results
results = [
MemorySearchResult(
path=chunk.path,
start_line=chunk.start_line,
end_line=chunk.end_line,
score=float(similarity),
snippet=chunk.text,
source=chunk.source,
raw_metric=1.0 - float(similarity),
)
for chunk, similarity in zip(candidates, similarities)
]
results.sort(key=lambda r: r.score, reverse=True)
return results[:limit]
async def keyword_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform keyword/full-text search via Python substring matching."""
if not self.fts_enabled or not query:
return []
words = query.split()
if not words:
return []
query_lower = query.lower()
words_lower = [w.lower() for w in words]
n_words = len(words)
results = []
for chunk in self._chunks.values():
if sources and chunk.source not in sources:
continue
text_lower = chunk.text.lower()
match_count = sum(1 for w in words_lower if w in text_lower)
if match_count == 0:
continue
base_score = match_count / n_words
# Bonus for full phrase match (multi-word queries only)
phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0
score = min(1.0, base_score + phrase_bonus)
results.append(
MemorySearchResult(
path=chunk.path,
start_line=chunk.start_line,
end_line=chunk.end_line,
score=score,
snippet=chunk.text,
source=chunk.source,
),
)
results.sort(key=lambda r: r.score, reverse=True)
return results[:limit]
async def hybrid_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
vector_weight: float = 0.7,
candidate_multiplier: float = 3.0,
) -> list[MemorySearchResult]:
"""Perform hybrid search combining vector and keyword search.
Args:
query: Search query text
limit: Maximum number of results
sources: Optional list of sources to filter
vector_weight: Weight for vector search results (0.0-1.0).
Keyword weight = 1.0 - vector_weight.
candidate_multiplier: Multiplier for candidate pool size.
Returns:
List of search results sorted by combined relevance score
"""
assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}"
candidates = min(200, max(1, int(limit * candidate_multiplier)))
text_weight = 1.0 - vector_weight
if self.vector_enabled and self.fts_enabled:
keyword_results = await self.keyword_search(query, candidates, sources)
vector_results = await self.vector_search(query, candidates, sources)
logger.info("\n=== Vector Search Results ===")
for i, r in enumerate(vector_results[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
logger.info("\n=== Keyword Search Results ===")
for i, r in enumerate(keyword_results[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
if not keyword_results:
return vector_results[:limit]
elif not vector_results:
return keyword_results[:limit]
else:
merged = self._merge_hybrid_results(
vector=vector_results,
keyword=keyword_results,
vector_weight=vector_weight,
text_weight=text_weight,
)
logger.info("\n=== Merged Hybrid Results ===")
for i, r in enumerate(merged[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
return merged[:limit]
elif self.vector_enabled:
return await self.vector_search(query, limit, sources)
elif self.fts_enabled:
return await self.keyword_search(query, limit, sources)
else:
return []
@staticmethod
def _merge_hybrid_results(
vector: list[MemorySearchResult],
keyword: list[MemorySearchResult],
vector_weight: float,
text_weight: float,
) -> list[MemorySearchResult]:
"""Merge vector and keyword search results with weighted scoring."""
merged: dict[str, MemorySearchResult] = {}
for result in vector:
result.metadata["_weighted_score"] = result.score * vector_weight
merged[result.merge_key] = result
for result in keyword:
key = result.merge_key
if key in merged:
merged[key].metadata["_weighted_score"] += result.score * text_weight
else:
result.metadata["_weighted_score"] = result.score * text_weight
merged[key] = result
results = list(merged.values())
for r in results:
r.score = r.metadata.pop("_weighted_score")
results.sort(key=lambda r: r.score, reverse=True)
return results
async def clear_all(self) -> None:
"""Clear all indexed data from memory and disk."""
self._chunks.clear()
self._files.clear()
await self._save_chunks()
await self._save_metadata()
logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'")

View file

@ -0,0 +1,978 @@
"""SQLite storage backend for file store."""
import json
import struct
import time
from loguru import logger
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
class SqliteFileStore(BaseFileStore):
"""SQLite file storage with vector and full-text search.
Inherits embedding methods from BaseFileStore:
- get_chunk_embedding / get_chunk_embeddings (async)
- get_chunk_embedding_sync / get_chunk_embeddings_sync (sync)
- get_embedding / get_embeddings (async)
Provides SQLite-backed persistent storage with:
- Vector similarity search (via sqlite-vec extension)
- Full-text search (via FTS5)
- Efficient chunk and file metadata management
"""
def __init__(self, vec_ext_path: str = "", **kwargs):
super().__init__(**kwargs)
self.vec_ext_path = vec_ext_path
import sqlite3
self.conn: sqlite3.Connection | None = None
@property
def vector_table_name(self) -> str:
"""Get the name of the vector table for this store."""
return f"chunks_vec_{self.store_name}"
@property
def fts_table_name(self) -> str:
"""Get the name of the FTS table for this store."""
return f"chunks_fts_{self.store_name}"
@property
def chunks_table_name(self) -> str:
"""Get the name of the chunks table for this store."""
return f"chunks_{self.store_name}"
@property
def files_table_name(self) -> str:
"""Get the name of the files table for this store."""
return f"files_{self.store_name}"
@staticmethod
def vector_to_blob(embedding: list[float]) -> bytes:
"""Convert vector to binary blob for sqlite-vec."""
return struct.pack(f"{len(embedding)}f", *embedding)
async def start(self) -> None:
"""Initialize database and load extensions."""
if self.conn is not None:
return
import sqlite3
self.conn = sqlite3.connect(self.db_path / "reme.db", check_same_thread=False)
# Only load sqlite-vec extension if vector search is enabled
if self.vector_enabled:
logger.warning(
"On macOS systems with version 14 or earlier, "
"loading the sqlite-vec vector extension carries a risk of crashes or hangs.",
)
self.conn.enable_load_extension(True)
# Load sqlite-vec extension
if self.vec_ext_path:
try:
self.conn.load_extension(self.vec_ext_path)
logger.info(f"Loaded sqlite-vec: {self.vec_ext_path}")
except Exception as e:
logger.warning(f"Failed to load sqlite-vec: {e}")
else:
try:
import sqlite_vec
ext_path = sqlite_vec.loadable_path()
self.conn.load_extension(ext_path)
logger.info(f"Loaded sqlite-vec from package: {ext_path}")
except Exception as e:
logger.warning(f"Failed to load sqlite-vec from package: {e}")
# Fallback: try common extension names
for name in ["vec0", "sqlite_vec", "vector0"]:
try:
self.conn.load_extension(name)
logger.info(f"Loaded sqlite-vec: {name}")
break
except Exception:
pass
self.conn.enable_load_extension(False)
else:
logger.info("Vector search disabled, skipping sqlite-vec extension loading")
await self._create_tables()
async def _create_tables(self) -> None:
"""Create database schema."""
cursor = self.conn.cursor()
try:
# Files
cursor.execute(
f"""
CREATE TABLE IF NOT EXISTS {self.files_table_name} (
path TEXT,
source TEXT,
hash TEXT,
mtime REAL,
size INTEGER,
PRIMARY KEY (path, source)
)
""",
)
# Chunks
cursor.execute(
f"""
CREATE TABLE IF NOT EXISTS {self.chunks_table_name} (
id TEXT PRIMARY KEY,
path TEXT,
source TEXT,
start_line INTEGER,
end_line INTEGER,
hash TEXT,
text TEXT,
embedding TEXT,
updated_at INTEGER
)
""",
)
# Vector table (sqlite-vec)
if self.vector_enabled:
cursor.execute(
f"""
CREATE VIRTUAL TABLE IF NOT EXISTS {self.vector_table_name} USING vec0(
id TEXT PRIMARY KEY,
embedding FLOAT[{self.embedding_dim}]
)
""",
)
logger.info(f"Created vector table (dims={self.embedding_dim})")
# FTS table
if self.fts_enabled:
cursor.execute(
f"""
CREATE VIRTUAL TABLE IF NOT EXISTS {self.fts_table_name} USING fts5(
text,
id UNINDEXED,
path UNINDEXED,
source UNINDEXED,
start_line UNINDEXED,
end_line UNINDEXED,
tokenize='trigram'
)
""",
)
logger.info("Created FTS5 table with trigram tokenizer")
self.conn.commit()
except Exception as e:
logger.error(f"Failed to create tables: {e}")
raise
finally:
cursor.close()
async def upsert_file(self, file_meta: FileMetadata, source: MemorySource, chunks: list[MemoryChunk]):
"""Insert or update file and its chunks."""
cursor = self.conn.cursor()
try:
cursor.execute("BEGIN")
# Insert file
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.files_table_name} (path, source, hash, mtime, size)
VALUES (?, ?, ?, ?, ?)
""",
(file_meta.path, source.value, file_meta.hash, file_meta.mtime_ms, file_meta.size),
)
# Insert chunks
now = int(time.time() * 1000)
for chunk in chunks:
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.chunks_table_name} (
id, path, source, start_line, end_line,
hash, text, embedding, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
chunk.id,
file_meta.path,
source.value,
chunk.start_line,
chunk.end_line,
chunk.hash,
chunk.text,
json.dumps(chunk.embedding) if chunk.embedding else None,
now,
),
)
# Insert vector (vec0 doesn't support OR REPLACE, use DELETE + INSERT)
if self.vector_enabled:
if not chunk.embedding:
logger.warning(
f"Chunk {chunk.id} missing embedding for vector insert, skipping vector indexing",
)
else:
# Delete existing vector first
cursor.execute(
f"DELETE FROM {self.vector_table_name} WHERE id = ?",
(chunk.id,),
)
# Then insert new vector
cursor.execute(
f"INSERT INTO {self.vector_table_name} (id, embedding) VALUES (?, ?)",
(chunk.id, self.vector_to_blob(chunk.embedding)),
)
# Insert FTS
if self.fts_enabled:
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.fts_table_name} (
text, id, path, source, start_line, end_line
) VALUES (?, ?, ?, ?, ?, ?)
""",
(
chunk.text,
chunk.id,
file_meta.path,
source.value,
chunk.start_line,
chunk.end_line,
),
)
cursor.execute("COMMIT")
except Exception as e:
cursor.execute("ROLLBACK")
logger.error(f"Failed to upsert file {file_meta.path}: {e}")
raise
finally:
cursor.close()
async def delete_file(self, path: str, source: MemorySource):
"""Delete file and all its chunks."""
cursor = self.conn.cursor()
try:
cursor.execute("BEGIN")
# Get chunk IDs for vector deletion
cursor.execute(
f"SELECT id FROM {self.chunks_table_name} WHERE path = ? AND source = ?",
(path, source.value),
)
chunk_ids = [row[0] for row in cursor.fetchall()]
# Delete vectors
if self.vector_enabled and chunk_ids:
for chunk_id in chunk_ids:
cursor.execute(
f"DELETE FROM {self.vector_table_name} WHERE id = ?",
(chunk_id,),
)
# Delete FTS entries
if self.fts_enabled:
cursor.execute(
f"DELETE FROM {self.fts_table_name} WHERE path = ? AND source = ?",
(path, source.value),
)
# Delete chunks and file
cursor.execute(
f"DELETE FROM {self.chunks_table_name} WHERE path = ? AND source = ?",
(path, source.value),
)
cursor.execute(
f"DELETE FROM {self.files_table_name} WHERE path = ? AND source = ?",
(path, source.value),
)
cursor.execute("COMMIT")
except Exception as e:
cursor.execute("ROLLBACK")
logger.error(f"Failed to delete file {path}: {e}")
raise
finally:
cursor.close()
async def delete_file_chunks(self, path: str, chunk_ids: list[str]):
"""Delete specific chunks for a file."""
if not chunk_ids:
return
cursor = self.conn.cursor()
try:
cursor.execute("BEGIN")
# Delete vectors
if self.vector_enabled:
for chunk_id in chunk_ids:
cursor.execute(
f"DELETE FROM {self.vector_table_name} WHERE id = ?",
(chunk_id,),
)
# Delete FTS entries
if self.fts_enabled:
placeholders = ",".join("?" * len(chunk_ids))
cursor.execute(
f"DELETE FROM {self.fts_table_name} WHERE id IN ({placeholders})",
chunk_ids,
)
# Delete chunks
placeholders = ",".join("?" * len(chunk_ids))
cursor.execute(
f"DELETE FROM {self.chunks_table_name} WHERE id IN ({placeholders})",
chunk_ids,
)
cursor.execute("COMMIT")
except Exception as e:
cursor.execute("ROLLBACK")
logger.error(f"Failed to delete chunks for {path}: {e}")
raise
finally:
cursor.close()
async def upsert_chunks(self, chunks: list[MemoryChunk], source: MemorySource):
"""Insert or update specific chunks without affecting other chunks."""
if not chunks:
return
cursor = self.conn.cursor()
try:
cursor.execute("BEGIN")
now = int(time.time() * 1000)
for chunk in chunks:
# Insert/update chunk
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.chunks_table_name} (
id, path, source, start_line, end_line,
hash, text, embedding, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
chunk.id,
chunk.path,
source.value,
chunk.start_line,
chunk.end_line,
chunk.hash,
chunk.text,
json.dumps(chunk.embedding) if chunk.embedding else None,
now,
),
)
# Insert/update vector (vec0 doesn't support OR REPLACE, use DELETE + INSERT)
if self.vector_enabled:
if not chunk.embedding:
logger.warning(
f"Chunk {chunk.id} missing embedding for vector insert, skipping vector indexing",
)
else:
# Delete existing vector first
cursor.execute(
f"DELETE FROM {self.vector_table_name} WHERE id = ?",
(chunk.id,),
)
# Then insert new vector
cursor.execute(
f"INSERT INTO {self.vector_table_name} (id, embedding) VALUES (?, ?)",
(chunk.id, self.vector_to_blob(chunk.embedding)),
)
# Insert/update FTS
if self.fts_enabled:
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.fts_table_name} (
text, id, path, source, start_line, end_line
) VALUES (?, ?, ?, ?, ?, ?)
""",
(
chunk.text,
chunk.id,
chunk.path,
source.value,
chunk.start_line,
chunk.end_line,
),
)
cursor.execute("COMMIT")
except Exception as e:
cursor.execute("ROLLBACK")
logger.error(f"Failed to upsert chunks: {e}")
raise
finally:
cursor.close()
async def list_files(self, source: MemorySource) -> list[str]:
"""List all indexed files."""
cursor = self.conn.cursor()
try:
cursor.execute(f"SELECT path FROM {self.files_table_name} WHERE source = ?", (source.value,))
paths = [row[0] for row in cursor.fetchall()]
return paths
except Exception as e:
logger.error(f"Failed to list files: {e}")
raise
finally:
cursor.close()
async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None:
"""Get file metadata with chunk count."""
cursor = self.conn.cursor()
try:
cursor.execute(
f"SELECT hash, mtime, size FROM {self.files_table_name} WHERE path = ? AND source = ?",
(path, source.value),
)
row = cursor.fetchone()
if not row:
return None
hash_val, mtime, size = row
cursor.execute(
f"SELECT COUNT(*) FROM {self.chunks_table_name} WHERE path = ? AND source = ?",
(path, source.value),
)
chunk_count = cursor.fetchone()[0]
return FileMetadata(
hash=hash_val,
mtime_ms=mtime,
size=size,
path=path,
chunk_count=chunk_count,
)
except Exception as e:
logger.error(f"Failed to get file metadata for {path}: {e}")
raise
finally:
cursor.close()
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
cursor = self.conn.cursor()
try:
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.files_table_name} (path, source, hash, mtime, size)
VALUES (?, ?, ?, ?, ?)
""",
(file_meta.path, source.value, file_meta.hash, file_meta.mtime_ms, file_meta.size),
)
self.conn.commit()
except Exception as e:
logger.error(f"Failed to update file metadata for {file_meta.path}: {e}")
raise
finally:
cursor.close()
async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]:
"""Get all chunks for a file."""
cursor = self.conn.cursor()
try:
cursor.execute(
f"""
SELECT id, path, source, start_line, end_line, text, hash, embedding
FROM {self.chunks_table_name} WHERE path = ? AND source = ?
ORDER BY start_line
""",
(path, source.value),
)
chunks = []
for row in cursor.fetchall():
chunk_id, path_val, source_val, start, end, text, hash_val, emb_str = row
# Parse embedding from JSON string
embedding = None
if emb_str:
try:
embedding = json.loads(emb_str)
except (json.JSONDecodeError, TypeError):
embedding = None
chunks.append(
MemoryChunk(
id=chunk_id,
path=path_val,
source=MemorySource(source_val),
start_line=start,
end_line=end,
text=text,
hash=hash_val,
embedding=embedding,
),
)
return chunks
except Exception as e:
logger.error(f"Failed to get file chunks for {path}: {e}")
raise
finally:
cursor.close()
async def vector_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform vector similarity search."""
if not self.vector_enabled or not query:
return []
# Get query embedding
query_embedding = await self.get_embedding(query)
if not query_embedding:
return []
cursor = self.conn.cursor()
source_filter = ""
params: list = []
if sources:
placeholders = ",".join("?" * len(sources))
source_filter = f" AND c.source IN ({placeholders})"
params = [s.value for s in sources]
try:
query_blob = self.vector_to_blob(query_embedding)
# Correct SQLite-vec syntax for vector search with limit
# vec0 requires 'k = ?' constraint for knn queries
query_sql = f"""
SELECT c.id, c.path, c.start_line, c.end_line, c.source, c.text, v.distance
FROM {self.vector_table_name} v
JOIN {self.chunks_table_name} c ON v.id = c.id
WHERE v.embedding MATCH ?
AND k = ?
"""
query_params: list = [query_blob, limit]
# Add source filter if specified
if source_filter:
query_sql += source_filter
query_params.extend(params)
# Order by distance (k constraint already limits results)
query_sql += " ORDER BY v.distance"
cursor.execute(query_sql, query_params)
results = []
for _, path, start, end, src, text, dist in cursor.fetchall():
# Convert L2 distance to similarity score
# For normalized vectors, L2 distance range is [0, 2]
# Map to [1, 0] score range (higher score = more similar)
score = max(0.0, 1.0 - dist / 2.0)
snippet = text
results.append(
MemorySearchResult(
path=path,
start_line=start,
end_line=end,
score=score,
snippet=snippet,
source=MemorySource(src),
raw_metric=dist,
),
)
results.sort(key=lambda r: r.score, reverse=True)
return results
except Exception as e:
logger.error(f"Vector search failed: {e}")
return []
finally:
cursor.close()
@staticmethod
def _sanitize_fts_query(query: str) -> str:
"""Sanitize query string for FTS5 search.
Removes or escapes special characters that have special meaning in FTS5:
- * (prefix match)
- ? (not used in FTS5, but can cause issues)
- " (phrase search, needs escaping)
- : (column filter)
- ^ (start of line anchor, not standard FTS5)
- ' (single quote, causes syntax errors)
- ` (backtick, can cause issues)
- | (pipe, OR operator)
- + (plus, can be used for required terms)
- - (minus, NOT operator)
- = (equals, can cause issues)
- < > (angle brackets, comparison operators)
- ! (exclamation, NOT operator variant)
- @ # $ % & (other special chars)
- "\"
- / (slash, can interfere)
- ; (semicolon, statement separator)
- , (comma, can interfere with phrase parsing)
Args:
query: Raw query string
Returns:
Sanitized query string safe for FTS5
"""
if not query:
return ""
# Remove FTS5 special characters that we don't want users to use
# Keep only alphanumeric, spaces, periods, and underscores
special_chars = [
"*",
"?",
":",
"^",
"(",
")",
"[",
"]",
"{",
"}",
"'",
'"',
"`",
"|",
"+",
"-",
"=",
"<",
">",
"!",
"@",
"#",
"$",
"%",
"&",
"\\",
"/",
";",
",",
]
cleaned = query
for char in special_chars:
cleaned = cleaned.replace(char, " ")
# Normalize whitespace
cleaned = " ".join(cleaned.split())
return cleaned
async def keyword_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform keyword search.
Strategy:
- FTS5 trigram (fast path): used when ALL terms >= 3 chars (trigram minimum).
- LIKE (universal fallback): used when any term < 3 chars, covering CJK
short words, single/double-char queries, and mixed-length queries.
"""
if not self.fts_enabled:
return []
cleaned = self._sanitize_fts_query(query)
if not cleaned:
return []
words = cleaned.split()
if not words:
return []
# FTS5 trigram requires every term >= 3 characters
if all(len(w) >= 3 for w in words):
results = await self._fts_trigram_search(words, limit, sources)
if results:
return results
# Universal fallback: LIKE-based substring search
return await self._like_search(cleaned, words, limit, sources)
async def _fts_trigram_search(
self,
words: list[str],
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""FTS5 trigram search. All terms must be >= 3 characters."""
escaped_words = [w.replace('"', '""') for w in words]
fts_query = " OR ".join(escaped_words)
cursor = self.conn.cursor()
source_filter = ""
params: list = [fts_query]
if sources:
placeholders = ",".join("?" * len(sources))
source_filter = f" AND fts.source IN ({placeholders})"
params.extend([s.value for s in sources])
params.append(limit)
try:
cursor.execute(
f"""
SELECT fts.id, fts.path, fts.start_line, fts.end_line,
fts.source, fts.text, rank
FROM {self.fts_table_name} fts
WHERE fts.text MATCH ?{source_filter}
ORDER BY rank
LIMIT ?
""",
params,
)
results = []
for _, path, start, end, src, text, rank in cursor.fetchall():
score = max(0.0, 1.0 / (1.0 + abs(rank)))
results.append(
MemorySearchResult(
path=path,
start_line=start,
end_line=end,
score=score,
snippet=text,
source=MemorySource(src),
raw_metric=rank,
),
)
results.sort(key=lambda r: r.score, reverse=True)
return results
except Exception as e:
logger.error(f"FTS trigram search failed: {e}")
return []
finally:
cursor.close()
async def _like_search(
self,
phrase: str,
words: list[str],
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""LIKE-based substring search with Python-side relevance scoring.
Handles any term length and all languages (CJK, Latin, etc.).
Scores results by: word-match ratio + full-phrase bonus.
"""
cursor = self.conn.cursor()
try:
# Build OR conditions: match any individual word
like_clauses = []
params: list = []
for word in words:
like_clauses.append("c.text LIKE ?")
params.append(f"%{word}%")
where_clause = " OR ".join(like_clauses)
source_filter = ""
if sources:
placeholders = ",".join("?" * len(sources))
source_filter = f" AND c.source IN ({placeholders})"
params.extend([s.value for s in sources])
# Fetch extra candidates for re-ranking in Python
fetch_limit = min(limit * 3, 200)
params.append(fetch_limit)
cursor.execute(
f"""
SELECT c.id, c.path, c.start_line, c.end_line, c.source, c.text
FROM {self.chunks_table_name} c
WHERE ({where_clause}){source_filter}
LIMIT ?
""",
params,
)
results = []
phrase_lower = phrase.lower()
words_lower = [w.lower() for w in words]
n_words = len(words)
for _, path, start, end, src, text in cursor.fetchall():
text_lower = text.lower()
# Base score: proportion of query words found in text
match_count = sum(1 for w in words_lower if w in text_lower)
base_score = match_count / n_words
# Bonus: full phrase appears as contiguous substring
phrase_bonus = 0.2 if n_words > 1 and phrase_lower in text_lower else 0.0
score = min(1.0, base_score * 0.8 + phrase_bonus)
results.append(
MemorySearchResult(
path=path,
start_line=start,
end_line=end,
score=score,
snippet=text,
source=MemorySource(src),
),
)
# Sort by score descending, return top `limit`
results.sort(key=lambda r: r.score, reverse=True)
return results[:limit]
except Exception as e:
logger.error(f"LIKE search failed: {e}")
return []
finally:
cursor.close()
async def hybrid_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
vector_weight: float = 0.7,
candidate_multiplier: float = 3.0,
) -> list[MemorySearchResult]:
"""Perform hybrid search combining vector and keyword search.
Args:
query: Search query text
limit: Maximum number of results
sources: Optional list of sources to filter
vector_weight: Weight for vector search results (0.0-1.0).
Keyword weight = 1.0 - vector_weight.
candidate_multiplier: Multiplier for candidate pool size.
Returns:
List of search results sorted by combined relevance score
"""
assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}"
candidates = min(200, max(1, int(limit * candidate_multiplier)))
text_weight = 1.0 - vector_weight
# Perform search based on enabled backends
if self.vector_enabled and self.fts_enabled:
keyword_results = await self.keyword_search(query, candidates, sources)
vector_results = await self.vector_search(query, candidates, sources)
# Log original vector results
logger.info("\n=== Vector Search Results ===")
for i, r in enumerate(vector_results[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
# Log original keyword results
logger.info("\n=== Keyword Search Results ===")
for i, r in enumerate(keyword_results[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
if not keyword_results:
return vector_results[:limit]
elif not vector_results:
return keyword_results[:limit]
else:
merged = self._merge_hybrid_results(
vector=vector_results,
keyword=keyword_results,
vector_weight=vector_weight,
text_weight=text_weight,
)
# Log merged results
logger.info("\n=== Merged Hybrid Results ===")
for i, r in enumerate(merged[:10], 1):
snippet_preview = (r.snippet[:100] + "...") if len(r.snippet) > 100 else r.snippet
logger.info(f"{i}. Score: {r.score:.4f} | Snippet: {snippet_preview}")
return merged[:limit]
elif self.vector_enabled:
vector_results = await self.vector_search(query, limit, sources)
return vector_results
elif self.fts_enabled:
keyword_results = await self.keyword_search(query, limit, sources)
return keyword_results
else:
return []
@staticmethod
def _merge_hybrid_results(
vector: list[MemorySearchResult],
keyword: list[MemorySearchResult],
vector_weight: float,
text_weight: float,
) -> list[MemorySearchResult]:
"""Merge vector and keyword search results with weighted scoring."""
merged: dict[str, MemorySearchResult] = {}
# Process vector results
for result in vector:
result.score = result.score * vector_weight
merged[result.merge_key] = result
# Process keyword results
for result in keyword:
key = result.merge_key
if key in merged:
merged[key].score += result.score * text_weight
else:
result.score = result.score * text_weight
merged[key] = result
# Sort by score and return
results = list(merged.values())
results.sort(key=lambda r: r.score, reverse=True)
return results
async def clear_all(self):
"""Clear all indexed data."""
cursor = self.conn.cursor()
try:
cursor.execute("BEGIN")
cursor.execute(f"DELETE FROM {self.files_table_name}")
cursor.execute(f"DELETE FROM {self.chunks_table_name}")
if self.vector_enabled:
cursor.execute(f"DELETE FROM {self.vector_table_name}")
if self.fts_enabled:
cursor.execute(f"DELETE FROM {self.fts_table_name}")
cursor.execute("COMMIT")
except Exception as e:
cursor.execute("ROLLBACK")
logger.error(f"Failed to clear all data: {e}")
raise
finally:
cursor.close()
async def close(self):
"""Close database connection."""
if self.conn:
self.conn.close()
self.conn = None
await super().close()

View file

@ -0,0 +1,19 @@
"""File watcher module for monitoring file system changes.
This module provides file watcher implementations for monitoring file changes
and updating memory stores accordingly.
"""
from .base_file_watcher import BaseFileWatcher
from .delta_file_watcher import DeltaFileWatcher
from .full_file_watcher import FullFileWatcher
from ..registry_factory import R
__all__ = [
"BaseFileWatcher",
"DeltaFileWatcher",
"FullFileWatcher",
]
R.file_watchers.register("full")(FullFileWatcher)
R.file_watchers.register("delta")(DeltaFileWatcher)

View file

@ -0,0 +1,240 @@
"""Base file watcher implementation.
This module provides the base class for file watcher implementations
that monitor file system changes and trigger callbacks.
"""
import asyncio
from collections.abc import Coroutine
from pathlib import Path
from typing import Any, Callable
from loguru import logger
from watchfiles import awatch, Change
from ..enumeration import MemorySource
from ..file_store import BaseFileStore
class BaseFileWatcher:
"""
Minimal file watcher base class
This base class provides basic file monitoring functionality that can be extended
to implement specific file monitoring requirements.
"""
def __init__(
self,
watch_paths: list[str] | str,
suffix_filters: list[str] | None = None,
recursive: bool = False,
debounce: int = 2000,
chunk_tokens: int = 400,
chunk_overlap: int = 80,
file_store: BaseFileStore | None = None,
callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None,
rebuild_index_on_start: bool = True,
poll_delay_ms: int = 2000,
**kwargs,
):
"""
Initialize the file watcher
Args:
watch_paths: Paths to watch for changes
suffix_filters: File suffix filters (e.g., ['.py', '.txt'])
recursive: Whether to watch directories recursively
debounce: Debounce time in milliseconds
chunk_tokens: Token size for chunking
chunk_overlap: Overlap size for chunks
file_store: File store instance
callback: Callback function for changes
rebuild_index_on_start: If True, clear all indexed data on start and rescan existing files.
If False, only monitor new changes without initialization.
poll_delay_ms: Polling delay in milliseconds. If > 300ms, force_polling will be enabled automatically.
**kwargs: Additional keyword arguments
"""
self.watch_paths: list[str] = [watch_paths] if isinstance(watch_paths, str) else watch_paths
self.suffix_filters: list[str] = suffix_filters or []
self.recursive: bool = recursive
self.debounce: int = debounce
self.chunk_tokens: int = chunk_tokens
self.chunk_overlap: int = chunk_overlap
self.file_store: BaseFileStore = file_store
self.callback = callback
self.rebuild_index_on_start: bool = rebuild_index_on_start
self.poll_delay_ms: int = poll_delay_ms
self.kwargs: dict = kwargs
self._stop_event = asyncio.Event()
self._watch_task: asyncio.Task | None = None
self._running = False
async def start(self):
"""Start the file watcher"""
if self._running:
return
self._running = True
async def _initialize_and_watch():
if self.rebuild_index_on_start:
await self.file_store.clear_all()
logger.info("Cleared all indexed data on start")
await self._scan_existing_files()
await self._watch_loop()
self._watch_task = asyncio.create_task(_initialize_and_watch())
logger.info(f"Started watching: {self.watch_paths}")
async def close(self):
"""Stop the file watcher"""
if not self._running:
return
self._stop_event.set()
if self._watch_task:
await self._watch_task
self._running = False
logger.info("Stopped watching")
def watch_filter(self, _change: Change, path: str) -> bool:
"""Filter function for file watching."""
# If no suffix filters are specified, watch all files
if not self.suffix_filters:
return True
# Check if the file has one of the allowed suffixes
for suffix in self.suffix_filters:
if path.endswith("." + suffix.strip(".")):
return True
return False
async def _scan_existing_files(self):
"""Scan existing files matching watch criteria and trigger on_changes with Change.added"""
existing_files: set[tuple[Change, str]] = set()
for watch_path_str in self.watch_paths:
watch_path = Path(watch_path_str)
if not watch_path.exists():
logger.warning(f"Watch path does not exist: {watch_path}")
continue
if watch_path.is_file():
# Single file
if self.watch_filter(Change.added, str(watch_path)):
existing_files.add((Change.added, str(watch_path)))
elif watch_path.is_dir():
# Directory
if self.recursive:
# Recursive scan
for file_path in watch_path.rglob("*"):
if file_path.is_file() and self.watch_filter(Change.added, str(file_path)):
existing_files.add((Change.added, str(file_path)))
else:
# Non-recursive scan (only immediate children)
for file_path in watch_path.iterdir():
if file_path.is_file() and self.watch_filter(Change.added, str(file_path)):
existing_files.add((Change.added, str(file_path)))
if existing_files:
logger.info(f"[SCAN_ON_START] Found {len(existing_files)} existing files matching watch criteria")
await self.on_changes(existing_files)
logger.info(f"[SCAN_ON_START] Added {len(existing_files)} files to memory store")
else:
logger.info("[SCAN_ON_START] No existing files found matching watch criteria")
if self.file_store is not None:
files: list[str] = await self.file_store.list_files(MemorySource.MEMORY)
for file_path in files:
chunks = await self.file_store.get_file_chunks(file_path, MemorySource.MEMORY)
logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks")
async def _interruptible_sleep(self, seconds: float):
"""Sleep that can be interrupted by stop_event."""
try:
await asyncio.wait_for(self._stop_event.wait(), timeout=seconds)
except asyncio.TimeoutError:
pass # Normal timeout, continue
async def _watch_loop(self):
"""Core monitoring loop with auto-restart on failure"""
if not self.watch_paths:
logger.warning("No watch paths specified")
return
while not self._stop_event.is_set():
# Filter out non-existent paths before each watch attempt
valid_paths = [p for p in self.watch_paths if Path(p).exists()]
if not valid_paths:
logger.warning("No valid watch paths exist, waiting 10 seconds before retry...")
await self._interruptible_sleep(10)
continue
invalid_paths = set(self.watch_paths) - set(valid_paths)
if invalid_paths:
logger.warning(f"Skipping non-existent paths: {invalid_paths}")
try:
logger.info(f"Starting watch on valid paths: {valid_paths}")
async for changes in awatch(
*valid_paths,
watch_filter=self.watch_filter,
recursive=self.recursive,
debounce=self.debounce,
poll_delay_ms=self.poll_delay_ms,
stop_event=self._stop_event,
):
if self._stop_event.is_set():
break
await self.on_changes(changes)
except FileNotFoundError as e:
# Watch path was deleted during monitoring
logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...")
if not self._stop_event.is_set():
await self._interruptible_sleep(10)
except Exception as e:
# Log other exceptions and restart
logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True)
if not self._stop_event.is_set():
await self._interruptible_sleep(10)
async def _on_changes(self, changes: set[tuple[Change, str]]):
"""Callback method to handle file changes"""
async def on_changes(self, changes: set[tuple[Change, str]]):
"""Hook method to handle file changes"""
if self.callback:
result = self.callback(changes)
if asyncio.iscoroutine(result):
await result
else:
await self._on_changes(changes)
logger.info(f"[{self.__class__.__name__}] on_changes: {changes}")
def is_running(self) -> bool:
"""Check if the watcher is running"""
return self._running
async def add_path(self, path: str):
"""Dynamically add a path to monitor"""
if path not in self.watch_paths:
self.watch_paths.append(path)
if self._running:
await self.close()
await self.start()
async def remove_path(self, path: str):
"""Remove a monitored path"""
if path in self.watch_paths:
self.watch_paths.remove(path)
if self._running:
await self.close()
await self.start()

View file

@ -0,0 +1,280 @@
"""Delta file watcher for incremental file synchronization.
This module provides a file watcher that detects append-only changes
and only processes newly added content, avoiding redundant operations.
"""
import asyncio
import os
from loguru import logger
from watchfiles import Change
from .base_file_watcher import BaseFileWatcher
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk
from ..utils import chunk_markdown, hash_text
class DeltaFileWatcher(BaseFileWatcher):
"""Delta file watcher implementation for incremental synchronization.
This watcher detects append-only changes (e.g., log files) and only processes
the newly added content, avoiding redundant embedding requests for unchanged content.
Strategy:
- Detect if file is append-only (new lines added at end)
- Find the safe cutoff point (considering chunk overlap)
- Only re-chunk and embed content from cutoff to end
- Delete affected old chunks and insert new chunks
"""
def __init__(self, overlap_lines: int = 2, **kwargs):
"""
Initialize delta file watcher.
Args:
chunk_tokens: Maximum tokens per chunk
chunk_overlap: Overlap tokens between chunks
"""
super().__init__(**kwargs)
self.overlap_lines = overlap_lines
self.dirty = False
@staticmethod
async def _build_file_metadata(path: str) -> FileMetadata:
"""Build file metadata from filesystem."""
def _read_file_sync():
stat_t = os.stat(path)
with open(path, "r", encoding="utf-8") as f:
content_t = f.read()
return stat_t, content_t
stat, content = await asyncio.to_thread(_read_file_sync)
return FileMetadata(
hash=hash_text(content),
mtime_ms=stat.st_mtime * 1000,
size=stat.st_size,
path=path,
content=content,
)
def _find_cutoff_line(
self,
old_chunks: list[MemoryChunk],
old_file_meta: FileMetadata,
new_file_meta: FileMetadata,
) -> int | None:
"""Find the safe cutoff line for incremental update.
Uses a heuristic approach: if file size increased and hash changed,
we verify by comparing content. For true append-only files (like logs),
the old content should be a prefix of new content.
Args:
old_chunks: Existing chunks sorted by start_line
old_file_meta: Previous file metadata
new_file_meta: Current file metadata (with content)
Returns:
Cutoff line number (1-indexed), or None if not append-only
"""
if not old_chunks:
return None
# File shrunk - definitely not append-only
if new_file_meta.size < old_file_meta.size:
logger.debug("File shrunk, not append-only")
return None
# File didn't grow much - might be a modification
size_growth = new_file_meta.size - old_file_meta.size
if size_growth < 10: # Less than 10 bytes growth
logger.debug("Minimal size growth, treating as modification")
return None
# Verify append-only by checking if old content is prefix
# We need to read old file content from chunks
old_chunks_sorted = sorted(old_chunks, key=lambda c: c.start_line)
# Simple heuristic: check if first few chunks' content matches
# This avoids reconstructing full old content
new_lines = new_file_meta.content.split("\n")
# Sample check: verify first chunk still matches
first_chunk = old_chunks_sorted[0]
first_chunk_lines = first_chunk.text.split("\n")
new_first_lines = new_lines[first_chunk.start_line - 1 : first_chunk.end_line]
# Compare (allowing for minor whitespace differences at boundaries)
if len(first_chunk_lines) > 0 and len(new_first_lines) > 0:
# Check if most of the lines match
matches = sum(1 for old, new in zip(first_chunk_lines, new_first_lines) if old == new)
if matches < len(first_chunk_lines) * 0.8: # Less than 80% match
logger.debug("First chunk content changed, not append-only")
return None
# File appears to be append-only
# Find the last chunk and set cutoff considering overlap
last_chunk = max(old_chunks_sorted, key=lambda c: c.end_line)
cutoff_line = max(1, last_chunk.end_line - self.overlap_lines)
logger.debug(
f"Append-only detected: size {old_file_meta.size} -> {new_file_meta.size}, "
f"cutoff at line {cutoff_line}",
)
return cutoff_line
@staticmethod
def _extract_content_from_line(content: str, start_line: int) -> str:
"""Extract content starting from a specific line number."""
lines = content.split("\n")
if start_line <= 1:
return content
if start_line > len(lines):
return ""
# start_line is 1-indexed, array is 0-indexed
return "\n".join(lines[start_line - 1 :])
async def _on_changes(self, changes: set[tuple[Change, str]]):
"""Handle file changes with incremental synchronization."""
self.dirty = True
for change_type, path in changes:
if change_type == Change.added:
# New file: process everything
file_meta = await self._build_file_metadata(path)
chunks = (
chunk_markdown(
file_meta.content,
file_meta.path,
MemorySource.MEMORY,
self.chunk_tokens,
self.chunk_overlap,
)
or []
)
if chunks:
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
logger.info(f"File added: {path} ({len(chunks)} chunks)")
else:
logger.warning(f"No chunks generated for new file {path}")
elif change_type == Change.modified:
# Get existing data
old_chunks = await self.file_store.get_file_chunks(path, MemorySource.MEMORY)
old_file_meta = await self.file_store.get_file_metadata(path, MemorySource.MEMORY)
# Read new file
file_meta = await self._build_file_metadata(path)
# If no old chunks, fallback to full update
if not old_chunks or not old_file_meta:
logger.debug(f"No existing chunks for {path}, doing full update")
chunks = (
chunk_markdown(
file_meta.content,
file_meta.path,
MemorySource.MEMORY,
self.chunk_tokens,
self.chunk_overlap,
)
or []
)
if chunks:
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.file_store.delete_file(path, MemorySource.MEMORY)
await self.file_store.upsert_file(
file_meta,
MemorySource.MEMORY,
chunks,
)
logger.info(f"File modified (full): {path} ({len(chunks)} chunks)")
continue
# Check if append-only and find cutoff line
old_chunks_sorted = sorted(old_chunks, key=lambda c: c.start_line)
cutoff_line = self._find_cutoff_line(old_chunks_sorted, old_file_meta, file_meta)
if cutoff_line is None:
# Not append-only, do full update
logger.debug(f"File {path} has modifications, doing full update")
chunks = (
chunk_markdown(
file_meta.content,
file_meta.path,
MemorySource.MEMORY,
self.chunk_tokens,
self.chunk_overlap,
)
or []
)
if chunks:
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.file_store.delete_file(path, MemorySource.MEMORY)
await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
logger.info(f"File modified (full): {path} ({len(chunks)} chunks)")
else:
# Append-only: incremental update
new_content_part = self._extract_content_from_line(file_meta.content, cutoff_line)
new_chunks = (
chunk_markdown(
new_content_part,
file_meta.path,
MemorySource.MEMORY,
self.chunk_tokens,
self.chunk_overlap,
)
or []
)
if not new_chunks:
logger.debug(f"No new chunks for {path}, skipping")
continue
for idx, chunk in enumerate(new_chunks):
chunk.start_line += cutoff_line - 1
chunk.end_line += cutoff_line - 1
chunk.id = hash_text(
f"{chunk.source}:{chunk.path}:{chunk.start_line}:" f"{chunk.end_line}:{chunk.hash}:{idx}",
)
new_chunks = await self.file_store.get_chunk_embeddings(new_chunks)
chunks_to_delete = [c.id for c in old_chunks_sorted if c.start_line >= cutoff_line]
# Apply incremental updates
if chunks_to_delete:
await self.file_store.delete_file_chunks(path, chunks_to_delete)
if new_chunks:
await self.file_store.upsert_chunks(new_chunks, MemorySource.MEMORY)
# Update file metadata to reflect the changes
# Calculate new chunk count: old chunks - deleted + new chunks
new_chunk_count = len(old_chunks) - len(chunks_to_delete) + len(new_chunks)
file_meta.chunk_count = new_chunk_count
await self.file_store.update_file_metadata(file_meta, MemorySource.MEMORY)
logger.info(
f"File modified (incremental): {path} "
f"(cutoff: line {cutoff_line}, "
f"+{len(new_chunks)} chunks, -{len(chunks_to_delete)} chunks)",
)
elif change_type == Change.deleted:
await self.file_store.delete_file(path, MemorySource.MEMORY)
logger.info(f"File deleted: {path}")
else:
logger.warning(f"Unknown change type: {change_type}")
self.dirty = False

View file

@ -0,0 +1,79 @@
"""Full file watcher for complete file synchronization.
This module provides a file watcher that processes entire files
on any change, ensuring complete synchronization.
"""
import asyncio
from pathlib import Path
from loguru import logger
from watchfiles import Change
from .base_file_watcher import BaseFileWatcher
from ..enumeration import MemorySource
from ..schema import FileMetadata
from ..utils import chunk_markdown, hash_text
class FullFileWatcher(BaseFileWatcher):
"""Full file watcher implementation for full synchronization"""
def __init__(self, **kwargs):
"""
Initialize full file watcher"""
super().__init__(**kwargs)
self.dirty = False
@staticmethod
async def _build_file_metadata(path: str) -> FileMetadata:
file_path = Path(path)
def _read_file_sync():
return file_path.stat(), file_path.read_text(encoding="utf-8")
stat, content = await asyncio.to_thread(_read_file_sync)
return FileMetadata(
hash=hash_text(content),
mtime_ms=stat.st_mtime * 1000,
size=stat.st_size,
path=str(file_path.absolute()),
content=content,
)
async def _on_changes(self, changes: set[tuple[Change, str]]):
"""Handle file changes with full synchronization"""
self.dirty = True
for change_type, path in changes:
if change_type in [Change.added, Change.modified]:
file_meta = await self._build_file_metadata(path)
chunks = (
chunk_markdown(
file_meta.content,
file_meta.path,
MemorySource.MEMORY,
self.chunk_tokens,
self.chunk_overlap,
)
or []
)
if chunks:
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.file_store.delete_file(file_meta.path, MemorySource.MEMORY)
logger.info(f"delete_file {file_meta.path}")
await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
logger.info(f"Upserted {file_meta.chunk_count} chunks for {file_meta.path}")
elif change_type == Change.deleted:
await self.file_store.delete_file(path, MemorySource.MEMORY)
logger.info(f"Deleted {path}")
else:
logger.warning(f"Unknown change type: {change_type}")
logger.info(f"File {change_type} changed: {path}")
self.dirty = False

View file

@ -0,0 +1,146 @@
"""Module for managing and formatting prompt templates from files or dictionaries."""
import json
from pathlib import Path
from string import Formatter
from typing import Any, Dict, Optional, Union
import yaml
from loguru import logger
from .base_dict import BaseDict
class PromptHandler(BaseDict):
"""A context-aware handler for loading, retrieving, and formatting prompt templates."""
def __init__(self, language: str = "", **kwargs):
super().__init__(**kwargs)
# Use object.__setattr__ to avoid storing 'language' in the dict
object.__setattr__(self, "language", language.strip())
def load_prompt_by_file(
self,
prompt_file_path: Optional[Union[Path, str]] = None,
overwrite: bool = True,
) -> "PromptHandler":
"""Load prompt configurations from a YAML or JSON file."""
if prompt_file_path is None:
return self
if isinstance(prompt_file_path, str):
prompt_file_path = Path(prompt_file_path)
if not prompt_file_path.exists():
return self
suffix = prompt_file_path.suffix.lower()
with prompt_file_path.open(encoding="utf-8") as f:
if suffix in [".yaml", ".yml"]:
prompt_dict = yaml.safe_load(f)
elif suffix == ".json":
prompt_dict = json.load(f)
else:
raise ValueError(f"Unsupported file format: {suffix}")
self.load_prompt_dict(prompt_dict, overwrite=overwrite)
return self
def load_prompt_dict(
self,
prompt_dict: Optional[Dict[str, Any]] = None,
overwrite: bool = True,
) -> "PromptHandler":
"""Merge a dictionary of prompt strings into the current context."""
if not prompt_dict:
return self
for key, value in prompt_dict.items():
if not isinstance(value, str):
continue
if key in self:
if overwrite:
logger.warning(f"Overwriting prompt '{key}'")
self[key] = value
else:
self[key] = value
return self
def get_prompt(self, prompt_name: str, fallback_to_base: bool = True) -> str:
"""Retrieve a prompt by name with automatic language suffix handling."""
if self.language and not prompt_name.endswith(f"_{self.language}"):
key_with_lang = f"{prompt_name}_{self.language}"
if key_with_lang in self:
return self[key_with_lang].strip()
if prompt_name in self:
return self[prompt_name].strip()
if fallback_to_base and self.language and prompt_name.endswith(f"_{self.language}"):
base_name = prompt_name[: -(len(self.language) + 1)]
if base_name in self:
return self[base_name].strip()
raise KeyError(f"Prompt '{prompt_name}' not found. Available: {list(self.keys())[:10]}")
def has_prompt(self, prompt_name: str) -> bool:
"""Check if a prompt exists."""
try:
self.get_prompt(prompt_name)
return True
except KeyError:
return False
def list_prompts(self, language_filter: Optional[str] = None) -> list[str]:
"""List all available prompt names."""
if language_filter is None:
return list(self.keys())
suffix = f"_{language_filter.strip()}"
return [key for key in self.keys() if key.endswith(suffix)]
@staticmethod
def _extract_format_fields(template: str) -> set[str]:
"""Extract all format field names from a template string."""
return {field_name for _, field_name, _, _ in Formatter().parse(template) if field_name is not None}
@staticmethod
def _filter_conditional_lines(prompt: str, flags: Dict[str, bool]) -> str:
"""Filter lines based on boolean flags."""
filtered_lines = []
for line in prompt.split("\n"):
matched_flag = None
for flag_name in flags:
if line.startswith(f"[{flag_name}]"):
matched_flag = flag_name
break
if matched_flag is None:
filtered_lines.append(line)
elif flags[matched_flag]:
filtered_lines.append(line[len(f"[{matched_flag}]") :])
return "\n".join(filtered_lines)
def prompt_format(self, prompt_name: str, validate: bool = True, **kwargs) -> str:
"""Format a prompt with conditional line filtering and variable substitution."""
prompt = self.get_prompt(prompt_name)
flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)}
format_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)}
if flag_kwargs:
prompt = self._filter_conditional_lines(prompt, flag_kwargs)
if validate:
required_fields = self._extract_format_fields(prompt)
missing_fields = required_fields - set(format_kwargs.keys())
if missing_fields:
raise ValueError(f"Missing format variables for '{prompt_name}': {sorted(missing_fields)}")
if format_kwargs:
prompt = prompt.format(**format_kwargs)
return prompt.strip()
def __repr__(self) -> str:
return f"PromptHandler(language='{self.language}', num_prompts={len(self)})"

View file

@ -0,0 +1,50 @@
"""Module providing a registry class for managing class-to-name mappings via decorators."""
import inspect
from typing import Callable, TypeVar
from .base_dict import BaseDict
from .utils import singleton
T = TypeVar("T")
class Registry(BaseDict):
"""A registry container that uses decorators to map and store class references."""
def register(self, name: str | type = "") -> Callable[[type[T]], type[T]] | type[T]:
"""Return a decorator that registers a class under a specific name in the registry."""
if inspect.isclass(name):
self[name.__name__] = name
return name
else:
def decorator(cls):
key: str = name if isinstance(name, str) and name else cls.__name__
self[key] = cls
return cls
return decorator
@singleton
class RegistryFactory:
"""A factory class for creating registries."""
def __init__(self):
self.llms = Registry()
self.as_llms = Registry()
self.as_llm_formatters = Registry()
self.as_token_counters = Registry()
self.embedding_models = Registry()
self.vector_stores = Registry()
self.file_stores = Registry()
self.ops = Registry()
self.flows = Registry()
self.services = Registry()
self.token_counters = Registry()
self.file_watchers = Registry()
R = RegistryFactory()

View file

0
reme_cli/op/__init__.py Normal file
View file

427
reme_cli/op/base_op.py Normal file
View file

@ -0,0 +1,427 @@
"""Base operator class for LLM workflow execution and composition."""
import asyncio
import copy
import inspect
from abc import ABCMeta
from pathlib import Path
from typing import Callable, Optional, Any
from agentscope.formatter import FormatterBase
from agentscope.model import ChatModelBase
from agentscope.token import HuggingFaceTokenCounter
from loguru import logger
from tqdm import tqdm
from ..embedding import BaseEmbeddingModel
from ..file_store import BaseFileStore
from ..llm import BaseLLM
from ..prompt_handler import PromptHandler
from ..runtime_context import RuntimeContext
from ..schema import Response, ServiceConfig
from ..schema.service_config import OpConfig
from ..service_context import ServiceContext
from ..token_counter import BaseTokenCounter
from ..utils import camel_to_snake, CacheHandler, timer
from ..vector_store import BaseVectorStore
class BaseOp(metaclass=ABCMeta):
"""Base operator class for LLM workflow execution and composition."""
__alias_name__: str = ""
def __new__(cls, *args, **kwargs):
"""Capture initialization arguments for object cloning."""
instance = super().__new__(cls)
instance._init_args = copy.copy(args)
instance._init_kwargs = copy.copy(kwargs)
return instance
def __init__(
self,
name: str = "",
async_mode: bool = True,
language: str = "",
prompt_name: str = "",
prompt_path: str = "",
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | HuggingFaceTokenCounter = "default",
llm: str | BaseLLM = "default",
embedding_model: str | BaseEmbeddingModel = "default",
vector_store: str | BaseVectorStore = "default",
file_store: str | BaseFileStore = "default",
token_counter: str | BaseTokenCounter = "default",
enable_cache: bool = False,
cache_path: str = "cache/op",
cache_expire_hours: float | None = None,
sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"] = None,
input_mapping: dict[str, str] | None = None,
output_mapping: dict[str, str] | None = None,
enable_parallel: bool = False,
max_retries: int = 1,
raise_exception: bool = False,
**kwargs,
):
"""Initialize operator configurations and internal state."""
self.name = name or self.__alias_name__ or camel_to_snake(self.__class__.__name__)
self.async_mode = async_mode
self.language = language
self.prompt = self._get_prompt_handler(prompt_name, prompt_path)
self._as_llm = as_llm
self._as_llm_formatter = as_llm_formatter
self._as_token_counter = as_token_counter
self._llm = llm
self._embedding_model = embedding_model
self._vector_store = vector_store
self._file_store = file_store
self._token_counter = token_counter
self.enable_cache = enable_cache
self.cache_path = cache_path
self.cache_expire_hours = cache_expire_hours
self.sub_ops: list["BaseOp"] = []
self.add_sub_ops(sub_ops)
self.input_mapping = input_mapping
self.output_mapping = output_mapping
self.enable_parallel = enable_parallel # Control whether to execute tasks in parallel
self.max_retries = max(1, max_retries)
self.raise_exception = raise_exception
self.op_params = kwargs
self._pending_tasks: list = []
self.context: RuntimeContext | None = None
self._cache: CacheHandler | None = None
def _get_prompt_handler(self, prompt_name: str, prompt_path: str) -> PromptHandler:
"""Load prompt configuration from the associated YAML file."""
if prompt_path:
path = Path(prompt_path)
else:
path = Path(inspect.getfile(self.__class__))
if prompt_name:
path = path.with_stem(prompt_name)
return PromptHandler(language=self.language).load_prompt_by_file(path.with_suffix(".yaml"))
def _handle_failure(self, e: Exception, attempt: int) -> str | None:
"""Log failures and handle final retry logic."""
message = f"[{self.__class__.__name__}] failed (attempt {attempt + 1}): {e}"
if attempt == self.max_retries - 1:
logger.exception(message)
if self.raise_exception:
raise e
return f"[{self.__class__.__name__}] failed: {e}"
else:
logger.warning(message)
return None
@property
def cache(self) -> CacheHandler:
"""Access the operator-specific cache handler."""
assert self.enable_cache, "Cache is disabled!"
if not self._cache:
self._cache = CacheHandler(f"{self.cache_path}/{self.name}")
return self._cache
@property
def service_context(self) -> ServiceContext:
"""Access the service context."""
assert self.context, "Service context is not initialized!"
return self.context.service_context
@property
def service_config(self) -> ServiceConfig:
"""Access the service configuration."""
return self.service_context.service_config
@property
def as_llm(self) -> ChatModelBase:
"""Get the AgentScope LLM instance from ServiceContext."""
if isinstance(self._as_llm, str):
self._as_llm = self.service_context.as_llms[self._as_llm]
return self._as_llm
@property
def as_llm_formatter(self) -> FormatterBase:
"""Get the AgentScope LLM formatter instance from ServiceContext."""
if isinstance(self._as_llm_formatter, str):
self._as_llm_formatter = self.service_context.as_llm_formatters[self._as_llm_formatter]
return self._as_llm_formatter
@property
def as_token_counter(self) -> HuggingFaceTokenCounter:
"""Get the token counter instance from ServiceContext."""
if isinstance(self._as_token_counter, str):
self._as_token_counter = self.service_context.as_token_counters[self._as_token_counter]
return self._as_token_counter
@property
def llm(self) -> BaseLLM:
"""Get the LLM instance from ServiceContext."""
if isinstance(self._llm, str):
self._llm = self.service_context.llms[self._llm]
return self._llm
@property
def embedding_model(self) -> BaseEmbeddingModel:
"""Get the embedding model instance from ServiceContext."""
if isinstance(self._embedding_model, str):
self._embedding_model = self.service_context.embedding_models[self._embedding_model]
return self._embedding_model
@property
def vector_store(self) -> BaseVectorStore:
"""Lazily initialize and return the vector store instance."""
if isinstance(self._vector_store, str):
self._vector_store = self.service_context.vector_stores[self._vector_store]
return self._vector_store
@property
def file_store(self) -> BaseFileStore:
"""Lazily initialize and return the file store instance."""
if isinstance(self._file_store, str):
self._file_store = self.service_context.file_stores[self._file_store]
return self._file_store
@property
def token_counter(self) -> BaseTokenCounter:
"""Get the token counter instance from ServiceContext."""
if isinstance(self._token_counter, str):
self._token_counter = self.service_context.token_counters[self._token_counter]
return self._token_counter
@property
def service_metadata(self) -> dict:
"""Get service configuration metadata."""
return self.service_context.service_config.metadata
@property
def response(self) -> Response:
"""Access the response object."""
return self.context.response
def before_execute_sync(self):
"""Prepare context and validate before sync execution.
This method performs the following steps:
1. Apply input mapping to transform context variables
2. Load operator-specific configuration from service config if available
3. Override operator parameters and prompts based on config
"""
self.context.apply_mapping(self.input_mapping)
if self.context.service_context is None:
return
service_config = self.service_context.service_config
if self.name not in service_config.ops:
return
op_config: OpConfig = service_config.ops[self.name]
# Override operator parameters from config
if op_config.params:
for k, v in op_config.params.items():
if hasattr(self, k):
setattr(self, k, v)
logger.info(f"[{self.__class__.__name__}] Set attribute '{k}' = {v}")
else:
self.op_params[k] = v
logger.info(f"[{self.__class__.__name__}] Set op_param '{k}' = {v}")
# Load custom prompt templates from config
if op_config.prompt_dict:
self.prompt.load_prompt_dict(op_config.prompt_dict)
logger.info(f"[{self.__class__.__name__}] Loaded prompt keys={list(op_config.prompt_dict.keys())}")
async def before_execute(self):
"""Prepare context and validate before async execution."""
self.before_execute_sync()
def execute_sync(self):
"""Define core sync logic in subclasses."""
async def execute(self):
"""Define core async logic in subclasses."""
def after_execute_sync(self, response: Any):
"""Finalize context and mappings after sync execution."""
self.context.apply_mapping(self.output_mapping)
if response is not None:
if isinstance(response, dict):
for k, v in response.items():
if k == "answer":
self.response.answer = v
elif k == "success":
self.response.success = v if isinstance(v, bool) else v.lower() == "true"
else:
self.response.metadata[k] = v
else:
self.response.answer = response
return response
async def after_execute(self, output: Any):
"""Finalize context and mappings after async execution."""
return self.after_execute_sync(output)
@timer
def call_sync(self, context: RuntimeContext = None, **kwargs):
"""Execute the operator synchronously with retry logic."""
self.context = RuntimeContext.from_context(context, **kwargs)
response = None
for i in range(self.max_retries):
try:
self.before_execute_sync()
response = self.execute_sync()
response = self.after_execute_sync(response)
break
except Exception as e:
response = self._handle_failure(e, i)
return response
@timer
async def call(self, context: RuntimeContext = None, **kwargs):
"""Execute the operator asynchronously with retry logic."""
self.context = RuntimeContext.from_context(context, **kwargs)
response = None
for i in range(self.max_retries):
try:
await self.before_execute()
response = await self.execute()
response = await self.after_execute(response)
break
except Exception as e:
response = self._handle_failure(e, i)
return response
def submit_sync_task(self, fn: Callable, *args, **kwargs) -> "BaseOp":
"""Submit a task to the thread pool or local queue."""
if self.enable_parallel and self.service_context.thread_pool is not None:
task = self.service_context.thread_pool.submit(fn, *args, **kwargs)
else:
task = (fn, args, kwargs)
self._pending_tasks.append(task)
return self
def submit_async_task(self, coro_fn: Callable, *args, **kwargs) -> "BaseOp":
"""Submit an async task to the pending tasks queue."""
task = coro_fn(*args, **kwargs)
self._pending_tasks.append(task)
return self
def join_sync_tasks(self, task_desc: str = None) -> list:
"""Wait for all pending sync tasks and return flattened results."""
results = []
for task in tqdm(self._pending_tasks, desc=task_desc or self.name):
if self.enable_parallel:
result = task.result()
else:
result = task[0](*task[1], **task[2])
if result:
if isinstance(result, list):
results.extend(result)
else:
results.append(result)
self._pending_tasks.clear()
return results
async def join_async_tasks(self, return_exceptions: bool = True) -> list:
"""Wait for all pending async tasks and aggregate results."""
if self.enable_parallel:
raw_results = await asyncio.gather(*self._pending_tasks, return_exceptions=return_exceptions)
else:
raw_results = []
for task in self._pending_tasks:
try:
result = await task
raw_results.append(result)
except Exception as e:
if return_exceptions:
raw_results.append(e)
else:
raise
results = []
for result in raw_results:
if isinstance(result, Exception):
logger.error(f"[{self.__class__.__name__}] Async task failed: {result}")
elif result:
if isinstance(result, list):
results.extend(result)
else:
results.append(result)
self._pending_tasks.clear()
return results
def add_sub_ops(self, sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"]):
"""Add child operators to this operator's sub_ops."""
if not sub_ops:
return
if isinstance(sub_ops, dict):
for name, op in sub_ops.items():
assert self.async_mode == op.async_mode, "Async mode mismatch!"
op.name = name
if self.language:
op.language = self.language
self.sub_ops.append(op)
elif isinstance(sub_ops, list):
for op in sub_ops:
assert self.async_mode == op.async_mode, "Async mode mismatch!"
if self.language:
op.language = self.language
self.sub_ops.append(op)
else:
assert self.async_mode == sub_ops.async_mode, "Async mode mismatch!"
if self.language:
sub_ops.language = self.language
self.sub_ops.append(sub_ops)
def add_sub_op(self, sub_op: "BaseOp"):
"""Add a single child operator to this operator's sub_ops."""
self.sub_ops.append(sub_op)
def __lshift__(self, ops):
"""Operator overload for adding sub-operators."""
self.add_sub_ops(ops)
return self
def __rshift__(self, op: "BaseOp"):
"""Operator overload for sequential execution composition."""
from .sequential_op import SequentialOp
seq = SequentialOp(sub_ops=[self], async_mode=self.async_mode)
seq.add_sub_ops(op.sub_ops if isinstance(op, SequentialOp) else op)
return seq
def __or__(self, op: "BaseOp"):
"""Operator overload for parallel execution composition."""
from .parallel_op import ParallelOp
par = ParallelOp(sub_ops=[self], async_mode=self.async_mode)
par.add_sub_ops(op.sub_ops if isinstance(op, ParallelOp) else op)
return par
def prompt_format(self, prompt_name: str, **kwargs) -> str:
"""Format a prompt template with provided keyword arguments."""
return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs)
def get_prompt(self, prompt_name: str) -> str:
"""Get a prompt template by name."""
return self.prompt.get_prompt(prompt_name=prompt_name)
def copy(self, **kwargs):
"""Create a copy of this operator with optional parameter overrides."""
copy_op = self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs})
if self.sub_ops:
copy_op.sub_ops.clear()
for op in self.sub_ops:
copy_op.add_sub_op(op.copy())
return copy_op

78
reme_cli/reme_cli.py Normal file
View file

@ -0,0 +1,78 @@
from agentscope.message import Msg
from reme_cli import Application
class ReMeCli(Application):
async def init(self, **kwargs) -> None:
"""Initialize the application."""
...
async def read(self, file: str, **kwargs) -> None:
"""Read a note file."""
# obsidian read file="My Note"
...
async def create(self, name: str, content: str, template: str, silent: bool, **kwargs) -> None:
"""Create a new note."""
# obsidian create name="New Note" content="# Hello" template="Template" silent
...
async def append(self, file: str, content: str, **kwargs) -> None:
"""Append content to a note."""
# obsidian append file="My Note" content="New line"
...
async def search(self, query: str, limit: int, **kwargs) -> None:
"""Search for notes."""
# obsidian search query="search term" limit=10
...
async def daily_read(self, **kwargs) -> None:
"""Read daily note."""
# obsidian daily:read
...
async def daily_append(self, content: str, **kwargs) -> None:
"""Append content to daily note."""
# obsidian daily:append content="- [ ] New task"
...
async def property_set(self, name: str, value: str, file: str, **kwargs) -> None:
"""Set a property on a note."""
# obsidian property:set name="status" value="done" file="My Note"
...
async def tasks(self, daily: bool, todo: bool, **kwargs) -> None:
"""Manage tasks."""
# obsidian tasks daily todo
...
async def tags(self, sort: str, counts: bool, **kwargs) -> None:
"""Manage tags."""
# obsidian tags sort=count counts
...
async def backlinks(self, file: str, **kwargs) -> None:
"""Get backlinks for a note."""
# obsidian backlinks file="My Note"
...
async def summary(self, messages: list[Msg], **kwargs):
...
async def dream(self) -> dict:
...
async def proactive(self, messages: list[Msg], **kwargs) -> dict:
...
def main():
"""Main entry point for running ReMe from command line."""
ReMeCli(*sys.argv[1:], config_path="service").run_service()
if __name__ == "__main__":
main()

View file

View file

@ -0,0 +1,144 @@
"""Configuration schemas for service components using Pydantic models."""
import os
from pydantic import BaseModel, Field, ConfigDict
from .tool_call import ToolCall
class MCPConfig(BaseModel):
"""Configuration for Model Context Protocol transport and network settings."""
model_config = ConfigDict(extra="allow")
transport: str = Field(default="stdio")
host: str = Field(default="0.0.0.0")
port: int = Field(default=8001)
class HttpConfig(BaseModel):
"""Configuration for the HTTP server interface and connection lifecycle."""
model_config = ConfigDict(extra="allow")
host: str = Field(default="0.0.0.0")
port: int = Field(default=8001)
timeout_keep_alive: int = Field(default=3600)
limit_concurrency: int = Field(default=1000)
class CmdConfig(BaseModel):
"""Configuration for command-line flow execution parameters."""
model_config = ConfigDict(extra="allow")
flow: str = Field(default="")
class OpConfig(BaseModel):
"""Configuration for op settings and parameters."""
model_config = ConfigDict(extra="allow")
prompt_dict: dict[str, str] = Field(default_factory=dict)
params: dict = Field(default_factory=dict)
class FlowConfig(ToolCall):
"""Configuration for workflow execution, caching, and error handling."""
model_config = ConfigDict(extra="allow")
flow_content: str = Field(default="")
stream: bool = Field(default=False)
raise_exception: bool = Field(default=True)
enable_cache: bool = Field(default=False)
cache_path: str = Field(default="cache/flow")
cache_expire_hours: float = Field(default=0.1)
class BasicConfig(BaseModel):
"""Configuration for basic service settings and parameters."""
model_config = ConfigDict(extra="allow")
backend: str = Field(default="")
class ModelConfig(BasicConfig):
"""Configuration for model-based services with backend and model name."""
model_name: str = Field(default="")
class LLMConfig(ModelConfig):
"""Configuration for Large Language Model backend and model identification."""
class EmbeddingModelConfig(ModelConfig):
"""Configuration for embedding model backends and identity."""
class TokenCounterConfig(ModelConfig):
"""Configuration for token counting services and model mapping."""
class StoreConfig(BasicConfig):
"""Configuration for storage services with embedding model support."""
embedding_model: str = Field(default="default")
class VectorStoreConfig(StoreConfig):
"""Configuration for vector database storage and associated embeddings."""
collection_name: str = Field(default="reme")
class FileStoreConfig(StoreConfig):
"""Configuration for file store database storage and associated embeddings."""
store_name: str = Field(default="reme")
class FileWatcherConfig(BasicConfig):
"""Configuration for file watcher service."""
file_store: str = Field(default="")
watch_paths: list[str] = Field(default_factory=list)
class ServiceConfig(BasicConfig):
"""Root configuration schema aggregating all service-level settings and components."""
app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"))
working_dir: str = Field(default=".reme")
enable_logo: bool = Field(default=True)
language: str = Field(default="")
thread_pool_max_workers: int = Field(
default=16,
description="Number of thread pool workers. Set to -1 to disable thread pool.",
)
ray_max_workers: int = Field(default=-1)
log_to_console: bool = Field(default=True)
disabled_flows: list[str] = Field(default_factory=list)
enabled_flows: list[str] = Field(default_factory=list)
mcp_servers: dict[str, dict] = Field(default_factory=dict)
mcp: MCPConfig = Field(default_factory=MCPConfig)
http: HttpConfig = Field(default_factory=HttpConfig)
cmd: CmdConfig = Field(default_factory=CmdConfig)
ops: dict[str, OpConfig] = Field(default_factory=dict)
flows: dict[str, FlowConfig] = Field(default_factory=dict)
as_llms: dict[str, BasicConfig] = Field(default_factory=dict)
as_llm_formatters: dict[str, BasicConfig] = Field(default_factory=dict)
as_token_counters: dict[str, BasicConfig] = Field(default_factory=dict)
llms: dict[str, LLMConfig] = Field(default_factory=dict)
embedding_models: dict[str, EmbeddingModelConfig] = Field(default_factory=dict)
vector_stores: dict[str, VectorStoreConfig] = Field(default_factory=dict)
file_stores: dict[str, FileStoreConfig] = Field(default_factory=dict)
token_counters: dict[str, TokenCounterConfig] = Field(default_factory=dict)
file_watchers: dict[str, FileWatcherConfig] = Field(default_factory=dict)
metadata: dict = Field(default_factory=dict)

View file

View file

@ -0,0 +1,59 @@
"""Logging configuration module for application-wide tracing."""
import os
import sys
from datetime import datetime
def init_logger(log_dir: str = "logs", level: str = "INFO", log_to_console: bool = True) -> None:
"""Initialize the logger with both file and console handlers.
Args:
log_dir: Directory path for log files
level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
log_to_console: Whether to print logs to console/screen
"""
from loguru import logger
# Remove default handler to avoid duplicate logs
logger.remove()
# Configure colorized standard output logging if enabled
if log_to_console:
logger.add(
sink=sys.stdout,
level=level,
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
colorize=True,
)
# Try to configure file-based logging (skip if permission denied)
try:
# Ensure the logging directory exists
os.makedirs(log_dir, exist_ok=True)
# Generate filename based on the current timestamp
# Use dashes instead of colons for Windows compatibility
current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
log_filename = f"{current_ts}.log"
log_filepath = os.path.join(log_dir, log_filename)
# Configure file-based logging with rotation and compression
logger.add(
log_filepath,
level=level,
rotation="00:00",
retention="7 days",
compression="zip",
encoding="utf-8",
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
)
except Exception as e:
logger.error(f"Error configuring file logging: {e}")
def get_logger():
"""Get a configured logger instance using loguru."""
from loguru import logger
return logger

View file

@ -0,0 +1,21 @@
"""Module providing a decorator to implement the Singleton design pattern."""
import threading
def singleton(cls):
"""A class decorator that ensures only one instance of a class exists."""
# Dictionary to cache the single instance of the class
_instance = {}
_lock = threading.Lock()
def _singleton(*args, **kwargs):
"""Return the existing instance or create a new one if it doesn't exist."""
with _lock:
if cls not in _instance:
# Create and store the instance if it's the first call
_instance[cls] = cls(*args, **kwargs)
return _instance[cls]
return _singleton