From 905c99411873f1494a50353ea85bac9ac115f480 Mon Sep 17 00:00:00 2001 From: jrauch713-svg Date: Fri, 10 Jul 2026 04:04:50 +0000 Subject: [PATCH 1/2] feat(memory): add optional Mem0 backend for skill iteration and reflection tracking Wires SkillMemory into ReflACTTrainer via non-breaking hooks: records each step's skill/score after the evaluation gate and each step's reflection patches after the accumulation loop. Degrades to a no-op when MEM0_API_KEY is unset. --- skillopt/engine/trainer.py | 12 ++ skillopt/memory/__init__.py | 4 + skillopt/memory/mem0_backend.py | 279 +++++++++++++++++++++++++++++++ skillopt/memory/trainer_hooks.py | 147 ++++++++++++++++ 4 files changed, 442 insertions(+) create mode 100644 skillopt/memory/__init__.py create mode 100644 skillopt/memory/mem0_backend.py create mode 100644 skillopt/memory/trainer_hooks.py diff --git a/skillopt/engine/trainer.py b/skillopt/engine/trainer.py index 1eae4209..aca5b34b 100644 --- a/skillopt/engine/trainer.py +++ b/skillopt/engine/trainer.py @@ -75,6 +75,7 @@ set_optimizer_deployment, ) from skillopt.utils import compute_score, skill_hash +from skillopt.memory.trainer_hooks import maybe_init_mem0, hook_post_evaluate, hook_post_reflect # ── Skill-aware reflection: appendix flush ─────────────────────────────────── @@ -1042,6 +1043,8 @@ def _persist_runtime_state(last_completed_step: int) -> None: # ── Training loop ──────────────────────────────────────────────── t_loop_start = time.time() + memory = maybe_init_mem0(cfg) + if resume_from > total_steps: print(f"\n [skip] all {total_steps} steps complete — jumping to evaluation") @@ -1204,6 +1207,11 @@ def _persist_runtime_state(last_completed_step: int) -> None: step_rec["timing"]["rollout_s"] = round(total_rollout_time, 1) step_rec["timing"]["reflect_s"] = round(total_reflect_time, 1) + hook_post_reflect( + memory, epoch, global_step, all_raw_patches, + scores={"hard": agg_hard, "soft": agg_soft}, + ) + n_total_patches = len(all_failure_patches) + len(all_success_patches) step_rec["n_patches"] = n_total_patches step_rec["n_failure_patches"] = len(all_failure_patches) @@ -1522,6 +1530,10 @@ def _persist_runtime_state(last_completed_step: int) -> None: ): best_origin = current_origin + hook_post_evaluate( + memory, epoch, global_step, current_skill, current_score, cfg, + ) + if use_skill_aware: current_skill = _flush_skill_aware_appendix( current_skill, all_raw_patches, step_rec, step_dir, cfg, diff --git a/skillopt/memory/__init__.py b/skillopt/memory/__init__.py new file mode 100644 index 00000000..4d6ce551 --- /dev/null +++ b/skillopt/memory/__init__.py @@ -0,0 +1,4 @@ +"""skillopt.memory — mem0-backed persistent memory for SkillOpt.""" +from skillopt.memory.mem0_backend import SkillMemory + +__all__ = ["SkillMemory"] diff --git a/skillopt/memory/mem0_backend.py b/skillopt/memory/mem0_backend.py new file mode 100644 index 00000000..ad59cc12 --- /dev/null +++ b/skillopt/memory/mem0_backend.py @@ -0,0 +1,279 @@ +"""mem0-backed persistent memory for SkillOpt. + +Stores skill iterations, reflection results, and experiment outcomes in mem0 +so that the ReflACT trainer can retrieve relevant historical context across +runs and identify the best skill versions discovered so far. + +Usage:: + + from skillopt.memory import SkillMemory + + m = SkillMemory(api_key="m0-...") + m.store_skill_iteration(epoch=1, step=3, skill_text="...", score=0.82) + ctx = m.retrieve_relevant_context("handling multi-step navigation") +""" +from __future__ import annotations + +import hashlib +import json +import os +from typing import Any + +try: + from mem0 import MemoryClient + _MEM0_AVAILABLE = True +except ImportError: + _MEM0_AVAILABLE = False + MemoryClient = None # type: ignore[assignment,misc] + + +class SkillMemory: + """Persistent memory backend for SkillOpt using mem0. + + Parameters + ---------- + api_key : str | None + mem0 API key. Falls back to ``MEM0_API_KEY`` env var, then to + ``MEM0_API_KEY`` set on the object at construction time. + user_id : str + Logical user/project identifier used to namespace memories in mem0. + """ + + def __init__( + self, + api_key: str | None = None, + user_id: str = "skillopt", + ) -> None: + if not _MEM0_AVAILABLE: + raise ImportError( + "mem0ai is not installed. Run: pip install mem0ai" + ) + + resolved_key = api_key or os.environ.get("MEM0_API_KEY", "") + if not resolved_key: + raise ValueError( + "No mem0 API key provided. Pass api_key= or set MEM0_API_KEY env var." + ) + + self.user_id = user_id + self._client = MemoryClient(api_key=resolved_key) + + # ── Internal helpers ────────────────────────────────────────────────── + + def _add(self, messages: list[dict], metadata: dict | None = None) -> Any: + """Low-level add wrapper — always tags with user_id.""" + kwargs: dict[str, Any] = {"user_id": self.user_id} + if metadata: + kwargs["metadata"] = metadata + return self._client.add(messages, **kwargs) + + def _search(self, query: str, limit: int = 5) -> list[dict]: + """Low-level search wrapper — scoped to this user_id.""" + results = self._client.search(query, user_id=self.user_id, limit=limit) + # mem0 returns a list of memory dicts + if isinstance(results, list): + return results + # Some versions wrap results in a dict + if isinstance(results, dict): + return results.get("results", []) + return [] + + @staticmethod + def _short_hash(text: str) -> str: + return hashlib.sha1(text.encode()).hexdigest()[:8] + + # ── Public API ──────────────────────────────────────────────────────── + + def store_skill_iteration( + self, + epoch: int, + step: int, + skill_text: str, + score: float, + metadata: dict | None = None, + ) -> Any: + """Store a skill version produced during training. + + Parameters + ---------- + epoch : int + Current training epoch. + step : int + Current training step within the epoch. + skill_text : str + Full text of the skill document at this point. + score : float + Evaluation score (0–1) for this skill version. + metadata : dict | None + Optional additional metadata to attach. + """ + skill_hash = self._short_hash(skill_text) + base_meta: dict[str, Any] = { + "event_type": "skill_iteration", + "epoch": epoch, + "step": step, + "score": round(float(score), 6), + "skill_hash": skill_hash, + "skill_length": len(skill_text), + } + if metadata: + base_meta.update(metadata) + + content = ( + f"[SkillOpt] Epoch {epoch} Step {step} — skill_hash={skill_hash} " + f"score={score:.4f}\n\n" + f"=== SKILL TEXT ===\n{skill_text[:4000]}" # cap to avoid huge payloads + ) + + messages = [{"role": "user", "content": content}] + return self._add(messages, metadata=base_meta) + + def store_reflection( + self, + epoch: int, + step: int, + patches: list[dict], + scores: dict | None = None, + ) -> Any: + """Store the result of a Reflect stage. + + Parameters + ---------- + epoch : int + Current training epoch. + step : int + Current training step. + patches : list[dict] + Raw patches produced by the reflection stage. + scores : dict | None + Optional dict of metric → value (e.g. ``{"hard": 0.7, "soft": 0.8}``). + """ + n_patches = len(patches) + scores_str = json.dumps(scores or {}, ensure_ascii=False) + patch_summary = json.dumps( + [ + {k: v for k, v in p.items() if k != "skill_text"} + for p in patches[:10] # only first 10 to keep it concise + ], + ensure_ascii=False, + ) + + base_meta: dict[str, Any] = { + "event_type": "reflection", + "epoch": epoch, + "step": step, + "n_patches": n_patches, + } + if scores: + base_meta.update({f"score_{k}": v for k, v in scores.items()}) + + content = ( + f"[SkillOpt] Reflection Epoch {epoch} Step {step} — " + f"{n_patches} patch(es) generated. Scores: {scores_str}\n\n" + f"Patch summary:\n{patch_summary}" + ) + + messages = [{"role": "user", "content": content}] + return self._add(messages, metadata=base_meta) + + def retrieve_relevant_context( + self, + query: str, + limit: int = 5, + ) -> list[dict]: + """Retrieve past memories relevant to *query*. + + Parameters + ---------- + query : str + Free-text query describing the context you want to retrieve. + limit : int + Maximum number of results to return. + + Returns + ------- + list[dict] + List of memory objects (each has at least a ``memory`` field). + """ + return self._search(query, limit=limit) + + def get_best_skill(self, user_id: str | None = None) -> dict | None: + """Return the memory record for the highest-scored skill iteration. + + Searches mem0 for skill_iteration records and picks the one with + the highest ``score`` in its metadata. + + Parameters + ---------- + user_id : str | None + Override the user_id for this query (defaults to ``self.user_id``). + + Returns + ------- + dict | None + The memory record with the highest score, or ``None`` if no skill + iterations have been stored yet. + """ + results = self._client.search( + "skill iteration score evaluation", + user_id=user_id or self.user_id, + limit=50, + ) + if isinstance(results, dict): + results = results.get("results", []) + if not results: + return None + + best: dict | None = None + best_score = -1.0 + for rec in results: + meta = rec.get("metadata") or {} + if meta.get("event_type") != "skill_iteration": + continue + try: + s = float(meta.get("score", -1)) + except (TypeError, ValueError): + continue + if s > best_score: + best_score = s + best = rec + + return best + + def store_experiment_result( + self, + config_name: str, + final_score: float, + skill_hash: str, + ) -> Any: + """Store the final outcome of an experiment run. + + Parameters + ---------- + config_name : str + Human-readable name / path of the config used for this run. + final_score : float + Best evaluation score achieved during the run. + skill_hash : str + Hash of the best skill version (from :meth:`store_skill_iteration`). + """ + base_meta: dict[str, Any] = { + "event_type": "experiment_result", + "config_name": config_name, + "final_score": round(float(final_score), 6), + "skill_hash": skill_hash, + } + + content = ( + f"[SkillOpt] Experiment complete — config={config_name!r} " + f"final_score={final_score:.4f} best_skill_hash={skill_hash}" + ) + + messages = [{"role": "user", "content": content}] + return self._add(messages, metadata=base_meta) + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/skillopt/memory/trainer_hooks.py b/skillopt/memory/trainer_hooks.py new file mode 100644 index 00000000..dcdd625f --- /dev/null +++ b/skillopt/memory/trainer_hooks.py @@ -0,0 +1,147 @@ +"""Lightweight integration hooks for injecting SkillMemory into ReflACT training. + +These hooks are designed to be **non-breaking**: if ``MEM0_API_KEY`` is not +present in the environment, :func:`maybe_init_mem0` returns ``None`` and +every ``hook_*`` function becomes a no-op. + +Typical usage inside ``trainer.py``:: + + from skillopt.memory.trainer_hooks import ( + maybe_init_mem0, + hook_post_evaluate, + hook_post_reflect, + ) + + memory = maybe_init_mem0(cfg) + + # ... inside training loop, after evaluation gate: + hook_post_evaluate(memory, epoch, step, current_skill, gate_score, cfg) + + # ... after reflection stage: + hook_post_reflect(memory, epoch, step, raw_patches) +""" +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from skillopt.memory.mem0_backend import SkillMemory + + +# ── Initialisation ──────────────────────────────────────────────────────────── + +def maybe_init_mem0(cfg: dict) -> "SkillMemory | None": + """Attempt to initialise a :class:`~skillopt.memory.SkillMemory` instance. + + Reads ``MEM0_API_KEY`` from the environment. If the key is absent or + mem0ai is not installed, returns ``None`` (graceful degradation). + + Parameters + ---------- + cfg : dict + Flat trainer config dict. The ``config_name`` key (if present) is used + to label experiment results. + + Returns + ------- + SkillMemory | None + A live ``SkillMemory`` instance, or ``None`` if mem0 is unavailable. + """ + api_key = os.environ.get("MEM0_API_KEY", "") + if not api_key: + return None + + try: + from skillopt.memory.mem0_backend import SkillMemory # local import to avoid hard dep + + user_id = str(cfg.get("config_name") or cfg.get("env") or "skillopt") + memory = SkillMemory(api_key=api_key, user_id=user_id) + print(f" [mem0] SkillMemory initialised — user_id={user_id!r}") + return memory + except Exception as exc: # pragma: no cover + print(f" [mem0] WARNING: could not initialise SkillMemory: {exc}") + return None + + +# ── Post-evaluate hook ──────────────────────────────────────────────────────── + +def hook_post_evaluate( + memory: "SkillMemory | None", + epoch: int, + step: int, + skill: str, + score: float, + cfg: dict, +) -> None: + """Called after the evaluation gate — stores the current skill + score. + + Parameters + ---------- + memory : SkillMemory | None + The memory backend. If ``None``, this function is a no-op. + epoch : int + Current training epoch. + step : int + Current training step. + skill : str + Full text of the current skill document. + score : float + Gate score for this skill version. + cfg : dict + Flat trainer config (used to extract optional metadata). + """ + if memory is None: + return + try: + meta = { + "env": cfg.get("env", ""), + "optimizer_model": cfg.get("optimizer_model", ""), + "target_model": cfg.get("target_model", ""), + } + memory.store_skill_iteration( + epoch=epoch, + step=step, + skill_text=skill, + score=score, + metadata=meta, + ) + except Exception as exc: # pragma: no cover + print(f" [mem0] WARNING: hook_post_evaluate failed: {exc}") + + +# ── Post-reflect hook ───────────────────────────────────────────────────────── + +def hook_post_reflect( + memory: "SkillMemory | None", + epoch: int, + step: int, + patches: list, + scores: dict | None = None, +) -> None: + """Called after the Reflect stage — stores patch metadata. + + Parameters + ---------- + memory : SkillMemory | None + The memory backend. If ``None``, this function is a no-op. + epoch : int + Current training epoch. + step : int + Current training step. + patches : list + Raw patches returned by the reflection stage (list of dicts). + scores : dict | None + Optional rollout scores from this step (e.g. ``{"hard": 0.7}``). + """ + if memory is None: + return + try: + memory.store_reflection( + epoch=epoch, + step=step, + patches=list(patches) if patches else [], + scores=scores, + ) + except Exception as exc: # pragma: no cover + print(f" [mem0] WARNING: hook_post_reflect failed: {exc}") From 27d6cd9501bb2ced711559bbc87b912d4d284bde Mon Sep 17 00:00:00 2001 From: jrauch713-svg Date: Mon, 27 Jul 2026 13:00:01 -0500 Subject: [PATCH 2/2] feat(memory): explicit opt-in, redaction, namespacing, and a retrieval path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the review on #118. The integration was previously write-only and enabled by the mere presence of MEM0_API_KEY; both are fixed here. 1. Explicit opt-in. Uploading now requires `mem0_enabled: true`. A key present in the environment for another application no longer enables anything — a disabled run does not even read the key. New `skillopt/memory/settings.py` resolves all `mem0_*` config in one place. 2. Redaction before any outbound request. New `skillopt/memory/redaction.py` is the single choke point: no payload reaches the client without passing through `redact_for_upload`, which strips vendor keys, bearer/basic tokens, JWTs, private keys, `key = value` secret assignments, the project root, and `/home/`-style prefixes. Credentials are scrubbed before paths are collapsed, so a key embedded in a home path cannot survive. Relative paths and filenames are preserved so stored memories remain useful. The patterns deliberately mirror `skillopt_sleep/staging.py` rather than importing it — pyproject keeps that package decoupled with zero research dependency. 3. Stable project-specific namespace. Memories are scoped to `skillopt::`, which is stable across runs of one project and distinct across projects. The raw path is hashed, never transmitted. Replaces the previous config-name/env user id that could mix unrelated projects. 4. One bounded retrieval call before reflection. `hook_pre_reflect` fetches relevant history and appends it to the reflection input, turning the integration from external logging into a memory loop. It is deliberately assigned to a separate `reflect_context` variable so the autonomous learning-rate decision and the rewrite prompt continue to see the unaugmented context. Also in this commit: - Every call is wall-clock bounded by `mem0_timeout_seconds` (default 5s), so an unreachable service costs one bounded pause instead of blocking a step. - Malformed patches (None, strings, non-lists) are filtered rather than raised inside a swallowing try. - The `mem0` optional extra is declared, deliberately outside `all`. - The `mem0_*` keys are mapped through `_FLATTEN_MAP` under `train.`. Without this they were silently dropped for structured configs — which is every config in the repo, since they all inherit `_base_/default.yaml` — leaving the feature inert with no error. - `configs/features/mem0_memory.yaml` is a documented opt-in example following the `soft_gate.yaml` convention. - The config reference documents every key and exactly what leaves the machine. Tests (`tests/test_mem0_memory.py`, 15 cases, no network) cover: a bare MEM0_API_KEY not enabling uploads, explicit opt-in writing, redaction of secrets and home paths, cap applied after redaction, namespace stability and project separation, structured-config keys surviving flattening, the shipped example config actually enabling the feature, retrieval reaching the actual reflection prompt (verified end-to-end by capturing the optimizer's user message), retrieval being independently disableable, graceful degradation on service failure, timeout bounding, and malformed-patch filtering. Not included, per review guidance: batching, async writes, ranking, and benchmark study are left as follow-up work. --- configs/features/mem0_memory.yaml | 71 +++++ docs/reference/config.md | 48 ++++ pyproject.toml | 4 + skillopt/config.py | 8 + skillopt/engine/trainer.py | 16 +- skillopt/memory/__init__.py | 19 +- skillopt/memory/mem0_backend.py | 345 ++++++++++++------------ skillopt/memory/redaction.py | 135 ++++++++++ skillopt/memory/settings.py | 123 +++++++++ skillopt/memory/trainer_hooks.py | 172 +++++++----- tests/test_mem0_memory.py | 417 ++++++++++++++++++++++++++++++ 11 files changed, 1106 insertions(+), 252 deletions(-) create mode 100644 configs/features/mem0_memory.yaml create mode 100644 skillopt/memory/redaction.py create mode 100644 skillopt/memory/settings.py create mode 100644 tests/test_mem0_memory.py diff --git a/configs/features/mem0_memory.yaml b/configs/features/mem0_memory.yaml new file mode 100644 index 00000000..6ad5505d --- /dev/null +++ b/configs/features/mem0_memory.yaml @@ -0,0 +1,71 @@ +# ───────────────────────────────────────────────────────────────────────────── +# Feature: optional mem0-backed persistent memory (community-contributed) +# ───────────────────────────────────────────────────────────────────────────── +# +# This is NOT a default SkillOpt setting and was NOT used to produce the +# numbers reported in the paper. It is off unless you turn it on here. +# +# IMPORTANT — this feature sends data to a third-party hosted service. +# Enabling it means skill text and reflection summaries leave your machine. +# Setting MEM0_API_KEY alone does *nothing*; `mem0_enabled: true` below is the +# only thing that turns uploading on, and a key set for some other application +# will never enable it by accident. +# +# What it does: +# - After each evaluation gate, stores the current skill and its score. +# - After each Reflect stage, stores a summary of the patches produced. +# - Before each Reflect stage, makes ONE bounded retrieval call and appends +# the relevant history to that stage's reflection input, so previous runs +# can inform the patches produced now. +# +# When to consider this: +# - You run the same project repeatedly and want later runs to benefit from +# what earlier ones discovered. +# - You are comfortable with the data-handling note above. +# +# When NOT to use this: +# - When reproducing the paper. +# - When the skill text or trajectories are sensitive. Payloads are redacted +# (see docs/reference/config.md → "What leaves the machine"), but redaction +# is a mitigation, not a guarantee about content you know to be private. +# +# Install the optional dependency first: +# pip install 'skillopt[mem0]' +# +# Then export your key and use this file: +# export MEM0_API_KEY=m0-... +# _base_: ../features/mem0_memory.yaml +# or copy the `train:` block below into your config. +# ───────────────────────────────────────────────────────────────────────────── + +_base_: ../_base_/default.yaml + +train: + # The master switch. Nothing is uploaded while this is false. + mem0_enabled: true + + # Optional: set the key in config instead of the MEM0_API_KEY env var. + # Leave empty to use the environment (recommended — keeps keys out of YAML). + mem0_api_key: "" + + # Optional: override the namespace. Default is derived per project as + # `skillopt::`, which is stable across + # runs and distinct across projects. Set this only if you deliberately want + # several checkouts to share one memory pool. + mem0_namespace: "" + + # Read memory back into reflection. Setting this false keeps the writes but + # removes the retrieval call — useful for measuring whether retrieval is what + # actually helps, rather than assuming it does. + mem0_retrieval_enabled: true + + # Records fetched per retrieval. + mem0_retrieval_limit: 5 + + # Hard per-call wall-clock bound. On timeout the step continues unchanged, so + # an unreachable service costs one bounded pause rather than a stalled run. + mem0_timeout_seconds: 5.0 + + # Cap on any single stored payload, applied AFTER redaction so the limit is + # measured on exactly the text that would be transmitted. + mem0_max_chars: 4000 diff --git a/docs/reference/config.md b/docs/reference/config.md index b8243c77..cbf6664b 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -138,6 +138,54 @@ defaults to `claude` and can be overridden with `CLAUDE_CLI_BIN`. Benchmark-specific `env` keys are passed through to the adapter. +## Persistent Memory (mem0) — optional, off by default + +Stores skill iterations and reflection summaries in [mem0](https://mem0.ai) and +reads a small amount of relevant history back into the Reflect stage. + +**Disabled unless `mem0_enabled` is explicitly true.** A `MEM0_API_KEY` present +in the environment for another application does *not* enable it. Install with +`pip install 'skillopt[mem0]'`. + +Set these under `train:` in a structured config (every shipped config is +structured), at the top level of a flat config, or via +`--cfg-options mem0_enabled=true`. A ready-made example ships at +`configs/features/mem0_memory.yaml`. + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `train.mem0_enabled` | bool | `false` | Master switch. Nothing is sent unless this is true | +| `train.mem0_api_key` | str | empty | Falls back to `MEM0_API_KEY`, read only when enabled | +| `train.mem0_namespace` | str | derived | Override the namespace; default is derived per project | +| `train.mem0_retrieval_enabled` | bool | `true` | Read memory back into reflection; writes continue if false | +| `train.mem0_retrieval_limit` | int | `5` | Max records fetched per retrieval | +| `train.mem0_timeout_seconds` | float | `5.0` | Hard per-call bound; on timeout training continues | +| `train.mem0_max_chars` | int | `4000` | Cap on any single stored payload, applied after redaction | + +### What leaves the machine + +When enabled, exactly two record types are sent, both redacted first: + +1. **Skill iteration** — epoch, step, score, skill hash and length, the `env` + and model names, and the skill text (capped at `mem0_max_chars`). +2. **Reflection summary** — epoch, step, patch count, rollout scores, and a + JSON summary of up to 10 patches with any `skill_text` field removed. + +Retrieval sends the first 600 characters of the current skill as a similarity +query. + +Before transmission every payload passes through +`skillopt.memory.redaction.redact_for_upload`, which strips vendor API keys, +bearer/basic tokens, JWTs, private keys, `key = value` secret assignments, the +project root, and `/home/`-style prefixes. Relative paths and filenames +are deliberately preserved so stored memories stay useful. + +### Namespacing + +Memories are scoped to `skillopt::`, where the digest is a SHA-256 +prefix of the absolute project root. Stable across runs of one project, distinct +across projects, and the raw path is never transmitted. + ## Credential Environment Variables ### Azure-family backend diff --git a/pyproject.toml b/pyproject.toml index 45544eb3..32dad45f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,10 @@ searchqa = ["datasets>=2.18.0"] docs = ["mkdocs-material>=9.5.0", "mkdocstrings[python]>=0.24.0"] # WebUI dashboard webui = ["gradio>=4.0.0"] +# Optional mem0-backed persistent memory. Deliberately excluded from `all`: +# installing it is harmless (uploads still require mem0_enabled: true), but +# a data-export integration should be an explicit choice at install time too. +mem0 = ["mem0ai>=0.1.0"] # Development tools dev = ["ruff>=0.4.0", "pytest>=8.0.0"] # All optional dependencies (except docs/dev/webui) diff --git a/skillopt/config.py b/skillopt/config.py index c47e6d9b..b1651217 100644 --- a/skillopt/config.py +++ b/skillopt/config.py @@ -105,6 +105,14 @@ "train.batch_size": "batch_size", "train.accumulation": "accumulation", "train.seed": "seed", + # Optional mem0 memory (off unless train.mem0_enabled is true). + "train.mem0_enabled": "mem0_enabled", + "train.mem0_api_key": "mem0_api_key", + "train.mem0_namespace": "mem0_namespace", + "train.mem0_retrieval_enabled": "mem0_retrieval_enabled", + "train.mem0_retrieval_limit": "mem0_retrieval_limit", + "train.mem0_timeout_seconds": "mem0_timeout_seconds", + "train.mem0_max_chars": "mem0_max_chars", "gradient.minibatch_size": "minibatch_size", "gradient.merge_batch_size": "merge_batch_size", "gradient.analyst_workers": "analyst_workers", diff --git a/skillopt/engine/trainer.py b/skillopt/engine/trainer.py index aca5b34b..77df5a08 100644 --- a/skillopt/engine/trainer.py +++ b/skillopt/engine/trainer.py @@ -75,7 +75,12 @@ set_optimizer_deployment, ) from skillopt.utils import compute_score, skill_hash -from skillopt.memory.trainer_hooks import maybe_init_mem0, hook_post_evaluate, hook_post_reflect +from skillopt.memory.trainer_hooks import ( + maybe_init_mem0, + hook_pre_reflect, + hook_post_evaluate, + hook_post_reflect, +) # ── Skill-aware reflection: appendix flush ─────────────────────────────────── @@ -1161,11 +1166,18 @@ def _persist_runtime_state(last_completed_step: int) -> None: # Build step context from buffer step_buffer_context = _format_step_buffer(step_buffer) + # Memory read: relevant history for THIS reflection only. + # Kept in a separate name so the LR decision and rewrite + # prompts below continue to see the unaugmented context. + reflect_context = hook_pre_reflect( + memory, current_skill, step_buffer_context, + ) + raw_patches = adapter.reflect( rollout_results, current_skill, batch_dir, prediction_dir=pred_dir, patches_dir=patches_dir, random_seed=batch_seed, - step_buffer_context=step_buffer_context, + step_buffer_context=reflect_context, meta_skill_context=active_meta_skill, ) failure_patches, success_patches = _normalise_patches( diff --git a/skillopt/memory/__init__.py b/skillopt/memory/__init__.py index 4d6ce551..069d9763 100644 --- a/skillopt/memory/__init__.py +++ b/skillopt/memory/__init__.py @@ -1,4 +1,17 @@ -"""skillopt.memory — mem0-backed persistent memory for SkillOpt.""" -from skillopt.memory.mem0_backend import SkillMemory +"""skillopt.memory — optional, opt-in mem0-backed persistent memory. -__all__ = ["SkillMemory"] +Disabled unless a config sets ``mem0_enabled: true``; see +:mod:`skillopt.memory.settings` for the resolution rules and +:mod:`skillopt.memory.redaction` for what is stripped before anything is sent. +""" +from skillopt.memory.mem0_backend import SkillMemory, mem0_available +from skillopt.memory.redaction import redact_for_upload +from skillopt.memory.settings import Mem0Settings, resolve_settings + +__all__ = [ + "Mem0Settings", + "SkillMemory", + "mem0_available", + "redact_for_upload", + "resolve_settings", +] diff --git a/skillopt/memory/mem0_backend.py b/skillopt/memory/mem0_backend.py index ad59cc12..f2c05d84 100644 --- a/skillopt/memory/mem0_backend.py +++ b/skillopt/memory/mem0_backend.py @@ -1,24 +1,43 @@ """mem0-backed persistent memory for SkillOpt. -Stores skill iterations, reflection results, and experiment outcomes in mem0 -so that the ReflACT trainer can retrieve relevant historical context across -runs and identify the best skill versions discovered so far. +Stores skill iterations and reflection outcomes, and reads a small amount of +relevant history back into the Reflect stage so the memory participates in +training rather than only recording it. + +Safety contract +--------------- +1. **Nothing is sent unless the run opted in.** The client is only constructed + from a :class:`~skillopt.memory.settings.Mem0Settings` whose ``usable`` is + true, which requires ``mem0_enabled`` to have been set explicitly. +2. **Everything outbound is redacted.** Every payload is built through + :meth:`SkillMemory._payload`, which applies + :func:`~skillopt.memory.redaction.redact_for_upload` and then truncates. + There is no code path that reaches ``self._client`` with unredacted text. +3. **Every call is time-bounded.** Network work runs through :meth:`_bounded`, + so a slow or unreachable service costs at most ``timeout_seconds`` per + training step instead of blocking it. +4. **Failures never break training.** Every public method returns a benign + value on error; the trainer hooks additionally swallow anything unexpected. Usage:: from skillopt.memory import SkillMemory + from skillopt.memory.settings import resolve_settings - m = SkillMemory(api_key="m0-...") - m.store_skill_iteration(epoch=1, step=3, skill_text="...", score=0.82) - ctx = m.retrieve_relevant_context("handling multi-step navigation") + settings = resolve_settings(cfg) + memory = SkillMemory.from_settings(settings) # None unless opted in """ from __future__ import annotations import hashlib import json -import os +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import TimeoutError as _FuturesTimeout from typing import Any +from skillopt.memory.redaction import redact_for_upload +from skillopt.memory.settings import Mem0Settings + try: from mem0 import MemoryClient _MEM0_AVAILABLE = True @@ -27,61 +46,120 @@ MemoryClient = None # type: ignore[assignment,misc] +def mem0_available() -> bool: + """Whether the optional ``mem0ai`` dependency is importable.""" + return _MEM0_AVAILABLE + + class SkillMemory: """Persistent memory backend for SkillOpt using mem0. - Parameters - ---------- - api_key : str | None - mem0 API key. Falls back to ``MEM0_API_KEY`` env var, then to - ``MEM0_API_KEY`` set on the object at construction time. - user_id : str - Logical user/project identifier used to namespace memories in mem0. + Prefer :meth:`from_settings`; the constructor stays explicit so tests can + inject a fake client without touching the network. """ - def __init__( - self, - api_key: str | None = None, - user_id: str = "skillopt", - ) -> None: - if not _MEM0_AVAILABLE: - raise ImportError( - "mem0ai is not installed. Run: pip install mem0ai" - ) + def __init__(self, settings: Mem0Settings, client: Any | None = None) -> None: + if client is None: + if not _MEM0_AVAILABLE: + raise ImportError("mem0ai is not installed. Run: pip install 'skillopt[mem0]'") + if not settings.usable: + raise ValueError("SkillMemory requires mem0_enabled=true and an API key") + client = MemoryClient(api_key=settings.api_key) + + self.settings = settings + self.namespace = settings.namespace + self._client = client + # One worker: calls are already serialised by the training loop, and a + # single thread keeps the timeout semantics easy to reason about. + self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mem0") + + # ── Construction ────────────────────────────────────────────────────── + + @classmethod + def from_settings(cls, settings: Mem0Settings, client: Any | None = None) -> "SkillMemory | None": + """Return a backend, or ``None`` when the run has not opted in. + + Returning ``None`` rather than raising keeps callers free of + conditionals: every hook already treats ``None`` as "do nothing". + """ + if not settings.usable and client is None: + return None + if client is None and not _MEM0_AVAILABLE: + return None + return cls(settings, client=client) - resolved_key = api_key or os.environ.get("MEM0_API_KEY", "") - if not resolved_key: - raise ValueError( - "No mem0 API key provided. Pass api_key= or set MEM0_API_KEY env var." - ) + # ── Internal helpers ────────────────────────────────────────────────── - self.user_id = user_id - self._client = MemoryClient(api_key=resolved_key) + def _bounded(self, fn, *args, **kwargs) -> Any: + """Run *fn* with a hard wall-clock bound; ``None`` on timeout or error. - # ── Internal helpers ────────────────────────────────────────────────── + The worker thread may keep running after a timeout — Python cannot + cancel a blocked socket read — but the *training step* is released on + schedule, which is the property that matters here. + """ + try: + future = self._pool.submit(fn, *args, **kwargs) + return future.result(timeout=self.settings.timeout_seconds) + except _FuturesTimeout: + print(f" [mem0] call exceeded {self.settings.timeout_seconds}s — continuing without it") + return None + except Exception as exc: # noqa: BLE001 - memory must never break training + print(f" [mem0] call failed ({type(exc).__name__}: {exc}) — continuing without it") + return None - def _add(self, messages: list[dict], metadata: dict | None = None) -> Any: - """Low-level add wrapper — always tags with user_id.""" - kwargs: dict[str, Any] = {"user_id": self.user_id} - if metadata: - kwargs["metadata"] = metadata - return self._client.add(messages, **kwargs) + def _payload(self, text: str) -> str: + """The only route by which text becomes outbound content. - def _search(self, query: str, limit: int = 5) -> list[dict]: - """Low-level search wrapper — scoped to this user_id.""" - results = self._client.search(query, user_id=self.user_id, limit=limit) - # mem0 returns a list of memory dicts + Redacts first, truncates second, so the cap is measured on exactly the + string that would be transmitted. + """ + redacted = redact_for_upload(text, self.settings.project_root) + if not isinstance(redacted, str): + redacted = str(redacted) + limit = max(0, int(self.settings.max_chars)) + if len(redacted) > limit: + return redacted[:limit] + "\n[...truncated]" + return redacted + + def _add(self, content: str, metadata: dict | None = None) -> Any: + messages = [{"role": "user", "content": self._payload(content)}] + kwargs: dict[str, Any] = {"user_id": self.namespace} + if metadata: + kwargs["metadata"] = redact_for_upload(metadata, self.settings.project_root) + return self._bounded(self._client.add, messages, **kwargs) + + def _search(self, query: str, limit: int) -> list[dict]: + results = self._bounded( + self._client.search, + self._payload(query), + user_id=self.namespace, + limit=limit, + ) if isinstance(results, list): return results - # Some versions wrap results in a dict if isinstance(results, dict): - return results.get("results", []) + found = results.get("results", []) + return found if isinstance(found, list) else [] return [] @staticmethod def _short_hash(text: str) -> str: return hashlib.sha1(text.encode()).hexdigest()[:8] + @staticmethod + def _valid_patches(patches: Any) -> list[dict]: + """Keep only genuine dict patches. + + The reflection stage may yield ``None`` for a minibatch that produced + nothing, and a malformed backend reply can yield a bare string. Both + are filtered here rather than raising inside a ``try`` that swallows + the error, so a malformed patch costs one skipped record instead of + the whole memory write. + """ + if not isinstance(patches, (list, tuple)): + return [] + return [p for p in patches if isinstance(p, dict)] + # ── Public API ──────────────────────────────────────────────────────── def store_skill_iteration( @@ -92,29 +170,15 @@ def store_skill_iteration( score: float, metadata: dict | None = None, ) -> Any: - """Store a skill version produced during training. - - Parameters - ---------- - epoch : int - Current training epoch. - step : int - Current training step within the epoch. - skill_text : str - Full text of the skill document at this point. - score : float - Evaluation score (0–1) for this skill version. - metadata : dict | None - Optional additional metadata to attach. - """ - skill_hash = self._short_hash(skill_text) + """Store the skill version produced at *epoch*/*step*.""" + skill_hash = self._short_hash(skill_text or "") base_meta: dict[str, Any] = { "event_type": "skill_iteration", "epoch": epoch, "step": step, "score": round(float(score), 6), "skill_hash": skill_hash, - "skill_length": len(skill_text), + "skill_length": len(skill_text or ""), } if metadata: base_meta.update(metadata) @@ -122,41 +186,29 @@ def store_skill_iteration( content = ( f"[SkillOpt] Epoch {epoch} Step {step} — skill_hash={skill_hash} " f"score={score:.4f}\n\n" - f"=== SKILL TEXT ===\n{skill_text[:4000]}" # cap to avoid huge payloads + f"=== SKILL TEXT ===\n{skill_text or ''}" ) - - messages = [{"role": "user", "content": content}] - return self._add(messages, metadata=base_meta) + return self._add(content, metadata=base_meta) def store_reflection( self, epoch: int, step: int, - patches: list[dict], + patches: Any, scores: dict | None = None, ) -> Any: - """Store the result of a Reflect stage. - - Parameters - ---------- - epoch : int - Current training epoch. - step : int - Current training step. - patches : list[dict] - Raw patches produced by the reflection stage. - scores : dict | None - Optional dict of metric → value (e.g. ``{"hard": 0.7, "soft": 0.8}``). - """ - n_patches = len(patches) + """Store a summary of the patches produced by one Reflect stage.""" + valid = self._valid_patches(patches) + n_patches = len(valid) scores_str = json.dumps(scores or {}, ensure_ascii=False) - patch_summary = json.dumps( - [ - {k: v for k, v in p.items() if k != "skill_text"} - for p in patches[:10] # only first 10 to keep it concise - ], - ensure_ascii=False, - ) + try: + patch_summary = json.dumps( + [{k: v for k, v in p.items() if k != "skill_text"} for p in valid[:10]], + ensure_ascii=False, + default=str, + ) + except (TypeError, ValueError): + patch_summary = "[unserialisable patch summary]" base_meta: dict[str, Any] = { "event_type": "reflection", @@ -172,108 +224,39 @@ def store_reflection( f"{n_patches} patch(es) generated. Scores: {scores_str}\n\n" f"Patch summary:\n{patch_summary}" ) + return self._add(content, metadata=base_meta) - messages = [{"role": "user", "content": content}] - return self._add(messages, metadata=base_meta) - - def retrieve_relevant_context( - self, - query: str, - limit: int = 5, - ) -> list[dict]: - """Retrieve past memories relevant to *query*. - - Parameters - ---------- - query : str - Free-text query describing the context you want to retrieve. - limit : int - Maximum number of results to return. - - Returns - ------- - list[dict] - List of memory objects (each has at least a ``memory`` field). - """ - return self._search(query, limit=limit) - - def get_best_skill(self, user_id: str | None = None) -> dict | None: - """Return the memory record for the highest-scored skill iteration. - - Searches mem0 for skill_iteration records and picks the one with - the highest ``score`` in its metadata. + def retrieve_relevant_context(self, query: str, limit: int | None = None) -> list[dict]: + """Return past memories relevant to *query* (empty on any failure).""" + if not self.settings.retrieval_enabled: + return [] + return self._search(query, limit=limit or self.settings.retrieval_limit) - Parameters - ---------- - user_id : str | None - Override the user_id for this query (defaults to ``self.user_id``). + def format_retrieved_context(self, memories: list[dict], max_chars: int = 2000) -> str: + """Render retrieved memories as text for inclusion in a prompt. - Returns - ------- - dict | None - The memory record with the highest score, or ``None`` if no skill - iterations have been stored yet. + Returns ``""`` when there is nothing useful, so callers can append + unconditionally without producing an empty heading. """ - results = self._client.search( - "skill iteration score evaluation", - user_id=user_id or self.user_id, - limit=50, - ) - if isinstance(results, dict): - results = results.get("results", []) - if not results: - return None - - best: dict | None = None - best_score = -1.0 - for rec in results: - meta = rec.get("metadata") or {} - if meta.get("event_type") != "skill_iteration": + lines: list[str] = [] + for rec in memories or []: + if not isinstance(rec, dict): continue - try: - s = float(meta.get("score", -1)) - except (TypeError, ValueError): + text = rec.get("memory") or rec.get("text") or "" + if not isinstance(text, str) or not text.strip(): continue - if s > best_score: - best_score = s - best = rec + lines.append(f"- {text.strip()}") + if not lines: + return "" - return best + body = "\n".join(lines) + if len(body) > max_chars: + body = body[:max_chars] + "\n[...truncated]" + return "## Relevant Memory From Previous Runs\n" + body - def store_experiment_result( - self, - config_name: str, - final_score: float, - skill_hash: str, - ) -> Any: - """Store the final outcome of an experiment run. - - Parameters - ---------- - config_name : str - Human-readable name / path of the config used for this run. - final_score : float - Best evaluation score achieved during the run. - skill_hash : str - Hash of the best skill version (from :meth:`store_skill_iteration`). - """ - base_meta: dict[str, Any] = { - "event_type": "experiment_result", - "config_name": config_name, - "final_score": round(float(final_score), 6), - "skill_hash": skill_hash, - } - - content = ( - f"[SkillOpt] Experiment complete — config={config_name!r} " - f"final_score={final_score:.4f} best_skill_hash={skill_hash}" - ) - - messages = [{"role": "user", "content": content}] - return self._add(messages, metadata=base_meta) + def close(self) -> None: + """Release the worker thread. Safe to call more than once.""" + self._pool.shutdown(wait=False) def __repr__(self) -> str: - return ( - f"" - ) + return f"" diff --git a/skillopt/memory/redaction.py b/skillopt/memory/redaction.py new file mode 100644 index 00000000..686eee65 --- /dev/null +++ b/skillopt/memory/redaction.py @@ -0,0 +1,135 @@ +"""Redaction applied to every payload before it leaves the process. + +Mem0 is a third-party hosted service, so anything the trainer stores can leave +the machine. This module is the single choke point: :func:`redact_for_upload` +is called on every string in every payload built by +:mod:`skillopt.memory.mem0_backend`, so there is exactly one place to audit. + +What gets removed +----------------- +* **Credentials** — vendor API keys, bearer/basic tokens, JWTs, private keys, + and ``key = value`` style assignments for api-key/token/password/secret. +* **Filesystem identity** — the project root, the user's home directory, and + ``/home/`` / ``/Users/`` / ``C:\\Users\\`` prefixes, all of + which carry the operating-system username. + +What deliberately stays +----------------------- +Relative paths and file *names* survive: a skill that says "edit +``configs/train.yaml``" is still useful after redaction, and the portion +removed is the machine-specific prefix rather than the structure. Redaction +that destroyed the text would defeat the point of storing it. + +The secret patterns intentionally mirror +``skillopt_sleep/staging.py::_SECRET_PATTERNS``. They are duplicated rather +than imported because ``skillopt_sleep`` is decoupled from the research +package by design (``pyproject.toml``: "the open-source Sleep tool (decoupled, +zero research dep)"); importing across that boundary to save a few lines would +couple two packages the project keeps apart on purpose. +""" +from __future__ import annotations + +import os +import re +from typing import Any + +# Credential patterns — mirrored from skillopt_sleep/staging.py (see module +# docstring for why this is a copy and not an import). +_SECRET_PATTERNS: tuple[tuple[re.Pattern[str], str], ...] = ( + (re.compile(r"sk-[A-Za-z0-9_-]{10,}"), "[REDACTED_API_KEY]"), + (re.compile(r"\bAKIA[0-9A-Z]{16}\b"), "[REDACTED_AWS_KEY]"), + (re.compile(r"\bgh[pousr]_[A-Za-z0-9]{20,}\b"), "[REDACTED_GITHUB_TOKEN]"), + (re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), "[REDACTED_SLACK_TOKEN]"), + (re.compile(r"\bAIza[0-9A-Za-z_-]{20,}\b"), "[REDACTED_GOOGLE_KEY]"), + (re.compile(r"\bm0-[A-Za-z0-9_-]{16,}\b"), "[REDACTED_MEM0_KEY]"), + ( + re.compile(r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\b"), + "[REDACTED_JWT]", + ), + (re.compile(r"(?i)(Authorization:\s*Bearer\s+)[^\s\"']+"), r"\1[REDACTED]"), + (re.compile(r"(?i)(Authorization:\s*Basic\s+)[^\s\"']+"), r"\1[REDACTED]"), + ( + re.compile(r"(?i)\b(api[_-]?key|token|password|secret)\b(\s*[:=]\s*)[^\s\"']+"), + r"\1\2[REDACTED]", + ), + ( + re.compile(r"(?i)\b(api[_-]?key|token|password|secret)\b(\s+)[^\s\"']+"), + r"\1\2[REDACTED]", + ), + ( + re.compile( + r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----", + re.DOTALL, + ), + "[REDACTED_PRIVATE_KEY]", + ), +) + +# Home-directory prefixes carry the OS username, which is personally +# identifying. Replaced with "~" so the trailing structure stays readable. +_HOME_PATTERNS: tuple[tuple[re.Pattern[str], str], ...] = ( + (re.compile(r"/home/[^/\s\"']+"), "~"), + (re.compile(r"/Users/[^/\s\"']+"), "~"), + (re.compile(r"(?i)\b[A-Z]:\\Users\\[^\\\s\"']+"), "~"), +) + +_PROJECT_PLACEHOLDER = "" + + +def redact_secrets(value: Any) -> Any: + """Scrub credential-looking substrings from every string leaf of *value*. + + Strings are rewritten; lists and dicts are walked recursively; other + scalars pass through unchanged. + """ + if isinstance(value, str): + out = value + for pattern, replacement in _SECRET_PATTERNS: + out = pattern.sub(replacement, out) + return out + if isinstance(value, list): + return [redact_secrets(v) for v in value] + if isinstance(value, dict): + return {k: redact_secrets(v) for k, v in value.items()} + return value + + +def redact_paths(value: Any, project_root: str = "") -> Any: + """Replace machine-identifying path prefixes in every string leaf. + + *project_root*, when given, is collapsed to ```` first so that a + path inside the run directory does not also match a home-directory rule. + """ + if isinstance(value, str): + out = value + root = (project_root or "").rstrip("/\\") + if root: + out = out.replace(root, _PROJECT_PLACEHOLDER) + for pattern, replacement in _HOME_PATTERNS: + out = pattern.sub(replacement, out) + return out + if isinstance(value, list): + return [redact_paths(v, project_root) for v in value] + if isinstance(value, dict): + return {k: redact_paths(v, project_root) for k, v in value.items()} + return value + + +def redact_for_upload(value: Any, project_root: str = "") -> Any: + """Full outbound redaction: credentials first, then path identity. + + This is the only function callers should need. Order matters — a secret + embedded in a path (``/home/alice/.config/sk-abc123``) must be scrubbed as + a credential before the home prefix is collapsed, or the key would survive + inside the shortened string. + """ + return redact_paths(redact_secrets(value), project_root) + + +def default_project_root(cfg: dict | None = None) -> str: + """Best-effort absolute project root, used to anchor path redaction.""" + if cfg: + root = cfg.get("out_root") + if root: + return os.path.abspath(str(root)) + return os.path.abspath(os.getcwd()) diff --git a/skillopt/memory/settings.py b/skillopt/memory/settings.py new file mode 100644 index 00000000..ff0efc3b --- /dev/null +++ b/skillopt/memory/settings.py @@ -0,0 +1,123 @@ +"""Resolution of the ``mem0_*`` trainer config keys. + +The single rule this module exists to enforce: **an API key present in the +environment must never, by itself, cause data to leave the machine.** Uploading +is opt-in via ``mem0_enabled``, and a key is only consulted once that flag is +explicitly true. A ``MEM0_API_KEY`` set for some unrelated application is +therefore inert here. + +Namespacing +----------- +Memories are scoped to a namespace derived from the *project root* rather than +from a config name, so two checkouts that happen to share a config label do not +read each other's memories. The project path is hashed rather than sent +verbatim: the namespace has to be stable and unique, and it does not need to be +legible to the remote service. +""" +from __future__ import annotations + +import hashlib +import os +from dataclasses import dataclass + +# Ceiling on any single stored payload. Applied after redaction, so the cap is +# measured on exactly the text that would be transmitted. +DEFAULT_MAX_CHARS = 4000 +DEFAULT_TIMEOUT_SECONDS = 5.0 +DEFAULT_RETRIEVAL_LIMIT = 5 + + +@dataclass(frozen=True) +class Mem0Settings: + """Fully resolved mem0 configuration for one training run.""" + + enabled: bool = False + api_key: str = "" + namespace: str = "skillopt" + retrieval_enabled: bool = True + retrieval_limit: int = DEFAULT_RETRIEVAL_LIMIT + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS + max_chars: int = DEFAULT_MAX_CHARS + project_root: str = "" + + @property + def usable(self) -> bool: + """True only when uploading was explicitly enabled *and* a key exists.""" + return bool(self.enabled and self.api_key) + + +def project_namespace(project_root: str, env_name: str = "") -> str: + """A stable, project-specific namespace. + + Stable across runs of the same project (same absolute root → same digest) + and distinct across projects, which is what keeps unrelated experiments + from sharing a memory pool. + """ + root = os.path.abspath(project_root or os.getcwd()) + digest = hashlib.sha256(root.encode("utf-8")).hexdigest()[:12] + env_part = (env_name or "default").strip().replace(" ", "_") + return f"skillopt:{env_part}:{digest}" + + +def _as_bool(value: object, default: bool = False) -> bool: + """YAML may hand us a real bool; ``--cfg-options`` hands us a string.""" + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +def _as_int(value: object, default: int) -> int: + try: + return int(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return default + + +def _as_float(value: object, default: float) -> float: + try: + return float(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return default + + +def resolve_settings(cfg: dict | None, env: dict | None = None) -> Mem0Settings: + """Build :class:`Mem0Settings` from a flat trainer config. + + Parameters + ---------- + cfg : dict | None + Flat trainer config. ``mem0_enabled`` is the master switch. + env : dict | None + Environment mapping, defaulting to :data:`os.environ`. Only consulted + for the API key, and only when ``mem0_enabled`` is true. + """ + cfg = cfg or {} + env = os.environ if env is None else env + + enabled = _as_bool(cfg.get("mem0_enabled"), False) + if not enabled: + # Return the inert default rather than reading the key at all, so a + # disabled run cannot even accidentally hold a credential. + return Mem0Settings(enabled=False) + + api_key = str(cfg.get("mem0_api_key") or env.get("MEM0_API_KEY", "") or "") + + project_root = os.path.abspath(str(cfg.get("out_root") or os.getcwd())) + namespace = str(cfg.get("mem0_namespace") or "").strip() + if not namespace: + namespace = project_namespace(project_root, str(cfg.get("env") or "")) + + return Mem0Settings( + enabled=True, + api_key=api_key, + namespace=namespace, + retrieval_enabled=_as_bool(cfg.get("mem0_retrieval_enabled"), True), + retrieval_limit=_as_int(cfg.get("mem0_retrieval_limit"), DEFAULT_RETRIEVAL_LIMIT), + timeout_seconds=_as_float(cfg.get("mem0_timeout_seconds"), DEFAULT_TIMEOUT_SECONDS), + max_chars=_as_int(cfg.get("mem0_max_chars"), DEFAULT_MAX_CHARS), + project_root=project_root, + ) diff --git a/skillopt/memory/trainer_hooks.py b/skillopt/memory/trainer_hooks.py index dcdd625f..b59f3dd3 100644 --- a/skillopt/memory/trainer_hooks.py +++ b/skillopt/memory/trainer_hooks.py @@ -1,67 +1,137 @@ -"""Lightweight integration hooks for injecting SkillMemory into ReflACT training. +"""Integration hooks binding :class:`SkillMemory` into the ReflACT loop. -These hooks are designed to be **non-breaking**: if ``MEM0_API_KEY`` is not -present in the environment, :func:`maybe_init_mem0` returns ``None`` and -every ``hook_*`` function becomes a no-op. +Every hook is a no-op when *memory* is ``None``, which is the state of any run +that did not set ``mem0_enabled: true``. That keeps the trainer free of +feature-flag branches: it calls the hooks unconditionally and they decide. + +The read/write pair is what makes this a memory rather than a log: + +* :func:`hook_pre_reflect` — **reads**. Runs before the Reflect stage and + appends relevant history to the reflection context, so past runs influence + the patches produced now. +* :func:`hook_post_reflect` / :func:`hook_post_evaluate` — **write**. Typical usage inside ``trainer.py``:: from skillopt.memory.trainer_hooks import ( - maybe_init_mem0, - hook_post_evaluate, - hook_post_reflect, + maybe_init_mem0, hook_pre_reflect, hook_post_evaluate, hook_post_reflect, ) memory = maybe_init_mem0(cfg) - # ... inside training loop, after evaluation gate: - hook_post_evaluate(memory, epoch, step, current_skill, gate_score, cfg) + # before the Reflect stage: + step_buffer_context = hook_pre_reflect(memory, current_skill, step_buffer_context) - # ... after reflection stage: - hook_post_reflect(memory, epoch, step, raw_patches) + # after reflection / after the evaluation gate: + hook_post_reflect(memory, epoch, step, raw_patches, scores=...) + hook_post_evaluate(memory, epoch, step, current_skill, gate_score, cfg) """ from __future__ import annotations -import os from typing import TYPE_CHECKING +from skillopt.memory.settings import resolve_settings + if TYPE_CHECKING: from skillopt.memory.mem0_backend import SkillMemory +# Longest slice of skill text used to build a retrieval query. The query is +# only a similarity probe, so a head slice is enough and keeps the outbound +# request small. +_QUERY_CHARS = 600 + # ── Initialisation ──────────────────────────────────────────────────────────── def maybe_init_mem0(cfg: dict) -> "SkillMemory | None": - """Attempt to initialise a :class:`~skillopt.memory.SkillMemory` instance. + """Initialise memory for this run, or return ``None``. + + Returns ``None`` — sending nothing — unless **all** of the following hold: + + 1. ``mem0_enabled`` is explicitly true in the config. A ``MEM0_API_KEY`` + present in the environment for some other application is *not* + sufficient and never has been sufficient since this check existed. + 2. An API key is resolvable (``mem0_api_key`` or ``MEM0_API_KEY``). + 3. The optional ``mem0ai`` dependency is installed. + """ + settings = resolve_settings(cfg) + if not settings.enabled: + return None + + if not settings.api_key: + print(" [mem0] mem0_enabled is set but no API key was found — memory disabled") + return None + + try: + from skillopt.memory.mem0_backend import SkillMemory, mem0_available + + if not mem0_available(): + print(" [mem0] mem0_enabled is set but mem0ai is not installed " + "(pip install 'skillopt[mem0]') — memory disabled") + return None + + memory = SkillMemory.from_settings(settings) + if memory is not None: + print(f" [mem0] memory enabled — namespace={memory.namespace!r} " + f"retrieval={'on' if settings.retrieval_enabled else 'off'}") + return memory + except Exception as exc: # noqa: BLE001 - memory must never break training + print(f" [mem0] WARNING: could not initialise memory: {exc}") + return None + - Reads ``MEM0_API_KEY`` from the environment. If the key is absent or - mem0ai is not installed, returns ``None`` (graceful degradation). +# ── Pre-reflect hook (the read side) ────────────────────────────────────────── + +def hook_pre_reflect( + memory: "SkillMemory | None", + skill_text: str, + step_buffer_context: str = "", + query: str | None = None, +) -> str: + """Return *step_buffer_context* with relevant past memories appended. + + This is the one retrieval call per step. It is bounded by + ``mem0_timeout_seconds`` inside the backend, and on any failure — timeout, + network error, empty result — the original context is returned unchanged, + so reflection proceeds exactly as it would with memory disabled. Parameters ---------- - cfg : dict - Flat trainer config dict. The ``config_name`` key (if present) is used - to label experiment results. + memory : SkillMemory | None + Backend, or ``None`` for a run without memory (returns input unchanged). + skill_text : str + Current skill document; its head is used as the similarity probe. + step_buffer_context : str + Context already assembled for this step. Retrieved memory is appended + to it rather than replacing it. + query : str | None + Explicit query, overriding the skill-derived one. Used by tests. Returns ------- - SkillMemory | None - A live ``SkillMemory`` instance, or ``None`` if mem0 is unavailable. + str + Possibly-augmented context, always safe to pass to ``adapter.reflect``. """ - api_key = os.environ.get("MEM0_API_KEY", "") - if not api_key: - return None + if memory is None: + return step_buffer_context try: - from skillopt.memory.mem0_backend import SkillMemory # local import to avoid hard dep + probe = query if query is not None else (skill_text or "")[:_QUERY_CHARS] + if not probe.strip(): + return step_buffer_context - user_id = str(cfg.get("config_name") or cfg.get("env") or "skillopt") - memory = SkillMemory(api_key=api_key, user_id=user_id) - print(f" [mem0] SkillMemory initialised — user_id={user_id!r}") - return memory - except Exception as exc: # pragma: no cover - print(f" [mem0] WARNING: could not initialise SkillMemory: {exc}") - return None + memories = memory.retrieve_relevant_context(probe) + block = memory.format_retrieved_context(memories) + if not block: + return step_buffer_context + + print(f" [mem0] retrieved {len(memories)} memory record(s) into reflection context") + if step_buffer_context and step_buffer_context.strip(): + return f"{step_buffer_context}\n\n{block}" + return block + except Exception as exc: # noqa: BLE001 - memory must never break training + print(f" [mem0] WARNING: hook_pre_reflect failed: {exc}") + return step_buffer_context # ── Post-evaluate hook ──────────────────────────────────────────────────────── @@ -74,23 +144,7 @@ def hook_post_evaluate( score: float, cfg: dict, ) -> None: - """Called after the evaluation gate — stores the current skill + score. - - Parameters - ---------- - memory : SkillMemory | None - The memory backend. If ``None``, this function is a no-op. - epoch : int - Current training epoch. - step : int - Current training step. - skill : str - Full text of the current skill document. - score : float - Gate score for this skill version. - cfg : dict - Flat trainer config (used to extract optional metadata). - """ + """Store the current skill and its gate score. No-op without memory.""" if memory is None: return try: @@ -106,7 +160,7 @@ def hook_post_evaluate( score=score, metadata=meta, ) - except Exception as exc: # pragma: no cover + except Exception as exc: # noqa: BLE001 - memory must never break training print(f" [mem0] WARNING: hook_post_evaluate failed: {exc}") @@ -119,29 +173,15 @@ def hook_post_reflect( patches: list, scores: dict | None = None, ) -> None: - """Called after the Reflect stage — stores patch metadata. - - Parameters - ---------- - memory : SkillMemory | None - The memory backend. If ``None``, this function is a no-op. - epoch : int - Current training epoch. - step : int - Current training step. - patches : list - Raw patches returned by the reflection stage (list of dicts). - scores : dict | None - Optional rollout scores from this step (e.g. ``{"hard": 0.7}``). - """ + """Store a summary of this step's patches. No-op without memory.""" if memory is None: return try: memory.store_reflection( epoch=epoch, step=step, - patches=list(patches) if patches else [], + patches=patches, scores=scores, ) - except Exception as exc: # pragma: no cover + except Exception as exc: # noqa: BLE001 - memory must never break training print(f" [mem0] WARNING: hook_post_reflect failed: {exc}") diff --git a/tests/test_mem0_memory.py b/tests/test_mem0_memory.py new file mode 100644 index 00000000..1f6f663f --- /dev/null +++ b/tests/test_mem0_memory.py @@ -0,0 +1,417 @@ +"""Mocked tests for the opt-in mem0 memory backend. + +Run directly (no pytest needed): + python tests/test_mem0_memory.py + +No network is touched: a fake client records what *would* have been sent, so +the assertions are about the exact bytes leaving the process. + +Covers: +1. Disabled by default — a ``MEM0_API_KEY`` in the environment must not, on its + own, enable uploads. This is the central safety property. +2. Explicit opt-in does enable, and writes reach the client. +3. Redaction — credentials and home-directory paths are stripped from the + payload, and the cap is applied to the redacted text. +4. Namespacing — stable per project, distinct across projects, and never the + raw filesystem path. +5. Retrieval reaches the *actual* reflection prompt, verified end-to-end by + capturing the user message handed to the optimizer. +6. Service failures (exceptions) degrade gracefully — training continues and + the reflection context is returned unchanged. +7. Slow services are bounded by ``mem0_timeout_seconds``. +8. Malformed patches (``None``, strings, non-lists) are filtered, not raised. +""" +from __future__ import annotations + +import json +import os +import sys +import tempfile +import time + +# Ensure THIS repo's skillopt is imported (not an installed copy) when the +# file is run directly: script mode puts tests/ on sys.path, not the repo root. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from skillopt.memory.mem0_backend import SkillMemory # noqa: E402 +from skillopt.memory.redaction import redact_for_upload # noqa: E402 +from skillopt.memory.settings import ( # noqa: E402 + project_namespace, + resolve_settings, +) +from skillopt.memory.trainer_hooks import ( # noqa: E402 + hook_post_reflect, + hook_pre_reflect, + maybe_init_mem0, +) + +FAKE_KEY = "m0-abcdefghijklmnopqrstuvwxyz012345" + + +class FakeClient: + """Stand-in for ``mem0.MemoryClient`` that records instead of sending.""" + + def __init__(self, search_results=None, fail: bool = False, delay: float = 0.0): + self.adds: list[tuple] = [] + self.searches: list[tuple] = [] + self._results = search_results or [] + self._fail = fail + self._delay = delay + + def add(self, messages, **kwargs): + if self._delay: + time.sleep(self._delay) + if self._fail: + raise RuntimeError("mem0 service unavailable") + self.adds.append((messages, kwargs)) + return {"ok": True} + + def search(self, query, **kwargs): + if self._delay: + time.sleep(self._delay) + if self._fail: + raise RuntimeError("mem0 service unavailable") + self.searches.append((query, kwargs)) + return list(self._results) + + +class _EnvPatch: + """Temporarily set environment variables.""" + + def __init__(self, **kv): + self._kv = kv + self._old: dict[str, str | None] = {} + + def __enter__(self): + for k, v in self._kv.items(): + self._old[k] = os.environ.get(k) + os.environ[k] = v + return self + + def __exit__(self, *exc): + for k, old in self._old.items(): + if old is None: + os.environ.pop(k, None) + else: + os.environ[k] = old + return False + + +def _enabled_settings(**overrides): + cfg = { + "mem0_enabled": True, + "mem0_api_key": FAKE_KEY, + "out_root": "/tmp/skillopt-test-project", + "env": "alfworld", + } + cfg.update(overrides) + return resolve_settings(cfg, env={}) + + +# ── 1. Disabled by default ──────────────────────────────────────────────────── + +def test_api_key_alone_does_not_enable(): + """The maintainer's primary objection: a stray key must export nothing.""" + cfg = {"out_root": "/tmp/proj", "env": "alfworld"} + settings = resolve_settings(cfg, env={"MEM0_API_KEY": FAKE_KEY}) + + assert settings.enabled is False, "a bare env key must not enable memory" + assert settings.usable is False + assert settings.api_key == "", "a disabled run must not even read the key" + + with _EnvPatch(MEM0_API_KEY=FAKE_KEY): + assert maybe_init_mem0(cfg) is None, "no backend may be constructed when disabled" + + print("ok 1. MEM0_API_KEY alone does not enable uploads") + + +def test_disabled_hooks_send_nothing(): + """With memory None, every hook is inert and returns context untouched.""" + original = "prior step context" + assert hook_pre_reflect(None, "skill", original) == original + hook_post_reflect(None, 0, 0, [{"patch": {}}]) # must not raise + print("ok 1b. hooks are no-ops when memory is disabled") + + +# ── 2. Explicit opt-in ──────────────────────────────────────────────────────── + +def test_explicit_opt_in_writes(): + settings = _enabled_settings() + assert settings.enabled and settings.usable + + client = FakeClient() + memory = SkillMemory(settings, client=client) + memory.store_skill_iteration(epoch=1, step=2, skill_text="do the thing", score=0.75) + + assert len(client.adds) == 1, "an opted-in run should write exactly once" + messages, kwargs = client.adds[0] + assert kwargs["user_id"] == settings.namespace + assert "do the thing" in messages[0]["content"] + memory.close() + print("ok 2. explicit opt-in writes to the backend") + + +# ── 3. Redaction ────────────────────────────────────────────────────────────── + +def test_redaction_strips_secrets_and_paths(): + settings = _enabled_settings(out_root="/home/alice/projects/run") + client = FakeClient() + memory = SkillMemory(settings, client=client) + + leaky = ( + "Use sk-abcdefghijklmnop1234 to authenticate.\n" + "Config lives at /home/alice/projects/run/configs/train.yaml\n" + "Also check /home/bob/other/notes.md\n" + "api_key = supersecretvalue\n" + ) + memory.store_skill_iteration(epoch=0, step=0, skill_text=leaky, score=1.0) + + sent = client.adds[0][0][0]["content"] + assert "sk-abcdefghijklmnop1234" not in sent, "API key leaked" + assert "supersecretvalue" not in sent, "assigned secret leaked" + assert "/home/alice" not in sent, "home path (username) leaked" + assert "/home/bob" not in sent, "third-party home path leaked" + # Structure that makes the memory useful must survive. + assert "configs/train.yaml" in sent, "redaction destroyed useful structure" + memory.close() + print("ok 3. secrets and home paths are stripped, structure preserved") + + +def test_truncation_measured_after_redaction(): + settings = _enabled_settings(mem0_max_chars=120) + client = FakeClient() + memory = SkillMemory(settings, client=client) + memory.store_skill_iteration(epoch=0, step=0, skill_text="x" * 5000, score=0.0) + + sent = client.adds[0][0][0]["content"] + assert len(sent) <= 120 + len("\n[...truncated]") + memory.close() + print("ok 3b. payload cap applied to the redacted text") + + +def test_redaction_helper_is_order_safe(): + """A key embedded in a home path must be scrubbed before the path collapses.""" + out = redact_for_upload("/home/alice/.config/sk-abcdefghijklmnop1234", "") + assert "sk-abcdefghijklmnop1234" not in out + assert "/home/alice" not in out + print("ok 3c. credential inside a home path is still redacted") + + +# ── 4. Namespacing ──────────────────────────────────────────────────────────── + +def test_namespace_is_stable_and_project_specific(): + a1 = project_namespace("/home/alice/projA", "alfworld") + a2 = project_namespace("/home/alice/projA", "alfworld") + b = project_namespace("/home/alice/projB", "alfworld") + other_env = project_namespace("/home/alice/projA", "searchqa") + + assert a1 == a2, "namespace must be stable across runs of one project" + assert a1 != b, "distinct projects must not share a namespace" + assert a1 != other_env, "distinct envs must not share a namespace" + assert "/home/alice" not in a1, "namespace must not carry the raw path" + assert a1.startswith("skillopt:alfworld:") + print("ok 4. namespace stable, project-specific, and path-free") + + +def test_explicit_namespace_override_wins(): + settings = _enabled_settings(mem0_namespace="team-shared-pool") + assert settings.namespace == "team-shared-pool" + print("ok 4b. explicit mem0_namespace overrides the derived one") + + +# ── 4c. Structured configs must actually reach the trainer ──────────────────── + +def test_structured_config_keys_survive_flattening(): + """Regression: every shipped config is structured (inherits _base_/default). + + Structured configs are flattened through an explicit key map, so a + top-level ``mem0_enabled`` would be silently dropped — the feature would + appear inert with no error. The keys must be mapped under ``train``. + """ + from skillopt.config import flatten_config + + flat = flatten_config({ + "train": { + "mem0_enabled": True, + "mem0_retrieval_limit": 3, + "mem0_timeout_seconds": 2.5, + }, + }) + assert flat.get("mem0_enabled") is True, "mem0_enabled dropped when flattening" + assert flat.get("mem0_retrieval_limit") == 3 + assert flat.get("mem0_timeout_seconds") == 2.5 + + settings = resolve_settings(flat, env={"MEM0_API_KEY": FAKE_KEY}) + assert settings.enabled and settings.retrieval_limit == 3 + print("ok 4c. structured-config mem0 keys survive flattening") + + +def test_shipped_example_config_enables_memory(): + """The opt-in smoke example must actually switch the feature on.""" + from skillopt.config import flatten_config, load_config + + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + example = os.path.join(repo_root, "configs", "features", "mem0_memory.yaml") + assert os.path.exists(example), "shipped example config is missing" + + flat = flatten_config(load_config(example)) + settings = resolve_settings(flat, env={"MEM0_API_KEY": FAKE_KEY}) + + assert settings.enabled, "the example config does not enable memory" + assert settings.usable, "the example config does not yield a usable backend" + assert settings.timeout_seconds == 5.0 + assert settings.max_chars == 4000 + print("ok 4d. shipped example config enables memory end to end") + + +# ── 5. Retrieval reaches the reflection prompt ──────────────────────────────── + +def test_retrieved_context_reaches_reflection_prompt(): + """End-to-end: stored memory → hook → the user message sent to the optimizer.""" + memo = "Open the drawer before searching it; searching first always failed." + settings = _enabled_settings() + client = FakeClient(search_results=[{"memory": memo}]) + memory = SkillMemory(settings, client=client) + + reflect_context = hook_pre_reflect(memory, "current skill text", "prior step context") + assert memo in reflect_context, "retrieved memory missing from reflection context" + assert "prior step context" in reflect_context, "existing context was dropped" + + # Now prove that context actually lands in the prompt the optimizer sees. + from skillopt.gradient import reflect as reflect_mod + + captured: dict[str, str] = {} + + def fake_chat_optimizer(system, user, **kwargs): + captured["user"] = user + return ('{"patch": {"edits": []}}', None) + + original = reflect_mod.chat_optimizer + reflect_mod.chat_optimizer = fake_chat_optimizer + try: + with tempfile.TemporaryDirectory() as tmp: + pred_dir = os.path.join(tmp, "predictions") + task_dir = os.path.join(pred_dir, "task-1") + os.makedirs(task_dir) + with open(os.path.join(task_dir, "conversation.json"), "w") as fh: + json.dump([{"role": "user", "content": "go to the kitchen"}], fh) + + reflect_mod.run_error_analyst_minibatch( + skill_content="current skill text", + items=[{ + "id": "task-1", + "task_description": "find the mug", + "task_type": "pick", + "fail_reason": "searched before opening", + }], + prediction_dir=pred_dir, + step_buffer_context=reflect_context, + ) + finally: + reflect_mod.chat_optimizer = original + + assert "user" in captured, "optimizer was never called — test did not exercise the prompt" + assert memo in captured["user"], "retrieved memory never reached the reflection prompt" + memory.close() + print("ok 5. retrieved memory reaches the actual reflection prompt") + + +def test_retrieval_can_be_disabled_independently(): + settings = _enabled_settings(mem0_retrieval_enabled=False) + client = FakeClient(search_results=[{"memory": "should not be used"}]) + memory = SkillMemory(settings, client=client) + + out = hook_pre_reflect(memory, "skill", "ctx") + assert out == "ctx" + assert client.searches == [], "no search should be issued when retrieval is off" + memory.close() + print("ok 5b. retrieval can be turned off while writes stay on") + + +# ── 6. Failure handling ─────────────────────────────────────────────────────── + +def test_service_failure_degrades_gracefully(): + settings = _enabled_settings() + client = FakeClient(fail=True) + memory = SkillMemory(settings, client=client) + + # Writes must not raise. + memory.store_skill_iteration(epoch=0, step=0, skill_text="s", score=0.0) + hook_post_reflect(memory, 0, 0, [{"patch": {}}]) + + # Reads must not raise, and must leave the context untouched. + out = hook_pre_reflect(memory, "skill", "prior context") + assert out == "prior context", "a failed retrieval must not alter reflection input" + memory.close() + print("ok 6. service failure degrades gracefully") + + +def test_slow_service_is_bounded(): + settings = _enabled_settings(mem0_timeout_seconds=0.2) + client = FakeClient(search_results=[{"memory": "late"}], delay=3.0) + memory = SkillMemory(settings, client=client) + + start = time.time() + out = hook_pre_reflect(memory, "skill", "prior context") + elapsed = time.time() - start + + assert elapsed < 2.0, f"retrieval blocked training for {elapsed:.1f}s despite a 0.2s bound" + assert out == "prior context", "timed-out retrieval must not alter reflection input" + memory.close() + print(f"ok 7. slow service bounded ({elapsed:.2f}s, limit 0.2s)") + + +# ── 8. Malformed patches ────────────────────────────────────────────────────── + +def test_malformed_patches_are_filtered_not_raised(): + settings = _enabled_settings() + client = FakeClient() + memory = SkillMemory(settings, client=client) + + for bad in (None, "a string", 42, [None, "junk", {"patch": {"edits": []}}], {}): + memory.store_reflection(epoch=0, step=0, patches=bad) + + # The mixed list contributes exactly one valid patch; nothing raised. + mixed = [c for c in client.adds if "1 patch(es)" in c[0][0]["content"]] + assert len(mixed) == 1, "the one valid dict patch should have been counted" + memory.close() + print("ok 8. malformed patches are filtered, not raised") + + +ALL_TESTS = [ + test_api_key_alone_does_not_enable, + test_disabled_hooks_send_nothing, + test_explicit_opt_in_writes, + test_redaction_strips_secrets_and_paths, + test_truncation_measured_after_redaction, + test_redaction_helper_is_order_safe, + test_namespace_is_stable_and_project_specific, + test_explicit_namespace_override_wins, + test_structured_config_keys_survive_flattening, + test_shipped_example_config_enables_memory, + test_retrieved_context_reaches_reflection_prompt, + test_retrieval_can_be_disabled_independently, + test_service_failure_degrades_gracefully, + test_slow_service_is_bounded, + test_malformed_patches_are_filtered_not_raised, +] + + +def main() -> int: + failures = 0 + for fn in ALL_TESTS: + try: + fn() + except AssertionError as exc: + failures += 1 + print(f"FAIL {fn.__name__}: {exc}") + except Exception as exc: # noqa: BLE001 - report, don't abort the suite + failures += 1 + print(f"ERROR {fn.__name__}: {type(exc).__name__}: {exc}") + print() + print(f"{len(ALL_TESTS) - failures}/{len(ALL_TESTS)} passed") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main())