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..84d81b70 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -138,6 +138,73 @@ 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. 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 | +| `train.mem0_timeout_seconds` | float | `5.0` | Hard per-call bound. On timeout the executor is retired and training continues; 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 | + +### 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. + +### Failure behaviour + +A call that exceeds `mem0_timeout_seconds` releases the training step on time and +retires the worker, so the next step does not queue behind a wedged request. After +three consecutive failures the backend stops calling out for the remainder of the run +and logs once. A persistently unreachable service therefore costs a bounded total +rather than a bounded amount on every step. + +Text retrieved from mem0 is redacted again before it is inserted into the reflection +prompt. Outbound redaction alone is not sufficient: the store is external and may hold +records written by an older version, another tool, or a shared namespace, and anything +it returns is 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 +current 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` +is still used to anchor path redaction, which is what it is right for. + ## 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 1eae4209..e361ea36 100644 --- a/skillopt/engine/trainer.py +++ b/skillopt/engine/trainer.py @@ -75,6 +75,12 @@ set_optimizer_deployment, ) 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 ─────────────────────────────────── @@ -331,6 +337,7 @@ def _build_longitudinal_pairs( "azure_api_key", "api_key", "openai_api_key", + "mem0_api_key", } @@ -1042,6 +1049,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") @@ -1158,11 +1167,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( @@ -1204,6 +1220,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 +1543,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, @@ -2413,4 +2438,10 @@ def _persist_runtime_state(last_completed_step: int) -> None: f"calls={t['calls']})" ) + # Release the memory worker thread. SkillMemory also registers an + # atexit hook, so an aborted run is covered too; this is the clean + # path so a long-lived process does not hold the thread until exit. + 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..8590576a --- /dev/null +++ b/skillopt/memory/mem0_backend.py @@ -0,0 +1,336 @@ +"""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]: + results = self._bounded( + self._client.search, + self._payload(query), + user_id=self.namespace, + limit=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..cf24f5e5 --- /dev/null +++ b/tests/test_mem0_memory.py @@ -0,0 +1,607 @@ +"""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 _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") + + +# ── 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_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())