mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-17 23:51:19 +00:00
``` (#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 ```
This commit is contained in:
parent
33b3aee540
commit
ec66fcf2ed
6 changed files with 957 additions and 168 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
"""
|
||||
|
|
|
|||
262
reme2/component/file_store/chroma_file_store.py
Normal file
262
reme2/component/file_store/chroma_file_store.py
Normal 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}'")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
564
reme2/component/file_store/sqlite_file_store.py
Normal file
564
reme2/component/file_store/sqlite_file_store.py
Normal 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}'")
|
||||
Loading…
Add table
Reference in a new issue