diff --git a/reme4/components/base_component.py b/reme4/components/base_component.py index 06e289e3..4c614d52 100644 --- a/reme4/components/base_component.py +++ b/reme4/components/base_component.py @@ -142,6 +142,12 @@ class BaseComponent(ABC): async def _close(self) -> None: """Subclass hook: close logic.""" + async def dump(self) -> None: + """Persist in-memory state to disk. Override in subclasses that need persistence.""" + + async def load(self) -> None: + """Restore in-memory state from disk. Override in subclasses that need persistence.""" + async def start(self) -> None: """Resolve bindings → start owned fallbacks → _start(). No-op if already started.""" async with self._lock: diff --git a/reme4/components/embedding/base_embedding_model.py b/reme4/components/embedding/base_embedding_model.py index 4fce0535..02e2c877 100644 --- a/reme4/components/embedding/base_embedding_model.py +++ b/reme4/components/embedding/base_embedding_model.py @@ -56,8 +56,7 @@ class BaseEmbeddingModel(BaseComponent): async def _start(self) -> None: """Load cache from disk on startup.""" - self._embedding_cache.clear() - self._load_cache() + await self.load() async def health_check(self, timeout: float = 2.0) -> bool: """Probe the provider; sets and returns is_healthy.""" @@ -78,7 +77,7 @@ class BaseEmbeddingModel(BaseComponent): async def _close(self) -> None: """Persist cache to disk on shutdown.""" - self._save_cache() + await self.dump() # -- Public API -- @@ -180,12 +179,10 @@ class BaseEmbeddingModel(BaseComponent): # -- Cache Persistence -- - def _load_cache(self) -> None: - """Load cached embeddings 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(): + async def load(self) -> None: + """Load cached embeddings from disk (npz format); replaces in-memory cache.""" + self._embedding_cache.clear() + if not self.enable_cache or not self.cache_path.exists(): return try: @@ -202,10 +199,11 @@ class BaseEmbeddingModel(BaseComponent): break self._embedding_cache[str(key)] = emb.astype(np.float16) - def _save_cache(self) -> None: + async def dump(self) -> None: """Persist in-memory cache to disk (npz format).""" if not self.enable_cache or not self._embedding_cache: return + self.cache_path.parent.mkdir(parents=True, exist_ok=True) keys = list(self._embedding_cache.keys()) embeddings = np.stack(list(self._embedding_cache.values())) try: diff --git a/reme4/components/file_graph/local_file_graph.py b/reme4/components/file_graph/local_file_graph.py index 8aed20d7..221b022a 100644 --- a/reme4/components/file_graph/local_file_graph.py +++ b/reme4/components/file_graph/local_file_graph.py @@ -22,7 +22,7 @@ class LocalFileGraph(BaseFileGraph): async def _start(self) -> None: await super()._start() - self._load() + await self.load() await self.rebuild_links() self.logger.info( f"LocalFileGraph '{self.graph_name}' ready: " @@ -31,22 +31,30 @@ class LocalFileGraph(BaseFileGraph): ) async def _close(self) -> None: - self._dump() + await self.dump() await super()._close() - def _load(self) -> None: - """Load nodes from JSONL file into memory.""" + async def load(self) -> None: + """Load nodes from JSONL file into memory; keep current state on failure.""" if not self._graph_file.exists(): return - with open(self._graph_file, "r", encoding="utf-8") as f: - self._nodes.update((n.path, n) for line in f if line.strip() for n in [FileNode.model_validate_json(line)]) + try: + with open(self._graph_file, "r", encoding="utf-8") as f: + self._nodes.update( + (n.path, n) for line in f if line.strip() for n in [FileNode.model_validate_json(line)] + ) + except Exception as e: + self.logger.exception(f"Failed to load {self._graph_file}: {e}") - def _dump(self) -> None: + async def dump(self) -> None: """Persist all nodes to JSONL via atomic rename.""" - tmp = self._graph_file.with_suffix(".tmp") - with open(tmp, "w", encoding="utf-8") as f: - f.writelines(f"{n.model_dump_json()}\n" for n in self._nodes.values()) - tmp.replace(self._graph_file) + try: + tmp = self._graph_file.with_suffix(".tmp") + with open(tmp, "w", encoding="utf-8") as f: + f.writelines(f"{n.model_dump_json()}\n" for n in self._nodes.values()) + tmp.replace(self._graph_file) + except Exception as e: + self.logger.exception(f"Failed to write {self._graph_file}: {e}") # -- Edge bookkeeping -------------------------------------------------- diff --git a/reme4/components/file_graph/nx_file_graph.py b/reme4/components/file_graph/nx_file_graph.py index 81695dec..dcb608d6 100644 --- a/reme4/components/file_graph/nx_file_graph.py +++ b/reme4/components/file_graph/nx_file_graph.py @@ -28,9 +28,7 @@ class NxFileGraph(BaseFileGraph): async def _start(self) -> None: await super()._start() - loaded = self._load() - if loaded is not None: - self._graph = loaded + await self.load() n_real = sum(1 for _, d in self._graph.nodes(data=True) if "node" in d) self.logger.info( f"NxFileGraph '{self.graph_name}' ready: " @@ -39,21 +37,20 @@ class NxFileGraph(BaseFileGraph): ) async def _close(self) -> None: - self._dump() + await self.dump() await super()._close() - def _load(self) -> nx.MultiDiGraph | None: - """Load graph from pickle file; return None on failure.""" + async def load(self) -> None: + """Load graph from pickle file; keep current graph on failure.""" if not self._graph_file.exists(): - return None + return try: with open(self._graph_file, "rb") as f: - return pickle.load(f) + self._graph = pickle.load(f) except Exception as e: self.logger.exception(f"Failed to load {self._graph_file}: {e}") - return None - def _dump(self) -> None: + async def dump(self) -> None: """Persist graph to pickle via atomic rename.""" try: tmp = self._graph_file.with_suffix(".tmp") diff --git a/reme4/components/file_store/local_file_store.py b/reme4/components/file_store/local_file_store.py index 63638dd6..052694e4 100644 --- a/reme4/components/file_store/local_file_store.py +++ b/reme4/components/file_store/local_file_store.py @@ -23,19 +23,30 @@ class LocalFileStore(BaseFileStore): async def _start(self) -> None: await super()._start() - if self.chunks_path.exists(): - try: - async with aiofiles.open(self.chunks_path, encoding=self.encoding) as f: - async for line in f: - line = line.strip() - if line: - chunk = FileChunk.model_validate_json(line) - self.file_chunks[chunk.id] = chunk - except Exception as e: - self.logger.exception(f"Failed to load {self.chunks_path}: {e}") + await self.load() self.logger.info(f"LocalFileStore '{self.store_name}' ready: {len(self.file_chunks)} chunks") async def _close(self) -> None: + await self.dump() + self.file_chunks.clear() + await super()._close() + + async def load(self) -> None: + """Load chunks from JSONL file into memory.""" + if not self.chunks_path.exists(): + return + try: + async with aiofiles.open(self.chunks_path, encoding=self.encoding) as f: + async for line in f: + line = line.strip() + if line: + chunk = FileChunk.model_validate_json(line) + self.file_chunks[chunk.id] = chunk + except Exception as e: + self.logger.exception(f"Failed to load {self.chunks_path}: {e}") + + async def dump(self) -> None: + """Persist chunks to JSONL via atomic rename.""" try: tmp = self.chunks_path.with_suffix(".tmp") async with aiofiles.open(tmp, "w", encoding=self.encoding) as f: @@ -43,8 +54,6 @@ class LocalFileStore(BaseFileStore): tmp.replace(self.chunks_path) except Exception as e: self.logger.exception(f"Failed to write {self.chunks_path}: {e}") - self.file_chunks.clear() - await super()._close() # Base class interface diff --git a/reme4/components/keyword_index/base_keyword_index.py b/reme4/components/keyword_index/base_keyword_index.py index 45681a4d..c577e878 100644 --- a/reme4/components/keyword_index/base_keyword_index.py +++ b/reme4/components/keyword_index/base_keyword_index.py @@ -24,9 +24,8 @@ class BaseKeywordIndex(BaseComponent): async def _start(self) -> None: """Load existing index from disk if available.""" - if self.index_file.exists(): - await self.load() - self.logger.info(f"Loaded index from {self.index_path}") + await self.load() + self.logger.info(f"Loaded index from {self.index_path}") async def _close(self) -> None: """Save index to disk on shutdown.""" @@ -59,14 +58,6 @@ class BaseKeywordIndex(BaseComponent): async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: """Search documents. Returns {doc_id: score} sorted descending.""" - @abstractmethod - async def dump(self) -> None: - """Persist index to disk.""" - - @abstractmethod - async def load(self) -> None: - """Load index from disk.""" - @abstractmethod async def clear(self) -> None: """Reset index to empty state.""" diff --git a/reme4/components/keyword_index/bm25_index.py b/reme4/components/keyword_index/bm25_index.py index c64c70f3..3c2591ac 100644 --- a/reme4/components/keyword_index/bm25_index.py +++ b/reme4/components/keyword_index/bm25_index.py @@ -127,22 +127,29 @@ class BM25Index(BaseKeywordIndex): return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:limit]) if scores else {} async def dump(self) -> None: - """Persist index to disk via pickle.""" - with open(self.index_file, "wb") as f: - pickle.dump( - { - "vocab": self.vocab, - "inverted_index": self.inverted_index, - "doc_meta": self.doc_meta, - "total_len": self.total_len, - "k1": self.k1, - "b": self.b, - }, - f, - ) + """Persist index to disk via pickle (atomic rename).""" + try: + tmp = self.index_file.with_suffix(".tmp") + with open(tmp, "wb") as f: + pickle.dump( + { + "vocab": self.vocab, + "inverted_index": self.inverted_index, + "doc_meta": self.doc_meta, + "total_len": self.total_len, + "k1": self.k1, + "b": self.b, + }, + f, + ) + tmp.replace(self.index_file) + except Exception as e: + self.logger.exception(f"Failed to write {self.index_file}: {e}") async def load(self) -> None: - """Load index from disk. Clears index on failure.""" + """Load index from disk. No-op if file missing; clears index on corruption.""" + if not self.index_file.exists(): + return try: with open(self.index_file, "rb") as f: data = pickle.load(f) diff --git a/reme4/config/default.yaml b/reme4/config/default.yaml index f06d6c5a..8bcfd5c0 100644 --- a/reme4/config/default.yaml +++ b/reme4/config/default.yaml @@ -40,6 +40,15 @@ jobs: steps: - backend: health_check_step + - backend: base + name: help + description: "list all registered jobs with their metadata" + parameters: + type: object + properties: {} + steps: + - backend: help_step + - backend: stream name: stream_demo description: "stream demo job: repeat query 10x and stream char-by-char" diff --git a/reme4/steps/common/__init__.py b/reme4/steps/common/__init__.py index fb67df6a..e31ba302 100644 --- a/reme4/steps/common/__init__.py +++ b/reme4/steps/common/__init__.py @@ -2,6 +2,7 @@ from .demo import DemoEchoStep1, DemoEchoStep2 from .health_check import HealthCheckStep +from .help import HelpStep from .stream_demo import StreamDemoStep1, StreamDemoStep2 from .version import VersionStep @@ -9,6 +10,7 @@ __all__ = [ "DemoEchoStep1", "DemoEchoStep2", "HealthCheckStep", + "HelpStep", "StreamDemoStep1", "StreamDemoStep2", "VersionStep", diff --git a/reme4/steps/common/help.py b/reme4/steps/common/help.py new file mode 100644 index 00000000..782ea877 --- /dev/null +++ b/reme4/steps/common/help.py @@ -0,0 +1,42 @@ +"""Return a one-line summary of every registered job for LLM consumption.""" + +from ..base_step import BaseStep +from ...components import R + + +def _format_params(parameters: dict) -> str: + props = (parameters or {}).get("properties") or {} + if not props: + return "no args" + required = set((parameters or {}).get("required") or []) + parts = [] + for pname, pschema in props.items(): + ptype = pschema.get("type", "any") + if pname in required: + parts.append(f"{pname}:{ptype}*") + elif "default" in pschema: + parts.append(f"{pname}:{ptype}={pschema['default']}") + else: + parts.append(f"{pname}:{ptype}") + return ", ".join(parts) + + +@R.register("help_step") +class HelpStep(BaseStep): + """List all registered jobs (excluding self) as compact one-liners for an LLM.""" + + async def execute(self): + assert self.context is not None + + lines = [] + if self.app_context is not None: + for name, job in self.app_context.jobs.items(): + if name == "help": + continue + lines.append(f"🛠️ `{name}` — {job.description} 📥 {_format_params(job.parameters)}") + + self.logger.info(f"[{self.name}] returning {len(lines)} jobs") + + self.context.response.answer = "\n".join(lines) + self.context.response.metadata["job_count"] = len(lines) + return self.context.response