From ef0b74b3e78a4fdb5cf9922c59c26640eca3144a Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 12 May 2026 20:46:40 +0800 Subject: [PATCH] up --- .../embedding/base_embedding_model.py | 7 +- reme2/component/file_store/base_file_store.py | 134 +++++++----------- .../component/file_store/local_file_store.py | 118 ++++++--------- reme2/memory/memory_io.py | 2 +- 4 files changed, 96 insertions(+), 165 deletions(-) diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index 8844e891..6bc19057 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/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: diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index 26240da1..bda59724 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/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 diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index 38d52f42..55ed3481 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/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 \ No newline at end of file diff --git a/reme2/memory/memory_io.py b/reme2/memory/memory_io.py index c9093d86..632b70ef 100644 --- a/reme2/memory/memory_io.py +++ b/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, }