up

reme2/component/embedding/base_embedding_model.py

@@ -1,8 +1,5 @@
-"""Base embedding model with caching and batching support."""
-
 import hashlib
 import os
-import time
 from abc import abstractmethod
 from collections import OrderedDict
 from pathlib import Path
@@ -15,13 +12,7 @@ from ...schema import EmbNode
 
 
 class BaseEmbeddingModel(BaseComponent):
-    """Abstract base class for embedding models with LRU cache.
-
-    Provides:
-        - LRU in-memory cache with disk persistence (npz)
-        - Automatic text truncation to max_input_length
-        - Batch embedding support
-    """
+    """Embedding model with LRU cache and disk persistence."""
 
     component_type = ComponentEnum.EMBEDDING_MODEL
 
@@ -34,14 +25,13 @@ class BaseEmbeddingModel(BaseComponent):
         pass_dimensions: bool = False,
         max_batch_size: int = 10,
         max_input_length: int = 8192,
-        max_cache_size: int = 2000,
+        max_cache_size: int = 5000,
         enable_cache: bool = True,
-        cache_name: str = "",
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.api_key: str = api_key or os.environ.get("EMBEDDING_API_KEY", "")
-        self.base_url: str = base_url or os.environ.get("EMBEDDING_BASE_URL", "")
+        self.api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "")
+        self.base_url = base_url or os.environ.get("EMBEDDING_BASE_URL", "")
         self.model_name = model_name
         self.dimensions = dimensions
         self.pass_dimensions = pass_dimensions
@@ -49,193 +39,128 @@ class BaseEmbeddingModel(BaseComponent):
         self.max_input_length = max_input_length
         self.max_cache_size = max_cache_size
         self.enable_cache = enable_cache
-
         self._embedding_cache: OrderedDict[str, list[float]] = OrderedDict()
-        self._cache_hits = 0
-        self._cache_misses = 0
 
-        self.working_dir = self.app_context.app_config.working_dir if self.app_context is not None else ""
-        self.cache_name: str = cache_name or self.name
-        self.cache_path: Path = Path(self.working_dir) / "embedding_cache" / f"{self.cache_name}.npz"
-
-    def clear_cache(self) -> None:
-        """Clear in-memory cache and reset statistics."""
-        self._embedding_cache.clear()
-        self._cache_hits = 0
-        self._cache_misses = 0
+    @property
+    def cache_path(self) -> Path:
+        working_dir = self.app_context.app_config.working_dir if self.app_context else ""
+        return Path(working_dir) / "embedding_cache" / f"{self.name}.npz"
 
     async def _start(self) -> None:
-        """Load cache on start."""
-        self.clear_cache()
+        self._embedding_cache.clear()
         self._load_cache()
 
     async def _close(self) -> None:
-        """Save cache on close."""
         self._save_cache()
 
-    def _validate_and_adjust_embedding(self, embedding: list[float]) -> list[float]:
-        """Adjust embedding dimensions to match expected dimensions."""
-        actual_len = len(embedding)
-        if actual_len == self.dimensions:
-            return embedding
-        if actual_len < self.dimensions:
-            self.logger.warning(f"Embedding dim {actual_len} < expected {self.dimensions}, padding")
-            return embedding + [0.0] * (self.dimensions - actual_len)
-        self.logger.warning(f"Embedding dim {actual_len} > expected {self.dimensions}, truncating")
-        return embedding[: self.dimensions]
-
     def _get_cache_key(self, text: str) -> str:
-        """Generate cache key from text + model_name + dimensions."""
         return hashlib.sha256(f"{text}|{self.model_name}|{self.dimensions}".encode()).hexdigest()
 
     def _load_cache(self) -> None:
-        """Load embedding cache from disk (npz format)."""
         if not self.enable_cache:
             return
 
         self.cache_path.parent.mkdir(parents=True, exist_ok=True)
         if not self.cache_path.exists():
-            self.logger.info(f"No cache file at {self.cache_path}, starting empty")
             return
 
-        load_start = time.time()
         try:
             data = np.load(self.cache_path)
         except Exception:
-            self.logger.exception(f"Failed to load cache from {self.cache_path}, deleting file")
             self.cache_path.unlink(missing_ok=True)
             return
 
-        loaded_count = 0
         for key, emb in zip(data["keys"], data["embeddings"]):
             emb_list = emb.tolist()
             if len(emb_list) != self.dimensions:
-                self.logger.warning(
-                    f"Cache dimension mismatch for {key}: expected {self.dimensions}, got {len(emb_list)}",
-                )
                 continue
             if len(self._embedding_cache) >= self.max_cache_size:
-                self.logger.info(f"Cache limit reached ({self.max_cache_size}), loaded {loaded_count}")
                 break
             self._embedding_cache[str(key)] = emb_list
-            loaded_count += 1
-
-        self.logger.info(f"Loaded {loaded_count} embeddings from {self.cache_path} in {time.time() - load_start:.2f}s")
 
     def _save_cache(self) -> None:
-        """Save embedding cache to disk (npz format)."""
         if not self.enable_cache or not self._embedding_cache:
             return
 
         keys, embeddings = [], []
-        for cache_key, embedding in self._embedding_cache.items():
-            keys.append(cache_key)
-            embeddings.append(embedding)
+        for k, v in self._embedding_cache.items():
+            keys.append(k)
+            embeddings.append(v)
 
         try:
             np.savez(self.cache_path, keys=np.array(keys, dtype=str), embeddings=np.array(embeddings, dtype=np.float32))
-        except Exception as e:
-            self.logger.error(f"Failed to save cache to {self.cache_path}: {e}")
-            return
-
-        self.logger.info(f"Saved {len(keys)} embeddings to {self.cache_path}")
+        except Exception:
+            pass
 
     def _get_from_cache(self, text: str) -> list[float] | None:
-        """Retrieve embedding from cache if available."""
         if not self.enable_cache:
             return None
 
-        cache_key = self._get_cache_key(text)
-        if cache_key not in self._embedding_cache:
-            self._cache_misses += 1
+        key = self._get_cache_key(text)
+        if key not in self._embedding_cache:
             return None
 
-        self._embedding_cache.move_to_end(cache_key)
-        self._cache_hits += 1
-        return self._embedding_cache[cache_key]
+        self._embedding_cache.move_to_end(key)
+        return self._embedding_cache[key]
 
     def _put_to_cache(self, text: str, embedding: list[float]) -> None:
-        """Store embedding in cache with LRU eviction."""
-        if not self.enable_cache or self.max_cache_size <= 0:
-            return
-        if len(embedding) != self.dimensions:
+        if not self.enable_cache or self.max_cache_size <= 0 or len(embedding) != self.dimensions:
             return
 
-        cache_key = self._get_cache_key(text)
-        if len(self._embedding_cache) >= self.max_cache_size and cache_key not in self._embedding_cache:
+        key = self._get_cache_key(text)
+        if len(self._embedding_cache) >= self.max_cache_size and key not in self._embedding_cache:
             self._embedding_cache.popitem(last=False)
 
-        self._embedding_cache[cache_key] = embedding
-        self._embedding_cache.move_to_end(cache_key)
-
-    def get_cache_stats(self) -> dict[str, int | float]:
-        """Return cache statistics: size, hits, misses, hit_rate."""
-        total = self._cache_hits + self._cache_misses
-        return {
-            "cache_size": len(self._embedding_cache),
-            "max_cache_size": self.max_cache_size,
-            "cache_hits": self._cache_hits,
-            "cache_misses": self._cache_misses,
-            "hit_rate": self._cache_hits / total if total > 0 else 0.0,
-        }
+        self._embedding_cache[key] = embedding
+        self._embedding_cache.move_to_end(key)
 
     @abstractmethod
-    async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]:
-        """Fetch embeddings for a batch of texts. Override in subclasses."""
+    async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: ...
 
     async def get_embedding(self, input_text: str, **kwargs) -> list[float] | None:
-        """Get embedding for a single text with cache. Returns None on failure."""
         results = await self.get_embeddings([input_text], **kwargs)
         return results[0] if results else None
 
     async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]:
-        """Get embeddings for multiple texts with cache and batching."""
-        # TODO change to bytes instead of str
-        truncated_texts = [t[: self.max_input_length] for t in input_text]
-        results: list[list[float] | None] = [None] * len(truncated_texts)
-        texts_to_compute: list[tuple[int, str]] = []
+        truncated = [t[: self.max_input_length] for t in input_text]
+        results: list[list[float] | None] = [None] * len(truncated)
+        to_compute: list[tuple[int, str]] = []
 
-        for idx, text in enumerate(truncated_texts):
+        for idx, text in enumerate(truncated):
             cached = self._get_from_cache(text)
             if cached is not None:
                 results[idx] = cached
             else:
-                texts_to_compute.append((idx, text))
+                to_compute.append((idx, text))
 
-        if texts_to_compute:
-            uncached_texts = [text for _, text in texts_to_compute]
-            for i in range(0, len(uncached_texts), self.max_batch_size):
-                batch = texts_to_compute[i: i + self.max_batch_size]
-                batch_indices = [idx for idx, _ in batch]
-                batch_texts = [text for _, text in batch]
+        if to_compute:
+            for i in range(0, len(to_compute), self.max_batch_size):
+                batch = to_compute[i: i + self.max_batch_size]
+                indices = [idx for idx, _ in batch]
+                texts = [text for _, text in batch]
 
                 try:
-                    batch_embeddings = await self._get_embeddings(batch_texts, **kwargs)
-                    if batch_embeddings and len(batch_embeddings) == len(batch_texts):
-                        for orig_idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
-                            adjusted = self._validate_and_adjust_embedding(embedding)
-                            results[orig_idx] = adjusted
-                            self._put_to_cache(text, adjusted)
-                    else:
-                        self.logger.warning(
-                            f"Batch returned {len(batch_embeddings) if batch_embeddings else 0} "
-                            f"results for {len(batch_texts)} inputs",
-                        )
-                except Exception as e:
-                    self.logger.error(f"Model {self.model_name} batch failed: {e}")
+                    embeddings = await self._get_embeddings(texts, **kwargs)
+                    if embeddings and len(embeddings) == len(texts):
+                        for orig_idx, text, emb in zip(indices, texts, embeddings):
+                            if emb is None:
+                                continue
+                            if len(emb) != self.dimensions:
+                                if len(emb) < self.dimensions:
+                                    emb = emb + [0.0] * (self.dimensions - len(emb))
+                                else:
+                                    emb = emb[: self.dimensions]
+                            results[orig_idx] = emb
+                            self._put_to_cache(text, emb)
+                except Exception:
+                    pass
 
         return results
 
     async def get_node_embeddings(self, nodes: list[EmbNode], **kwargs) -> list[EmbNode]:
-        """Get embeddings for a list of nodes and assign to node.embedding."""
-        texts = [node.text for node in nodes]
-        embeddings = await self.get_embeddings(texts, **kwargs)
+        embeddings = await self.get_embeddings([n.text for n in nodes], **kwargs)
         if len(embeddings) == len(nodes):
             for node, vec in zip(nodes, embeddings):
                 if vec is not None:
                     node.embedding = vec
-                else:
-                    self.logger.warning(f"Embedding failed for node, skipping assignment")
-        else:
-            self.logger.warning(f"Mismatch: {len(embeddings)} vectors for {len(nodes)} nodes, skipping assignment")
-        return nodes
+        return nodes

reme2/component/embedding/openai_embedding_model.py

@@ -27,7 +27,11 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel):
         if self._client is None:
             raise RuntimeError("Client not initialized. Call _start() first.")
 
-        create_kwargs: dict = {"model": self.model_name, "input": input_text, **kwargs}
+        create_kwargs: dict = {
+            "model": self.model_name,
+            "input": input_text,
+            **kwargs,
+        }
         if self.pass_dimensions:
             create_kwargs["dimensions"] = self.dimensions
 
@@ -35,8 +39,8 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel):
 
         result: list[list[float] | None] = [None] * len(input_text)
         for emb in completion.data:
-            vec = getattr(emb, "embedding", None) or getattr(emb, "dense_embedding", None)
             if 0 <= emb.index < len(input_text):
+                vec = emb.embedding or getattr(emb, "dense_embedding", None)
                 if vec is not None:
                     result[emb.index] = list(vec)
                 else:

reme2/component/file_store/base_file_store.py

@@ -1,9 +1,10 @@
-"""Abstract base class for file stores — persistence + search for nodes & chunks."""
+"""Abstract base class for file stores."""
 
 import re
 from abc import abstractmethod
 from pathlib import Path
 
+from .bm25_lite import BM25Lite
 from ..base_component import BaseComponent
 from ..embedding import BaseEmbeddingModel
 from ...enumeration import ComponentEnum
@@ -11,7 +12,7 @@ from ...schema import FileChunk, FileNode
 
 
 class BaseFileStore(BaseComponent):
-    """Abstract file-store engine: persistence + search for nodes & chunks."""
+    """Abstract file-store engine."""
 
     component_type = ComponentEnum.FILE_STORE
 
@@ -19,14 +20,17 @@ class BaseFileStore(BaseComponent):
             self,
             store_name: str,
             embedding_model: str = "default",
+            tokenizer: str = "default",
             fts_enabled: bool = True,
             **kwargs,
     ):
+        """Initialize the file store."""
         super().__init__(**kwargs)
         if not re.match(r"^[a-zA-Z0-9_]+$", store_name):
             raise ValueError(f"Invalid store name '{store_name}'. Only alphanumeric and underscores allowed.")
         self.store_name = store_name or self.name
         self._embedding_model_name = embedding_model
+        self.tokenizer = tokenizer
         self.fts_enabled = fts_enabled
 
         self.embedding_model: BaseEmbeddingModel | None = None
@@ -37,9 +41,10 @@ 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.")
 
-    # Lifecycle
+        self.bm25: BM25Lite | None = None
 
     async def _start(self) -> None:
+        """Start the file store."""
         if self.vector_enabled and self.app_context is not None:
             model_dict = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {})
             if self._embedding_model_name not in model_dict:
@@ -49,66 +54,32 @@ class BaseFileStore(BaseComponent):
                 raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}")
             self.embedding_model = model
 
+        if self.fts_enabled:
+            self.bm25 = BM25Lite(index_dir=self.store_path, tokenizer=self.tokenizer, app_context=self.app_context)
+            if self.bm25 is not None:
+                await self.bm25.start()
+
     async def _close(self) -> None:
+        """Close the file store and release resources."""
         self.embedding_model = None
+        if self.fts_enabled and self.bm25 is not None:
+            await self.bm25.close()
 
     # Composite operations
 
-    async def upsert(self, node: FileNode, chunks: list[FileChunk]) -> None:
-        await self.upsert_node(node)
-        await self.upsert_chunks(node.path, chunks)
+    async def upsert_file(self, node: FileNode, chunks: list[FileChunk]) -> None:
+        """Upsert a node and its associated chunks."""
 
     async def delete(self, path: str) -> None:
-        await self.delete_chunks(path)
-        await self.delete_node(path)
+        """Delete a node and all its associated chunks by path."""
 
-    # Abstract: node operations
-
-    @abstractmethod
-    async def upsert_node(self, node: FileNode) -> None:
-        """Persist a node, replacing any prior entry for `node.path`."""
-
-    @abstractmethod
-    async def get_node(self, path: str) -> FileNode | None:
-        """Fetch a node by path, or None if absent."""
-
-    @abstractmethod
-    async def delete_node(self, path: str) -> None:
-        """Delete the node entry for `path`."""
-
-    # Abstract: chunk operations
-
-    @abstractmethod
-    async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None:
-        """Insert or replace all chunks for `path`. Handles embedding internally."""
-
-    @abstractmethod
-    async def get_chunks(self, path: str) -> list[FileChunk]:
-        """All chunks for `path`."""
-
-    @abstractmethod
-    async def delete_chunks(self, path: str) -> None:
-        """Delete all chunks for `path`."""
-
-    # Abstract: search
+    async def reindex(self) -> None:
+        """Re-index all nodes and chunks in the store."""
 
     @abstractmethod
     async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
-        """Vector similarity search."""
+        """Perform vector similarity search."""
 
     @abstractmethod
     async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
-        """Full-text / keyword search."""
-
-    # Internal helpers
-
-    async def get_embeddings(self, texts: list[str]) -> list[list[float] | None] | None:
-        """Embed texts. Returns None if vector search disabled or on API error."""
-        if not self.vector_enabled or not texts or not self.embedding_model:
-            return None
-        try:
-            return await self.embedding_model.get_embeddings(texts)
-        except Exception as e:
-            self.logger.warning(f"[{self.store_name}] Disabling vector search: {e}")
-            self.vector_enabled = False
-            return None
+        """Perform full-text keyword search."""

reme2/component/file_store/bm25_lite.py

@@ -212,12 +212,12 @@ class BM25Lite(BaseComponent):
         self._idf_cache[token_id] = math.log(1 + (self.n_docs - df + 0.5) / (df + 0.5)) if df else 0.0
         return self._idf_cache[token_id]
 
-    def retrieve(self, query: str, k: int = 3) -> dict[str, float]:
+    def retrieve(self, query: str, limit: int = 3) -> dict[str, float]:
         """Search for documents matching the query.
 
         Args:
             query: Search query string.
-            k: Maximum number of results to return.
+            limit: Maximum number of results to return.
 
         Returns:
             Dict mapping document IDs to BM25 scores, sorted by score descending.
@@ -238,7 +238,7 @@ class BM25Lite(BaseComponent):
                 tf_score = tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * doc_len / avg_len))
                 scores[doc_id] = scores.get(doc_id, 0.0) + idf * tf_score
 
-        return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:k]) if scores else {}
+        return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:limit]) if scores else {}
 
     def dump(self):
         """Persist index to disk via pickle."""

reme2/component/file_store/local_file_store.py

@@ -2,13 +2,16 @@
 
 from pathlib import Path
 
+import aiofiles
 import numpy as np
 from pydantic import BaseModel
 
 from .base_file_store import BaseFileStore
 from ..component_registry import R
 from ...schema import FileChunk, FileNode
-from ...utils import batch_cosine_similarity
+from ...utils import batch_cosine_similarity, get_logger
+
+logger = get_logger()
 
 
 @R.register("local")
@@ -26,60 +29,62 @@ class LocalFileStore(BaseFileStore):
     # Lifecycle
 
     async def _start(self) -> None:
-        self._load(self._nodes_file, self._nodes, FileNode, "path")
-        self._load(self._chunks_file, self._chunks, FileChunk, "id")
         await super()._start()
-        self.logger.info(
-            f"LocalFileStore '{self.store_name}' ready: "
-            f"{len(self._nodes)} nodes, {len(self._chunks)} chunks"
-        )
+        await self._load(self._nodes_file, self._nodes, FileNode, "path")
+        await self._load(self._chunks_file, self._chunks, FileChunk, "id")
+        self.logger.info(f"LocalFileStore '{self.store_name}' ready: "
+                         f"{len(self._nodes)} nodes, {len(self._chunks)} chunks")
 
     async def _close(self) -> None:
-        self._dump(self._nodes_file, self._nodes.values())
-        self._dump(self._chunks_file, self._chunks.values())
+        await self._dump(self._nodes_file, list(self._nodes.values()))
+        await self._dump(self._chunks_file, list(self._chunks.values()))
         self._nodes.clear()
         self._chunks.clear()
         await super()._close()
 
-    def _load(self, file: Path, target: dict, model: type[BaseModel], key: str) -> None:
+    async 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 line.strip():
-                    obj = model.model_validate_json(line)
-                    target[getattr(obj, key)] = obj
+            async with aiofiles.open(file, encoding=self._encoding) as f:
+                async for line in f:
+                    line = line.strip()
+                    if line:
+                        obj = model.model_validate_json(line)
+                        target[getattr(obj, key)] = obj
         except Exception as e:
-            self.logger.warning(f"Failed to load {file}: {e}")
+            self.logger.exception(f"Failed to load {file}: {e}")
 
-    def _dump(self, file: Path, items) -> None:
+    async def _dump(self, file: Path, items: list[BaseModel]) -> 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)
+            async with aiofiles.open(tmp, "w", encoding=self._encoding) as f:
+                await f.write(content)
             tmp.replace(file)
         except Exception as e:
-            self.logger.error(f"Failed to write {file}: {e}")
+            self.logger.exception(f"Failed to write {file}: {e}")
 
     # Node operations
 
     async def upsert_node(self, node: FileNode) -> None:
         self._nodes[node.path] = node
 
-    async def get_node(self, path: str) -> FileNode | None:
+    async def get_node_by_path(self, path: str) -> FileNode | None:
         return self._nodes.get(path)
 
-    async def delete_node(self, path: str) -> None:
-        self._nodes.pop(path, None)
+    async def delete_node_by_path(self, path: str) -> FileNode | None:
+        return self._nodes.pop(path, None)
 
     # Chunk operations
 
-    async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None:
-        existing = await self.get_chunks(path)
-        cached = {c.hash: c.embedding for c in existing if c.embedding}
+    async def upsert_chunks_with_path(self, path: str, chunks: list[FileChunk]) -> None:
+        existing = await self.get_chunks_by_path(path)
+        cached = {c.id: c.embedding for c in existing if c.embedding}
 
-        await self.delete_chunks(path)
+        await self.delete_chunks_by_path(path)
         if not chunks:
             return
 
@@ -87,8 +92,8 @@ class LocalFileStore(BaseFileStore):
         for c in chunks:
             if c.embedding:
                 continue
-            if c.hash in cached:
-                c.embedding = cached[c.hash]
+            if c.id in cached:
+                c.embedding = cached[c.id]
             elif c.text:
                 needs_embed.append(c)
 
@@ -102,33 +107,35 @@ class LocalFileStore(BaseFileStore):
         for c in chunks:
             self._chunks[c.id] = c
 
-    async def get_chunks(self, path: str) -> list[FileChunk]:
+    async def get_chunks_by_path(self, path: str) -> list[FileChunk]:
         chunks = [c for c in self._chunks.values() if c.path == path]
         chunks.sort(key=lambda c: c.start_line)
         return chunks
 
-    async def delete_chunks(self, path: str) -> None:
+    async def delete_chunks_by_path(self, path: str) -> None:
         stale = [cid for cid, c in self._chunks.items() if c.path == path]
         for cid in stale:
             del self._chunks[cid]
 
     # Search
 
-    async def vector_search(
-            self, query: str, limit: int, search_filter: dict,
-    ) -> list[FileChunk]:
-        if not self.vector_enabled or not query:
+    async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
+        if self.embedding_model is None or not query:
             return []
-        embeddings = await self.get_embeddings([query])
-        if not embeddings or not embeddings[0]:
+
+        query_embedding = await self.embedding_model.get_embedding(query)
+        if not query_embedding:
             return []
 
         candidates = [c for c in self._chunks.values() if c.embedding]
+        emb_missing = [c for c in self._chunks.values() if not c.embedding]
+        if emb_missing:
+            logger.warning(f"Embedding missing for {len(emb_missing)} chunks")
         if not candidates:
             return []
 
         chunk_embs = np.array([c.embedding for c in candidates])
-        similarities = batch_cosine_similarity(np.array([embeddings[0]]), chunk_embs)[0]
+        similarities = batch_cosine_similarity(np.array([query_embedding]), chunk_embs)[0]
 
         results = [
             c.model_copy(update={"scores": {"vector": float(s), "score": float(s)}})
@@ -137,33 +144,32 @@ class LocalFileStore(BaseFileStore):
         results.sort(key=lambda r: r.score, reverse=True)
         return results[:limit]
 
-    async def keyword_search(
-            self, query: str, limit: int, search_filter: dict,
-    ) -> list[FileChunk]:
-        if not self.fts_enabled or not query.split():
+    async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
+        if not self.fts_enabled or self.bm25 is None:
             return []
 
-        results = [
-            c.model_copy(update={"scores": {"keyword": s, "score": s}})
-            for c in self._chunks.values()
-            if (s := self._keyword_score(query, c.text)) > 0
-        ]
-        results.sort(key=lambda r: r.score, reverse=True)
-        return results[:limit]
+        query = query.strip()
+        if not query:
+            return []
 
-    # Helpers
+        doc_id_score_dict = self.bm25.retrieve(query, limit=limit)
+        results = []
+        for doc_id, score in doc_id_score_dict.items():
+            chunk = self._chunks.get(doc_id)
+            if chunk:
+                results.append(chunk.model_copy(update={"scores": {"keyword": score, "score": score}}))
+
+        return results
+
+    # Reindex
+
+    async def reindex_vector(self) -> None:
+        if not self.vector_enabled or self.embedding_model is None:
+            return
+
+        await self.embedding_model.get_node_embeddings(list(self._chunks.values()))
+
+    async def reindex_keyword(self) -> None:
+        if not self.fts_enabled or self.bm25 is None:
+            return
 
-    @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

reme2/component/file_watcher/base_file_watcher.py

@@ -1,427 +1,38 @@
-"""Base file watcher with watchfiles integration.
-
-Per change → single-step pipeline:
-  added/modified → node, chunks = parser.parse(path)
-                   file_store.upsert(node, chunks)
-  deleted        → file_store.delete_chunks + delete_node
-
-The parser owns: chunking and edge extraction (text-only, no embedding).
-The file_store owns: persistence of node + chunks AND the embedding
-pipeline — `upsert` hash-diffs incoming chunks against its persisted
-ones, reuses cached embeddings for unchanged blocks, and calls the
-embedding API only for new-hash blocks.
-The watcher owns: scheduling, cancel-on-modify, retry/timeout, and
-startup recovery (re-parsing files whose mtime drifted while offline).
-"""
-
-import asyncio
-import os
-import time
-from pathlib import Path
-
-from watchfiles import Change, awatch
-
 from ..base_component import BaseComponent
-from ..file_parser import BaseFileParser
-from ..file_store import BaseFileStore
 from ...enumeration import ComponentEnum
+from ..file_parser import BaseFileParser
 
 
 class BaseFileWatcher(BaseComponent):
-    """Watches a directory and feeds the file_store via single-step parse+upsert."""
-
     component_type = ComponentEnum.FILE_WATCHER
 
-    _META_DIR = ".reme"
-
     def __init__(
-        self,
-        recursive: bool = False,
-        debounce: int = 2000,
-        file_store: str = "default",
-        default_parser: str | None = None,
-        rebuild_index_on_start: bool = False,
-        poll_delay_ms: int = 2000,
-        parse_max_attempts: int = 3,
-        parse_retry_backoff: float = 2.0,
-        parse_task_timeout: float = 300.0,
-        **kwargs,
+            self,
+            watch_path: list[str],
+            recursive: bool = False,
+            file_parser: str = "default",
+            file_store: str = "default",
+            debounce: int = 2000,
+            poll_delay_ms: int = 2000,
+            **kwargs,
     ):
         super().__init__(**kwargs)
-        self._file_store_name: str = file_store
-        self._default_parser_name: str | None = default_parser
-        self.file_store: BaseFileStore | None = None
-        self._suffix_to_parser: dict[str, BaseFileParser] = {}
-        self._default_parser: BaseFileParser | None = None
+        self.watch_path: list[str] = watch_path
         self.recursive: bool = recursive
+        self.file_parser_name: str = file_parser
+        self.file_store: str = file_store
         self.debounce: int = debounce
-        self.rebuild_index_on_start: bool = rebuild_index_on_start
         self.poll_delay_ms: int = poll_delay_ms
 
-        self._stop_event = asyncio.Event()
-        self._watch_task: asyncio.Task | None = None
+        self.file_parser: BaseFileParser | None = None
 
-        # Parse pipeline configuration.
-        self._parse_max_attempts: int = max(1, int(parse_max_attempts))
-        self._parse_retry_backoff: float = float(parse_retry_backoff)
-        self._parse_task_timeout: float = float(parse_task_timeout)
-        self._tasks: dict[str, asyncio.Task] = {}
-        self._mgmt_lock: asyncio.Lock | None = None
-        self._run_lock: asyncio.Lock | None = None
-        self._last_failure: dict[str, float] = {}
+    async def _start(self) -> None:
+        await super()._start()
+        if self.app_context is not None:
+            file_parser_dict = self.app_context.components.get(ComponentEnum.FILE_PARSER, {})
+            if self.file_parser not in file_parser_dict:
+                raise ValueError(f"File parser {self.file_parser} not found")
+            self.file_parser = file_parser_dict[self.file_parser]
 
-    @property
-    def watch_path(self) -> str:
-        """Vault root the watcher is bound to (== file_store.working_dir).
-
-        Available after _start. Reads the live value from the bound
-        store so config has a single source of truth.
-        """
-        if self.file_store is None or self.file_store.working_dir is None:
-            raise RuntimeError(
-                "watch_path requires a started file_store with working_dir set",
-            )
-        return str(self.file_store.working_dir)
-
-    @property
-    def _meta_path(self) -> Path:
-        return (Path(self.watch_path) / self._META_DIR).resolve()
-
-    # -- Lifecycle ----------------------------------------------------------
-
-    async def _start(self):
-        """Resolve file_store + parsers, sync any disk drift, start watch loop."""
-        if self._file_store_name:
-            assert self.app_context is not None, "app_context must be provided"
-
-            stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {})
-            if self._file_store_name not in stores:
-                raise ValueError(f"File store '{self._file_store_name}' not found.")
-            store = stores[self._file_store_name]
-            if not isinstance(store, BaseFileStore):
-                raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}")
-            if store.working_dir is None:
-                raise ValueError(
-                    f"File store '{self._file_store_name}' has no working_dir; "
-                    f"set components.file_store.{self._file_store_name}.working_dir in config",
-                )
-            self.file_store = store
-
-            parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {})
-            for parser in parsers.values():
-                if isinstance(parser, BaseFileParser):
-                    for suffix in parser.suffixes:
-                        self._suffix_to_parser[suffix] = parser
-
-            if self._default_parser_name and self._default_parser_name in parsers:
-                parser = parsers[self._default_parser_name]
-                if isinstance(parser, BaseFileParser):
-                    self._default_parser = parser
-
-            if not self._suffix_to_parser and not self._default_parser:
-                self.logger.warning("No file parsers registered")
-
-        async def _initialize_and_watch():
-            await self._initial_sync_and_recovery()
-            await self._watch_loop()
-
-        self._mgmt_lock = asyncio.Lock()
-        self._run_lock = asyncio.Lock()
-        self._last_failure.clear()
-        self._stop_event.clear()
-        self._watch_task = asyncio.create_task(_initialize_and_watch())
-        self.logger.info(f"Started watching: {self.watch_path}")
-
-    async def _close(self):
-        """Stop watching and release resources."""
-        self._stop_event.set()
-        if self._watch_task and not self._watch_task.done():
-            self._watch_task.cancel()
-            try:
-                await self._watch_task
-            except asyncio.CancelledError:
-                pass
-
-        # Cancel all in-flight parse tasks before tearing down state.
-        tasks = list(self._tasks.values())
-        for t in tasks:
-            if not t.done():
-                t.cancel()
-        if tasks:
-            await asyncio.gather(*tasks, return_exceptions=True)
-        self._tasks.clear()
-        self._last_failure.clear()
-
-        self._watch_task = None
-        self._stop_event.clear()
-        self._mgmt_lock = None
-        self._run_lock = None
-        self.file_store = None
-        self._suffix_to_parser.clear()
-        self._default_parser = None
-        self.logger.info("Stopped watching")
-
-    # -- Startup sync -------------------------------------------------------
-
-    async def _initial_sync_and_recovery(self) -> None:
-        """At startup: file_store has its persisted state loaded already; we
-        diff against current disk for any changes that happened while we
-        were offline, then re-enqueue them through the parse pipeline."""
-        await self._sync_with_disk()
-
-    async def _sync_with_disk(self) -> bool:
-        """Diff cached graph (in file_store) vs disk state; reindex changes."""
-        if self.file_store is None:
-            return False
-
-        watch_path = Path(self.watch_path)
-        if not watch_path.exists():
-            self.logger.warning(f"Watch path does not exist: {watch_path}")
-            return False
-
-        on_disk: dict[str, float] = {}
-        if watch_path.is_file():
-            on_disk[str(watch_path.resolve())] = watch_path.stat().st_mtime
-        elif watch_path.is_dir():
-            iterator = watch_path.rglob("*") if self.recursive else watch_path.iterdir()
-            for fp in iterator:
-                if not fp.is_file():
-                    continue
-                abs_path = str(fp.resolve())
-                if not self._watch_filter(abs_path):
-                    continue
-                if not self._get_parser(abs_path):
-                    continue
-                try:
-                    on_disk[abs_path] = fp.stat().st_mtime
-                except OSError:
-                    continue
-
-        cached_paths = set(self.file_store.nodes.keys())
-        disk_paths = set(on_disk.keys())
-
-        added = disk_paths - cached_paths
-        deleted = cached_paths - disk_paths
-        modified: set[str] = set()
-        for p in disk_paths & cached_paths:
-            cached_node = self.file_store.get_node(p)
-            if cached_node is None or cached_node.st_mtime != on_disk[p]:
-                modified.add(p)
-
-        unchanged = len(disk_paths & cached_paths) - len(modified)
-        if not (added or deleted or modified):
-            self.logger.info(f"Index up-to-date ({unchanged} files cached, 0 changes)")
-            return False
-
-        self.logger.info(
-            f"Incremental sync: +{len(added)} ~{len(modified)} -{len(deleted)} (unchanged {unchanged})",
-        )
-
-        changes: set[tuple[Change, str]] = set()
-        for p in deleted:
-            changes.add((Change.deleted, p))
-        for p in added:
-            changes.add((Change.added, p))
-        for p in modified:
-            changes.add((Change.modified, p))
-        await self.on_changes(changes)
-        return True
-
-    # -- Watch loop ---------------------------------------------------------
-
-    async def _interruptible_sleep(self, seconds: float):
-        try:
-            await asyncio.wait_for(self._stop_event.wait(), timeout=seconds)
-        except asyncio.TimeoutError:
-            pass
-
-    async def _watch_loop(self):
-        while not self._stop_event.is_set():
-            if not Path(self.watch_path).exists():
-                self.logger.warning(f"Watch path does not exist: {self.watch_path}, waiting 10 seconds...")
-                await self._interruptible_sleep(10)
-                continue
-
-            try:
-                self.logger.info(f"Starting watch on: {self.watch_path}")
-                async for changes in awatch(
-                    self.watch_path,
-                    watch_filter=lambda _, p: self._watch_filter(p),
-                    recursive=self.recursive,
-                    debounce=self.debounce,
-                    poll_delay_ms=self.poll_delay_ms,
-                    stop_event=self._stop_event,
-                ):
-                    if self._stop_event.is_set():
-                        break
-                    await self.on_changes(changes)
-
-            except FileNotFoundError as e:
-                self.logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...")
-                if not self._stop_event.is_set():
-                    await self._interruptible_sleep(10)
-
-            except Exception as e:
-                self.logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True)
-                if not self._stop_event.is_set():
-                    await self._interruptible_sleep(10)
-
-    def _watch_filter(self, path: str) -> bool:
-        resolved = Path(path).resolve()
-        if resolved == self._meta_path:
-            return False
-        return self._meta_path not in resolved.parents
-
-    def _get_parser(self, path: str) -> BaseFileParser | None:
-        suffix = Path(path).suffix.lower()
-        return self._suffix_to_parser.get(suffix, self._default_parser)
-
-    # -- Change dispatch ----------------------------------------------------
-
-    async def on_changes(self, changes: set[tuple[Change, str]]) -> None:
-        if self.file_store is None:
-            self.logger.warning("File store not initialized, skipping changes")
-            return
-
-        for change_type, path in changes:
-            try:
-                if change_type in (Change.added, Change.modified):
-                    await self._on_modified(path)
-                elif change_type == Change.deleted:
-                    await self._on_deleted(path)
-            except FileNotFoundError:
-                self.logger.warning(f"File not found: {path}, skipping")
-            except PermissionError:
-                self.logger.warning(f"Permission denied: {path}, skipping")
-            except Exception as e:
-                self.logger.opt(exception=True).error("Error processing {p}: {err}", p=path, err=str(e))
-
-    async def _on_modified(self, path: str) -> None:
-        if not Path(path).is_file():
-            return
-        parser = self._get_parser(path)
-        if parser is None:
-            self.logger.debug(f"No parser for {path}, skipping")
-            return
-        await self._submit_parse_task(path, parser)
-
-    async def _on_deleted(self, path: str) -> None:
-        # Cancel any in-flight parse task FIRST so a delayed cancellation
-        # can't race the cleanup writes below.
-        assert self.file_store is not None, "_on_deleted requires file_store"
-        targets = [path, *self._descendant_indexed_paths(path)]
-        for p in targets:
-            await self._cancel_parse_task(p)
-            await self.file_store.delete_chunks(p)
-            await self.file_store.delete_node(p)
-        if len(targets) == 1:
-            self.logger.info(f"Deleted {path}")
-        else:
-            self.logger.info(f"Deleted directory {path} ({len(targets) - 1} indexed children)")
-
-    def _descendant_indexed_paths(self, path: str) -> list[str]:
-        """Indexed file paths that live beneath `path` as a directory.
-
-        watchfiles emits a single Change.deleted for a removed directory
-        (no per-file events), so we have to find the orphans ourselves.
-        Resolves both sides to handle symlinks and trailing-separator
-        drift between the watcher and the file_store keys.
-        """
-        assert self.file_store is not None
-        try:
-            prefix = str(Path(path).resolve()) + os.sep
-        except OSError:
-            prefix = path.rstrip(os.sep) + os.sep
-        return [p for p in self.file_store.nodes if p.startswith(prefix)]
-
-    # -- Parse task pipeline ------------------------------------------------
-
-    async def _submit_parse_task(self, path: str, parser: BaseFileParser) -> None:
-        """Cancel any prior task + spawn a new one. Atomic under _mgmt_lock."""
-        assert self._mgmt_lock is not None, "_submit_parse_task before _start()"
-        async with self._mgmt_lock:
-            await self._cancel_locked(path)
-            self._tasks[path] = asyncio.create_task(
-                self._run_parse_task(path, parser),
-            )
-
-    async def _cancel_parse_task(self, path: str) -> None:
-        """Cancel + await any in-flight parse task for `path`."""
-        assert self._mgmt_lock is not None, "_cancel_parse_task before _start()"
-        async with self._mgmt_lock:
-            await self._cancel_locked(path)
-
-    async def _cancel_locked(self, path: str) -> None:
-        """Cancel + await the task for `path`. Caller holds `_mgmt_lock`."""
-        task = self._tasks.pop(path, None)
-        if task is not None and not task.done():
-            task.cancel()
-            try:
-                await task
-            except (asyncio.CancelledError, Exception):
-                pass
-
-    async def flush_parse_tasks(self) -> None:
-        """Wait until all parse tasks finish (loops because retries spawn new ones)."""
-        while self._tasks:
-            await asyncio.gather(*list(self._tasks.values()), return_exceptions=True)
-
-    def pending_parse_count(self) -> int:
-        return sum(1 for t in self._tasks.values() if not t.done())
-
-    def failed_parse_paths(self) -> dict[str, float]:
-        return dict(self._last_failure)
-
-    async def _run_parse_task(self, path: str, parser: BaseFileParser) -> None:
-        """One per-path parse task: parse → upsert.
-
-        The store owns hash-diff + embedding: it inspects its own
-        existing chunks for `path`, reuses cached embeddings for
-        unchanged blocks, and only calls the embedding API for new-hash
-        blocks. The whole upsert is one call from this side.
-        """
-        try:
-            for attempt in range(1, self._parse_max_attempts + 1):
-                try:
-                    assert self.file_store is not None
-                    node, chunks = await asyncio.wait_for(
-                        parser.parse(path),
-                        timeout=self._parse_task_timeout,
-                    )
-
-                    assert self._run_lock is not None
-                    async with self._run_lock:
-                        await asyncio.wait_for(
-                            self.file_store.upsert(node, chunks),
-                            timeout=self._parse_task_timeout,
-                        )
-                        self._last_failure.pop(path, None)
-
-                    self.logger.info(
-                        f"indexed {path}: chunks={len(chunks)} edges={len(node.edges)}",
-                    )
-                    return
-                except asyncio.CancelledError:
-                    raise
-                except FileNotFoundError:
-                    self.logger.debug(f"parse task {path}: file gone, skipping")
-                    self._last_failure.pop(path, None)
-                    return
-                except Exception as e:
-                    self._last_failure[path] = time.time()
-                    if attempt < self._parse_max_attempts:
-                        backoff = self._parse_retry_backoff * attempt
-                        self.logger.warning(
-                            f"parse task {path} attempt {attempt} failed "
-                            f"({type(e).__name__}: {e}); retry in {backoff:.1f}s",
-                        )
-                        await asyncio.sleep(backoff)
-                        continue
-                    self.logger.error(
-                        f"parse task {path} giving up after {attempt} attempts: "
-                        f"{type(e).__name__}: {e}",
-                    )
-                    return
-        finally:
-            if self._tasks.get(path) is asyncio.current_task():
-                self._tasks.pop(path, None)
+    async def _close(self) -> None:
+        await super()._close()

reme2/component/file_watcher/base_file_watcher_old.py

@@ -0,0 +1,427 @@
+"""Base file watcher with watchfiles integration.
+
+Per change → single-step pipeline:
+  added/modified → node, chunks = parser.parse(path)
+                   file_store.upsert(node, chunks)
+  deleted        → file_store.delete_chunks + delete_node
+
+The parser owns: chunking and edge extraction (text-only, no embedding).
+The file_store owns: persistence of node + chunks AND the embedding
+pipeline — `upsert` hash-diffs incoming chunks against its persisted
+ones, reuses cached embeddings for unchanged blocks, and calls the
+embedding API only for new-hash blocks.
+The watcher owns: scheduling, cancel-on-modify, retry/timeout, and
+startup recovery (re-parsing files whose mtime drifted while offline).
+"""
+
+import asyncio
+import os
+import time
+from pathlib import Path
+
+from watchfiles import Change, awatch
+
+from ..base_component import BaseComponent
+from ..file_parser import BaseFileParser
+from ..file_store import BaseFileStore
+from ...enumeration import ComponentEnum
+
+
+class BaseFileWatcher(BaseComponent):
+    """Watches a directory and feeds the file_store via single-step parse+upsert."""
+
+    component_type = ComponentEnum.FILE_WATCHER
+
+    _META_DIR = ".reme"
+
+    def __init__(
+        self,
+        recursive: bool = False,
+        debounce: int = 2000,
+        file_store: str = "default",
+        default_parser: str | None = None,
+        rebuild_index_on_start: bool = False,
+        poll_delay_ms: int = 2000,
+        parse_max_attempts: int = 3,
+        parse_retry_backoff: float = 2.0,
+        parse_task_timeout: float = 300.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._file_store_name: str = file_store
+        self._default_parser_name: str | None = default_parser
+        self.file_store: BaseFileStore | None = None
+        self._suffix_to_parser: dict[str, BaseFileParser] = {}
+        self._default_parser: BaseFileParser | None = None
+        self.recursive: bool = recursive
+        self.debounce: int = debounce
+        self.rebuild_index_on_start: bool = rebuild_index_on_start
+        self.poll_delay_ms: int = poll_delay_ms
+
+        self._stop_event = asyncio.Event()
+        self._watch_task: asyncio.Task | None = None
+
+        # Parse pipeline configuration.
+        self._parse_max_attempts: int = max(1, int(parse_max_attempts))
+        self._parse_retry_backoff: float = float(parse_retry_backoff)
+        self._parse_task_timeout: float = float(parse_task_timeout)
+        self._tasks: dict[str, asyncio.Task] = {}
+        self._mgmt_lock: asyncio.Lock | None = None
+        self._run_lock: asyncio.Lock | None = None
+        self._last_failure: dict[str, float] = {}
+
+    @property
+    def watch_path(self) -> str:
+        """Vault root the watcher is bound to (== file_store.working_dir).
+
+        Available after _start. Reads the live value from the bound
+        store so config has a single source of truth.
+        """
+        if self.file_store is None or self.file_store.working_dir is None:
+            raise RuntimeError(
+                "watch_path requires a started file_store with working_dir set",
+            )
+        return str(self.file_store.working_dir)
+
+    @property
+    def _meta_path(self) -> Path:
+        return (Path(self.watch_path) / self._META_DIR).resolve()
+
+    # -- Lifecycle ----------------------------------------------------------
+
+    async def _start(self):
+        """Resolve file_store + parsers, sync any disk drift, start watch loop."""
+        if self._file_store_name:
+            assert self.app_context is not None, "app_context must be provided"
+
+            stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {})
+            if self._file_store_name not in stores:
+                raise ValueError(f"File store '{self._file_store_name}' not found.")
+            store = stores[self._file_store_name]
+            if not isinstance(store, BaseFileStore):
+                raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}")
+            if store.working_dir is None:
+                raise ValueError(
+                    f"File store '{self._file_store_name}' has no working_dir; "
+                    f"set components.file_store.{self._file_store_name}.working_dir in config",
+                )
+            self.file_store = store
+
+            parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {})
+            for parser in parsers.values():
+                if isinstance(parser, BaseFileParser):
+                    for suffix in parser.suffixes:
+                        self._suffix_to_parser[suffix] = parser
+
+            if self._default_parser_name and self._default_parser_name in parsers:
+                parser = parsers[self._default_parser_name]
+                if isinstance(parser, BaseFileParser):
+                    self._default_parser = parser
+
+            if not self._suffix_to_parser and not self._default_parser:
+                self.logger.warning("No file parsers registered")
+
+        async def _initialize_and_watch():
+            await self._initial_sync_and_recovery()
+            await self._watch_loop()
+
+        self._mgmt_lock = asyncio.Lock()
+        self._run_lock = asyncio.Lock()
+        self._last_failure.clear()
+        self._stop_event.clear()
+        self._watch_task = asyncio.create_task(_initialize_and_watch())
+        self.logger.info(f"Started watching: {self.watch_path}")
+
+    async def _close(self):
+        """Stop watching and release resources."""
+        self._stop_event.set()
+        if self._watch_task and not self._watch_task.done():
+            self._watch_task.cancel()
+            try:
+                await self._watch_task
+            except asyncio.CancelledError:
+                pass
+
+        # Cancel all in-flight parse tasks before tearing down state.
+        tasks = list(self._tasks.values())
+        for t in tasks:
+            if not t.done():
+                t.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        self._tasks.clear()
+        self._last_failure.clear()
+
+        self._watch_task = None
+        self._stop_event.clear()
+        self._mgmt_lock = None
+        self._run_lock = None
+        self.file_store = None
+        self._suffix_to_parser.clear()
+        self._default_parser = None
+        self.logger.info("Stopped watching")
+
+    # -- Startup sync -------------------------------------------------------
+
+    async def _initial_sync_and_recovery(self) -> None:
+        """At startup: file_store has its persisted state loaded already; we
+        diff against current disk for any changes that happened while we
+        were offline, then re-enqueue them through the parse pipeline."""
+        await self._sync_with_disk()
+
+    async def _sync_with_disk(self) -> bool:
+        """Diff cached graph (in file_store) vs disk state; reindex changes."""
+        if self.file_store is None:
+            return False
+
+        watch_path = Path(self.watch_path)
+        if not watch_path.exists():
+            self.logger.warning(f"Watch path does not exist: {watch_path}")
+            return False
+
+        on_disk: dict[str, float] = {}
+        if watch_path.is_file():
+            on_disk[str(watch_path.resolve())] = watch_path.stat().st_mtime
+        elif watch_path.is_dir():
+            iterator = watch_path.rglob("*") if self.recursive else watch_path.iterdir()
+            for fp in iterator:
+                if not fp.is_file():
+                    continue
+                abs_path = str(fp.resolve())
+                if not self._watch_filter(abs_path):
+                    continue
+                if not self._get_parser(abs_path):
+                    continue
+                try:
+                    on_disk[abs_path] = fp.stat().st_mtime
+                except OSError:
+                    continue
+
+        cached_paths = set(self.file_store.nodes.keys())
+        disk_paths = set(on_disk.keys())
+
+        added = disk_paths - cached_paths
+        deleted = cached_paths - disk_paths
+        modified: set[str] = set()
+        for p in disk_paths & cached_paths:
+            cached_node = self.file_store.get_node_by_path(p)
+            if cached_node is None or cached_node.st_mtime != on_disk[p]:
+                modified.add(p)
+
+        unchanged = len(disk_paths & cached_paths) - len(modified)
+        if not (added or deleted or modified):
+            self.logger.info(f"Index up-to-date ({unchanged} files cached, 0 changes)")
+            return False
+
+        self.logger.info(
+            f"Incremental sync: +{len(added)} ~{len(modified)} -{len(deleted)} (unchanged {unchanged})",
+        )
+
+        changes: set[tuple[Change, str]] = set()
+        for p in deleted:
+            changes.add((Change.deleted, p))
+        for p in added:
+            changes.add((Change.added, p))
+        for p in modified:
+            changes.add((Change.modified, p))
+        await self.on_changes(changes)
+        return True
+
+    # -- Watch loop ---------------------------------------------------------
+
+    async def _interruptible_sleep(self, seconds: float):
+        try:
+            await asyncio.wait_for(self._stop_event.wait(), timeout=seconds)
+        except asyncio.TimeoutError:
+            pass
+
+    async def _watch_loop(self):
+        while not self._stop_event.is_set():
+            if not Path(self.watch_path).exists():
+                self.logger.warning(f"Watch path does not exist: {self.watch_path}, waiting 10 seconds...")
+                await self._interruptible_sleep(10)
+                continue
+
+            try:
+                self.logger.info(f"Starting watch on: {self.watch_path}")
+                async for changes in awatch(
+                    self.watch_path,
+                    watch_filter=lambda _, p: self._watch_filter(p),
+                    recursive=self.recursive,
+                    debounce=self.debounce,
+                    poll_delay_ms=self.poll_delay_ms,
+                    stop_event=self._stop_event,
+                ):
+                    if self._stop_event.is_set():
+                        break
+                    await self.on_changes(changes)
+
+            except FileNotFoundError as e:
+                self.logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...")
+                if not self._stop_event.is_set():
+                    await self._interruptible_sleep(10)
+
+            except Exception as e:
+                self.logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True)
+                if not self._stop_event.is_set():
+                    await self._interruptible_sleep(10)
+
+    def _watch_filter(self, path: str) -> bool:
+        resolved = Path(path).resolve()
+        if resolved == self._meta_path:
+            return False
+        return self._meta_path not in resolved.parents
+
+    def _get_parser(self, path: str) -> BaseFileParser | None:
+        suffix = Path(path).suffix.lower()
+        return self._suffix_to_parser.get(suffix, self._default_parser)
+
+    # -- Change dispatch ----------------------------------------------------
+
+    async def on_changes(self, changes: set[tuple[Change, str]]) -> None:
+        if self.file_store is None:
+            self.logger.warning("File store not initialized, skipping changes")
+            return
+
+        for change_type, path in changes:
+            try:
+                if change_type in (Change.added, Change.modified):
+                    await self._on_modified(path)
+                elif change_type == Change.deleted:
+                    await self._on_deleted(path)
+            except FileNotFoundError:
+                self.logger.warning(f"File not found: {path}, skipping")
+            except PermissionError:
+                self.logger.warning(f"Permission denied: {path}, skipping")
+            except Exception as e:
+                self.logger.opt(exception=True).error("Error processing {p}: {err}", p=path, err=str(e))
+
+    async def _on_modified(self, path: str) -> None:
+        if not Path(path).is_file():
+            return
+        parser = self._get_parser(path)
+        if parser is None:
+            self.logger.debug(f"No parser for {path}, skipping")
+            return
+        await self._submit_parse_task(path, parser)
+
+    async def _on_deleted(self, path: str) -> None:
+        # Cancel any in-flight parse task FIRST so a delayed cancellation
+        # can't race the cleanup writes below.
+        assert self.file_store is not None, "_on_deleted requires file_store"
+        targets = [path, *self._descendant_indexed_paths(path)]
+        for p in targets:
+            await self._cancel_parse_task(p)
+            await self.file_store.delete_chunks_by_path(p)
+            await self.file_store.delete_node_by_path(p)
+        if len(targets) == 1:
+            self.logger.info(f"Deleted {path}")
+        else:
+            self.logger.info(f"Deleted directory {path} ({len(targets) - 1} indexed children)")
+
+    def _descendant_indexed_paths(self, path: str) -> list[str]:
+        """Indexed file paths that live beneath `path` as a directory.
+
+        watchfiles emits a single Change.deleted for a removed directory
+        (no per-file events), so we have to find the orphans ourselves.
+        Resolves both sides to handle symlinks and trailing-separator
+        drift between the watcher and the file_store keys.
+        """
+        assert self.file_store is not None
+        try:
+            prefix = str(Path(path).resolve()) + os.sep
+        except OSError:
+            prefix = path.rstrip(os.sep) + os.sep
+        return [p for p in self.file_store.nodes if p.startswith(prefix)]
+
+    # -- Parse task pipeline ------------------------------------------------
+
+    async def _submit_parse_task(self, path: str, parser: BaseFileParser) -> None:
+        """Cancel any prior task + spawn a new one. Atomic under _mgmt_lock."""
+        assert self._mgmt_lock is not None, "_submit_parse_task before _start()"
+        async with self._mgmt_lock:
+            await self._cancel_locked(path)
+            self._tasks[path] = asyncio.create_task(
+                self._run_parse_task(path, parser),
+            )
+
+    async def _cancel_parse_task(self, path: str) -> None:
+        """Cancel + await any in-flight parse task for `path`."""
+        assert self._mgmt_lock is not None, "_cancel_parse_task before _start()"
+        async with self._mgmt_lock:
+            await self._cancel_locked(path)
+
+    async def _cancel_locked(self, path: str) -> None:
+        """Cancel + await the task for `path`. Caller holds `_mgmt_lock`."""
+        task = self._tasks.pop(path, None)
+        if task is not None and not task.done():
+            task.cancel()
+            try:
+                await task
+            except (asyncio.CancelledError, Exception):
+                pass
+
+    async def flush_parse_tasks(self) -> None:
+        """Wait until all parse tasks finish (loops because retries spawn new ones)."""
+        while self._tasks:
+            await asyncio.gather(*list(self._tasks.values()), return_exceptions=True)
+
+    def pending_parse_count(self) -> int:
+        return sum(1 for t in self._tasks.values() if not t.done())
+
+    def failed_parse_paths(self) -> dict[str, float]:
+        return dict(self._last_failure)
+
+    async def _run_parse_task(self, path: str, parser: BaseFileParser) -> None:
+        """One per-path parse task: parse → upsert.
+
+        The store owns hash-diff + embedding: it inspects its own
+        existing chunks for `path`, reuses cached embeddings for
+        unchanged blocks, and only calls the embedding API for new-hash
+        blocks. The whole upsert is one call from this side.
+        """
+        try:
+            for attempt in range(1, self._parse_max_attempts + 1):
+                try:
+                    assert self.file_store is not None
+                    node, chunks = await asyncio.wait_for(
+                        parser.parse(path),
+                        timeout=self._parse_task_timeout,
+                    )
+
+                    assert self._run_lock is not None
+                    async with self._run_lock:
+                        await asyncio.wait_for(
+                            self.file_store.upsert(node, chunks),
+                            timeout=self._parse_task_timeout,
+                        )
+                        self._last_failure.pop(path, None)
+
+                    self.logger.info(
+                        f"indexed {path}: chunks={len(chunks)} edges={len(node.edges)}",
+                    )
+                    return
+                except asyncio.CancelledError:
+                    raise
+                except FileNotFoundError:
+                    self.logger.debug(f"parse task {path}: file gone, skipping")
+                    self._last_failure.pop(path, None)
+                    return
+                except Exception as e:
+                    self._last_failure[path] = time.time()
+                    if attempt < self._parse_max_attempts:
+                        backoff = self._parse_retry_backoff * attempt
+                        self.logger.warning(
+                            f"parse task {path} attempt {attempt} failed "
+                            f"({type(e).__name__}: {e}); retry in {backoff:.1f}s",
+                        )
+                        await asyncio.sleep(backoff)
+                        continue
+                    self.logger.error(
+                        f"parse task {path} giving up after {attempt} attempts: "
+                        f"{type(e).__name__}: {e}",
+                    )
+                    return
+        finally:
+            if self._tasks.get(path) is asyncio.current_task():
+                self._tasks.pop(path, None)

reme2/component/file_watcher/light_file_watcher.py

@@ -1,16 +1,7 @@
-"""Lightweight file watcher for Markdown files only."""
-
-from pathlib import Path
-
 from .base_file_watcher import BaseFileWatcher
 from ..component_registry import R
 
-_MD_SUFFIXES = {".md", ".markdown"}
-
 
 @R.register("light")
 class LightFileWatcher(BaseFileWatcher):
-    """Watches only Markdown files and delegates parsing to registered parsers."""
-
-    def _watch_filter(self, path: str) -> bool:
-        return super()._watch_filter(path) and Path(path).suffix.lower() in _MD_SUFFIXES
+    ...

reme2/mcp/steps/memory_retriever.py

@@ -84,7 +84,7 @@ def _serialize_chunk(chunk, file_store, extras: dict | None = None) -> dict:
     `extras` lets the caller attach step-specific fields (e.g. `graph_hop`).
     """
     item = chunk.model_dump(exclude_none=True, exclude={"embedding"})
-    node = file_store.get_node(chunk.path)
+    node = file_store.get_node_by_path(chunk.path)
     if node is not None:
         item["file_metadata"] = node.metadata
         item["file_st_mtime"] = node.st_mtime

reme2/memory/memory_io.py

@@ -62,7 +62,7 @@ async def get_file(
     On-disk frontmatter is the source of truth — the file_store cache may
     lag a write that hasn't been picked up by the watcher yet.
     """
-    node = file_store.get_node(path)
+    node = file_store.get_node_by_path(path)
     result: dict = {"path": path, "exists": False}
     if node is not None:
         result.update({
@@ -80,7 +80,7 @@ async def get_file(
         result["metadata"] = dict(post.metadata)
 
     if include_chunks:
-        chunks = await file_store.get_chunks(path)
+        chunks = await file_store.get_chunks_by_path(path)
         result["chunks"] = [c.model_dump(exclude_none=True) for c in chunks]
     return result
 

reme2/schema/file_edge.py

@@ -10,10 +10,7 @@ maintainer's job, not the schema's.
 
 ## Inline forms recognised by `parse_wikilinks`
 
-    [[X]]                           bare wikilink         → predicate=None
-    extends:: [[X]]                 line-level Dataview   → predicate="extends"
-    [extends:: [[X]]]               inline-bracketed      → predicate="extends"
-    extends:: [[A]], [[B]]          multi-target          → 2 edges
+
 """
 
 from __future__ import annotations

reme2/schema/file_node.py

@@ -2,6 +2,12 @@ from pydantic import BaseModel, Field, ConfigDict
 
 
 class FileEdge(BaseModel):
+    """ Format:
+    [[X]]                           bare wikilink         → predicate=None
+    extends:: [[X]]                 line-level Dataview   → predicate="extends"
+    [extends:: [[X]]]               inline-bracketed      → predicate="extends"
+    extends:: [[A]], [[B]]          multi-target          → 2 edges
+    """
     link: str = Field(default=...)
     predicate: str | None = Field(default=None)
 
@@ -31,4 +37,5 @@ class FileNode(BaseModel):
     path: str = Field(default=...)
     st_mtime: float = Field(default=...)
     edges: list[FileEdge] = Field(default_factory=list)
+    chunk_ids: list[str] = Field(default_factory=list)
     front_matter: FileFrontMatter = Field(default_factory=FileFrontMatter)

reme2/utils/similarity_utils.py

@@ -63,14 +63,6 @@ def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np.
 
     Raises:
         ValueError: If embedding dimensions don't match between arrays.
-
-    Examples:
-        >>> import numpy as np
-        >>> arr1 = np.array([[1.0, 0.0], [0.0, 1.0]])
-        >>> arr2 = np.array([[1.0, 0.0], [1.0, 1.0]])
-        >>> batch_cosine_similarity(arr1, arr2)
-        array([[1.        , 0.70710678],
-               [0.        , 0.70710678]])
     """
     if nd_array1.shape[1] != nd_array2.shape[1]:
         raise ValueError(

test/reme2/test_bm25_lite.py

@@ -114,7 +114,7 @@ def test_retrieve_basic():
             }
             bm25.add_docs(docs)
 
-            results = bm25.retrieve("python", k=3)
+            results = bm25.retrieve("python", limit=3)
             assert len(results) <= 3
             assert "doc1" in results or "doc3" in results
 
@@ -136,10 +136,10 @@ def test_retrieve_with_limit():
             }
             bm25.add_docs(docs)
 
-            results = bm25.retrieve("python", k=3)
+            results = bm25.retrieve("python", limit=3)
             assert len(results) == 3
 
-            results = bm25.retrieve("python", k=5)
+            results = bm25.retrieve("python", limit=5)
             assert len(results) == 5
 
             await bm25.close()
@@ -158,10 +158,10 @@ def test_retrieve_empty_query():
             docs = {"doc1": "hello world"}
             bm25.add_docs(docs)
 
-            results = bm25.retrieve("", k=3)
+            results = bm25.retrieve("", limit=3)
             assert results == {}
 
-            results = bm25.retrieve("unknownxyz", k=3)
+            results = bm25.retrieve("unknownxyz", limit=3)
             assert results == {}
 
             await bm25.close()
@@ -177,7 +177,7 @@ def test_retrieve_empty_index():
         with tempfile.TemporaryDirectory() as tmpdir:
             bm25 = await create_bm25(Path(tmpdir))
 
-            results = bm25.retrieve("python", k=3)
+            results = bm25.retrieve("python", limit=3)
             assert results == {}
 
             await bm25.close()
@@ -200,7 +200,7 @@ def test_update_doc():
             assert bm25.n_docs == 1
             assert bm25.total_len != old_len
 
-            results = bm25.retrieve("java", k=1)
+            results = bm25.retrieve("java", limit=1)
             assert "doc1" in results
 
             await bm25.close()
@@ -227,7 +227,7 @@ def test_remove_doc():
             assert bm25.n_docs == 1
             assert "doc1" not in bm25.doc_meta
 
-            results = bm25.retrieve("hello", k=2)
+            results = bm25.retrieve("hello", limit=2)
             assert "doc1" not in results
             assert "doc2" in results
 
@@ -325,7 +325,7 @@ def test_reindex_with_docs():
             assert "doc2" in bm25.doc_meta
             assert len(bm25.vocab) < len(old_vocab)
 
-            results = bm25.retrieve("hello", k=1)
+            results = bm25.retrieve("hello", limit=1)
             assert "doc2" in results
 
             await bm25.close()
@@ -362,7 +362,7 @@ def test_persistence():
             for doc_id, meta in old_doc_meta.items():
                 assert doc_id in bm25_new.doc_meta
 
-            results = bm25_new.retrieve("hello", k=2)
+            results = bm25_new.retrieve("hello", limit=2)
             assert "doc1" in results or "doc2" in results
 
             await bm25_new.close()
@@ -382,7 +382,7 @@ def test_custom_params():
             assert bm25.b == 0.5
 
             bm25.add_docs({"doc1": "test document"})
-            results = bm25.retrieve("test", k=1)
+            results = bm25.retrieve("test", limit=1)
             assert "doc1" in results
 
             await bm25.close()
@@ -405,7 +405,7 @@ def test_chinese_text():
             }
             bm25.add_docs(docs)
 
-            results = bm25.retrieve("北京", k=2)
+            results = bm25.retrieve("北京", limit=2)
             assert len(results) <= 2
             assert "doc1" in results or "doc2" in results
 
@@ -429,10 +429,10 @@ def test_mixed_chinese_english():
             }
             bm25.add_docs(docs)
 
-            results = bm25.retrieve("Python", k=3)
+            results = bm25.retrieve("Python", limit=3)
             assert len(results) > 0
 
-            results = bm25.retrieve("编程", k=2)
+            results = bm25.retrieve("编程", limit=2)
             assert len(results) > 0
 
             await bm25.close()
@@ -503,7 +503,7 @@ def test_score_ordering():
             }
             bm25.add_docs(docs)
 
-            results = bm25.retrieve("python", k=3)
+            results = bm25.retrieve("python", limit=3)
             scores = list(results.values())
 
             for i in range(len(scores) - 1):

test/test_cache_memory_usage.py

@@ -158,10 +158,7 @@ def print_usage_guidelines():
     print("   • ~420 MB per 50,000 entries")
     print("   • Scale linearly: ~8.4 MB per 1,000 entries\n")
 
-    print("3️⃣  Monitor cache statistics:")
-    print("   • Use model.get_cache_stats() to check hit rate")
-    print("   • If hit rate < 50%, consider reducing cache size")
-    print("   • If hit rate > 90%, you might benefit from larger cache\n")
+    print("3️⃣  Monitor cache size:\n")
 
     print("4️⃣  Configuration examples:\n")
 
@@ -193,10 +190,8 @@ def print_usage_guidelines():
     print("       max_cache_size=1000,  # ~8.4 MB")
     print("   )")
     print()
-    print("   # Monitor cache performance")
-    print("   stats = model.get_cache_stats()")
-    print("   print(f'Hit rate: {stats[\"hit_rate\"]:.1%}')")
-    print("   print(f'Memory used: ~{stats[\"cache_size\"] * 8.4:.1f} KB')")
+    print("   # Monitor cache")
+    print("   print(f'Cache entries: {len(model._embedding_cache)}')")
 
 
 if __name__ == "__main__":

test/test_embedding_cache.py

@@ -1,427 +0,0 @@
-"""
-Async unit tests for embedding cache functionality.
-
-Tests cover:
-- Cache hit/miss tracking
-- LRU eviction policy
-- Cache statistics
-- Performance improvements with repeated queries
-- Cache clearing
-
-Usage:
-    python test_embedding_cache.py
-"""
-
-# flake8: noqa: E402
-# pylint: disable=C0413
-
-import asyncio
-import shutil
-import tempfile
-from typing import List
-
-from reme.core.utils import load_env
-
-load_env()
-
-from reme.core.embedding import OpenAIEmbeddingModel
-
-
-def get_test_texts() -> List[str]:
-    """Create test texts for embedding cache testing."""
-    return [
-        "What is machine learning?",
-        "How does neural network work?",
-        "Explain artificial intelligence",
-        "Define deep learning",
-        "What is data science?",
-    ]
-
-
-async def test_cache_basic_functionality():
-    """Test basic cache hit/miss functionality."""
-    print(f"\n{'='*60}")
-    print("Test 1: Basic Cache Functionality")
-    print(f"{'='*60}")
-
-    temp_dir = tempfile.mkdtemp()
-    try:
-        model = OpenAIEmbeddingModel(
-            model_name="text-embedding-v4",
-            dimensions=1024,
-            max_cache_size=100,
-            max_retries=2,
-            raise_exception=True,
-            cache_dir=temp_dir,
-        )
-
-        test_text = "Hello, this is a test sentence for embedding cache."
-
-        print(f"Input text: {test_text}")
-        print(f"Cache directory: {temp_dir}")
-
-        # First call - should be a cache miss
-        print("\n1️⃣ First embedding call (cold cache):")
-        embedding1 = await model.get_embedding(test_text)
-        stats1 = model.get_cache_stats()
-
-        print(f"   Embedding dimension: {len(embedding1)}")
-        print(f"   Cache size: {stats1['cache_size']}")
-        print(f"   Cache hits: {stats1['cache_hits']}")
-        print(f"   Cache misses: {stats1['cache_misses']}")
-        print(f"   Hit rate: {stats1['hit_rate']:.2%}")
-
-        assert len(embedding1) == 1024, "Embedding dimension mismatch"
-        assert stats1["cache_misses"] == 1, "Should have 1 cache miss"
-        assert stats1["cache_hits"] == 0, "Should have 0 cache hits"
-        assert stats1["cache_size"] == 1, "Cache should have 1 entry"
-
-        # Second call with same text - should be a cache hit
-        print("\n2️⃣ Second embedding call (same text):")
-        embedding2 = await model.get_embedding(test_text)
-        stats2 = model.get_cache_stats()
-
-        print(f"   Cache hits: {stats2['cache_hits']}")
-        print(f"   Cache misses: {stats2['cache_misses']}")
-        print(f"   Hit rate: {stats2['hit_rate']:.2%}")
-
-        assert embedding1 == embedding2, "Cached embedding should be identical"
-        assert stats2["cache_hits"] == 1, "Should have 1 cache hit"
-        assert stats2["cache_misses"] == 1, "Should still have 1 cache miss"
-        assert stats2["hit_rate"] == 0.5, "Hit rate should be 50%"
-
-        await model.close()
-        print("\n✓ PASSED: Basic cache functionality works correctly")
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
-async def test_batch_cache_efficiency():
-    """Test cache efficiency with batch embeddings including duplicates."""
-    print(f"\n{'='*60}")
-    print("Test 2: Batch Cache Efficiency")
-    print(f"{'='*60}")
-
-    temp_dir = tempfile.mkdtemp()
-    try:
-        model = OpenAIEmbeddingModel(
-            model_name="text-embedding-v4",
-            dimensions=1024,
-            max_cache_size=1000,
-            max_retries=2,
-            raise_exception=True,
-            cache_dir=temp_dir,
-        )
-
-        texts = get_test_texts()
-
-        # Create a list with duplicates
-        texts_with_duplicates = texts + texts[:3]  # 5 unique + 3 duplicates = 8 total
-
-        print(f"Processing {len(texts_with_duplicates)} texts (5 unique + 3 duplicates)")
-
-        # First batch
-        print("\n1️⃣ First batch (cold cache):")
-        embeddings1 = await model.get_embeddings(texts)
-        stats1 = model.get_cache_stats()
-
-        print(f"   Embeddings generated: {len(embeddings1)}")
-        print(f"   Cache size: {stats1['cache_size']}")
-        print(f"   Cache misses: {stats1['cache_misses']}")
-        print(f"   Cache hits: {stats1['cache_hits']}")
-
-        assert len(embeddings1) == len(texts), "Embeddings count mismatch"
-        assert stats1["cache_size"] == len(texts), f"Cache should have {len(texts)} entries"
-        assert stats1["cache_misses"] == len(texts), "All should be cache misses"
-
-        # Second batch with duplicates
-        print("\n2️⃣ Second batch (with duplicates):")
-        embeddings2 = await model.get_embeddings(texts_with_duplicates)
-        stats2 = model.get_cache_stats()
-
-        print(f"   Embeddings generated: {len(embeddings2)}")
-        print(f"   Cache hits: {stats2['cache_hits']}")
-        print(f"   Cache misses: {stats2['cache_misses']}")
-        print(f"   Hit rate: {stats2['hit_rate']:.2%}")
-
-        assert len(embeddings2) == len(texts_with_duplicates), "Embeddings count mismatch"
-        assert stats2["cache_hits"] >= 3, "Should have at least 3 cache hits from duplicates"
-
-        # Verify embeddings are identical for duplicated texts
-        for i in range(3):
-            assert embeddings2[i] == embeddings2[len(texts) + i], f"Duplicate {i} should have identical embedding"
-
-        await model.close()
-        print("\n✓ PASSED: Batch cache efficiently handles duplicates")
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
-async def test_cache_lru_eviction():
-    """Test LRU cache eviction policy."""
-    print(f"\n{'='*60}")
-    print("Test 3: LRU Cache Eviction")
-    print(f"{'='*60}")
-
-    temp_dir = tempfile.mkdtemp()
-    try:
-        # Create model with small cache size
-        model = OpenAIEmbeddingModel(
-            model_name="text-embedding-v4",
-            dimensions=1024,
-            max_cache_size=3,  # Small cache for testing eviction
-            max_retries=2,
-            raise_exception=True,
-            cache_dir=temp_dir,
-        )
-
-        texts = get_test_texts()[:5]  # Use 5 texts, cache size is 3
-
-        print(f"Cache size limit: {model.max_cache_size}")
-        print(f"Number of unique texts: {len(texts)}")
-
-        # Fill cache beyond capacity
-        print("\n1️⃣ Filling cache with 5 texts (capacity = 3):")
-        for i, text in enumerate(texts):
-            await model.get_embedding(text)
-            stats = model.get_cache_stats()
-            print(
-                f"   After text {i+1}: cache_size={stats['cache_size']}, "
-                f"hits={stats['cache_hits']}, misses={stats['cache_misses']}",
-            )
-
-        final_stats = model.get_cache_stats()
-        assert final_stats["cache_size"] <= 3, "Cache size should not exceed max_cache_size"
-        assert final_stats["cache_misses"] == 5, "Should have 5 cache misses for 5 unique texts"
-
-        # Access the most recent entries - should be cache hits
-        print("\n2️⃣ Accessing recent entries (should be cached):")
-        recent_texts = texts[-3:]  # Last 3 texts should still be in cache
-
-        for i, text in enumerate(recent_texts):
-            await model.get_embedding(text)
-            stats = model.get_cache_stats()
-            print(f"   Text {len(texts) - 3 + i + 1}: hits={stats['cache_hits']}")
-
-        final_stats = model.get_cache_stats()
-        assert final_stats["cache_hits"] == 3, "Should have 3 cache hits for recent entries"
-
-        # Access oldest entries - should be cache misses (evicted)
-        print("\n3️⃣ Accessing oldest entries (should be evicted):")
-        old_texts = texts[:2]  # First 2 texts should have been evicted
-
-        before_misses = final_stats["cache_misses"]
-        for i, text in enumerate(old_texts):
-            await model.get_embedding(text)
-            stats = model.get_cache_stats()
-            print(f"   Text {i + 1}: misses={stats['cache_misses']}")
-
-        final_stats = model.get_cache_stats()
-        assert final_stats["cache_misses"] == before_misses + 2, "Should have 2 more cache misses for evicted entries"
-
-        await model.close()
-        print("\n✓ PASSED: LRU eviction works correctly")
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
-async def test_cache_stats_and_clear():
-    """Test cache statistics tracking and clearing."""
-    print(f"\n{'='*60}")
-    print("Test 4: Cache Statistics and Clearing")
-    print(f"{'='*60}")
-
-    temp_dir = tempfile.mkdtemp()
-    try:
-        model = OpenAIEmbeddingModel(
-            model_name="text-embedding-v4",
-            dimensions=1024,
-            max_cache_size=100,
-            max_retries=2,
-            raise_exception=True,
-            cache_dir=temp_dir,
-        )
-
-        texts = get_test_texts()
-
-        # Generate some cache activity
-        print("\n1️⃣ Generating cache activity:")
-        await model.get_embeddings(texts)
-        await model.get_embeddings(texts[:3])  # Repeat first 3
-
-        stats = model.get_cache_stats()
-        print(f"   Cache size: {stats['cache_size']}")
-        print(f"   Max cache size: {stats['max_cache_size']}")
-        print(f"   Cache hits: {stats['cache_hits']}")
-        print(f"   Cache misses: {stats['cache_misses']}")
-        print(f"   Hit rate: {stats['hit_rate']:.2%}")
-
-        assert stats["cache_size"] > 0, "Cache should not be empty"
-        assert stats["cache_hits"] >= 3, "Should have at least 3 cache hits"
-        assert "hit_rate" in stats, "Stats should include hit_rate"
-
-        # Clear cache
-        print("\n2️⃣ Clearing cache:")
-        model.clear_cache()
-        stats_after_clear = model.get_cache_stats()
-
-        print(f"   Cache size after clear: {stats_after_clear['cache_size']}")
-        print(f"   Hits after clear: {stats_after_clear['cache_hits']}")
-        print(f"   Misses after clear: {stats_after_clear['cache_misses']}")
-        print(f"   Hit rate after clear: {stats_after_clear['hit_rate']:.2%}")
-
-        assert stats_after_clear["cache_size"] == 0, "Cache should be empty after clear"
-        assert stats_after_clear["cache_hits"] == 0, "Hits should be reset"
-        assert stats_after_clear["cache_misses"] == 0, "Misses should be reset"
-        assert stats_after_clear["hit_rate"] == 0.0, "Hit rate should be 0"
-
-        await model.close()
-        print("\n✓ PASSED: Cache statistics and clearing work correctly")
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
-async def test_cache_disabled():
-    """Test behavior when cache is disabled (max_cache_size=0)."""
-    print(f"\n{'='*60}")
-    print("Test 5: Cache Disabled")
-    print(f"{'='*60}")
-
-    temp_dir = tempfile.mkdtemp()
-    try:
-        model = OpenAIEmbeddingModel(
-            model_name="text-embedding-v4",
-            dimensions=1024,
-            max_cache_size=0,  # Disable cache
-            max_retries=2,
-            raise_exception=True,
-            cache_dir=temp_dir,
-        )
-
-        test_text = "Test text with cache disabled"
-
-        print(f"Cache size limit: {model.max_cache_size} (disabled)")
-        print(f"Input text: {test_text}")
-
-        # Call twice with same text
-        print("\n1️⃣ First call:")
-        embedding1 = await model.get_embedding(test_text)
-        stats1 = model.get_cache_stats()
-        print(f"   Cache size: {stats1['cache_size']}")
-        print(f"   Cache misses: {stats1['cache_misses']}")
-
-        print("\n2️⃣ Second call (same text):")
-        embedding2 = await model.get_embedding(test_text)
-        stats2 = model.get_cache_stats()
-        print(f"   Cache size: {stats2['cache_size']}")
-        print(f"   Cache misses: {stats2['cache_misses']}")
-        print(f"   Cache hits: {stats2['cache_hits']}")
-
-        assert stats2["cache_size"] == 0, "Cache should remain empty when disabled"
-        assert stats2["cache_misses"] == 2, "Both calls should be cache misses"
-        assert stats2["cache_hits"] == 0, "Should have no cache hits when disabled"
-        assert embedding1 == embedding2, "Embeddings should still be consistent"
-
-        await model.close()
-        print("\n✓ PASSED: Cache correctly disabled when max_cache_size=0")
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
-async def test_cache_performance_demo():
-    """Demonstrate cache performance improvements."""
-    print(f"\n{'='*60}")
-    print("Test 6: Cache Performance Demo")
-    print(f"{'='*60}")
-
-    temp_dir = tempfile.mkdtemp()
-    try:
-        model = OpenAIEmbeddingModel(
-            model_name="text-embedding-v4",
-            dimensions=1024,
-            max_cache_size=1000,
-            max_retries=2,
-            raise_exception=True,
-            cache_dir=temp_dir,
-        )
-
-        texts = get_test_texts()
-
-        # Create a realistic workload with many repeated queries
-        workload = texts * 3  # 15 queries total, 5 unique
-
-        print(f"\nProcessing {len(workload)} queries ({len(texts)} unique texts)")
-        print("This simulates a realistic scenario with repeated queries\n")
-
-        # Process all queries
-        for i, text in enumerate(workload, 1):
-            await model.get_embedding(text)
-            if i % 5 == 0:  # Report every 5 queries
-                stats = model.get_cache_stats()
-                print(
-                    f"After {i:2d} queries: hits={stats['cache_hits']:2d}, "
-                    f"misses={stats['cache_misses']:2d}, "
-                    f"hit_rate={stats['hit_rate']:5.1%}",
-                )
-
-        final_stats = model.get_cache_stats()
-        total_requests = final_stats["cache_hits"] + final_stats["cache_misses"]
-
-        print(f"\n{'─'*60}")
-        print("📊 Final Statistics:")
-        print(f"{'─'*60}")
-        print(f"  Total queries: {total_requests}")
-        print(f"  Unique texts: {len(texts)}")
-        print(f"  Cache hits: {final_stats['cache_hits']}")
-        print(f"  Cache misses: {final_stats['cache_misses']}")
-        print(f"  Hit rate: {final_stats['hit_rate']:.1%}")
-        print(f"  Cache size: {final_stats['cache_size']}/{final_stats['max_cache_size']}")
-        print(f"{'─'*60}")
-        print(
-            f"💰 API calls saved: {final_stats['cache_hits']} out of {total_requests} "
-            f"({final_stats['cache_hits']/total_requests*100:.1f}%)",
-        )
-        print(f"{'─'*60}")
-
-        assert final_stats["cache_hits"] == 10, "Should have 10 cache hits (2 repeats × 5 texts)"
-        assert final_stats["cache_misses"] == 5, "Should have 5 cache misses (5 unique texts)"
-        assert final_stats["hit_rate"] > 0.6, "Hit rate should be > 60%"
-
-        await model.close()
-        print("\n✓ PASSED: Cache provides significant performance improvement")
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
-async def main():
-    """Run all cache tests."""
-    print("\n" + "#" * 60)
-    print("# EMBEDDING CACHE TESTS")
-    print("#" * 60)
-
-    try:
-        await test_cache_basic_functionality()
-        await test_batch_cache_efficiency()
-        await test_cache_lru_eviction()
-        await test_cache_stats_and_clear()
-        await test_cache_disabled()
-        await test_cache_performance_demo()
-
-        print("\n" + "=" * 60)
-        print("✅ ALL CACHE TESTS PASSED")
-        print("=" * 60)
-        print("\nKey takeaways:")
-        print("  • Cache correctly tracks hits/misses")
-        print("  • LRU eviction works as expected")
-        print("  • Duplicate queries are efficiently cached")
-        print("  • Cache can be disabled or cleared")
-        print("  • Significant performance improvement with realistic workloads")
-        print("=" * 60 + "\n")
-
-    except Exception as e:
-        print(f"\n✗ TEST FAILED: {type(e).__name__}: {e}")
-        raise
-
-
-if __name__ == "__main__":
-    asyncio.run(main())