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