up

reme2/component/embedding/base_embedding_model.py

@@ -52,7 +52,9 @@ class BaseEmbeddingModel(BaseComponent):
         self._embedding_cache: OrderedDict[str, list[float]] = OrderedDict()
         self._cache_hits = 0
         self._cache_misses = 0
-        self.cache_path: Path = Path()
+
+        self.working_dir = self.app_context.app_config.working_dir if self.app_context is not None else ""
+        self.cache_path: Path = Path(self.working_dir) / "embedding_cache" / f"{self.name}.npz"
 
     def clear_cache(self) -> None:
         """Clear in-memory cache and reset statistics."""
@@ -62,10 +64,7 @@ class BaseEmbeddingModel(BaseComponent):
 
     async def _start(self) -> None:
         """Load cache on start."""
-        assert self.app_context is not None, "app_context must be provided"
         self.clear_cache()
-        working_path = Path(self.app_context.app_config.working_dir)
-        self.cache_path = working_path / "embedding_cache" / f"{self.name}.npz"
         self._load_cache()
 
     async def _close(self) -> None:

reme2/component/file_store/base_file_store.py

@@ -1,16 +1,4 @@
-"""Abstract base class for file stores — minimal engine surface.
-
-The store owns persistence + search for the (file → chunks) graph.
-Subclasses implement every read/write/search verb; the base only
-resolves shared infrastructure:
-
-  * `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
+"""Abstract base class for file stores — persistence + search for nodes & chunks."""
 
 import re
 from abc import abstractmethod
@@ -28,40 +16,33 @@ class BaseFileStore(BaseComponent):
     component_type = ComponentEnum.FILE_STORE
 
     def __init__(
-        self,
-        store_name: str,
-        store_path: str | Path,
-        embedding_model: str = "default",
-        fts_enabled: bool = True,
-        **kwargs,
+            self,
+            store_name: str,
+            store_path: str | Path,
+            embedding_model: str = "default",
+            fts_enabled: bool = True,
+            **kwargs,
     ):
         super().__init__(**kwargs)
         if not re.match(r"^[a-zA-Z0-9_]+$", store_name):
-            raise ValueError(
-                f"Invalid store name '{store_name}'. "
-                f"Only alphanumeric characters and underscores are allowed.",
-            )
-        self.store_name: str = store_name
-        self.store_path: Path = Path(store_path)
+            raise ValueError(f"Invalid store name '{store_name}'. Only alphanumeric and underscores allowed.")
+        self.store_name = store_name
+        self.store_path = Path(store_path)
         self.store_path.mkdir(parents=True, exist_ok=True)
-        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.working_dir = self.app_context.app_config.working_dir if self.app_context else ""
+        self._embedding_model_name = embedding_model
         self.embedding_model: BaseEmbeddingModel | None = None
-        self.vector_enabled: bool = bool(embedding_model)
-        self.fts_enabled: bool = fts_enabled
+        self.vector_enabled = bool(embedding_model)
+        self.fts_enabled = fts_enabled
         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 ---------------------------------------------------------
+    # Lifecycle
 
     async def _start(self) -> None:
         if not self._embedding_model_name:
             return
-        assert self.app_context is not None, "app_context must be provided"
+        assert self.app_context is not None
         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.")
@@ -73,82 +54,63 @@ class BaseFileStore(BaseComponent):
     async def _close(self) -> None:
         self.embedding_model = None
 
-    # -- 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 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
-        try:
-            return await self.embedding_model.get_embeddings(texts)
-        except Exception as e:
-            self._disable_vector_search(str(e))
-            return None
-
-    def _disable_vector_search(self, reason: str = "embedding API error") -> None:
-        if self.vector_enabled:
-            self.logger.warning(f"[{self.store_name}] Disabling vector search: {reason}")
-            self.vector_enabled = False
-
-    # -- Composite write entries (concrete, in terms of the four abstracts) -
+    # Composite operations
 
     async def upsert(self, node: FileNode, chunks: list[FileChunk]) -> None:
-        """Persist a node and its chunks together. Node first, then chunks."""
         await self.upsert_node(node)
         await self.upsert_chunks(node.path, chunks)
 
     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) ---------------------------
+    # Abstract: node operations
 
     @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."""
+        """Persist a node, replacing any prior entry for `node.path`."""
 
     @abstractmethod
     async def get_node(self, path: str) -> FileNode | None:
-        """Fetch a single node by path. None if absent."""
+        """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`.
-
-        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`."""
+        """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 vector_search(
-        self, query: str, limit: int, search_filter: dict,
-    ) -> list[FileChunk]:
+    async def delete_chunks(self, path: str) -> None:
+        """Delete all chunks for `path`."""
+
+    # Abstract: search
+
+    @abstractmethod
+    async def vector_search(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,
-    ) -> list[FileChunk]:
+    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

reme2/component/file_store/local_file_store.py

@@ -1,16 +1,4 @@
-"""Pure-Python in-memory store with on-close JSONL persistence.
-
-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.
-"""
+"""In-memory file store with JSONL persistence on close."""
 
 from __future__ import annotations
 
@@ -31,21 +19,21 @@ class LocalFileStore(BaseFileStore):
 
     def __init__(self, encoding: str = "utf-8", **kwargs):
         super().__init__(**kwargs)
-        self._encoding: str = encoding
+        self._encoding = 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"
+        self._nodes_file = self.store_path / f"{self.store_name}_nodes.jsonl"
+        self._chunks_file = self.store_path / f"{self.store_name}_chunks.jsonl"
 
-    # -- Lifecycle ---------------------------------------------------------
+    # Lifecycle
 
     async def _start(self) -> None:
-        self._load(self._nodes_file, self._nodes, FileNode, key="path")
-        self._load(self._chunks_file, self._chunks, FileChunk, key="id")
+        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",
+            f"{len(self._nodes)} nodes, {len(self._chunks)} chunks"
         )
 
     async def _close(self) -> None:
@@ -55,18 +43,15 @@ class LocalFileStore(BaseFileStore):
         self._chunks.clear()
         await super()._close()
 
-    def _load(
-        self, file: Path, target: dict, model: type[BaseModel], key: str,
-    ) -> None:
+    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
+                if line.strip():
+                    obj = model.model_validate_json(line)
+                    target[getattr(obj, key)] = obj
         except Exception as e:
             self.logger.warning(f"Failed to load {file}: {e}")
 
@@ -79,28 +64,22 @@ class LocalFileStore(BaseFileStore):
         except Exception as e:
             self.logger.error(f"Failed to write {file}: {e}")
 
-    # -- Node CRUD ---------------------------------------------------------
+    # Node 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 get_node(self, path: str) -> FileNode | None:
         return self._nodes.get(path)
 
-    # -- Chunk CRUD --------------------------------------------------------
+    async def delete_node(self, path: str) -> None:
+        self._nodes.pop(path, None)
+
+    # Chunk operations
 
     async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None:
-        """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}
+        cached = {c.hash: c.embedding for c in existing if c.embedding}
 
         await self.delete_chunks(path)
         if not chunks:
@@ -108,61 +87,55 @@ class LocalFileStore(BaseFileStore):
 
         needs_embed: list[FileChunk] = []
         for c in chunks:
-            if c.embedding is not None:
+            if c.embedding:
                 continue
-            cached_emb = cached.get(c.hash)
-            if cached_emb is not None:
-                c.embedding = cached_emb
+            if c.hash in cached:
+                c.embedding = cached[c.hash]
             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:
+            embeddings = await self.get_embeddings([c.text for c in needs_embed])
+            if embeddings:
                 for c, emb in zip(needs_embed, embeddings):
-                    if emb is not None:
+                    if emb:
                         c.embedding = emb
 
         for c in chunks:
             self._chunks[c.id] = c
 
-    async def delete_chunks(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]
-
     async def get_chunks(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
 
-    # -- Search ------------------------------------------------------------
+    async def delete_chunks(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:
             return []
-        embeddings = await self.embed([query])
-        if not embeddings or embeddings[0] is None:
+        embeddings = await self.get_embeddings([query])
+        if not embeddings or not embeddings[0]:
             return []
-        query_emb = embeddings[0]
 
-        # TODO: honor `search_filter` (paths / exclude_paths / tags / ...).
         candidates = [c for c in self._chunks.values() if c.embedding]
         if not candidates:
             return []
 
         chunk_embs = np.array([c.embedding for c in candidates])
-        similarities = batch_cosine_similarity(
-            np.array([query_emb]), chunk_embs,
-        )[0]
+        similarities = batch_cosine_similarity(np.array([embeddings[0]]), 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 = [
+            c.model_copy(update={"scores": {"vector": float(s), "score": float(s)}})
+            for c, s in zip(candidates, similarities)
+        ]
         results.sort(key=lambda r: r.score, reverse=True)
         return results[:limit]
 
@@ -172,18 +145,15 @@ class LocalFileStore(BaseFileStore):
         if not self.fts_enabled or not query.split():
             return []
 
-        # 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 = [
+            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]
 
-    # -- Helpers -----------------------------------------------------------
+    # Helpers
 
     @staticmethod
     def _keyword_score(query: str, text: str) -> float:
@@ -198,4 +168,4 @@ class LocalFileStore(BaseFileStore):
         base = matches / len(words)
         if len(words) > 1 and query.lower() in text_lower:
             base = min(1.0, base + 0.2)
-        return base
+        return base

reme2/memory/memory_io.py

@@ -120,7 +120,7 @@ def _edge_to_dict(node, edge) -> dict:
         "predicate": edge.predicate,
         "anchor": edge.anchor,
         "alias": edge.alias,
-        "embed": edge.embed,
+        "embed": edge.get_embeddings,
     }