From 701f8f6e3aaaa76000cc910ec591ccef734f43b3 Mon Sep 17 00:00:00 2001 From: huangsen Date: Mon, 11 May 2026 20:57:10 +0800 Subject: [PATCH] refactor(file_store): simplify file store architecture and remove unused features MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove SqliteFileStore backend and simplify the base file store interface. The file store now focuses on in-memory + JSONL persistence for the (file → chunks) graph with reduced complexity. BREAKING CHANGE: SqliteFileStore is removed, only LocalFileStore remains available. --- reme2/component/file_store/__init__.py | 18 +- reme2/component/file_store/base_file_store.py | 363 ++++-------------- .../component/file_store/local_file_store.py | 281 ++++++-------- 3 files changed, 199 insertions(+), 463 deletions(-) diff --git a/reme2/component/file_store/__init__.py b/reme2/component/file_store/__init__.py index 42299077..7739fb99 100644 --- a/reme2/component/file_store/__init__.py +++ b/reme2/component/file_store/__init__.py @@ -1,26 +1,14 @@ """File store module. -Unified storage for file metadata (frontmatter + mtime), the wikilink -graph (edges per path), and chunks (text + embeddings) with vector / -keyword / hybrid search. - -The store accepts `(node, chunks)` from the watcher's parsing pass via -`upsert(node, chunks)`, dispatching node + chunks into the appropriate -persistence pipelines (and attaching embeddings via hash-diff). - -Two backends: - LocalFileStore — pure-Python with JSONL persistence (default, - zero deps, fine for small vaults). - SqliteFileStore — SQLite + FTS5 + sqlite-vec for keyword/vector - at scale; nodes live in a relational table. +In-memory + JSONL backend for the (file → chunks) graph. Subclass +`BaseFileStore` to add other backends; only `LocalFileStore` is +shipped today. """ from .base_file_store import BaseFileStore from .local_file_store import LocalFileStore -from .sqlite_file_store import SqliteFileStore __all__ = [ "BaseFileStore", "LocalFileStore", - "SqliteFileStore", ] diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index 1a8dc019..d8331a7c 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/file_store/base_file_store.py @@ -1,41 +1,29 @@ -"""Abstract base class for file stores — minimalist engine surface. +"""Abstract base class for file stores — minimal engine surface. -A `FileStore` manages two things, no more: +The store owns persistence + search for the (file → chunks) graph. +Subclasses implement every read/write/search verb; the base only +resolves shared infrastructure: - * **graph** — `dict[path, FileNode]` in memory, JSONL-backed by default - (sqlite backend overrides for relational storage). Edges - live on `FileNode.edges`; there is no separate edge index. - * **chunks** — `dict[path, list[FileChunk]]`-shaped reads/writes plus - vector / keyword search; subclasses implement. - -Read APIs (`get_node`, `get_links`, `get_backlinks`, …) are **synchronous** -— they hit the in-memory index rebuilt from persistence on `_start`. -Writes are **async** — they touch persistence first, then update the -in-memory index, so on crash the on-disk view is the source of truth. + * `working_dir` pulled from `app_config` for path-relative work. + * `embedding_model` resolved on `_start`. + * `embed(texts)` — single async entry that wraps the model and + auto-disables vector search on persistent failure. """ from __future__ import annotations import re from abc import abstractmethod -from collections import defaultdict from pathlib import Path -from typing import Iterable, Mapping from ..base_component import BaseComponent from ..embedding import BaseEmbeddingModel from ...enumeration import ComponentEnum -from ...schema import ChunkFilter, FileChunk, FileEdge, FileNode - - -def _target_stem(raw: str) -> str: - """Stem of a raw wikilink target — `"topics/X/X.md"` → `"X"`.""" - target = raw.strip().removesuffix(".md") - return target.rsplit("/", 1)[-1] +from ...schema import FileChunk, FileNode class BaseFileStore(BaseComponent): - """Engine-level store: graph (in-memory + JSON) + chunks (db) + search.""" + """Abstract file-store engine: persistence + search for nodes & chunks.""" component_type = ComponentEnum.FILE_STORE @@ -56,9 +44,10 @@ class BaseFileStore(BaseComponent): self.store_name: str = store_name self.store_path: Path = Path(store_path) self.store_path.mkdir(parents=True, exist_ok=True) - - self.working_dir = self.app_context.app_config.working_dir if self.app_context is not None else "" - + self.working_dir: str = ( + self.app_context.app_config.working_dir + if self.app_context is not None else "" + ) self._embedding_model_name: str = embedding_model self.embedding_model: BaseEmbeddingModel | None = None @@ -67,57 +56,33 @@ class BaseFileStore(BaseComponent): if not self.vector_enabled and not self.fts_enabled: raise ValueError("At least one of embedding_model or fts_enabled must be set.") - # In-memory state (rebuilt from persistence on _start). - self._nodes: dict[str, FileNode] = {} - self._backlinks: dict[str, set[str]] = defaultdict(set) - # -- Lifecycle --------------------------------------------------------- async def _start(self) -> None: - if self._embedding_model_name: - assert self.app_context is not None, "app_context must be provided" - models = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {}) - if self._embedding_model_name not in models: - raise ValueError(f"Embedding model '{self._embedding_model_name}' not found.") - model = models[self._embedding_model_name] - if not isinstance(model, BaseEmbeddingModel): - raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}") - self.embedding_model = model - # Subclass `_start` opens persistence FIRST; this base `_start` - # runs LAST so persistence is ready by the time we iterate it. - # Subclasses should `await super()._start()` at the END of their _start. - await self._reload_nodes() + if not self._embedding_model_name: + return + assert self.app_context is not None, "app_context must be provided" + models = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {}) + if self._embedding_model_name not in models: + raise ValueError(f"Embedding model '{self._embedding_model_name}' not found.") + model = models[self._embedding_model_name] + if not isinstance(model, BaseEmbeddingModel): + raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}") + self.embedding_model = model async def _close(self) -> None: self.embedding_model = None - self._nodes.clear() - self._backlinks.clear() - async def _reload_nodes(self) -> None: - """Rebuild in-memory indices from persisted nodes. Called on _start.""" - self._nodes.clear() - self._backlinks.clear() - node_count = 0 - edge_count = 0 - for node in self._iter_persisted_nodes(): - self._index_node(node) - node_count += 1 - edge_count += len(node.edges) - self.logger.info( - f"[{self.store_name}] Loaded {node_count} nodes, " - f"{edge_count} edges into in-memory index", - ) - - # -- Embedding (single public entrypoint) ------------------------------ + # -- Embedding (shared helper) ----------------------------------------- async def embed(self, texts: list[str]) -> list[list[float] | None] | None: """Embed a batch of texts. None on disabled / API failure. Returns a list parallel to `texts` whose entries may individually be None if the embedding model dropped them. The whole call - returns None if the store has vector search disabled or the - embedding API errors out (which also auto-disables vector - search for the rest of the process). + returns None if vector search is disabled or the API errors out + (which also auto-disables vector search for the rest of the + process). """ if not self.vector_enabled or not texts or self.embedding_model is None: return None @@ -132,250 +97,58 @@ class BaseFileStore(BaseComponent): self.logger.warning(f"[{self.store_name}] Disabling vector search: {reason}") self.vector_enabled = False - # -- Graph CRUD -------------------------------------------------------- - - async def upsert_node(self, node: FileNode) -> None: - """Persist + reindex a node. Replaces any prior entry for `node.path`.""" - # await self._persist_upsert_node(node) - self._unindex_node(node.path) - self._index_node(node) - - # async def patch_node(self, path: str, **fields) -> FileNode | None: - # """Convenience: `upsert_node(existing.model_copy(update=fields))`.""" - # existing = self._nodes.get(path) - # if existing is None: - # return None - # updated = existing.model_copy(update=fields) - # await self.upsert_node(updated) - # return updated - - async def delete_node(self, path: str) -> FileNode | None: - """Persist deletion + drop from indices. Returns prior node if any.""" - await self._persist_delete_node(path) - prior = self._unindex_node(path) - return prior - - async def read_node(self, path: str) -> FileNode | None: - return self._nodes.get(path) - - # def get_edges(self, path: str) -> list[FileEdge]: - # node = self._nodes.get(path) - # return list(node.edges) if node is not None else [] - # - # @property - # def nodes(self) -> Mapping[str, FileNode]: - # return self._nodes - # - # def get_paths_by_stem(self, stem: str) -> list[str]: - # return sorted(self._stems.get(stem, set())) - - # -- Index maintenance (private) --------------------------------------- - - def _index_node(self, node: FileNode) -> None: - self._nodes[node.path] = node - for edge in node.edges: - self._backlinks[_target_stem(edge.target)].add(node.path) - - def _unindex_node(self, path: str) -> FileNode | None: - prior = self._nodes.pop(path, None) - if prior is None: - return None - stem = Path(prior.path).stem - for edge in prior.edges: - tgt = _target_stem(edge.target) - self._backlinks.get(tgt, set()).discard(prior.path) - if not self._backlinks.get(tgt): - self._backlinks.pop(tgt, None) - return prior - - # -- Link queries (resolve via reme2.utils.wikilink_resolver) ---------- - - # def get_links(self, path: str) -> list[tuple[FileNode, FileEdge]]: - # """Files `path` links TO (resolved). One pair per edge, dedup by caller.""" - # from ...utils.wikilink_resolver import resolve_wikilink - # - # node = self._nodes.get(path) - # if node is None: - # return [] - # out: list[tuple[FileNode, FileEdge]] = [] - # for edge in node.edges: - # hit = resolve_wikilink(self, edge.target) - # if hit is not None and hit in self._nodes: - # out.append((self._nodes[hit], edge)) - # return out - - # def get_backlinks(self, path: str) -> list[tuple[FileNode, FileEdge]]: - # """Files that link TO `path` (resolved). One pair per qualifying edge.""" - # from ...utils.wikilink_resolver import resolve_wikilink - # - # if path not in self._nodes: - # return [] - # stem = Path(path).stem - # out: list[tuple[FileNode, FileEdge]] = [] - # for src in self._backlinks.get(stem, set()): - # src_node = self._nodes.get(src) - # if src_node is None: - # continue - # for edge in src_node.edges: - # if resolve_wikilink(self, edge.target) == path: - # out.append((src_node, edge)) - # return out - - # -- Hot write entry (called by watcher per file change) --------------- + # -- Composite write entries (concrete, in terms of the four abstracts) - async def upsert(self, node: FileNode, chunks: list[FileChunk]) -> None: - """Single fan-out from the watcher's parse pass. - - The parser hands over a fresh node + text-only chunks; the store - owns the embedding pipeline: - 1. fetch persisted chunks for `node.path` - 2. attach cached embeddings to incoming chunks whose hash matches - 3. embed only the dirty (new-hash) subset - 4. persist node, then chunks - """ - existing = await self.get_chunks(node.path) - dirty = self._hash_diff_attach(chunks, existing) - if dirty: - await self._embed_chunks(dirty) - + """Persist a node and its chunks together. Node first, then chunks.""" await self.upsert_node(node) await self.upsert_chunks(node.path, chunks) - @staticmethod - def _hash_diff_attach( - chunks: list[FileChunk], - existing_chunks: list[FileChunk] | None, - ) -> list[FileChunk]: - """Attach cached embeddings to chunks whose hash already exists. - - Mutates input chunks in place; returns the dirty subset still - needing embeddings. - """ - if not existing_chunks: - return list(chunks) - cached = {c.hash: c.embedding for c in existing_chunks if c.embedding} - if not cached: - return list(chunks) - dirty: list[FileChunk] = [] - for c in chunks: - cached_emb = cached.get(c.hash) - if cached_emb is not None: - c.embedding = cached_emb - else: - dirty.append(c) - return dirty - - async def _embed_chunks(self, chunks: list[FileChunk]) -> None: - """In-place attach embeddings via the configured model.""" - if not chunks or self.embedding_model is None or not self.vector_enabled: - return - try: - await self.embedding_model.get_node_embeddings(chunks) - except Exception as e: - self.logger.warning(f"embedding chunks failed: {e}") - self._disable_vector_search(str(e)) - - # -- Default node persistence (single sidecar JSONL) ------------------- - - @property - def _nodes_path(self) -> Path: - return self.store_path / f"{self.store_name}_nodes.jsonl" - - async def _persist_upsert_node(self, node: FileNode) -> None: - """Default: full rewrite. Backends with relational storage override.""" - snapshot = {**self._nodes, node.path: node} - self._write_nodes_jsonl(snapshot.values()) - - async def _persist_delete_node(self, path: str) -> None: - """Default: full rewrite, omitting `path`.""" - snapshot = {p: n for p, n in self._nodes.items() if p != path} - self._write_nodes_jsonl(snapshot.values()) - - def _iter_persisted_nodes(self) -> Iterable[FileNode]: - path = self._nodes_path - if not path.exists(): - return [] - out: list[FileNode] = [] - for line in path.read_text(encoding="utf-8").splitlines(): - if not line.strip(): - continue - try: - out.append(FileNode.model_validate_json(line)) - except Exception as e: - self.logger.warning(f"Bad row in {path}: {e}") - return out - - def _write_nodes_jsonl(self, nodes: Iterable[FileNode]) -> None: - path = self._nodes_path - content = "\n".join(n.model_dump_json() for n in nodes) - tmp = path.with_suffix(".tmp") - try: - tmp.write_text(content, encoding="utf-8") - tmp.replace(path) - except Exception as e: - self.logger.error(f"Failed to write {path}: {e}") - raise - finally: - if tmp.exists(): - tmp.unlink() - - # -- Chunk APIs (subclasses implement) --------------------------------- - - @abstractmethod - async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: - """Insert or replace all chunks for a file path. Embeddings pre-attached.""" - - @abstractmethod - async def delete_chunks(self, path: str) -> None: - """Delete all chunks for a file path.""" - - @abstractmethod - async def get_chunks(self, path: str) -> list[FileChunk]: - """All chunks for a file path.""" - - @abstractmethod - async def get_chunks_by_paths(self, paths: Iterable[str]) -> list[FileChunk]: - """Batch fetch chunks across many paths.""" - - - async def upsert_node(self, node: FileNode): - ... - - async def delete_node(self, path: str): - ... - - async def get_node(self, path: str) -> FileNode: - ... - - async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: - ... - - async def delete_chunks(self, path: str) -> None: - ... - - async def get_chunks(self, path: str) -> list[FileChunk]: - ... - - async def upsert(self, node: FileNode, chunks: list[FileChunk]) -> None: - ... - async def delete(self, path: str) -> None: - ... + """Delete chunks then node — chunks-first avoids orphan rows.""" + await self.delete_chunks(path) + await self.delete_node(path) + + # -- Abstract surface (subclasses implement) --------------------------- + + @abstractmethod + async def upsert_node(self, node: FileNode) -> None: + """Persist a single node, replacing any prior entry for `node.path`.""" + + @abstractmethod + async def delete_node(self, path: str) -> None: + """Delete the node entry for `path`. Chunks are managed separately.""" + + @abstractmethod + async def read_node(self, path: str) -> FileNode | None: + """Fetch a single node by path. None if absent.""" + + @abstractmethod + async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: + """Insert or replace all chunks for `path`. + + The store owns the embedding pipeline: it should hash-diff + incoming chunks against persisted ones, reuse cached embeddings + for unchanged blocks, and only call the embedding API for new + hashes. + """ + + @abstractmethod + async def delete_chunks(self, path: str) -> None: + """Delete all chunks for `path`.""" + + @abstractmethod + async def get_chunks(self, path: str) -> list[FileChunk]: + """All chunks for `path`.""" @abstractmethod async def vector_search( - self, - query: str, - limit: int, - search_filter: dict, + self, query: str, limit: int, search_filter: dict, ) -> list[FileChunk]: """Vector similarity search.""" @abstractmethod async def keyword_search( - self, - query: str, - limit: int, - search_filter: dict, + self, query: str, limit: int, search_filter: dict, ) -> list[FileChunk]: - """Full-text/keyword search.""" + """Full-text / keyword search.""" diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index 58acd205..e1d56e57 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/file_store/local_file_store.py @@ -1,226 +1,201 @@ """Pure-Python in-memory store with on-close JSONL persistence. -Design (per project clarification): - - During runtime, ALL state lives in memory — nodes + chunks. - Per-write disk I/O is suppressed; `_persist_upsert_node` is a no-op - so each upsert costs only a dict mutation. - - On `_start`, load the JSONL sidecars into memory. - - On `_close`, flush the full in-memory snapshot back to JSONL. +Runtime model: + * All state lives in two dicts (`_nodes`, `_chunks`) — every read / + write is a dict op, no per-call I/O. + * `_start` rehydrates from `{store_name}_nodes.jsonl` and + `{store_name}_chunks.jsonl` under `store_path`. + * `_close` flushes the full in-memory snapshot back to the same + sidecar files (atomic via tmp + replace). Trade-off: lower write latency, but a hard crash drops anything since -the last clean shutdown. For larger / write-critical workloads, use -`SqliteFileStore` instead (per-write transaction). +the last clean shutdown. """ -import json +from __future__ import annotations + from pathlib import Path -from typing import Iterable import numpy as np +from pydantic import BaseModel from .base_file_store import BaseFileStore from ..component_registry import R -from ...schema import ChunkFilter, FileChunk, FileNode +from ...schema import FileChunk, FileNode from ...utils import batch_cosine_similarity -from ...utils.chunk_search import filter_chunks, keyword_score @R.register("local") class LocalFileStore(BaseFileStore): - """In-memory chunk + node store with deferred JSONL persistence.""" + """In-memory file store with deferred JSONL persistence.""" def __init__(self, encoding: str = "utf-8", **kwargs): super().__init__(**kwargs) self._encoding: str = encoding + self._nodes: dict[str, FileNode] = {} self._chunks: dict[str, FileChunk] = {} + self._nodes_file: Path = self.store_path / f"{self.store_name}_nodes.jsonl" self._chunks_file: Path = self.store_path / f"{self.store_name}_chunks.jsonl" - async d - - - # -- Persistence helpers ------------------------------------------------ - - async def _load_chunks(self) -> None: - if not self._chunks_file.exists(): - return - try: - data = self._chunks_file.read_text(encoding=self._encoding) - self._chunks = {} - for line in data.strip().split("\n"): - if not line: - continue - chunk = FileChunk.model_validate(json.loads(line)) - self._chunks[chunk.id] = chunk - except Exception as e: - self.logger.warning(f"Failed to load chunks: {e}") - - async def _save_chunks(self) -> None: - lines = [json.dumps(c.model_dump(mode="json"), ensure_ascii=False) for c in self._chunks.values()] - content = "\n".join(lines) - temp_path = self._chunks_file.with_suffix(".tmp") - try: - temp_path.write_text(content, encoding=self._encoding) - temp_path.replace(self._chunks_file) - except Exception as e: - self.logger.error(f"Failed to save chunks: {e}") - raise - finally: - if temp_path.exists(): - temp_path.unlink() - - # -- Lifecycle ---------------------------------------------------------- + # -- Lifecycle --------------------------------------------------------- async def _start(self) -> None: - await self._load_chunks() + self._load(self._nodes_file, self._nodes, FileNode, key="path") + self._load(self._chunks_file, self._chunks, FileChunk, key="id") await super()._start() - edge_count = sum(len(n.edges) for n in self._nodes.values()) self.logger.info( f"LocalFileStore '{self.store_name}' ready: " - f"{len(self._nodes)} files, {edge_count} edges, " - f"{len(self._chunks)} chunks", + f"{len(self._nodes)} nodes, {len(self._chunks)} chunks", ) async def _close(self) -> None: - # Flush full in-memory snapshot before tearing down state. - try: - await self._save_chunks() - self._write_nodes_jsonl(self._nodes.values()) - except Exception as e: - self.logger.error(f"Failed to flush LocalFileStore '{self.store_name}': {e}") + self._dump(self._nodes_file, self._nodes.values()) + self._dump(self._chunks_file, self._chunks.values()) + self._nodes.clear() self._chunks.clear() await super()._close() - # -- Persistence overrides: no-op (we flush on close) ------------------- + def _load( + self, file: Path, target: dict, model: type[BaseModel], key: str, + ) -> None: + if not file.exists(): + return + target.clear() + try: + for line in file.read_text(encoding=self._encoding).splitlines(): + if not line.strip(): + continue + obj = model.model_validate_json(line) + target[getattr(obj, key)] = obj + except Exception as e: + self.logger.warning(f"Failed to load {file}: {e}") - async def _persist_upsert_node(self, node: FileNode) -> None: - pass + def _dump(self, file: Path, items) -> None: + try: + content = "\n".join(o.model_dump_json() for o in items) + tmp = file.with_suffix(".tmp") + tmp.write_text(content, encoding=self._encoding) + tmp.replace(file) + except Exception as e: + self.logger.error(f"Failed to write {file}: {e}") - async def _persist_delete_node(self, path: str) -> None: - pass + # -- Node CRUD --------------------------------------------------------- - # -- Write operations --------------------------------------------------- + async def upsert_node(self, node: FileNode) -> None: + self._nodes[node.path] = node + + async def delete_node(self, path: str) -> None: + self._nodes.pop(path, None) + + async def read_node(self, path: str) -> FileNode | None: + return self._nodes.get(path) + + # -- Chunk CRUD -------------------------------------------------------- async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: - """Insert or update a file's chunks. Chunks arrive embedded.""" + """Replace all chunks for `path`. + + Hash-diff: chunks whose `hash` matches a persisted one inherit + the cached embedding; only the new-hash subset hits the embedding + API. + """ + existing = await self.get_chunks(path) + cached = {c.hash: c.embedding for c in existing if c.embedding is not None} + await self.delete_chunks(path) if not chunks: return - for chunk in chunks: - self._chunks[chunk.id] = chunk + + needs_embed: list[FileChunk] = [] + for c in chunks: + if c.embedding is not None: + continue + cached_emb = cached.get(c.hash) + if cached_emb is not None: + c.embedding = cached_emb + elif c.text: + needs_embed.append(c) + + if needs_embed: + embeddings = await self.embed([c.text for c in needs_embed]) + if embeddings is not None: + for c, emb in zip(needs_embed, embeddings): + if emb is not None: + c.embedding = emb + + for c in chunks: + self._chunks[c.id] = c async def delete_chunks(self, path: str) -> None: - to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path] - for cid in to_delete: + stale = [cid for cid, c in self._chunks.items() if c.path == path] + for cid in stale: del self._chunks[cid] - # -- Read operations ---------------------------------------------------- - async def get_chunks(self, path: str) -> list[FileChunk]: - chunks = [chunk for chunk in self._chunks.values() if chunk.path == path] + chunks = [c for c in self._chunks.values() if c.path == path] chunks.sort(key=lambda c: c.start_line) return chunks - async def get_chunks_by_paths(self, paths: Iterable[str]) -> list[FileChunk]: - wanted = set(paths) - if not wanted: - return [] - chunks = [c for c in self._chunks.values() if c.path in wanted] - chunks.sort(key=lambda c: (c.path, c.start_line)) - return chunks - - # -- Search operations -------------------------------------------------- + # -- Search ------------------------------------------------------------ async def vector_search( - self, - query: str, - limit: int, - chunk_filter: ChunkFilter | None = None, + self, query: str, limit: int, search_filter: dict, ) -> list[FileChunk]: if not self.vector_enabled or not query: return [] - embeddings = await self.embed([query]) - query_embedding = embeddings[0] if embeddings else None - if not query_embedding: + if not embeddings or embeddings[0] is None: return [] + query_emb = embeddings[0] - candidates = filter_chunks( - [c for c in self._chunks.values() if c.embedding], - chunk_filter, - ) + # TODO: honor `search_filter` (paths / exclude_paths / tags / ...). + candidates = [c for c in self._chunks.values() if c.embedding] if not candidates: return [] - expected_dim = self.embedding_model.dimensions if self.embedding_model else len(query_embedding) - valid_embeddings = [] - for chunk in candidates: - emb = chunk.embedding - if emb is None: - continue - emb_len = len(emb) - if emb_len != expected_dim: - emb = (emb + [0.0] * (expected_dim - emb_len)) if emb_len < expected_dim else emb[:expected_dim] - valid_embeddings.append(emb) - - query_array = np.array([query_embedding]) - chunk_embeddings = np.array(valid_embeddings) - similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] - - results = [] - for chunk, sim in zip(candidates, similarities): - results.append( - FileChunk( - id=chunk.id, - path=chunk.path, - start_line=chunk.start_line, - end_line=chunk.end_line, - hash=chunk.hash, - text=chunk.text, - embedding=chunk.embedding, - scores={"vector": float(sim), "score": float(sim)}, - ), - ) + chunk_embs = np.array([c.embedding for c in candidates]) + similarities = batch_cosine_similarity( + np.array([query_emb]), chunk_embs, + )[0] + results: list[FileChunk] = [] + for c, sim in zip(candidates, similarities): + results.append(c.model_copy( + update={"scores": {"vector": float(sim), "score": float(sim)}}, + )) results.sort(key=lambda r: r.score, reverse=True) return results[:limit] async def keyword_search( - self, - query: str, - limit: int, - chunk_filter: ChunkFilter | None = None, + self, query: str, limit: int, search_filter: dict, ) -> list[FileChunk]: - if not self.fts_enabled or not query or not query.split(): + if not self.fts_enabled or not query.split(): return [] - candidates = filter_chunks(list(self._chunks.values()), chunk_filter) - - results = [] - for chunk in candidates: - score = keyword_score(query, chunk.text) - if score == 0.0: - continue - results.append( - FileChunk( - id=chunk.id, - path=chunk.path, - start_line=chunk.start_line, - end_line=chunk.end_line, - hash=chunk.hash, - text=chunk.text, - scores={"keyword": score, "score": score}, - ), - ) - + # TODO: honor `search_filter` (paths / exclude_paths / tags / ...). + results: list[FileChunk] = [] + for c in self._chunks.values(): + score = self._keyword_score(query, c.text) + if score > 0: + results.append(c.model_copy( + update={"scores": {"keyword": score, "score": score}}, + )) results.sort(key=lambda r: r.score, reverse=True) return results[:limit] - async def clear_all(self) -> None: - """Clear all in-memory state and on-disk JSONL sidecars.""" - self._chunks.clear() - self._nodes.clear() - self._stems.clear() - self._backlinks.clear() - await self._save_chunks() - self._write_nodes_jsonl([]) - self.logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'") + # -- Helpers ----------------------------------------------------------- + + @staticmethod + def _keyword_score(query: str, text: str) -> float: + """Word-overlap score with phrase bonus. Range [0, 1].""" + words = query.split() + if not words or not text: + return 0.0 + text_lower = text.lower() + matches = sum(1 for w in words if w.lower() in text_lower) + if matches == 0: + return 0.0 + base = matches / len(words) + if len(words) > 1 and query.lower() in text_lower: + base = min(1.0, base + 0.2) + return base