```
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled

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:
huangsen 2026-04-28 11:26:32 +08:00
parent 98632dfc24
commit c2e43c398b
16 changed files with 308 additions and 454 deletions

View file

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

View file

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

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -26,7 +26,7 @@ class ComponentEnum(str, Enum):
FILE_PARSER = "file_parser"
FILE_STORE = "file_store"
CHUNK_STORE = "chunk_store"
FILE_WATCHER = "file_watcher"

View file

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

View file

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

View 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

View file

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

View file

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