From 709cd52e86b48ea343785a02f3e601daf8c1a5f0 Mon Sep 17 00:00:00 2001 From: jrauch713-svg Date: Sat, 5 Sep 2026 14:18:03 -0500 Subject: [PATCH] feat(memory): optional opt-in Mem0 backend with retrieval, rebased onto main MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebuilt on current main. The branch was 203 commits behind and trainer.py had been substantially rewritten upstream, so replaying the four prior commits conflicted across essentially the whole file. The integration is reapplied to the current trainer.py instead; the resulting head is what matters for review. Prior history remains at a16783e for reference. Adds an optional, opt-in mem0 backend that stores skill iterations and reflection summaries and reads relevant history back into the Reflect stage. Safety and correctness properties, each with a test: - Opt-in only. Uploading requires `train.mem0_enabled: true`; a MEM0_API_KEY present for another application enables nothing, and a disabled run does not read the key at all. - Single redaction choke point. No payload reaches the client without passing `redact_for_upload`, which strips vendor keys, bearer/basic tokens, JWTs, private keys, `key = value` secret assignments, the project root, and `/home/` prefixes. Text retrieved *from* mem0 is redacted again before entering the reflection prompt, since it is forwarded to the optimizer's model provider. - `mem0_api_key` is in `trainer._SECRET_KEYS`, so a YAML-supplied key is redacted from both `config.json` and the run summary. - Namespace derives from a stable project identity (git repository root, else cwd), never from `out_root` — the CLIs default that to a timestamped directory, which would mint a new namespace every run. - Timeout enforced on the injected `httpx.Client` rather than only at the caller, calls run on daemon threads so a blocked request cannot delay interpreter exit, and the backend disables itself after three consecutive failures. Mem0 API contract ----------------- The extra now pins `mem0ai>=2.0.0,<3` and the code targets the Mem0 2.x client contract, verified against 2.0.14: add(messages, user_id=, ...) # top-level user_id search(query, filters={"user_id": }, top_k=) # scoping via filters The previous `>=0.1.0` range was wrong at both ends: the injected client used to enforce timeouts first appears in 0.1.98, and 2.x raises ValueError for a top-level `user_id` in `search()` before issuing a request. Because failures degrade to an empty result, the wrong shape was silent — writes kept succeeding while retrieval always returned nothing. Retrieval is therefore tested against `_StrictMem0V2Client`, whose `search(query, *, filters=None, top_k=None)` signature mirrors the real client. The permissive `FakeClient` accepts arbitrary kwargs and could not catch this; reverting the call shape makes the new test fail, which was checked. Tests: 24 in tests/test_mem0_memory.py, no network. Full suite 1521 passed, 11 skipped. ruff clean on all new files. The API mismatch was identified by @WODE25500, who also supplied the corrected call and the strict-signature double adopted here. Co-authored-by: WODE25500 --- configs/features/mem0_memory.yaml | 71 ++++ docs/reference/config.md | 80 ++++ pyproject.toml | 7 + skillopt/config.py | 8 + skillopt/engine/trainer.py | 32 +- skillopt/memory/__init__.py | 17 + skillopt/memory/mem0_backend.py | 342 +++++++++++++++ skillopt/memory/redaction.py | 135 ++++++ skillopt/memory/settings.py | 166 ++++++++ skillopt/memory/trainer_hooks.py | 187 ++++++++ tests/test_mem0_memory.py | 683 ++++++++++++++++++++++++++++++ 11 files changed, 1727 insertions(+), 1 deletion(-) create mode 100644 configs/features/mem0_memory.yaml create mode 100644 skillopt/memory/__init__.py create mode 100644 skillopt/memory/mem0_backend.py create mode 100644 skillopt/memory/redaction.py create mode 100644 skillopt/memory/settings.py create mode 100644 skillopt/memory/trainer_hooks.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 6371f69d..66cb0c3c 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -210,6 +210,86 @@ forwarded consistently to both SDK and CLI execution paths. 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. Redacted from `config.json` and the run summary; prefer the env var | +| `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 (sent as `top_k`) | +| `train.mem0_timeout_seconds` | float | `5.0` | Hard per-call bound, enforced on the HTTP client. After 3 consecutive failures memory disables itself for the run | +| `train.mem0_max_chars` | int | `4000` | Cap on any single stored payload, applied after redaction | + +### Supported Mem0 versions + +The extra pins `mem0ai>=2.0.0,<3`, and the integration is written against the Mem0 2.x +client contract, verified against 2.0.14: + +- `add(messages, user_id=, metadata=...)` — top-level `user_id`. +- `search(query, filters={"user_id": }, top_k=)` — entity scoping goes in + `filters`, and the count is `top_k`. + +The two calls differ on purpose. Mem0 2.x raises `ValueError` for a top-level `user_id` +in `search()` *before* issuing a request, and because failures degrade to an empty +result, an incorrect call shape is silent — writes keep succeeding while retrieval +returns nothing. `tests/test_mem0_memory.py` therefore exercises retrieval against a +signature-strict client double, not a permissive mock. + +0.1.x is not supported: it rejects the injected `httpx` client used to enforce timeouts +(added in 0.1.98) and uses the older `user_id`/`limit` search shape. + +### 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. Text *retrieved* from mem0 is redacted again +before it enters the reflection prompt, since the store is external and its contents are +forwarded to the optimizer's model provider. + +### Namespacing + +Memories are scoped to `skillopt::`, where the digest is a SHA-256 prefix of +a **stable project identity** — the enclosing git repository root, or the working +directory when there is no repository. Stable across runs of one project, distinct across +projects, and the raw path is never transmitted. + +Identity is deliberately *not* derived from `out_root`. The train/eval CLIs default that +to `outputs/skillopt___`, so deriving from it would mint a new +namespace on every run and cross-run retrieval would never return anything. `out_root` +still anchors path redaction, which is what it is right for. + +### Failure behaviour + +A call exceeding `mem0_timeout_seconds` releases the training step on time. The timeout is +enforced on the injected `httpx.Client`, so the request aborts rather than merely being +abandoned, and calls run on daemon threads so a stuck request can never delay interpreter +exit. After three consecutive failures the backend stops calling out for the remainder of +the run and logs once. + ## Credential Environment Variables ### Azure-family backend diff --git a/pyproject.toml b/pyproject.toml index 5d50b8fe..2791eaab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,13 @@ docs = ["mkdocs-material>=9.5.0", "mkdocstrings[python]>=0.24.0"] # is 5.50.0: earlier gradio 5.x (5.0-5.49) still import HfFolder or don't apply # the Blocks theme the same way, so they are not valid supported versions. webui = ["gradio>=5.50.0,<7"] +# 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. +# Pinned to 2.x: the retrieval path targets the Mem0 2.x search contract +# (filters={"user_id": ...}, top_k=...), which 0.1.x does not accept, and the +# injected httpx client used to enforce timeouts predates 0.1.98. +mem0 = ["mem0ai>=2.0.0,<3"] # 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 e1fdec04..4bf262a8 100644 --- a/skillopt/config.py +++ b/skillopt/config.py @@ -129,6 +129,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 520a90eb..4989e583 100644 --- a/skillopt/engine/trainer.py +++ b/skillopt/engine/trainer.py @@ -79,6 +79,12 @@ ) from skillopt.model.common import normalize_backend_name from skillopt.utils import compute_score, skill_hash +from skillopt.memory.trainer_hooks import ( + maybe_init_mem0, + hook_pre_reflect, + hook_post_evaluate, + hook_post_reflect, +) # ── Skill-aware reflection: appendix flush ─────────────────────────────────── @@ -335,6 +341,7 @@ def _build_longitudinal_pairs( "azure_api_key", "api_key", "openai_api_key", + "mem0_api_key", } @@ -1125,6 +1132,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") @@ -1241,11 +1250,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( @@ -1287,6 +1303,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) @@ -1605,6 +1626,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, @@ -2496,4 +2521,9 @@ def _persist_runtime_state(last_completed_step: int) -> None: f"calls={t['calls']})" ) + # Release the memory worker. SkillMemory calls run on daemon threads, + # so this is tidiness rather than a shutdown requirement. + if memory is not None: + memory.close() + return summary diff --git a/skillopt/memory/__init__.py b/skillopt/memory/__init__.py new file mode 100644 index 00000000..069d9763 --- /dev/null +++ b/skillopt/memory/__init__.py @@ -0,0 +1,17 @@ +"""skillopt.memory — optional, opt-in mem0-backed persistent memory. + +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 new file mode 100644 index 00000000..3be18cf4 --- /dev/null +++ b/skillopt/memory/mem0_backend.py @@ -0,0 +1,342 @@ +"""mem0-backed persistent memory for SkillOpt. + +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 + + settings = resolve_settings(cfg) + memory = SkillMemory.from_settings(settings) # None unless opted in +""" +from __future__ import annotations + +import hashlib +import json +import threading +from typing import Any + +from skillopt.memory.redaction import redact_for_upload +from skillopt.memory.settings import Mem0Settings + +# After this many consecutive timeouts/errors the backend stops calling out for +# the rest of the run, so a persistently unreachable service costs a bounded +# total rather than a bounded amount on every step. +MAX_CONSECUTIVE_FAILURES = 3 + +try: + import httpx + from mem0 import MemoryClient + _MEM0_AVAILABLE = True +except ImportError: + _MEM0_AVAILABLE = False + httpx = None # type: ignore[assignment] + 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. + + Prefer :meth:`from_settings`; the constructor stays explicit so tests can + inject a fake client without touching the network. + """ + + 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") + # Enforce the timeout where the request actually happens. mem0's own + # default client uses timeout=300, and no amount of waiting on the + # caller's side can abort a socket read that the HTTP layer is still + # willing to wait 5 minutes for. Injecting the client is supported: + # mem0 only overrides base_url and headers on it. + client = MemoryClient( + api_key=settings.api_key, + client=httpx.Client(timeout=settings.timeout_seconds), + ) + + self.settings = settings + self.namespace = settings.namespace + self._client = client + self._consecutive_failures = 0 + self._degraded = False + self._closed = False + + # ── 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) + + # ── Internal helpers ────────────────────────────────────────────────── + + def _note_failure(self, reason: str) -> None: + self._consecutive_failures += 1 + if self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES and not self._degraded: + self._degraded = True + print(f" [mem0] disabled for the rest of this run after " + f"{self._consecutive_failures} consecutive failures ({reason})") + + def _bounded(self, fn, *args, **kwargs) -> Any: + """Run *fn* with a hard wall-clock bound; ``None`` on timeout or error. + + Two layers, because neither alone is sufficient: + + * The injected ``httpx.Client`` carries the same timeout, so the request + itself aborts rather than merely being abandoned. This is the layer + that actually stops the work. + * The call runs on a **daemon** thread, so even a pathologically stuck + request can never delay interpreter exit. A + ``ThreadPoolExecutor`` cannot provide this: its workers are + non-daemon and ``concurrent.futures`` installs an ``atexit`` hook that + joins them, so one blocked read would hang the process at shutdown no + matter what ``shutdown(wait=False)`` was told. + + After :data:`MAX_CONSECUTIVE_FAILURES` the backend stops calling out + entirely. + """ + if self._degraded or self._closed: + return None + + box: dict[str, Any] = {} + + def _runner() -> None: + try: + box["value"] = fn(*args, **kwargs) + except BaseException as exc: # noqa: BLE001 - reported on the caller's thread + box["error"] = exc + + thread = threading.Thread(target=_runner, daemon=True, name="mem0-call") + thread.start() + thread.join(self.settings.timeout_seconds) + + if thread.is_alive(): + print(f" [mem0] call exceeded {self.settings.timeout_seconds}s — continuing without it") + self._note_failure("timeout") + return None + if "error" in box: + exc = box["error"] + print(f" [mem0] call failed ({type(exc).__name__}: {exc}) — continuing without it") + self._note_failure(type(exc).__name__) + return None + + self._consecutive_failures = 0 + return box.get("value") + + def _payload(self, text: str) -> str: + """The only route by which text becomes outbound content. + + 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]: + # Mem0 2.x search contract: entity scoping goes in `filters`, and the + # result count is `top_k`. Passing a top-level `user_id` raises + # ValueError before any request is made, which `_bounded` would turn + # into a silent empty result — writes would keep succeeding while + # retrieval never returned anything. `add()` still takes `user_id` at + # the top level, so the two calls differ on purpose. + results = self._bounded( + self._client.search, + self._payload(query), + filters={"user_id": self.namespace}, + top_k=limit, + ) + if isinstance(results, list): + return results + if isinstance(results, dict): + 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( + self, + epoch: int, + step: int, + skill_text: str, + score: float, + metadata: dict | None = None, + ) -> Any: + """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 or ""), + } + 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 or ''}" + ) + return self._add(content, metadata=base_meta) + + def store_reflection( + self, + epoch: int, + step: int, + patches: Any, + scores: dict | None = None, + ) -> Any: + """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) + 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", + "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}" + ) + return self._add(content, metadata=base_meta) + + 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) + + def format_retrieved_context(self, memories: list[dict], max_chars: int = 2000) -> str: + """Render retrieved memories as text for inclusion in a prompt. + + Retrieved text is redacted **again** on the way in. Outbound redaction + alone is not enough: mem0 is an external store that may hold records + written by an older version of this code, by another tool, or by a + deliberately shared namespace. Anything returned from it flows straight + into the reflection prompt and on to the optimizer's model provider, so + it is untrusted input to a *second* third party and gets the same + treatment as anything leaving the process. + + Returns ``""`` when there is nothing useful, so callers can append + unconditionally without producing an empty heading. + """ + lines: list[str] = [] + for rec in memories or []: + if not isinstance(rec, dict): + continue + text = rec.get("memory") or rec.get("text") or "" + if not isinstance(text, str) or not text.strip(): + continue + clean = redact_for_upload(text, self.settings.project_root) + if not isinstance(clean, str) or not clean.strip(): + continue + lines.append(f"- {clean.strip()}") + if not lines: + return "" + + body = "\n".join(lines) + if len(body) > max_chars: + body = body[:max_chars] + "\n[...truncated]" + return "## Relevant Memory From Previous Runs\n" + body + + def close(self) -> None: + """Stop issuing calls, and close the underlying HTTP client. + + Idempotent. There is no worker pool to join: calls run on daemon + threads, so nothing here is load-bearing for interpreter exit. + """ + if self._closed: + return + self._closed = True + inner = getattr(self._client, "client", None) + close_fn = getattr(inner, "close", None) + if callable(close_fn): + try: + close_fn() + except Exception: # noqa: BLE001 - cleanup must never raise + pass + + def __repr__(self) -> str: + 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..061c29e5 --- /dev/null +++ b/skillopt/memory/settings.py @@ -0,0 +1,166 @@ +"""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 a *stable project identity*, so +that a later run of the same project reads what earlier runs wrote. That +identity must not come from ``out_root``: the train/eval CLIs default it to +``outputs/skillopt___``, so deriving from it would mint a +fresh namespace on every run and cross-run retrieval would never return +anything. Identity is resolved as: an explicit ``mem0_namespace``, else the +enclosing git repository root, else the current working directory. + +The 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 +import subprocess +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 git_repo_root(start: str | None = None) -> str: + """Absolute path of the enclosing git work tree, or ``""`` if there is none. + + Preferred identity source: it is the same for every run of a checkout no + matter which output directory a run happens to write to, and no matter which + subdirectory the command was invoked from. + """ + try: + out = subprocess.run( + ["git", "rev-parse", "--show-toplevel"], + cwd=start or os.getcwd(), + capture_output=True, + text=True, + timeout=5, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return "" + if out.returncode != 0: + return "" + return os.path.abspath(out.stdout.strip()) if out.stdout.strip() else "" + + +def project_identity(start: str | None = None) -> str: + """The stable path this project is identified by. + + Never derived from ``out_root`` — see the module docstring for why. + """ + return git_repo_root(start) or os.path.abspath(start or os.getcwd()) + + +def project_namespace(project_root: str, env_name: str = "") -> str: + """Hash *project_root* into a namespace. + + Stable across runs of the same project (same absolute path → same digest) + and distinct across projects, which is what keeps unrelated experiments from + sharing a memory pool. Callers should pass :func:`project_identity`, not a + run output directory. + """ + 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 "") + + # Identity for namespacing: deliberately NOT out_root, which is per-run. + identity = project_identity() + namespace = str(cfg.get("mem0_namespace") or "").strip() + if not namespace: + namespace = project_namespace(identity, str(cfg.get("env") or "")) + + # out_root still anchors *path redaction* — collapsing the run directory out + # of stored text is exactly what it is right for. + project_root = os.path.abspath(str(cfg.get("out_root") or identity)) + + 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 new file mode 100644 index 00000000..b59f3dd3 --- /dev/null +++ b/skillopt/memory/trainer_hooks.py @@ -0,0 +1,187 @@ +"""Integration hooks binding :class:`SkillMemory` into the ReflACT loop. + +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_pre_reflect, hook_post_evaluate, hook_post_reflect, + ) + + memory = maybe_init_mem0(cfg) + + # before the Reflect stage: + step_buffer_context = hook_pre_reflect(memory, current_skill, step_buffer_context) + + # 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 + +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": + """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 + + +# ── 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 + ---------- + 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 + ------- + str + Possibly-augmented context, always safe to pass to ``adapter.reflect``. + """ + if memory is None: + return step_buffer_context + + try: + probe = query if query is not None else (skill_text or "")[:_QUERY_CHARS] + if not probe.strip(): + return step_buffer_context + + 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 ──────────────────────────────────────────────────────── + +def hook_post_evaluate( + memory: "SkillMemory | None", + epoch: int, + step: int, + skill: str, + score: float, + cfg: dict, +) -> None: + """Store the current skill and its gate score. No-op without memory.""" + 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: # noqa: BLE001 - memory must never break training + 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: + """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=patches, + scores=scores, + ) + 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..789bf42c --- /dev/null +++ b/tests/test_mem0_memory.py @@ -0,0 +1,683 @@ +"""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 subprocess +import sys +import tempfile +import textwrap +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_identity, + 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 _StrictMem0V2Client: + """A stand-in that enforces the real Mem0 2.x ``search`` signature. + + ``FakeClient`` accepts arbitrary keyword arguments, so it happily swallowed + the 0.1.x call shape (``user_id=`` / ``limit=``) and every test passed while + real Mem0 2.x raised ``ValueError`` before sending a request. ``_bounded`` + then converted that into a silent empty result, so writes kept succeeding + while retrieval returned nothing. This double makes that regression fail + loudly instead. + + Contract mirrored from mem0 2.0.x: ``search(query, *, filters=None, + top_k=None)``; ``add()`` still accepts a top-level ``user_id``. + + Thanks to @WODE25500, who identified the mismatch and supplied this shape. + """ + + def __init__(self, search_results=None): + self.searches: list[tuple] = [] + self.adds: list[tuple] = [] + self._results = search_results or [] + + def add(self, messages, **kwargs): + self.adds.append((messages, kwargs)) + return {"ok": True} + + def search(self, query, *, filters=None, top_k=None): + self.searches.append((query, filters, top_k)) + 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)") + + +# ── 9. Copilot review follow-ups ────────────────────────────────────────────── + +def test_retrieved_text_is_redacted_before_entering_the_prompt(): + """Inbound redaction: mem0 is an external store and may hold raw secrets. + + Anything it returns flows into the reflection prompt and on to the + optimizer's model provider, so it must be scrubbed on the way *in* as well + as on the way out. + """ + poisoned = "Earlier run used sk-abcdefghijklmnop1234 from /home/alice/creds.txt" + settings = _enabled_settings() + client = FakeClient(search_results=[{"memory": poisoned}]) + memory = SkillMemory(settings, client=client) + + ctx = hook_pre_reflect(memory, "skill", "prior context") + assert "sk-abcdefghijklmnop1234" not in ctx, "secret from mem0 reached the prompt" + assert "/home/alice" not in ctx, "home path from mem0 reached the prompt" + assert "Earlier run used" in ctx, "redaction destroyed the whole record" + memory.close() + print("ok 9. retrieved text is redacted before entering the prompt") + + +def test_repeated_timeouts_disable_the_backend(): + """A wedged service must cost a bounded total, not a bounded amount per step. + + With a single worker, a stuck call would otherwise make every later call + queue behind it and time out again, so the cost would recur every step. + """ + settings = _enabled_settings(mem0_timeout_seconds=0.1) + client = FakeClient(search_results=[{"memory": "late"}], delay=5.0) + memory = SkillMemory(settings, client=client) + + start = time.time() + for _ in range(6): + hook_pre_reflect(memory, "skill", "ctx") + elapsed = time.time() - start + + assert memory._degraded, "backend should disable itself after repeated timeouts" + # 3 timeouts at 0.1s, then short-circuited: well under 6 x 0.1s. + assert elapsed < 1.0, f"repeated timeouts cost {elapsed:.2f}s — not bounded" + memory.close() + print(f"ok 9b. repeated timeouts disable the backend ({elapsed:.2f}s for 6 calls)") + + +def test_blocked_request_cannot_prevent_interpreter_exit(): + """Process-level proof, not caller-level. + + A caller-side timing assertion is not enough: with a + ``ThreadPoolExecutor`` the caller returned on schedule while + ``concurrent.futures``' atexit hook still joined the non-daemon worker, so + the process hung at shutdown forever. Calls now run on daemon threads, so + the only honest test is to start a real interpreter, wedge a request, and + require the process to exit on its own. + """ + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + child = textwrap.dedent( + """ + import sys, time + sys.path.insert(0, %r) + from skillopt.memory.mem0_backend import SkillMemory + from skillopt.memory.settings import resolve_settings + + class Wedged: + def search(self, q, **kw): time.sleep(600) + def add(self, m, **kw): time.sleep(600) + + s = resolve_settings( + {"mem0_enabled": True, "mem0_api_key": "m0-" + "x" * 24, + "mem0_timeout_seconds": 0.2}, + env={}, + ) + m = SkillMemory(s, client=Wedged()) + m.retrieve_relevant_context("probe") # times out, thread stays blocked + m.close() + print("child-done", flush=True) + """ + ) % repo_root + + with tempfile.TemporaryDirectory() as tmp: + script = os.path.join(tmp, "exit_probe.py") + with open(script, "w") as fh: + fh.write(child) + + start = time.time() + try: + proc = subprocess.run( + [sys.executable, script], + capture_output=True, text=True, timeout=60, + ) + except subprocess.TimeoutExpired: + raise AssertionError( + "interpreter did not exit within 60s — a blocked request is pinning the process" + ) from None + elapsed = time.time() - start + + assert "child-done" in proc.stdout, f"child did not reach the end: {proc.stdout!r} {proc.stderr[-400:]!r}" + assert proc.returncode == 0, f"child exited {proc.returncode}: {proc.stderr[-400:]}" + print(f"ok 9c. a permanently blocked request cannot prevent interpreter exit ({elapsed:.1f}s)") + + +def test_namespace_is_stable_across_different_out_roots(): + """The run output directory must not change the namespace. + + ``scripts/train.py`` defaults ``out_root`` to + ``outputs/skillopt___``, so deriving identity from it + minted a fresh namespace every run and cross-run retrieval could never + return anything. This is the regression test for that. + """ + a = resolve_settings({ + "mem0_enabled": True, "mem0_api_key": FAKE_KEY, + "env": "alfworld", "out_root": "/tmp/outputs/skillopt_alfworld_gpt_20260101-000000", + }, env={}) + b = resolve_settings({ + "mem0_enabled": True, "mem0_api_key": FAKE_KEY, + "env": "alfworld", "out_root": "/tmp/outputs/skillopt_alfworld_gpt_20260729-235959", + }, env={}) + + assert a.namespace == b.namespace, ( + f"two runs of the same project resolved to different namespaces:\n" + f" {a.namespace}\n {b.namespace}" + ) + # And the run directory must not leak into the namespace at all. + assert "20260101" not in a.namespace and "outputs" not in a.namespace + + # Identity itself must come from somewhere other than the run directory. + identity = project_identity() + assert identity, "project identity should never be empty" + assert not identity.startswith("/tmp/outputs"), ( + f"identity is still derived from the run directory: {identity}" + ) + print("ok 9f. namespace is stable across different out_roots") + + +def test_mem0_api_key_is_redacted_from_config_artifacts(): + """A YAML-supplied key must not reach config.json or the run summary. + + Both artifacts serialise through ``trainer._redact_cfg``, so covering it + covers both call sites (``config.json`` and ``summary["config"]``). + """ + from skillopt.engine.trainer import _SECRET_KEYS, _redact_cfg + + assert "mem0_api_key" in _SECRET_KEYS, "mem0_api_key missing from the central redaction set" + + cfg = {"mem0_enabled": True, "mem0_api_key": FAKE_KEY, "env": "alfworld"} + redacted = _redact_cfg(cfg) + + assert FAKE_KEY not in json.dumps(redacted), "key survived config redaction" + assert redacted["mem0_api_key"] != FAKE_KEY + assert redacted["env"] == "alfworld", "redaction damaged unrelated config" + # The original dict must not be mutated — the trainer keeps using it. + assert cfg["mem0_api_key"] == FAKE_KEY, "_redact_cfg mutated the live config" + print("ok 9g. mem0_api_key is redacted from config.json and the run summary") + + +def test_close_is_idempotent_and_unregisters(): + settings = _enabled_settings() + memory = SkillMemory(settings, client=FakeClient()) + memory.close() + memory.close() # must not raise + assert memory._closed + # A closed backend refuses further calls rather than resurrecting the pool. + assert memory.retrieve_relevant_context("anything") == [] + print("ok 9d. close() is idempotent and stops further calls") + + +def test_successful_call_resets_the_failure_counter(): + settings = _enabled_settings() + client = FakeClient(fail=True) + memory = SkillMemory(settings, client=client) + + memory.store_skill_iteration(epoch=0, step=0, skill_text="s", score=0.0) + assert memory._consecutive_failures == 1 + client._fail = False + memory.store_skill_iteration(epoch=0, step=1, skill_text="s", score=0.0) + assert memory._consecutive_failures == 0, "a success should reset the counter" + assert not memory._degraded + memory.close() + print("ok 9e. a successful call resets the failure counter") + + +# ── 10. Mem0 API contract ───────────────────────────────────────────────────── + +def test_search_uses_the_mem0_v2_contract(): + """Retrieval must call Mem0 2.x's actual search signature. + + Mem0 2.x rejects a top-level ``user_id``/``limit`` in ``search()`` with a + ``ValueError`` raised *before* any network call, requiring + ``filters={"user_id": ...}`` and ``top_k=``. Because ``_bounded`` degrades + failures to ``[]``, that mismatch is silent: writes keep succeeding while + retrieval quietly returns nothing. Hence a signature-strict double rather + than the permissive one, which cannot catch this. + """ + memo = "Open the drawer before searching it." + settings = _enabled_settings() + client = _StrictMem0V2Client(search_results=[{"memory": memo}]) + memory = SkillMemory(settings, client=client) + + ctx = hook_pre_reflect(memory, "current skill text", "prior step context") + + assert client.searches, "search was never called against the strict client" + _query, filters, top_k = client.searches[0] + assert filters == {"user_id": settings.namespace}, ( + f"scoping must go through filters, got {filters!r}" + ) + assert top_k == settings.retrieval_limit, f"expected top_k, got {top_k!r}" + assert memo in ctx, "retrieval produced nothing under the 2.x contract" + memory.close() + print("ok 10. search uses the Mem0 2.x contract (filters/top_k)") + + +def test_write_path_still_uses_top_level_user_id(): + """``add()`` keeps a top-level ``user_id`` in 2.x — the two calls differ.""" + settings = _enabled_settings() + client = _StrictMem0V2Client() + memory = SkillMemory(settings, client=client) + memory.store_skill_iteration(epoch=0, step=0, skill_text="s", score=1.0) + + assert client.adds, "write never reached the client" + _messages, kwargs = client.adds[0] + assert kwargs.get("user_id") == settings.namespace + memory.close() + print("ok 10b. add() still scopes with a top-level user_id") + + +# ── 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_retrieved_text_is_redacted_before_entering_the_prompt, + test_repeated_timeouts_disable_the_backend, + test_blocked_request_cannot_prevent_interpreter_exit, + test_namespace_is_stable_across_different_out_roots, + test_mem0_api_key_is_redacted_from_config_artifacts, + test_close_is_idempotent_and_unregisters, + test_successful_call_resets_the_failure_counter, + test_search_uses_the_mem0_v2_contract, + test_write_path_still_uses_top_level_user_id, + 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())