From ec66fcf2edf688839722ed7e9b1d9b34ff261feb Mon Sep 17 00:00:00 2001 From: Sen Huang <48879559+ployts@users.noreply.github.com> Date: Thu, 23 Apr 2026 10:59:44 +0800 Subject: [PATCH] ``` (#216) feat(file-store): add Chroma and SQLite file store implementations - Add ChromaFileStore and SqliteFileStore classes to support additional storage backends - Export new store classes in __init__.py module refactor(file-store): enhance BaseFileStore with hybrid search capabilities - Add keyword scoring utility with word-match ratio and phrase bonus - Implement hybrid search method that combines vector and keyword results - Add merge logic for combining vector and keyword search results with weighted scoring - Move abstract methods to separate section for better organization - Remove redundant search filter parameter from local implementation refactor(file-store): simplify LocalFileStore implementation - Remove unused delete_file_chunks and upsert_chunks methods - Remove redundant update_file_metadata method - Update chunk counting logic to use file_meta directly - Simplify keyword search to use new base class scoring utility - Remove duplicate hybrid search implementation since it's now in base class ``` --- .../estimate_token_counter.py | 4 +- reme2/component/file_store/__init__.py | 4 + reme2/component/file_store/base_file_store.py | 160 +++-- .../component/file_store/chroma_file_store.py | 262 ++++++++ .../component/file_store/local_file_store.py | 131 +--- .../component/file_store/sqlite_file_store.py | 564 ++++++++++++++++++ 6 files changed, 957 insertions(+), 168 deletions(-) create mode 100644 reme2/component/file_store/chroma_file_store.py create mode 100644 reme2/component/file_store/sqlite_file_store.py diff --git a/reme2/component/as_token_counter/estimate_token_counter.py b/reme2/component/as_token_counter/estimate_token_counter.py index b3ed881b..790dca10 100644 --- a/reme2/component/as_token_counter/estimate_token_counter.py +++ b/reme2/component/as_token_counter/estimate_token_counter.py @@ -12,7 +12,9 @@ class EstimatedTokenCounter(TokenCounterBase): """ def __init__( - self, estimate_divisor: float = 4, encoding: str = "utf-8" + self, + estimate_divisor: float = 4, + encoding: str = "utf-8", ): """Initialize the estimated token counter. diff --git a/reme2/component/file_store/__init__.py b/reme2/component/file_store/__init__.py index 8fac7640..28616593 100644 --- a/reme2/component/file_store/__init__.py +++ b/reme2/component/file_store/__init__.py @@ -5,9 +5,13 @@ 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", ] diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index b1a84cfc..19f78e01 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/file_store/base_file_store.py @@ -13,8 +13,9 @@ from ...schema import FileChunk, FileMetadata, SearchFilter class BaseFileStore(BaseComponent): """Abstract base class for file storage backends. - Provides embedding resolution, validation, and safe embedding retrieval - with automatic fallback on failure. + 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. """ component_type = ComponentEnum.FILE_STORE @@ -109,41 +110,94 @@ class BaseFileStore(BaseComponent): chunk.embedding = None return chunks - @abstractmethod - async def clear_all(self): - """Clear all indexed data.""" + # -- Keyword scoring utility -------------------------------------------- - @abstractmethod - async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]): - """Insert or update a file and its chunks.""" + @staticmethod + def _score_keyword_match(query: str, text: str) -> float: + """Score a keyword match using word-match ratio + phrase bonus.""" + words = query.split() + if not words: + return 0.0 + query_lower = query.lower() + words_lower = [w.lower() for w in words] + text_lower = text.lower() + n_words = len(words) - @abstractmethod - async def delete_file(self, path: str): - """Delete a file and all its chunks.""" + match_count = sum(1 for w in words_lower if w in text_lower) + if match_count == 0: + return 0.0 - @abstractmethod - async def delete_file_chunks(self, path: str, chunk_ids: list[str]): - """Delete specific chunks for a file.""" + base_score = match_count / n_words + phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0 + return min(1.0, base_score + phrase_bonus) - @abstractmethod - async def upsert_chunks(self, chunks: list[FileChunk]): - """Insert or update specific chunks without affecting others.""" + # -- Hybrid search (concrete, delegates to abstract vector/keyword) ----- - @abstractmethod - async def list_files(self) -> list[str]: - """List all indexed file paths.""" + async def hybrid_search( + self, + query: str, + limit: int, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + search_filter: SearchFilter | None = None, + ) -> list[FileChunk]: + """Perform hybrid search combining vector and keyword results.""" + assert 0.0 <= vector_weight <= 1.0 - @abstractmethod - async def get_file_metadata(self, path: str) -> FileMetadata | None: - """Get file metadata.""" + candidates = min(200, max(1, int(limit * candidate_multiplier))) + text_weight = 1.0 - vector_weight - @abstractmethod - async def update_file_metadata(self, file_meta: FileMetadata) -> None: - """Update file metadata without affecting chunks.""" + 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) - @abstractmethod - async def get_file_chunks(self, path: str) -> list[FileChunk]: - """Get all chunks for a file.""" + if not keyword_results: + return vector_results[:limit] + if not vector_results: + return keyword_results[:limit] + + merged = self._merge_hybrid_results( + vector_results, + keyword_results, + vector_weight, + text_weight, + ) + return merged[:limit] + elif self.vector_enabled: + return await self.vector_search(query, limit, search_filter) + elif self.fts_enabled: + return await self.keyword_search(query, limit, search_filter) + return [] + + @staticmethod + def _merge_hybrid_results( + vector: list[FileChunk], + keyword: list[FileChunk], + vector_weight: float, + text_weight: float, + ) -> list[FileChunk]: + """Merge vector and keyword results with weighted scoring.""" + merged: dict[str, FileChunk] = {} + + for result in vector: + v_score = result.scores.get("vector", 0) + result.scores["score"] = v_score * vector_weight + merged[result.merge_key] = result + + for result in keyword: + key = result.merge_key + k_score = result.scores.get("keyword", 0) + if key in merged: + merged[key].scores["score"] += k_score * text_weight + else: + result.scores["score"] = k_score * text_weight + merged[key] = result + + results = list(merged.values()) + results.sort(key=lambda r: r.score, reverse=True) + return results + + # -- Filter utility ----------------------------------------------------- def _apply_filter( self, @@ -164,6 +218,32 @@ class BaseFileStore(BaseComponent): 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)] + # -- Abstract methods --------------------------------------------------- + + @abstractmethod + async def clear_all(self): + """Clear all indexed data.""" + + @abstractmethod + async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]): + """Insert or update a file and its chunks.""" + + @abstractmethod + async def delete_file(self, path: str): + """Delete a file and all its chunks.""" + + @abstractmethod + async def list_files(self) -> list[str]: + """List all indexed file paths.""" + + @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.""" @@ -176,25 +256,3 @@ class BaseFileStore(BaseComponent): search_filter: SearchFilter | None = None, ) -> list[FileChunk]: """Perform full-text/keyword search.""" - - @abstractmethod - async def hybrid_search( - self, - query: str, - limit: int, - vector_weight: float = 0.7, - candidate_multiplier: float = 3.0, - search_filter: SearchFilter | None = None, - ) -> list[FileChunk]: - """Perform hybrid search combining vector and keyword results. - - Args: - query: Search query text. - limit: Maximum number of results. - vector_weight: Weight for vector scores (0.0-1.0). - candidate_multiplier: Multiplier for candidate pool size. - search_filter: Optional filter for paths/tags. - - Returns: - FileChunk list with score populated, sorted by relevance. - """ diff --git a/reme2/component/file_store/chroma_file_store.py b/reme2/component/file_store/chroma_file_store.py new file mode 100644 index 00000000..cb709be2 --- /dev/null +++ b/reme2/component/file_store/chroma_file_store.py @@ -0,0 +1,262 @@ +"""ChromaDB storage backend for file store.""" + +import json +import time +from pathlib import Path + +from .base_file_store import BaseFileStore +from ..component_registry import R +from ...schema import FileChunk, FileMetadata, SearchFilter + +try: + import chromadb + from chromadb.config import Settings + + _CHROMADB_IMPORT_ERROR: Exception | None = None +except Exception as e: + _CHROMADB_IMPORT_ERROR = e + chromadb = None + Settings = None + + +@R.register("chroma") +class ChromaFileStore(BaseFileStore): + """ChromaDB file 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. + """ + + def __init__(self, **kwargs): + if _CHROMADB_IMPORT_ERROR is not None: + raise _CHROMADB_IMPORT_ERROR + super().__init__(**kwargs) + self.client: "chromadb.ClientAPI | None" = None + self.chunks_collection: "chromadb.Collection | None" = None + 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(exclude={"content"}, 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: + self.client = chromadb.PersistentClient( + path=str(self.db_path), + settings=Settings(allow_reset=True, anonymized_telemetry=False), + ) + self.chunks_collection = self.client.get_or_create_collection( + 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}", + ) + await super()._start(app_context) + + async def _close(self) -> None: + await self._save_metadata() + self.client = None + self.chunks_collection = None + await super()._close() + + # -- 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) + + if chunks: + chunks = await self.get_chunk_embeddings(chunks) + + ids, documents, embeddings, metadatas = [], [], [], [] + now = int(time.time() * 1000) + for chunk in chunks: + ids.append(chunk.id) + documents.append(chunk.text) + embeddings.append(chunk.embedding 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, + }) + + self.chunks_collection.upsert( + ids=ids, + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ) + + file_meta.chunk_count = len(chunks) + if file_meta.path: + self._metadata_cache[file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + metadata=file_meta.metadata, + ) + + async def delete_file(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) + + # -- Search helpers ----------------------------------------------------- + + def _chunk_from_chroma(self, chunk_id: str, md: dict, text: str, embedding=None) -> FileChunk: + return FileChunk( + id=chunk_id, + path=md["path"], + start_line=md["start_line"], + end_line=md["end_line"], + text=text, + hash=md["hash"], + embedding=embedding, + ) + + # -- Search operations -------------------------------------------------- + + async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: + if not self.vector_enabled or not query: + return [] + + query_embedding = await self.get_embedding(query) + if not query_embedding: + return [] + + try: + results = self.chunks_collection.query( + query_embeddings=[query_embedding], + n_results=limit, + include=["documents", "metadatas", "distances"], + ) + except Exception as e: + self.logger.error(f"Vector search failed: {e}") + return [] + + chunks = [] + if results["ids"] and results["ids"][0]: + for i, cid in enumerate(results["ids"][0]): + md = results["metadatas"][0][i] + distance = results["distances"][0][i] + score = max(0.0, 1.0 - distance / 2.0) + + chunk = self._chunk_from_chroma(cid, md, results["documents"][0][i]) + 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] + + async def keyword_search( + self, + query: str, + limit: int, + search_filter: SearchFilter | None = None, + ) -> list[FileChunk]: + """Keyword search via ChromaDB $contains with case variants.""" + if not self.fts_enabled or not query: + return [] + + words = query.split() + if not words: + return [] + + # Generate case variants for case-insensitive matching + word_variants = set() + for word in words: + word_variants.add(word) + word_variants.add(word.lower()) + word_variants.add(word.capitalize()) + word_variants.add(word.upper()) + variants_list = list(word_variants) + + if len(variants_list) == 1: + where_document: dict = {"$contains": variants_list[0]} + else: + where_document = {"$or": [{"$contains": w} for w in variants_list]} + + results = self.chunks_collection.get( + where_document=where_document, + include=["documents", "metadatas"], + ) + + chunks = [] + for i, cid in enumerate(results["ids"]): + md = results["metadatas"][i] + text = results["documents"][i] + score = self._score_keyword_match(query, text) + if score == 0.0: + continue + + chunk = self._chunk_from_chroma(cid, md, text) + 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] + + # -- Clear -------------------------------------------------------------- + + async def clear_all(self) -> None: + self.client.delete_collection(name=self.collection_name) + self.chunks_collection = self.client.get_or_create_collection( + name=self.collection_name, + metadata={"hnsw:space": "cosine"}, + ) + self._metadata_cache = {} + await self._save_metadata() + self.logger.info(f"Cleared all data from ChromaFileStore '{self.store_name}'") diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index fa410006..98b785b8 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/file_store/local_file_store.py @@ -116,14 +116,16 @@ class LocalFileStore(BaseFileStore): for chunk in chunks: self._chunks[chunk.id] = chunk - self._files[file_meta.path] = FileMetadata( - hash=file_meta.hash, - mtime_ms=file_meta.mtime_ms, - size=file_meta.size, - path=file_meta.path, - chunk_count=len(chunks), - metadata=file_meta.metadata, - ) + file_meta.chunk_count = len(chunks) + if file_meta.path: + self._files[file_meta.path] = FileMetadata( + hash=file_meta.hash, + mtime_ms=file_meta.mtime_ms, + size=file_meta.size, + path=file_meta.path, + chunk_count=file_meta.chunk_count, + metadata=file_meta.metadata, + ) async def delete_file(self, path: str) -> None: """Delete a file and all its chunks.""" @@ -132,23 +134,6 @@ class LocalFileStore(BaseFileStore): del self._chunks[cid] self._files.pop(path, None) - async def delete_file_chunks(self, path: str, chunk_ids: list[str]) -> None: - """Delete specific chunks for a file.""" - if not chunk_ids: - return - for cid in chunk_ids: - self._chunks.pop(cid, None) - if path in self._files: - self._files[path].chunk_count = sum(1 for chunk in self._chunks.values() if chunk.path == path) - - async def upsert_chunks(self, chunks: list[FileChunk]) -> None: - """Insert or update specific chunks without affecting others.""" - if not chunks: - return - chunks = await self.get_chunk_embeddings(chunks) - for chunk in chunks: - self._chunks[chunk.id] = chunk - # -- Read operations ---------------------------------------------------- async def list_files(self) -> list[str]: @@ -156,27 +141,16 @@ class LocalFileStore(BaseFileStore): return list(self._files.keys()) async def get_file_metadata(self, path: str) -> FileMetadata | None: - """Get file metadata.""" + """Get metadata for a specific file.""" return self._files.get(path) - async def update_file_metadata(self, file_meta: FileMetadata) -> None: - """Update file metadata without affecting chunks.""" - self._files[file_meta.path] = FileMetadata( - hash=file_meta.hash, - mtime_ms=file_meta.mtime_ms, - size=file_meta.size, - path=file_meta.path, - chunk_count=file_meta.chunk_count, - metadata=file_meta.metadata, - ) - async def get_file_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) return chunks - # -- Search ------------------------------------------------------------- + # -- Search operations -------------------------------------------------- async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: """Cosine-similarity vector search over in-memory embeddings.""" @@ -190,7 +164,6 @@ class LocalFileStore(BaseFileStore): candidates = self._apply_filter( [c for c in self._chunks.values() if c.embedding], search_filter, - self._files, ) if not candidates: return [] @@ -242,23 +215,13 @@ class LocalFileStore(BaseFileStore): if not words: return [] - query_lower = query.lower() - words_lower = [w.lower() for w in words] - n_words = len(words) - - filtered_chunks = self._apply_filter(list(self._chunks.values()), search_filter, self._files) + filtered_chunks = self._apply_filter(list(self._chunks.values()), search_filter) results = [] for chunk in filtered_chunks: - text_lower = chunk.text.lower() - match_count = sum(1 for w in words_lower if w in text_lower) - if match_count == 0: + score = self._score_keyword_match(query, chunk.text) + if score == 0.0: continue - - base_score = match_count / n_words - phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0 - score = min(1.0, base_score + phrase_bonus) - results.append( FileChunk( id=chunk.id, @@ -274,70 +237,6 @@ class LocalFileStore(BaseFileStore): results.sort(key=lambda r: r.score, reverse=True) return results[:limit] - async def hybrid_search( - self, - query: str, - limit: int, - vector_weight: float = 0.7, - candidate_multiplier: float = 3.0, - search_filter: SearchFilter | None = None, - ) -> list[FileChunk]: - """Hybrid search combining vector and keyword results.""" - assert 0.0 <= vector_weight <= 1.0 - - candidates = min(200, max(1, int(limit * candidate_multiplier))) - text_weight = 1.0 - vector_weight - - if self.vector_enabled and self.fts_enabled: - keyword_results = await self.keyword_search(query, candidates, search_filter) - vector_results = await self.vector_search(query, candidates, search_filter) - - if not keyword_results: - return vector_results[:limit] - if not vector_results: - return keyword_results[:limit] - - merged = self._merge_hybrid_results( - vector=vector_results, - keyword=keyword_results, - vector_weight=vector_weight, - text_weight=text_weight, - ) - return merged[:limit] - elif self.vector_enabled: - return await self.vector_search(query, limit, search_filter) - elif self.fts_enabled: - return await self.keyword_search(query, limit, search_filter) - return [] - - @staticmethod - def _merge_hybrid_results( - vector: list[FileChunk], - keyword: list[FileChunk], - vector_weight: float, - text_weight: float, - ) -> list[FileChunk]: - """Merge vector and keyword results with weighted scoring.""" - merged: dict[str, FileChunk] = {} - - for result in vector: - v_score = result.scores.get("vector", 0) - result.scores["score"] = v_score * vector_weight - merged[result.merge_key] = result - - for result in keyword: - key = result.merge_key - k_score = result.scores.get("keyword", 0) - if key in merged: - merged[key].scores["score"] += k_score * text_weight - else: - result.scores["score"] = k_score * text_weight - merged[key] = result - - results = list(merged.values()) - results.sort(key=lambda r: r.score, reverse=True) - return results - async def clear_all(self) -> None: """Clear all indexed data from memory and disk.""" self._chunks.clear() diff --git a/reme2/component/file_store/sqlite_file_store.py b/reme2/component/file_store/sqlite_file_store.py new file mode 100644 index 00000000..1e97a7d6 --- /dev/null +++ b/reme2/component/file_store/sqlite_file_store.py @@ -0,0 +1,564 @@ +"""SQLite storage backend for file store.""" + +import json +import struct +import time + +import sqlite3 + +from .base_file_store import BaseFileStore +from ..component_registry import R +from ...schema import FileChunk, FileMetadata, SearchFilter + + +@R.register("sqlite") +class SqliteFileStore(BaseFileStore): + """SQLite file 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 + search for short query terms. + """ + + def __init__(self, vec_ext_path: str = "", **kwargs): + super().__init__(**kwargs) + self.vec_ext_path = vec_ext_path + self.conn: sqlite3.Connection | None = None + + # -- Table names -------------------------------------------------------- + + @property + 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}" + + @property + def fts_table(self) -> str: + return f"chunks_fts_{self.store_name}" + + @staticmethod + def vector_to_blob(embedding: list[float]) -> bytes: + return struct.pack(f"{len(embedding)}f", *embedding) + + # -- Lifecycle ---------------------------------------------------------- + + async def _start(self, app_context=None) -> None: + self.conn = sqlite3.connect(self.db_path / "reme.db", check_same_thread=False) + + if self.vector_enabled: + self.conn.enable_load_extension(True) + if self.vec_ext_path: + try: + self.conn.load_extension(self.vec_ext_path) + self.logger.info(f"Loaded sqlite-vec: {self.vec_ext_path}") + except Exception as e: + self.logger.warning(f"Failed to load sqlite-vec: {e}") + self._disable_vector_search(f"extension load failed: {e}") + else: + loaded = False + try: + import sqlite_vec + + ext_path = sqlite_vec.loadable_path() + self.conn.load_extension(ext_path) + self.logger.info(f"Loaded sqlite-vec from package: {ext_path}") + loaded = True + except Exception: + pass + + if not loaded: + for name in ["vec0", "sqlite_vec", "vector0"]: + try: + self.conn.load_extension(name) + self.logger.info(f"Loaded sqlite-vec: {name}") + loaded = True + break + except Exception: + pass + + if not loaded: + self._disable_vector_search("sqlite-vec extension not available") + + 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'}", + ) + 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, + hash TEXT, + mtime REAL, + size INTEGER, + metadata TEXT, + chunk_count INTEGER DEFAULT 0 + ) + """) + + cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self.chunks_table} ( + id TEXT PRIMARY KEY, + path TEXT, + start_line INTEGER, + end_line INTEGER, + hash TEXT, + text TEXT, + embedding TEXT, + updated_at INTEGER + ) + """) + + if self.vector_enabled: + cursor.execute(f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self.vector_table} USING vec0( + id TEXT PRIMARY KEY, + embedding FLOAT[{self.embedding_dim}] + ) + """) + self.logger.info(f"Created vector table (dims={self.embedding_dim})") + + if self.fts_enabled: + cursor.execute(f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self.fts_table} USING fts5( + text, + id UNINDEXED, + path UNINDEXED, + start_line UNINDEXED, + end_line UNINDEXED, + tokenize='trigram' + ) + """) + self.logger.info("Created FTS5 table with trigram tokenizer") + + self.conn.commit() + except Exception as e: + self.logger.error(f"Failed to create tables: {e}") + raise + finally: + cursor.close() + + async def _close(self) -> None: + if self.conn: + self.conn.close() + 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, hash, mtime, size, metadata FROM {self.files_table}") + result = {} + for path, hash_val, mtime, size, meta_str in cursor.fetchall(): + metadata = json.loads(meta_str) if meta_str else {} + result[path] = FileMetadata( + hash=hash_val, + mtime_ms=mtime, + size=size, + 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: + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + + # Upsert file metadata + cursor.execute( + f"""INSERT OR REPLACE INTO {self.files_table} + (path, hash, mtime, size, metadata, chunk_count) + VALUES (?, ?, ?, ?, ?, ?)""", + ( + file_meta.path, + file_meta.hash, + file_meta.mtime_ms, + file_meta.size, + json.dumps(file_meta.metadata, ensure_ascii=False) if file_meta.metadata else None, + len(chunks), + ), + ) + + # 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,), + ).fetchall() + ] + if old_ids: + placeholders = ",".join("?" * len(old_ids)) + cursor.execute(f"DELETE FROM {self.chunks_table} WHERE id IN ({placeholders})", old_ids) + if self.vector_enabled: + 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,)) + + # Insert new chunks + if chunks: + chunks = await self.get_chunk_embeddings(chunks) + now = int(time.time() * 1000) + for chunk in chunks: + cursor.execute( + f"""INSERT INTO {self.chunks_table} + (id, path, start_line, end_line, hash, text, embedding, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", + ( + chunk.id, + file_meta.path, + chunk.start_line, + chunk.end_line, + chunk.hash, + chunk.text, + json.dumps(chunk.embedding) if chunk.embedding else None, + now, + ), + ) + + if self.vector_enabled and chunk.embedding: + cursor.execute( + f"INSERT INTO {self.vector_table} (id, embedding) VALUES (?, ?)", + (chunk.id, self.vector_to_blob(chunk.embedding)), + ) + + if self.fts_enabled: + cursor.execute( + 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), + ) + + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + self.logger.error(f"Failed to upsert file {file_meta.path}: {e}") + raise + finally: + cursor.close() + + async def delete_file(self, path: str) -> None: + cursor = self.conn.cursor() + try: + cursor.execute("BEGIN") + + chunk_ids = [ + row[0] for row in cursor.execute( + f"SELECT id FROM {self.chunks_table} WHERE path = ?", (path,), + ).fetchall() + ] + + if self.vector_enabled and chunk_ids: + for cid in chunk_ids: + cursor.execute(f"DELETE FROM {self.vector_table} WHERE id = ?", (cid,)) + + if self.fts_enabled: + 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}") + raise + finally: + cursor.close() + + # -- Read operations (SQL overrides) ------------------------------------ + + 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 hash, mtime, size, metadata, chunk_count FROM {self.files_table} WHERE path = ?", + (path,), + ) + row = cursor.fetchone() + if not row: + return None + hash_val, mtime, size, meta_str, chunk_count = row + metadata = json.loads(meta_str) if meta_str else {} + return FileMetadata( + hash=hash_val, + mtime_ms=mtime, + size=size, + path=path, + chunk_count=chunk_count, + 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]: + cursor = self.conn.cursor() + try: + cursor.execute( + f"""SELECT id, path, start_line, end_line, text, hash, embedding + FROM {self.chunks_table} WHERE path = ? + ORDER BY start_line""", + (path,), + ) + chunks = [] + for row in cursor.fetchall(): + chunk_id, path_val, start, end, text, hash_val, emb_str = row + embedding = None + if emb_str: + try: + embedding = json.loads(emb_str) + except (json.JSONDecodeError, TypeError): + pass + chunks.append(FileChunk( + id=chunk_id, + path=path_val, + start_line=start, + end_line=end, + text=text, + hash=hash_val, + embedding=embedding, + )) + return chunks + except Exception as e: + self.logger.error(f"Failed to get file chunks for {path}: {e}") + return [] + finally: + cursor.close() + + # -- Search helpers ----------------------------------------------------- + + @staticmethod + def _sanitize_fts_query(query: str) -> str: + if not query: + return "" + special_chars = list('*?:^()[]{}\'"`|+-=<>!@#$%&\\/,;') + cleaned = query + for ch in special_chars: + cleaned = cleaned.replace(ch, " ") + return " ".join(cleaned.split()) + + # -- Search operations -------------------------------------------------- + + async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: + if not self.vector_enabled or not query: + return [] + + query_embedding = await self.get_embedding(query) + if not query_embedding: + return [] + + cursor = self.conn.cursor() + try: + query_blob = self.vector_to_blob(query_embedding) + cursor.execute( + f""" + SELECT c.id, c.path, c.start_line, c.end_line, c.text, v.distance + FROM {self.vector_table} v + JOIN {self.chunks_table} c ON v.id = c.id + WHERE v.embedding MATCH ? AND k = ? + ORDER BY v.distance + """, + [query_blob, limit], + ) + + chunks = [] + for cid, path, start, end, text, dist in cursor.fetchall(): + score = max(0.0, 1.0 - dist / 2.0) + chunk = FileChunk( + id=cid, + path=path, + start_line=start, + end_line=end, + 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.sort(key=lambda c: c.score, reverse=True) + return chunks[:limit] + except Exception as e: + self.logger.error(f"Vector search failed: {e}") + return [] + finally: + cursor.close() + + async def keyword_search( + self, + query: str, + limit: int, + search_filter: SearchFilter | None = None, + ) -> list[FileChunk]: + if not self.fts_enabled or not query: + return [] + + cleaned = self._sanitize_fts_query(query) + if not cleaned: + return [] + + words = cleaned.split() + 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( + await self._like_search(cleaned, words, limit), + search_filter, + file_meta, + )[:limit] + + async def _fts_trigram_search(self, words: list[str], limit: int) -> list[FileChunk]: + escaped = [w.replace('"', '""') for w in words] + fts_query = " OR ".join(escaped) + + cursor = self.conn.cursor() + try: + cursor.execute( + f""" + SELECT fts.id, fts.path, fts.start_line, fts.end_line, fts.text, rank + FROM {self.fts_table} fts + WHERE fts.text MATCH ? + ORDER BY rank + LIMIT ? + """, + [fts_query, limit], + ) + + chunks = [] + for cid, path, start, end, text, rank in cursor.fetchall(): + score = max(0.0, 1.0 / (1.0 + abs(rank))) + chunk = FileChunk( + id=cid, + path=path, + start_line=start, + end_line=end, + 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: + self.logger.error(f"FTS trigram search failed: {e}") + return [] + finally: + cursor.close() + + async def _like_search(self, phrase: str, words: list[str], limit: int) -> list[FileChunk]: + cursor = self.conn.cursor() + try: + like_clauses = [] + params: list = [] + for word in words: + like_clauses.append("text LIKE ?") + params.append(f"%{word}%") + + where_clause = " OR ".join(like_clauses) + fetch_limit = min(limit * 3, 200) + params.append(fetch_limit) + + cursor.execute( + f""" + SELECT id, path, start_line, end_line, text + FROM {self.chunks_table} + WHERE ({where_clause}) + LIMIT ? + """, + params, + ) + + chunks = [] + for cid, path, start, end, text in cursor.fetchall(): + score = self._score_keyword_match(phrase, text) + if score == 0.0: + continue + + chunk = FileChunk( + id=cid, + path=path, + start_line=start, + end_line=end, + text=text, + hash="", + scores={"keyword": score, "score": score}, + ) + chunks.append(chunk) + + chunks.sort(key=lambda c: c.score, reverse=True) + return chunks[:limit] + except Exception as e: + self.logger.error(f"LIKE search failed: {e}") + return [] + finally: + cursor.close() + + # -- Clear -------------------------------------------------------------- + + async def clear_all(self) -> None: + 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}") + if self.fts_enabled: + cursor.execute(f"DELETE FROM {self.fts_table}") + cursor.execute("COMMIT") + except Exception as e: + cursor.execute("ROLLBACK") + self.logger.error(f"Failed to clear all data: {e}") + raise + finally: + cursor.close() + self.logger.info(f"Cleared all data from SqliteFileStore '{self.store_name}'")