From 64326a5468ae1332784fe744d3fb3496a5232b76 Mon Sep 17 00:00:00 2001 From: jrauch713-svg Date: Fri, 10 Jul 2026 04:04:50 +0000 Subject: [PATCH] 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. Co-Authored-By: Claude Sonnet 5 --- skillopt/engine/trainer.py | 52 +++--- skillopt/memory/__init__.py | 4 + skillopt/memory/mem0_backend.py | 279 +++++++++++++++++++++++++++++++ skillopt/memory/trainer_hooks.py | 147 ++++++++++++++++ 4 files changed, 462 insertions(+), 20 deletions(-) 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 85aae53c..78640acd 100644 --- a/skillopt/engine/trainer.py +++ b/skillopt/engine/trainer.py @@ -74,6 +74,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 ─────────────────────────────────── @@ -709,26 +710,26 @@ def _build_eval_env(split: str, env_num: int, seed: int): effort=cfg.get("claude_code_exec_effort", cfg.get("reasoning_effort", "medium")), max_thinking_tokens=cfg.get("claude_code_exec_max_thinking_tokens", 16384), ) - configure_qwen_chat( - base_url=cfg.get("qwen_chat_base_url") or None, - api_key=cfg.get("qwen_chat_api_key") or None, - temperature=cfg.get("qwen_chat_temperature"), - timeout_seconds=cfg.get("qwen_chat_timeout_seconds"), - max_tokens=cfg.get("qwen_chat_max_tokens"), - enable_thinking=cfg.get("qwen_chat_enable_thinking"), - optimizer_base_url=cfg.get("optimizer_qwen_chat_base_url") or None, - optimizer_api_key=cfg.get("optimizer_qwen_chat_api_key") or None, - optimizer_temperature=cfg.get("optimizer_qwen_chat_temperature"), - optimizer_timeout_seconds=cfg.get("optimizer_qwen_chat_timeout_seconds"), - optimizer_max_tokens=cfg.get("optimizer_qwen_chat_max_tokens"), - optimizer_enable_thinking=cfg.get("optimizer_qwen_chat_enable_thinking"), - target_base_url=cfg.get("target_qwen_chat_base_url") or None, - target_api_key=cfg.get("target_qwen_chat_api_key") or None, - target_temperature=cfg.get("target_qwen_chat_temperature"), - target_timeout_seconds=cfg.get("target_qwen_chat_timeout_seconds"), - target_max_tokens=cfg.get("target_qwen_chat_max_tokens"), - target_enable_thinking=cfg.get("target_qwen_chat_enable_thinking"), - ) + configure_qwen_chat( + base_url=cfg.get("qwen_chat_base_url") or None, + api_key=cfg.get("qwen_chat_api_key") or None, + temperature=cfg.get("qwen_chat_temperature"), + timeout_seconds=cfg.get("qwen_chat_timeout_seconds"), + max_tokens=cfg.get("qwen_chat_max_tokens"), + enable_thinking=cfg.get("qwen_chat_enable_thinking"), + optimizer_base_url=cfg.get("optimizer_qwen_chat_base_url") or None, + optimizer_api_key=cfg.get("optimizer_qwen_chat_api_key") or None, + optimizer_temperature=cfg.get("optimizer_qwen_chat_temperature"), + optimizer_timeout_seconds=cfg.get("optimizer_qwen_chat_timeout_seconds"), + optimizer_max_tokens=cfg.get("optimizer_qwen_chat_max_tokens"), + optimizer_enable_thinking=cfg.get("optimizer_qwen_chat_enable_thinking"), + target_base_url=cfg.get("target_qwen_chat_base_url") or None, + target_api_key=cfg.get("target_qwen_chat_api_key") or None, + target_temperature=cfg.get("target_qwen_chat_temperature"), + target_timeout_seconds=cfg.get("target_qwen_chat_timeout_seconds"), + target_max_tokens=cfg.get("target_qwen_chat_max_tokens"), + target_enable_thinking=cfg.get("target_qwen_chat_enable_thinking"), + ) configure_minimax_chat( base_url=cfg.get("minimax_base_url") or None, api_key=cfg.get("minimax_api_key") or None, @@ -1019,6 +1020,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") @@ -1181,6 +1184,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) @@ -1492,6 +1500,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}")