mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
refactor(component): rename file_store to chunk_store and update interfaces - Rename BaseFileStore to BaseChunkStore and update component type - Replace file_store property with chunk_store in BaseStep - Add file_graph property to access file metadata from FileWatcher - Update all storage backends (Chroma, Local, SQLite) to use chunk-focused APIs - Remove file metadata handling from chunk stores (moved to FileGraph) - Update search methods to use ChunkFilter instead of SearchFilter - Remove file_store imports and add chunk_store imports ```
This commit is contained in:
parent
98632dfc24
commit
c2e43c398b
16 changed files with 308 additions and 454 deletions
|
|
@ -2,10 +2,10 @@
|
|||
|
||||
from . import as_llm
|
||||
from . import as_llm_formatter
|
||||
from . import chunk_store
|
||||
from . import client
|
||||
from . import embedding
|
||||
from . import file_parser
|
||||
from . import file_store
|
||||
from . import file_watcher
|
||||
from . import job
|
||||
from . import service
|
||||
|
|
@ -27,10 +27,10 @@ __all__ = [
|
|||
# base components
|
||||
"as_llm",
|
||||
"as_llm_formatter",
|
||||
"chunk_store",
|
||||
"client",
|
||||
"embedding",
|
||||
"file_parser",
|
||||
"file_store",
|
||||
"file_watcher",
|
||||
"job",
|
||||
"service",
|
||||
|
|
|
|||
|
|
@ -8,11 +8,12 @@ from agentscope.model import ChatModelBase
|
|||
from agentscope.token import TokenCounterBase
|
||||
|
||||
from .base_component import BaseComponent
|
||||
from .chunk_store import BaseChunkStore
|
||||
from .embedding import BaseEmbeddingModel
|
||||
from .file_store import BaseFileStore
|
||||
from .prompt_handler import PromptHandler
|
||||
from .runtime_context import RuntimeContext
|
||||
from ..enumeration import ComponentEnum
|
||||
from ..schema.file_graph import FileGraph
|
||||
|
||||
|
||||
class BaseStep(BaseComponent):
|
||||
|
|
@ -83,9 +84,15 @@ class BaseStep(BaseComponent):
|
|||
"token_counter")
|
||||
|
||||
@property
|
||||
def file_store(self) -> BaseFileStore:
|
||||
name = self.kwargs.get("file_store", "default")
|
||||
return name if isinstance(name, BaseFileStore) else self._get_component(ComponentEnum.FILE_STORE, name)
|
||||
def chunk_store(self) -> BaseChunkStore:
|
||||
name = self.kwargs.get("chunk_store", "default")
|
||||
return name if isinstance(name, BaseChunkStore) else self._get_component(ComponentEnum.CHUNK_STORE, name)
|
||||
|
||||
@property
|
||||
def file_graph(self) -> FileGraph:
|
||||
name = self.kwargs.get("file_watcher", "default")
|
||||
watcher = self._get_component(ComponentEnum.FILE_WATCHER, name)
|
||||
return watcher.file_graph
|
||||
|
||||
@property
|
||||
def embedding(self) -> BaseEmbeddingModel:
|
||||
|
|
|
|||
17
reme2/component/chunk_store/__init__.py
Normal file
17
reme2/component/chunk_store/__init__.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
"""Chunk store module.
|
||||
|
||||
Storage backends for FileChunks with vector and full-text search.
|
||||
File metadata is managed by FileGraph, not by ChunkStore.
|
||||
"""
|
||||
|
||||
from .base_chunk_store import BaseChunkStore
|
||||
from .chroma_chunk_store import ChromaChunkStore
|
||||
from .local_chunk_store import LocalChunkStore
|
||||
from .sqlite_chunk_store import SqliteChunkStore
|
||||
|
||||
__all__ = [
|
||||
"BaseChunkStore",
|
||||
"ChromaChunkStore",
|
||||
"LocalChunkStore",
|
||||
"SqliteChunkStore",
|
||||
]
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
"""Abstract base class for file storage backends."""
|
||||
"""Abstract base class for chunk storage backends."""
|
||||
|
||||
import re
|
||||
from abc import abstractmethod
|
||||
|
|
@ -7,18 +7,18 @@ from pathlib import Path
|
|||
from ..base_component import BaseComponent
|
||||
from ..embedding import BaseEmbeddingModel
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import FileChunk, FileMetadata, SearchFilter
|
||||
from ...schema import ChunkFilter, FileChunk
|
||||
|
||||
|
||||
class BaseFileStore(BaseComponent):
|
||||
"""Abstract base class for file storage backends.
|
||||
class BaseChunkStore(BaseComponent):
|
||||
"""Abstract base class for chunk storage backends.
|
||||
|
||||
Provides embedding resolution, validation, safe embedding retrieval,
|
||||
metadata caching, hybrid search, and keyword scoring utilities.
|
||||
Subclasses must implement the storage-specific CRUD and search methods.
|
||||
Handles chunk persistence and retrieval (vector / keyword / hybrid search).
|
||||
File-level metadata and search filter resolution live in FileGraph; this
|
||||
layer only consumes a compiled ChunkFilter (path set) for restricting search.
|
||||
"""
|
||||
|
||||
component_type = ComponentEnum.FILE_STORE
|
||||
component_type = ComponentEnum.CHUNK_STORE
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -45,7 +45,6 @@ class BaseFileStore(BaseComponent):
|
|||
raise ValueError("At least one of embedding_model or fts_enabled must be set.")
|
||||
|
||||
async def _start(self):
|
||||
"""Resolve embedding model from app_context."""
|
||||
if not self._embedding_model_name:
|
||||
return
|
||||
assert self.app_context is not None, "app_context must be provided"
|
||||
|
|
@ -139,7 +138,7 @@ class BaseFileStore(BaseComponent):
|
|||
limit: int,
|
||||
vector_weight: float = 0.7,
|
||||
candidate_multiplier: float = 3.0,
|
||||
search_filter: SearchFilter | None = None,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Perform hybrid search combining vector and keyword results."""
|
||||
assert 0.0 <= vector_weight <= 1.0
|
||||
|
|
@ -148,8 +147,8 @@ class BaseFileStore(BaseComponent):
|
|||
text_weight = 1.0 - vector_weight
|
||||
|
||||
if self.vector_enabled and self.fts_enabled:
|
||||
keyword_results = await self.keyword_search(query, candidates, search_filter)
|
||||
vector_results = await self.vector_search(query, candidates, search_filter)
|
||||
keyword_results = await self.keyword_search(query, candidates, chunk_filter)
|
||||
vector_results = await self.vector_search(query, candidates, chunk_filter)
|
||||
|
||||
if not keyword_results:
|
||||
return vector_results[:limit]
|
||||
|
|
@ -164,9 +163,9 @@ class BaseFileStore(BaseComponent):
|
|||
)
|
||||
return merged[:limit]
|
||||
elif self.vector_enabled:
|
||||
return await self.vector_search(query, limit, search_filter)
|
||||
return await self.vector_search(query, limit, chunk_filter)
|
||||
elif self.fts_enabled:
|
||||
return await self.keyword_search(query, limit, search_filter)
|
||||
return await self.keyword_search(query, limit, chunk_filter)
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -199,24 +198,11 @@ class BaseFileStore(BaseComponent):
|
|||
|
||||
# -- Filter utility -----------------------------------------------------
|
||||
|
||||
def _apply_filter(
|
||||
self,
|
||||
chunks: list[FileChunk],
|
||||
search_filter: SearchFilter | None,
|
||||
file_metadata: dict[str, FileMetadata] | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Apply search filter to a list of chunks.
|
||||
|
||||
Args:
|
||||
chunks: Candidate chunks to filter.
|
||||
search_filter: Filter conditions.
|
||||
file_metadata: File-level metadata lookup (path -> FileMetadata).
|
||||
Used for tag filtering since tags are file-level, not chunk-level.
|
||||
"""
|
||||
if not search_filter or search_filter.is_empty():
|
||||
@staticmethod
|
||||
def _apply_filter(chunks: list[FileChunk], chunk_filter: ChunkFilter | None) -> list[FileChunk]:
|
||||
if chunk_filter is None or chunk_filter.resolved_paths is None:
|
||||
return chunks
|
||||
fm = file_metadata or {}
|
||||
return [c for c in chunks if search_filter.match(c.path, fm[c.path].metadata if c.path in fm else None)]
|
||||
return [c for c in chunks if chunk_filter.match_path(c.path)]
|
||||
|
||||
# -- Abstract methods ---------------------------------------------------
|
||||
|
||||
|
|
@ -225,34 +211,31 @@ class BaseFileStore(BaseComponent):
|
|||
"""Clear all indexed data."""
|
||||
|
||||
@abstractmethod
|
||||
async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]):
|
||||
"""Insert or update a file and its chunks."""
|
||||
async def upsert_chunks(self, path: str, chunks: list[FileChunk]):
|
||||
"""Insert or update all chunks for a file path."""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_file(self, path: str):
|
||||
"""Delete a file and all its chunks."""
|
||||
async def delete_chunks(self, path: str):
|
||||
"""Delete all chunks for a file path."""
|
||||
|
||||
@abstractmethod
|
||||
async def list_files(self) -> list[str]:
|
||||
"""List all indexed file paths."""
|
||||
async def get_chunks(self, path: str) -> list[FileChunk]:
|
||||
"""Get all chunks for a file path."""
|
||||
|
||||
@abstractmethod
|
||||
async def get_file_chunks(self, path: str) -> list[FileChunk]:
|
||||
"""Get all chunks for a file."""
|
||||
|
||||
@abstractmethod
|
||||
async def get_file_metadata(self, path: str) -> FileMetadata | None:
|
||||
"""Get metadata for a specific file."""
|
||||
|
||||
@abstractmethod
|
||||
async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]:
|
||||
"""Perform vector similarity search."""
|
||||
async def vector_search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Perform vector similarity search, optionally restricted by chunk_filter."""
|
||||
|
||||
@abstractmethod
|
||||
async def keyword_search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
search_filter: SearchFilter | None = None,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Perform full-text/keyword search."""
|
||||
"""Perform full-text/keyword search, optionally restricted by chunk_filter."""
|
||||
|
|
@ -1,12 +1,10 @@
|
|||
"""ChromaDB storage backend for file store."""
|
||||
"""ChromaDB chunk storage backend."""
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from .base_chunk_store import BaseChunkStore
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileMetadata, SearchFilter
|
||||
from ...schema import ChunkFilter, FileChunk
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
|
|
@ -20,12 +18,11 @@ except Exception as e:
|
|||
|
||||
|
||||
@R.register("chroma")
|
||||
class ChromaFileStore(BaseFileStore):
|
||||
"""ChromaDB file storage with vector and full-text search.
|
||||
class ChromaChunkStore(BaseChunkStore):
|
||||
"""ChromaDB chunk storage with vector and full-text search.
|
||||
|
||||
Uses ChromaDB's native vector search and `where_document` $contains
|
||||
for keyword matching. File metadata is persisted to a JSON file
|
||||
alongside the ChromaDB database.
|
||||
for keyword matching.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
|
|
@ -34,42 +31,11 @@ class ChromaFileStore(BaseFileStore):
|
|||
super().__init__(**kwargs)
|
||||
self.client: "chromadb.ClientAPI | None" = None
|
||||
self.chunks_collection: "chromadb.Collection | None" = None
|
||||
self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json"
|
||||
self._metadata_cache: dict[str, FileMetadata] = {}
|
||||
|
||||
@property
|
||||
def collection_name(self) -> str:
|
||||
return f"chunks_{self.store_name}"
|
||||
|
||||
# -- Persistence helpers ------------------------------------------------
|
||||
|
||||
async def _load_metadata(self) -> None:
|
||||
if not self._metadata_file.exists():
|
||||
return
|
||||
try:
|
||||
data = self._metadata_file.read_text(encoding="utf-8")
|
||||
raw = json.loads(data)
|
||||
self._metadata_cache = {path: FileMetadata(**meta) for path, meta in raw.items()}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to load metadata: {e}")
|
||||
|
||||
async def _save_metadata(self) -> None:
|
||||
try:
|
||||
raw = {
|
||||
path: meta.model_dump(mode="json")
|
||||
for path, meta in self._metadata_cache.items()
|
||||
}
|
||||
data = json.dumps(raw, indent=2, ensure_ascii=False)
|
||||
temp = self._metadata_file.with_suffix(".tmp")
|
||||
temp.write_text(data, encoding="utf-8")
|
||||
temp.replace(self._metadata_file)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to save metadata: {e}")
|
||||
raise
|
||||
finally:
|
||||
if temp.exists():
|
||||
temp.unlink()
|
||||
|
||||
# -- Lifecycle ----------------------------------------------------------
|
||||
|
||||
async def _start(self, app_context=None) -> None:
|
||||
|
|
@ -81,76 +47,77 @@ class ChromaFileStore(BaseFileStore):
|
|||
name=self.collection_name,
|
||||
metadata={"hnsw:space": "cosine"},
|
||||
)
|
||||
await self._load_metadata()
|
||||
self.logger.info(
|
||||
f"ChromaFileStore '{self.store_name}' ready: "
|
||||
f"collection={self.collection_name}, metadata at {self._metadata_file}",
|
||||
)
|
||||
self.logger.info(f"ChromaChunkStore '{self.store_name}' ready: collection={self.collection_name}")
|
||||
await super()._start(app_context)
|
||||
|
||||
async def _close(self) -> None:
|
||||
await self._save_metadata()
|
||||
self.client = None
|
||||
self.chunks_collection = None
|
||||
await super()._close()
|
||||
|
||||
# -- Filter helper ------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _path_where(chunk_filter: ChunkFilter | None) -> dict | None:
|
||||
if chunk_filter is None or chunk_filter.resolved_paths is None:
|
||||
return None
|
||||
paths = chunk_filter.resolved_paths
|
||||
if not paths:
|
||||
return {"path": "__nonexistent__"}
|
||||
if len(paths) == 1:
|
||||
return {"path": next(iter(paths))}
|
||||
return {"path": {"$in": list(paths)}}
|
||||
|
||||
# -- Write operations ---------------------------------------------------
|
||||
|
||||
async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]) -> None:
|
||||
# Always delete existing data for this file first
|
||||
if file_meta.path:
|
||||
await self.delete_file(file_meta.path)
|
||||
async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None:
|
||||
await self.delete_chunks(path)
|
||||
|
||||
if chunks:
|
||||
chunks = await self.get_chunk_embeddings(chunks)
|
||||
if not chunks:
|
||||
return
|
||||
|
||||
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 if chunk.embedding else [0.0] * self.embedding_dim)
|
||||
metadatas.append({
|
||||
"path": file_meta.path,
|
||||
"start_line": chunk.start_line,
|
||||
"end_line": chunk.end_line,
|
||||
"hash": chunk.hash,
|
||||
"updated_at": now,
|
||||
})
|
||||
chunks = await self.get_chunk_embeddings(chunks)
|
||||
|
||||
self.chunks_collection.upsert(
|
||||
ids=ids,
|
||||
documents=documents,
|
||||
embeddings=embeddings,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
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 if chunk.embedding else [0.0] * self.embedding_dim)
|
||||
metadatas.append({
|
||||
"path": path,
|
||||
"start_line": chunk.start_line,
|
||||
"end_line": chunk.end_line,
|
||||
"hash": chunk.hash,
|
||||
"updated_at": now,
|
||||
})
|
||||
|
||||
if file_meta.path:
|
||||
self._metadata_cache[file_meta.path] = FileMetadata(
|
||||
modified_time=file_meta.modified_time,
|
||||
path=file_meta.path,
|
||||
metadata=file_meta.metadata,
|
||||
)
|
||||
self.chunks_collection.upsert(
|
||||
ids=ids,
|
||||
documents=documents,
|
||||
embeddings=embeddings,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
|
||||
async def delete_file(self, path: str) -> None:
|
||||
async def delete_chunks(self, path: str) -> None:
|
||||
results = self.chunks_collection.get(where={"path": path}, include=[])
|
||||
if results["ids"]:
|
||||
self.chunks_collection.delete(ids=results["ids"])
|
||||
self._metadata_cache.pop(path, None)
|
||||
|
||||
# -- Read operations ----------------------------------------------------
|
||||
|
||||
async def list_files(self) -> list[str]:
|
||||
"""List all indexed file paths."""
|
||||
return list(self._metadata_cache.keys())
|
||||
|
||||
async def get_file_metadata(self, path: str) -> FileMetadata | None:
|
||||
"""Get metadata for a specific file."""
|
||||
return self._metadata_cache.get(path)
|
||||
async def get_chunks(self, path: str) -> list[FileChunk]:
|
||||
results = self.chunks_collection.get(where={"path": path}, include=["documents", "metadatas"])
|
||||
chunks: list[FileChunk] = []
|
||||
for cid, md, text in zip(results["ids"], results["metadatas"], results["documents"]):
|
||||
chunks.append(self._chunk_from_chroma(cid, md, text))
|
||||
chunks.sort(key=lambda c: c.start_line)
|
||||
return chunks
|
||||
|
||||
# -- Search helpers -----------------------------------------------------
|
||||
|
||||
def _chunk_from_chroma(self, chunk_id: str, md: dict, text: str, embedding=None) -> FileChunk:
|
||||
@staticmethod
|
||||
def _chunk_from_chroma(chunk_id: str, md: dict, text: str, embedding=None) -> FileChunk:
|
||||
return FileChunk(
|
||||
id=chunk_id,
|
||||
path=md["path"],
|
||||
|
|
@ -163,7 +130,12 @@ class ChromaFileStore(BaseFileStore):
|
|||
|
||||
# -- Search operations --------------------------------------------------
|
||||
|
||||
async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]:
|
||||
async def vector_search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
if not self.vector_enabled or not query:
|
||||
return []
|
||||
|
||||
|
|
@ -175,6 +147,7 @@ class ChromaFileStore(BaseFileStore):
|
|||
results = self.chunks_collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=limit,
|
||||
where=self._path_where(chunk_filter),
|
||||
include=["documents", "metadatas", "distances"],
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -192,7 +165,6 @@ class ChromaFileStore(BaseFileStore):
|
|||
chunk.scores = {"vector": score, "score": score}
|
||||
chunks.append(chunk)
|
||||
|
||||
chunks = self._apply_filter(chunks, search_filter)
|
||||
chunks.sort(key=lambda c: c.score, reverse=True)
|
||||
return chunks[:limit]
|
||||
|
||||
|
|
@ -200,9 +172,8 @@ class ChromaFileStore(BaseFileStore):
|
|||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
search_filter: SearchFilter | None = None,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Keyword search via ChromaDB $contains with case variants."""
|
||||
if not self.fts_enabled or not query:
|
||||
return []
|
||||
|
||||
|
|
@ -210,7 +181,6 @@ class ChromaFileStore(BaseFileStore):
|
|||
if not words:
|
||||
return []
|
||||
|
||||
# Generate case variants for case-insensitive matching
|
||||
word_variants = set()
|
||||
for word in words:
|
||||
word_variants.add(word)
|
||||
|
|
@ -225,6 +195,7 @@ class ChromaFileStore(BaseFileStore):
|
|||
where_document = {"$or": [{"$contains": w} for w in variants_list]}
|
||||
|
||||
results = self.chunks_collection.get(
|
||||
where=self._path_where(chunk_filter),
|
||||
where_document=where_document,
|
||||
include=["documents", "metadatas"],
|
||||
)
|
||||
|
|
@ -241,7 +212,6 @@ class ChromaFileStore(BaseFileStore):
|
|||
chunk.scores = {"keyword": score, "score": score}
|
||||
chunks.append(chunk)
|
||||
|
||||
chunks = self._apply_filter(chunks, search_filter)
|
||||
chunks.sort(key=lambda c: c.score, reverse=True)
|
||||
return chunks[:limit]
|
||||
|
||||
|
|
@ -253,6 +223,4 @@ class ChromaFileStore(BaseFileStore):
|
|||
name=self.collection_name,
|
||||
metadata={"hnsw:space": "cosine"},
|
||||
)
|
||||
self._metadata_cache = {}
|
||||
await self._save_metadata()
|
||||
self.logger.info(f"Cleared all data from ChromaFileStore '{self.store_name}'")
|
||||
self.logger.info(f"Cleared all data from ChromaChunkStore '{self.store_name}'")
|
||||
|
|
@ -1,31 +1,25 @@
|
|||
"""Pure-Python file storage with JSONL persistence."""
|
||||
"""Pure-Python chunk storage with JSONL persistence."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from .base_chunk_store import BaseChunkStore
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileMetadata, SearchFilter
|
||||
from ...schema import ChunkFilter, FileChunk
|
||||
from ...utils import batch_cosine_similarity
|
||||
|
||||
|
||||
@R.register("local")
|
||||
class LocalFileStore(BaseFileStore):
|
||||
"""In-memory file storage with JSONL disk persistence.
|
||||
|
||||
No external database required. All data lives in Python dicts;
|
||||
writes are flushed to JSONL files on disk and survive restarts.
|
||||
"""
|
||||
class LocalChunkStore(BaseChunkStore):
|
||||
"""In-memory chunk storage with JSONL disk persistence."""
|
||||
|
||||
def __init__(self, encoding: str = "utf-8", **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._encoding: str = encoding
|
||||
self._chunks: dict[str, FileChunk] = {}
|
||||
self._files: dict[str, FileMetadata] = {}
|
||||
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 ------------------------------------------------
|
||||
|
||||
|
|
@ -59,88 +53,40 @@ class LocalFileStore(BaseFileStore):
|
|||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
|
||||
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=self._encoding)
|
||||
raw: dict = json.loads(data)
|
||||
self._files = {path: FileMetadata(**meta) for path, meta in raw.items()}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to load metadata: {e}")
|
||||
|
||||
async def _save_metadata(self) -> None:
|
||||
"""Persist file metadata to JSON file with atomic write."""
|
||||
raw = {path: meta.model_dump(mode="json") for path, meta in self._files.items()}
|
||||
content = json.dumps(raw, indent=2, ensure_ascii=False)
|
||||
temp_path = self._metadata_file.with_suffix(".tmp")
|
||||
try:
|
||||
temp_path.write_text(content, encoding=self._encoding)
|
||||
temp_path.replace(self._metadata_file)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to save metadata: {e}")
|
||||
raise
|
||||
finally:
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
|
||||
# -- Lifecycle ----------------------------------------------------------
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Load persisted data into memory."""
|
||||
await self._load_metadata()
|
||||
await self._load_chunks()
|
||||
self.logger.info(
|
||||
f"LocalFileStore '{self.store_name}' ready: "
|
||||
f"{len(self._chunks)} chunks, metadata at {self._metadata_file}",
|
||||
)
|
||||
self.logger.info(f"LocalChunkStore '{self.store_name}' ready: {len(self._chunks)} chunks")
|
||||
await super()._start()
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Flush state to disk and clear memory."""
|
||||
await self._save_metadata()
|
||||
await self._save_chunks()
|
||||
self._chunks.clear()
|
||||
self._files.clear()
|
||||
await super()._close()
|
||||
|
||||
# -- Write operations ---------------------------------------------------
|
||||
|
||||
async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]) -> None:
|
||||
async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None:
|
||||
"""Insert or update a file and its chunks."""
|
||||
await self.delete_file(file_meta.path)
|
||||
await self.delete_chunks(path)
|
||||
if not chunks:
|
||||
return
|
||||
chunks = await self.get_chunk_embeddings(chunks)
|
||||
for chunk in chunks:
|
||||
self._chunks[chunk.id] = chunk
|
||||
|
||||
if chunks:
|
||||
chunks = await self.get_chunk_embeddings(chunks)
|
||||
for chunk in chunks:
|
||||
self._chunks[chunk.id] = chunk
|
||||
|
||||
if file_meta.path:
|
||||
self._files[file_meta.path] = FileMetadata(
|
||||
modified_time=file_meta.modified_time,
|
||||
path=file_meta.path,
|
||||
metadata=file_meta.metadata,
|
||||
)
|
||||
|
||||
async def delete_file(self, path: str) -> None:
|
||||
async def delete_chunks(self, path: str) -> None:
|
||||
"""Delete a file and all its chunks."""
|
||||
to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path]
|
||||
for cid in to_delete:
|
||||
del self._chunks[cid]
|
||||
self._files.pop(path, None)
|
||||
|
||||
# -- Read operations ----------------------------------------------------
|
||||
|
||||
async def list_files(self) -> list[str]:
|
||||
"""List all indexed file paths."""
|
||||
return list(self._files.keys())
|
||||
|
||||
async def get_file_metadata(self, path: str) -> FileMetadata | None:
|
||||
"""Get metadata for a specific file."""
|
||||
return self._files.get(path)
|
||||
|
||||
async def get_file_chunks(self, path: str) -> list[FileChunk]:
|
||||
async def get_chunks(self, path: str) -> list[FileChunk]:
|
||||
"""Get all chunks for a file, sorted by start_line."""
|
||||
chunks = [chunk for chunk in self._chunks.values() if chunk.path == path]
|
||||
chunks.sort(key=lambda c: c.start_line)
|
||||
|
|
@ -148,7 +94,12 @@ class LocalFileStore(BaseFileStore):
|
|||
|
||||
# -- Search operations --------------------------------------------------
|
||||
|
||||
async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]:
|
||||
async def vector_search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Cosine-similarity vector search over in-memory embeddings."""
|
||||
if not self.vector_enabled or not query:
|
||||
return []
|
||||
|
|
@ -159,14 +110,12 @@ class LocalFileStore(BaseFileStore):
|
|||
|
||||
candidates = self._apply_filter(
|
||||
[c for c in self._chunks.values() if c.embedding],
|
||||
search_filter,
|
||||
chunk_filter,
|
||||
)
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
expected_dim = self.embedding_dim
|
||||
|
||||
# Validate and align embedding dimensions
|
||||
valid_embeddings = []
|
||||
for chunk in candidates:
|
||||
emb = chunk.embedding
|
||||
|
|
@ -201,17 +150,16 @@ class LocalFileStore(BaseFileStore):
|
|||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
search_filter: SearchFilter | None = None,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
"""Keyword search via substring matching."""
|
||||
if not self.fts_enabled or not query:
|
||||
return []
|
||||
|
||||
words = query.split()
|
||||
if not words:
|
||||
if not query.split():
|
||||
return []
|
||||
|
||||
filtered_chunks = self._apply_filter(list(self._chunks.values()), search_filter)
|
||||
filtered_chunks = self._apply_filter(list(self._chunks.values()), chunk_filter)
|
||||
|
||||
results = []
|
||||
for chunk in filtered_chunks:
|
||||
|
|
@ -236,7 +184,5 @@ class LocalFileStore(BaseFileStore):
|
|||
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()
|
||||
self.logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'")
|
||||
self.logger.info(f"Cleared all data from LocalChunkStore '{self.store_name}'")
|
||||
|
|
@ -1,18 +1,18 @@
|
|||
"""SQLite storage backend for file store."""
|
||||
"""SQLite chunk storage backend."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import struct
|
||||
import time
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from .base_chunk_store import BaseChunkStore
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileMetadata, SearchFilter
|
||||
from ...schema import ChunkFilter, FileChunk
|
||||
|
||||
|
||||
@R.register("sqlite")
|
||||
class SqliteFileStore(BaseFileStore):
|
||||
"""SQLite file storage with vector and full-text search.
|
||||
class SqliteChunkStore(BaseChunkStore):
|
||||
"""SQLite chunk storage with vector and full-text search.
|
||||
|
||||
Uses sqlite-vec for vector similarity search and FTS5 with trigram
|
||||
tokenizer for keyword search. Falls back to LIKE-based substring
|
||||
|
|
@ -30,10 +30,6 @@ class SqliteFileStore(BaseFileStore):
|
|||
def chunks_table(self) -> str:
|
||||
return f"chunks_{self.store_name}"
|
||||
|
||||
@property
|
||||
def files_table(self) -> str:
|
||||
return f"files_{self.store_name}"
|
||||
|
||||
@property
|
||||
def vector_table(self) -> str:
|
||||
return f"chunks_vec_{self.store_name}"
|
||||
|
|
@ -88,23 +84,12 @@ class SqliteFileStore(BaseFileStore):
|
|||
self.conn.enable_load_extension(False)
|
||||
|
||||
await self._create_tables()
|
||||
self.logger.info(
|
||||
f"SqliteFileStore '{self.store_name}' ready: "
|
||||
f"db={self.db_path / 'reme.db'}",
|
||||
)
|
||||
self.logger.info(f"SqliteChunkStore '{self.store_name}' ready: db={self.db_path / 'reme.db'}")
|
||||
await super()._start(app_context)
|
||||
|
||||
async def _create_tables(self) -> None:
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute(f"""
|
||||
CREATE TABLE IF NOT EXISTS {self.files_table} (
|
||||
path TEXT PRIMARY KEY,
|
||||
mtime REAL,
|
||||
metadata TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
cursor.execute(f"""
|
||||
CREATE TABLE IF NOT EXISTS {self.chunks_table} (
|
||||
id TEXT PRIMARY KEY,
|
||||
|
|
@ -117,6 +102,10 @@ class SqliteFileStore(BaseFileStore):
|
|||
updated_at INTEGER
|
||||
)
|
||||
""")
|
||||
cursor.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_{self.chunks_table}_path "
|
||||
f"ON {self.chunks_table}(path)",
|
||||
)
|
||||
|
||||
if self.vector_enabled:
|
||||
cursor.execute(f"""
|
||||
|
|
@ -153,52 +142,17 @@ class SqliteFileStore(BaseFileStore):
|
|||
self.conn = None
|
||||
await super()._close()
|
||||
|
||||
# -- Metadata helpers ---------------------------------------------------
|
||||
|
||||
async def _get_all_metadata(self) -> dict[str, FileMetadata]:
|
||||
"""Load all file metadata from SQL into a dict for tag filtering."""
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute(f"SELECT path, mtime, metadata FROM {self.files_table}")
|
||||
result = {}
|
||||
for path, mtime, meta_str in cursor.fetchall():
|
||||
metadata = json.loads(meta_str) if meta_str else {}
|
||||
result[path] = FileMetadata(
|
||||
modified_time=mtime,
|
||||
path=path,
|
||||
metadata=metadata,
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to load all metadata: {e}")
|
||||
return {}
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
# -- Write operations ---------------------------------------------------
|
||||
|
||||
async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]) -> None:
|
||||
async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None:
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute("BEGIN")
|
||||
|
||||
# Upsert file metadata
|
||||
cursor.execute(
|
||||
f"""INSERT OR REPLACE INTO {self.files_table}
|
||||
(path, mtime, metadata)
|
||||
VALUES (?, ?, ?)""",
|
||||
(
|
||||
file_meta.path,
|
||||
file_meta.modified_time,
|
||||
json.dumps(file_meta.metadata, ensure_ascii=False) if file_meta.metadata else None,
|
||||
),
|
||||
)
|
||||
|
||||
# Delete old chunks and vectors/fts for this file
|
||||
old_ids = [
|
||||
row[0] for row in cursor.execute(
|
||||
f"SELECT id FROM {self.chunks_table} WHERE path = ?",
|
||||
(file_meta.path,),
|
||||
(path,),
|
||||
).fetchall()
|
||||
]
|
||||
if old_ids:
|
||||
|
|
@ -208,9 +162,8 @@ class SqliteFileStore(BaseFileStore):
|
|||
for oid in old_ids:
|
||||
cursor.execute(f"DELETE FROM {self.vector_table} WHERE id = ?", (oid,))
|
||||
if self.fts_enabled:
|
||||
cursor.execute(f"DELETE FROM {self.fts_table} WHERE path = ?", (file_meta.path,))
|
||||
cursor.execute(f"DELETE FROM {self.fts_table} WHERE path = ?", (path,))
|
||||
|
||||
# Insert new chunks
|
||||
if chunks:
|
||||
chunks = await self.get_chunk_embeddings(chunks)
|
||||
now = int(time.time() * 1000)
|
||||
|
|
@ -221,7 +174,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
(
|
||||
chunk.id,
|
||||
file_meta.path,
|
||||
path,
|
||||
chunk.start_line,
|
||||
chunk.end_line,
|
||||
chunk.hash,
|
||||
|
|
@ -242,18 +195,18 @@ class SqliteFileStore(BaseFileStore):
|
|||
f"""INSERT INTO {self.fts_table}
|
||||
(text, id, path, start_line, end_line)
|
||||
VALUES (?, ?, ?, ?, ?)""",
|
||||
(chunk.text, chunk.id, file_meta.path, chunk.start_line, chunk.end_line),
|
||||
(chunk.text, chunk.id, path, chunk.start_line, chunk.end_line),
|
||||
)
|
||||
|
||||
cursor.execute("COMMIT")
|
||||
except Exception as e:
|
||||
cursor.execute("ROLLBACK")
|
||||
self.logger.error(f"Failed to upsert file {file_meta.path}: {e}")
|
||||
self.logger.error(f"Failed to upsert chunks for {path}: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
async def delete_file(self, path: str) -> None:
|
||||
async def delete_chunks(self, path: str) -> None:
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute("BEGIN")
|
||||
|
|
@ -272,53 +225,18 @@ class SqliteFileStore(BaseFileStore):
|
|||
cursor.execute(f"DELETE FROM {self.fts_table} WHERE path = ?", (path,))
|
||||
|
||||
cursor.execute(f"DELETE FROM {self.chunks_table} WHERE path = ?", (path,))
|
||||
cursor.execute(f"DELETE FROM {self.files_table} WHERE path = ?", (path,))
|
||||
|
||||
cursor.execute("COMMIT")
|
||||
except Exception as e:
|
||||
cursor.execute("ROLLBACK")
|
||||
self.logger.error(f"Failed to delete file {path}: {e}")
|
||||
self.logger.error(f"Failed to delete chunks for {path}: {e}")
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
# -- Read operations (SQL overrides) ------------------------------------
|
||||
# -- Read operations ----------------------------------------------------
|
||||
|
||||
async def list_files(self) -> list[str]:
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute(f"SELECT path FROM {self.files_table}")
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to list files: {e}")
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
async def get_file_metadata(self, path: str) -> FileMetadata | None:
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
f"SELECT mtime, metadata FROM {self.files_table} WHERE path = ?",
|
||||
(path,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if not row:
|
||||
return None
|
||||
mtime, meta_str = row
|
||||
metadata = json.loads(meta_str) if meta_str else {}
|
||||
return FileMetadata(
|
||||
modified_time=mtime,
|
||||
path=path,
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to get file metadata for {path}: {e}")
|
||||
return None
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
async def get_file_chunks(self, path: str) -> list[FileChunk]:
|
||||
async def get_chunks(self, path: str) -> list[FileChunk]:
|
||||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
|
|
@ -347,7 +265,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
))
|
||||
return chunks
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to get file chunks for {path}: {e}")
|
||||
self.logger.error(f"Failed to get chunks for {path}: {e}")
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
|
|
@ -364,9 +282,25 @@ class SqliteFileStore(BaseFileStore):
|
|||
cleaned = cleaned.replace(ch, " ")
|
||||
return " ".join(cleaned.split())
|
||||
|
||||
@staticmethod
|
||||
def _path_filter_clause(chunk_filter: ChunkFilter | None, column: str = "path") -> tuple[str, list]:
|
||||
"""Build a WHERE clause fragment for ChunkFilter; returns (sql_fragment, params)."""
|
||||
if chunk_filter is None or chunk_filter.resolved_paths is None:
|
||||
return "", []
|
||||
paths = chunk_filter.resolved_paths
|
||||
if not paths:
|
||||
return "0", [] # filter excludes everything
|
||||
placeholders = ",".join("?" * len(paths))
|
||||
return f"{column} IN ({placeholders})", list(paths)
|
||||
|
||||
# -- Search operations --------------------------------------------------
|
||||
|
||||
async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]:
|
||||
async def vector_search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
if not self.vector_enabled or not query:
|
||||
return []
|
||||
|
||||
|
|
@ -391,7 +325,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
chunks = []
|
||||
for cid, path, start, end, text, dist in cursor.fetchall():
|
||||
score = max(0.0, 1.0 - dist / 2.0)
|
||||
chunk = FileChunk(
|
||||
chunks.append(FileChunk(
|
||||
id=cid,
|
||||
path=path,
|
||||
start_line=start,
|
||||
|
|
@ -399,12 +333,9 @@ class SqliteFileStore(BaseFileStore):
|
|||
text=text,
|
||||
hash="",
|
||||
scores={"vector": score, "score": score},
|
||||
)
|
||||
chunks.append(chunk)
|
||||
))
|
||||
|
||||
# Apply filter with SQL-based metadata for tag support
|
||||
file_meta = await self._get_all_metadata() if (search_filter and search_filter.tags) else None
|
||||
chunks = self._apply_filter(chunks, search_filter, file_meta)
|
||||
chunks = self._apply_filter(chunks, chunk_filter)
|
||||
chunks.sort(key=lambda c: c.score, reverse=True)
|
||||
return chunks[:limit]
|
||||
except Exception as e:
|
||||
|
|
@ -417,7 +348,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
self,
|
||||
query: str,
|
||||
limit: int,
|
||||
search_filter: SearchFilter | None = None,
|
||||
chunk_filter: ChunkFilter | None = None,
|
||||
) -> list[FileChunk]:
|
||||
if not self.fts_enabled or not query:
|
||||
return []
|
||||
|
|
@ -430,18 +361,14 @@ class SqliteFileStore(BaseFileStore):
|
|||
if not words:
|
||||
return []
|
||||
|
||||
file_meta = await self._get_all_metadata() if (search_filter and search_filter.tags) else None
|
||||
|
||||
# FTS5 trigram requires all terms >= 3 chars
|
||||
if all(len(w) >= 3 for w in words):
|
||||
results = await self._fts_trigram_search(words, limit)
|
||||
if results:
|
||||
return self._apply_filter(results, search_filter, file_meta)[:limit]
|
||||
return self._apply_filter(results, chunk_filter)[:limit]
|
||||
|
||||
return self._apply_filter(
|
||||
await self._like_search(cleaned, words, limit),
|
||||
search_filter,
|
||||
file_meta,
|
||||
chunk_filter,
|
||||
)[:limit]
|
||||
|
||||
async def _fts_trigram_search(self, words: list[str], limit: int) -> list[FileChunk]:
|
||||
|
|
@ -464,7 +391,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
chunks = []
|
||||
for cid, path, start, end, text, rank in cursor.fetchall():
|
||||
score = max(0.0, 1.0 / (1.0 + abs(rank)))
|
||||
chunk = FileChunk(
|
||||
chunks.append(FileChunk(
|
||||
id=cid,
|
||||
path=path,
|
||||
start_line=start,
|
||||
|
|
@ -472,8 +399,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
text=text,
|
||||
hash="",
|
||||
scores={"keyword": score, "score": score},
|
||||
)
|
||||
chunks.append(chunk)
|
||||
))
|
||||
chunks.sort(key=lambda c: c.score, reverse=True)
|
||||
return chunks
|
||||
except Exception as e:
|
||||
|
|
@ -511,7 +437,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
if score == 0.0:
|
||||
continue
|
||||
|
||||
chunk = FileChunk(
|
||||
chunks.append(FileChunk(
|
||||
id=cid,
|
||||
path=path,
|
||||
start_line=start,
|
||||
|
|
@ -519,8 +445,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
text=text,
|
||||
hash="",
|
||||
scores={"keyword": score, "score": score},
|
||||
)
|
||||
chunks.append(chunk)
|
||||
))
|
||||
|
||||
chunks.sort(key=lambda c: c.score, reverse=True)
|
||||
return chunks[:limit]
|
||||
|
|
@ -536,7 +461,6 @@ class SqliteFileStore(BaseFileStore):
|
|||
cursor = self.conn.cursor()
|
||||
try:
|
||||
cursor.execute("BEGIN")
|
||||
cursor.execute(f"DELETE FROM {self.files_table}")
|
||||
cursor.execute(f"DELETE FROM {self.chunks_table}")
|
||||
if self.vector_enabled:
|
||||
cursor.execute(f"DELETE FROM {self.vector_table}")
|
||||
|
|
@ -549,4 +473,4 @@ class SqliteFileStore(BaseFileStore):
|
|||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
self.logger.info(f"Cleared all data from SqliteFileStore '{self.store_name}'")
|
||||
self.logger.info(f"Cleared all data from SqliteChunkStore '{self.store_name}'")
|
||||
|
|
@ -1,17 +0,0 @@
|
|||
"""File store module for persistent memory management.
|
||||
|
||||
Provides storage backends for memory chunks and file metadata with
|
||||
vector and full-text search capabilities.
|
||||
"""
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from .chroma_file_store import ChromaFileStore
|
||||
from .local_file_store import LocalFileStore
|
||||
from .sqlite_file_store import SqliteFileStore
|
||||
|
||||
__all__ = [
|
||||
"BaseFileStore",
|
||||
"ChromaFileStore",
|
||||
"LocalFileStore",
|
||||
"SqliteFileStore",
|
||||
]
|
||||
|
|
@ -6,10 +6,10 @@ from pathlib import Path
|
|||
from watchfiles import Change, awatch
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..chunk_store import BaseChunkStore
|
||||
from ..file_parser import BaseFileParser
|
||||
from ..file_store import BaseFileStore
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...file_graph import FileGraph
|
||||
from ...schema.file_graph import FileGraph
|
||||
|
||||
|
||||
class BaseFileWatcher(BaseComponent):
|
||||
|
|
@ -31,16 +31,16 @@ class BaseFileWatcher(BaseComponent):
|
|||
debounce: int = 2000,
|
||||
chunk_tokens: int = 400,
|
||||
chunk_overlap: int = 80,
|
||||
file_store: str = "default",
|
||||
chunk_store: str = "default",
|
||||
default_parser: str | None = None,
|
||||
rebuild_index_on_start: bool = False,
|
||||
poll_delay_ms: int = 2000,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self._file_store_name: str = file_store
|
||||
self._chunk_store_name: str = chunk_store
|
||||
self._default_parser_name: str | None = default_parser
|
||||
self.file_store: BaseFileStore | None = None
|
||||
self.chunk_store: BaseChunkStore | None = None
|
||||
self._suffix_to_parser: dict[str, BaseFileParser] = {}
|
||||
self._default_parser: BaseFileParser | None = None
|
||||
self.watch_path: str = watch_path
|
||||
|
|
@ -66,17 +66,17 @@ class BaseFileWatcher(BaseComponent):
|
|||
return self._meta_path / "file_graph.json"
|
||||
|
||||
async def _start(self):
|
||||
"""Resolve file_store, load or build file_graph, and start watching."""
|
||||
if self._file_store_name:
|
||||
"""Resolve chunk_store, load or build file_graph, and start watching."""
|
||||
if self._chunk_store_name:
|
||||
assert self.app_context is not None, "app_context must be provided"
|
||||
|
||||
stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {})
|
||||
if self._file_store_name not in stores:
|
||||
raise ValueError(f"File store '{self._file_store_name}' not found.")
|
||||
store = stores[self._file_store_name]
|
||||
if not isinstance(store, BaseFileStore):
|
||||
raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}")
|
||||
self.file_store = store
|
||||
stores = self.app_context.components.get(ComponentEnum.CHUNK_STORE, {})
|
||||
if self._chunk_store_name not in stores:
|
||||
raise ValueError(f"Chunk store '{self._chunk_store_name}' not found.")
|
||||
store = stores[self._chunk_store_name]
|
||||
if not isinstance(store, BaseChunkStore):
|
||||
raise TypeError(f"Expected BaseChunkStore, got {type(store).__name__}")
|
||||
self.chunk_store = store
|
||||
|
||||
parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {})
|
||||
for parser in parsers.values():
|
||||
|
|
@ -127,13 +127,13 @@ class BaseFileWatcher(BaseComponent):
|
|||
|
||||
self._watch_task = None
|
||||
self._stop_event.clear()
|
||||
self.file_store = None
|
||||
self.chunk_store = None
|
||||
self._suffix_to_parser.clear()
|
||||
self._default_parser = None
|
||||
self.logger.info("Stopped watching")
|
||||
|
||||
async def _scan_existing_files(self) -> None:
|
||||
if not self.file_store:
|
||||
if not self.chunk_store:
|
||||
return
|
||||
|
||||
watch_path = Path(self.watch_path)
|
||||
|
|
@ -202,7 +202,7 @@ class BaseFileWatcher(BaseComponent):
|
|||
return self._suffix_to_parser.get(suffix, self._default_parser)
|
||||
|
||||
async def on_changes(self, changes: set[tuple[Change, str]]) -> None:
|
||||
if not self.file_store:
|
||||
if not self.chunk_store:
|
||||
self.logger.warning("File store not initialized, skipping changes")
|
||||
return
|
||||
|
||||
|
|
@ -227,8 +227,8 @@ class BaseFileWatcher(BaseComponent):
|
|||
self.logger.debug(f"No parser for {path}, skipping")
|
||||
return
|
||||
file_meta, chunks = await parser.parse(path)
|
||||
await self.file_store.upsert_file(file_meta, chunks)
|
||||
self.file_graph.add(file_meta)
|
||||
await self.chunk_store.upsert_chunks(path, chunks)
|
||||
self.file_graph.create(file_meta)
|
||||
self.logger.info(f"Added {path} ({len(chunks)} chunks)")
|
||||
|
||||
async def _on_modified(self, path: str) -> None:
|
||||
|
|
@ -237,11 +237,11 @@ class BaseFileWatcher(BaseComponent):
|
|||
self.logger.debug(f"No parser for {path}, skipping")
|
||||
return
|
||||
file_meta, chunks = await parser.parse(path)
|
||||
await self.file_store.upsert_file(file_meta, chunks)
|
||||
self.file_graph.update(file_meta)
|
||||
await self.chunk_store.upsert_chunks(path, chunks)
|
||||
self.file_graph.create(file_meta)
|
||||
self.logger.info(f"Modified {path} ({len(chunks)} chunks)")
|
||||
|
||||
async def _on_deleted(self, path: str) -> None:
|
||||
await self.file_store.delete_file(path)
|
||||
self.file_graph.remove(path)
|
||||
await self.chunk_store.delete_chunks(path)
|
||||
self.file_graph.delete(path)
|
||||
self.logger.info(f"Deleted {path}")
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ components:
|
|||
default:
|
||||
backend: default
|
||||
|
||||
file_store:
|
||||
chunk_store:
|
||||
default:
|
||||
backend: local
|
||||
embedding_model: default
|
||||
|
|
@ -58,7 +58,7 @@ components:
|
|||
file_watcher:
|
||||
default:
|
||||
backend: full
|
||||
file_store: default
|
||||
chunk_store: default
|
||||
default_parser: default
|
||||
watch_paths: [ "./test_data" ]
|
||||
recursive: true
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ class ComponentEnum(str, Enum):
|
|||
|
||||
FILE_PARSER = "file_parser"
|
||||
|
||||
FILE_STORE = "file_store"
|
||||
CHUNK_STORE = "chunk_store"
|
||||
|
||||
FILE_WATCHER = "file_watcher"
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ import json
|
|||
|
||||
from ..component import R
|
||||
from ..component.base_step import BaseStep
|
||||
from ..schema import SearchFilter
|
||||
|
||||
|
||||
@R.register("memory_search")
|
||||
|
|
@ -38,22 +37,20 @@ class MemorySearch(BaseStep):
|
|||
isinstance(max_results, int) and max_results > 0
|
||||
), f"max_results must be a positive integer, got {max_results}"
|
||||
|
||||
filter_paths: list[str] | None = self.context.get("paths") or None
|
||||
filter_tags: list[str] | None = self.context.get("tags") or None
|
||||
exclude_paths: list[str] | None = self.context.get("exclude_paths") or None
|
||||
search_filter = None
|
||||
if filter_paths or filter_tags or exclude_paths:
|
||||
search_filter = SearchFilter(paths=filter_paths, tags=filter_tags, exclude_paths=exclude_paths)
|
||||
chunk_filter = self.file_graph.filter(
|
||||
paths=self.context.get("paths") or None,
|
||||
tags=self.context.get("tags") or None,
|
||||
exclude_paths=self.context.get("exclude_paths") or None,
|
||||
)
|
||||
|
||||
results = await self.file_store.hybrid_search(
|
||||
results = await self.chunk_store.hybrid_search(
|
||||
query=query,
|
||||
limit=max_results,
|
||||
vector_weight=self.vector_weight,
|
||||
candidate_multiplier=self.candidate_multiplier,
|
||||
search_filter=search_filter,
|
||||
chunk_filter=chunk_filter,
|
||||
)
|
||||
|
||||
# Filter by min_score
|
||||
results = [r for r in results if r.score >= min_score]
|
||||
|
||||
return json.dumps([result.model_dump(exclude_none=True) for result in results], indent=2, ensure_ascii=False)
|
||||
|
|
|
|||
|
|
@ -3,11 +3,11 @@
|
|||
from .application_config import ApplicationConfig, ComponentConfig, JobConfig
|
||||
from .as_msg_stat import AsBlockStat, AsMsgStat
|
||||
from .base_node import BaseNode
|
||||
from .chunk_filter import ChunkFilter
|
||||
from .file_chunk import FileChunk
|
||||
from .file_metadata import FileMetadata
|
||||
from .request import Request
|
||||
from .response import Response
|
||||
from .search_filter import SearchFilter
|
||||
from .stream_chunk import StreamChunk
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -17,10 +17,10 @@ __all__ = [
|
|||
"AsBlockStat",
|
||||
"AsMsgStat",
|
||||
"BaseNode",
|
||||
"ChunkFilter",
|
||||
"FileChunk",
|
||||
"FileMetadata",
|
||||
"Request",
|
||||
"Response",
|
||||
"SearchFilter",
|
||||
"StreamChunk",
|
||||
]
|
||||
|
|
|
|||
49
reme2/schema/chunk_filter.py
Normal file
49
reme2/schema/chunk_filter.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
"""Filter for chunk-store search.
|
||||
|
||||
User-facing fields (paths/tags/exclude_paths) describe metadata-level intent.
|
||||
FileGraph compiles them into `resolved_paths` (the concrete path set that
|
||||
ChunkStore actually consumes for filtering).
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ChunkFilter(BaseModel):
|
||||
"""User-facing search filter, resolved by FileGraph into a path set.
|
||||
|
||||
User input fields:
|
||||
paths: include only paths starting with any of these prefixes
|
||||
tags: include only files whose metadata contains ALL these tags
|
||||
exclude_paths: exclude paths starting with any of these prefixes
|
||||
|
||||
Compiled field (set by FileGraph.filter):
|
||||
resolved_paths: concrete path set to restrict ChunkStore search.
|
||||
None = no restriction (all chunks).
|
||||
Empty set = no chunks match (search returns empty).
|
||||
"""
|
||||
|
||||
paths: list[str] | None = Field(default=None)
|
||||
tags: list[str] | None = Field(default=None)
|
||||
exclude_paths: list[str] | None = Field(default=None)
|
||||
resolved_paths: set[str] | None = Field(default=None)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return not self.paths and not self.tags and not self.exclude_paths
|
||||
|
||||
def match_metadata(self, path: str, metadata: dict | None = None) -> bool:
|
||||
"""Match a single path+metadata against user-input conditions."""
|
||||
if self.paths and not any(path.startswith(p) for p in self.paths):
|
||||
return False
|
||||
if self.exclude_paths and any(path.startswith(p) for p in self.exclude_paths):
|
||||
return False
|
||||
if self.tags:
|
||||
file_tags = set((metadata or {}).get("tags", []))
|
||||
if not all(t in file_tags for t in self.tags):
|
||||
return False
|
||||
return True
|
||||
|
||||
def match_path(self, path: str) -> bool:
|
||||
"""Match a path against the resolved path set (used by ChunkStore)."""
|
||||
if self.resolved_paths is None:
|
||||
return True
|
||||
return path in self.resolved_paths
|
||||
|
|
@ -2,7 +2,8 @@ import json
|
|||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
from reme2.schema.file_metadata import FileMetadata
|
||||
from .chunk_filter import ChunkFilter
|
||||
from .file_metadata import FileMetadata
|
||||
|
||||
|
||||
class FileGraph:
|
||||
|
|
@ -13,7 +14,7 @@ class FileGraph:
|
|||
|
||||
# -- CRUD ----------------------------------------------------------------
|
||||
|
||||
def add(self, *metadatas: FileMetadata) -> None:
|
||||
def create(self, *metadatas: FileMetadata) -> None:
|
||||
for metadata in metadatas:
|
||||
path = metadata.path
|
||||
if path in self._nodes:
|
||||
|
|
@ -22,6 +23,9 @@ class FileGraph:
|
|||
for metadata in metadatas:
|
||||
self._add_forward(metadata)
|
||||
|
||||
def read(self, path: str) -> FileMetadata | None:
|
||||
return self._nodes.get(path)
|
||||
|
||||
def update(self, path: str, **fields) -> FileMetadata | None:
|
||||
metadata = self._nodes.get(path)
|
||||
if metadata is None:
|
||||
|
|
@ -32,7 +36,7 @@ class FileGraph:
|
|||
self._add_forward(updated)
|
||||
return updated
|
||||
|
||||
def remove(self, path: str) -> FileMetadata | None:
|
||||
def delete(self, path: str) -> FileMetadata | None:
|
||||
metadata = self._nodes.pop(path, None)
|
||||
if metadata is None:
|
||||
return None
|
||||
|
|
@ -40,9 +44,6 @@ class FileGraph:
|
|||
self._backlinks.pop(path, None)
|
||||
return metadata
|
||||
|
||||
def get(self, path: str) -> FileMetadata | None:
|
||||
return self._nodes.get(path)
|
||||
|
||||
# -- Link queries --------------------------------------------------------
|
||||
|
||||
def get_links(self, path: str) -> list[FileMetadata]:
|
||||
|
|
@ -58,6 +59,21 @@ class FileGraph:
|
|||
if src in self._nodes
|
||||
]
|
||||
|
||||
def filter(
|
||||
self,
|
||||
paths: list[str] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
exclude_paths: list[str] | None = None,
|
||||
) -> ChunkFilter:
|
||||
cf = ChunkFilter(paths=paths, tags=tags, exclude_paths=exclude_paths)
|
||||
if cf.is_empty():
|
||||
return cf
|
||||
cf.resolved_paths = {
|
||||
path for path, meta in self._nodes.items()
|
||||
if cf.match_metadata(path, meta.metadata)
|
||||
}
|
||||
return cf
|
||||
|
||||
# -- Index helpers -------------------------------------------------------
|
||||
|
||||
def _add_forward(self, metadata: FileMetadata) -> None:
|
||||
|
|
@ -87,7 +103,7 @@ class FileGraph:
|
|||
return graph
|
||||
raw: dict = json.loads(path.read_text(encoding="utf-8"))
|
||||
nodes = [FileMetadata(**meta) for meta in raw.values()]
|
||||
graph.add(*nodes)
|
||||
graph.create(*nodes)
|
||||
return graph
|
||||
|
||||
# -- Dunder --------------------------------------------------------------
|
||||
|
|
@ -1,36 +0,0 @@
|
|||
"""Search filter schema for constraining search results."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchFilter(BaseModel):
|
||||
"""Filter conditions for search operations.
|
||||
|
||||
All specified conditions are combined with AND logic.
|
||||
Within paths/exclude_paths, items are combined with OR logic.
|
||||
Within tags, items are combined with AND logic (all must match).
|
||||
"""
|
||||
|
||||
paths: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Include only chunks whose path starts with any of these prefixes",
|
||||
)
|
||||
tags: list[str] | None = Field(default=None, description="Include only chunks containing ALL specified tags")
|
||||
exclude_paths: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Exclude chunks whose path starts with any of these prefixes",
|
||||
)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return not self.paths and not self.tags and not self.exclude_paths
|
||||
|
||||
def match(self, path: str, metadata: dict | None = None) -> bool:
|
||||
if self.paths and not any(path.startswith(p) for p in self.paths):
|
||||
return False
|
||||
if self.exclude_paths and any(path.startswith(p) for p in self.exclude_paths):
|
||||
return False
|
||||
if self.tags:
|
||||
chunk_tags = set((metadata or {}).get("tags", []))
|
||||
if not all(t in chunk_tags for t in self.tags):
|
||||
return False
|
||||
return True
|
||||
Loading…
Add table
Reference in a new issue