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
```
This commit is contained in:
Sen Huang 2026-04-23 10:59:44 +08:00 committed by GitHub
parent 33b3aee540
commit ec66fcf2ed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 957 additions and 168 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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