mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-17 23:51:19 +00:00
init
This commit is contained in:
parent
935e886af3
commit
16d139583f
28 changed files with 4195 additions and 0 deletions
6
reme_cli/__init__.py
Normal file
6
reme_cli/__init__.py
Normal 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
20
reme_cli/application.py
Normal 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
|
||||
|
||||
5
reme_cli/component/__init__.py
Normal file
5
reme_cli/component/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
from .base_component import BaseComponent
|
||||
|
||||
__all__ = [
|
||||
"BaseComponent",
|
||||
]
|
||||
9
reme_cli/component/as_llm/__init__.py
Normal file
9
reme_cli/component/as_llm/__init__.py
Normal 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)
|
||||
9
reme_cli/component/as_llm_formatter/__init__.py
Normal file
9
reme_cli/component/as_llm_formatter/__init__.py
Normal 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)
|
||||
|
|
@ -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
|
||||
25
reme_cli/component/base_component.py
Normal file
25
reme_cli/component/base_component.py
Normal 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()
|
||||
41
reme_cli/component/base_dict.py
Normal file
41
reme_cli/component/base_dict.py
Normal 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__()
|
||||
23
reme_cli/component/file_store/__init__.py
Normal file
23
reme_cli/component/file_store/__init__.py
Normal 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)
|
||||
227
reme_cli/component/file_store/base_file_store.py
Normal file
227
reme_cli/component/file_store/base_file_store.py
Normal 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."""
|
||||
633
reme_cli/component/file_store/chroma_file_store.py
Normal file
633
reme_cli/component/file_store/chroma_file_store.py
Normal 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()
|
||||
461
reme_cli/component/file_store/local_file_store.py
Normal file
461
reme_cli/component/file_store/local_file_store.py
Normal 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}'")
|
||||
978
reme_cli/component/file_store/sqlite_file_store.py
Normal file
978
reme_cli/component/file_store/sqlite_file_store.py
Normal 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()
|
||||
19
reme_cli/component/file_watcher/__init__.py
Normal file
19
reme_cli/component/file_watcher/__init__.py
Normal 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)
|
||||
240
reme_cli/component/file_watcher/base_file_watcher.py
Normal file
240
reme_cli/component/file_watcher/base_file_watcher.py
Normal 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()
|
||||
280
reme_cli/component/file_watcher/delta_file_watcher.py
Normal file
280
reme_cli/component/file_watcher/delta_file_watcher.py
Normal 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
|
||||
79
reme_cli/component/file_watcher/full_file_watcher.py
Normal file
79
reme_cli/component/file_watcher/full_file_watcher.py
Normal 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
|
||||
146
reme_cli/component/prompt_handler.py
Normal file
146
reme_cli/component/prompt_handler.py
Normal 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)})"
|
||||
50
reme_cli/component/registry_factory.py
Normal file
50
reme_cli/component/registry_factory.py
Normal 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()
|
||||
0
reme_cli/config/__init__.py
Normal file
0
reme_cli/config/__init__.py
Normal file
0
reme_cli/op/__init__.py
Normal file
0
reme_cli/op/__init__.py
Normal file
427
reme_cli/op/base_op.py
Normal file
427
reme_cli/op/base_op.py
Normal 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
78
reme_cli/reme_cli.py
Normal 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()
|
||||
0
reme_cli/schema/__init__.py
Normal file
0
reme_cli/schema/__init__.py
Normal file
144
reme_cli/schema/service_config.py
Normal file
144
reme_cli/schema/service_config.py
Normal 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)
|
||||
0
reme_cli/utils/__init__.py
Normal file
0
reme_cli/utils/__init__.py
Normal file
59
reme_cli/utils/logger_utils.py
Normal file
59
reme_cli/utils/logger_utils.py
Normal 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
|
||||
21
reme_cli/utils/singleton.py
Normal file
21
reme_cli/utils/singleton.py
Normal 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
|
||||
Loading…
Add table
Reference in a new issue