default/openevolve_world_knowledge

"""OpenEvolve + World Knowledge — the same island MAP-Elites search, with an open world.

**All six components are OpenEvolve's own**, inherited rather than rebuilt: the World Knowledge Layer
is orthogonal to the component vocabulary, bolted onto the scaffold rather than swapped into a slot,
which is what makes the two methods differ by exactly one thing.

The layer is one object (``self.world_knowledge``, built only when the config enables it) and three
call sites — the shape the reference implementation uses on its controller: :meth:`augment_prompt`
inserts the retrieved evidence ``S`` before the task, :meth:`after_step` writes ``Delta_t`` back (``UpdateDelta``),
and :meth:`setup` / :meth:`_resume` preflight the backend and seed ``D_search``. Its state rides in
``_checkpoint_scaffold_state``, the hook for search state the scaffold owns rather than a component.
"""
from __future__ import annotations

import logging
from dataclasses import replace

from ...records import Genome
from ...world_knowledge import WorldKnowledgeLayer
from ..openevolve.scaffold import OpenEvolveScaffold
from ..registry import register_scaffold

log = logging.getLogger(__name__)


@register_scaffold("openevolve_world_knowledge")
class OpenEvolveWorldKnowledgeScaffold(OpenEvolveScaffold):
    name = "openevolve_world_knowledge"

    def __init__(self, **kw):
        super().__init__(**kw)
        config = self.config.world_knowledge
        # The layer runs on the SAME model as the proposer: the method's single frozen M_theta plays
        # the mutation, gating, and query roles, separated by the role prompt alone. Disabled -> the
        # attribute is None and every call site below is a no-op, so this scaffold is exactly
        # OpenEvolve.
        self.world_knowledge = (
            WorldKnowledgeLayer(config, self.model, seed=int(self.seed)) if config.enabled else None
        )
        if self.world_knowledge is not None:
            steps = int(config.query_evolution_steps)
            refinement = ("LCLM result reranking" if steps == 1
                          else f"query evolution ({steps} full rounds)")
            log.info("World Knowledge Layer enabled (gate: %s, retrieval: %s)",
                     config.gate, refinement)

    # ---- the three call sites --------------------------------------------------------------------
    def augment_prompt(self, prompt, selection):
        """Insert retrieved evidence between the current program and its mutation task, if any.

        With the layer off, or on any iteration the gate declines, the prompt is returned untouched —
        byte-identical to OpenEvolve's.
        """
        if self.world_knowledge is None:
            return prompt
        try:
            evidence = self.world_knowledge.step(
                parent=selection.parent,
                previous_parent=self._immediate_parent(selection.parent),
                history=self._rendered_history(selection),
                ctx=self.ctx,
            )
        except Exception:   # noqa: BLE001 - the layer is optional; the base method is the fallback
            self.world_knowledge.discard_credit(self.ctx.iteration)
            log.exception("World Knowledge Layer failed; this iteration evolves closed-world")
            return prompt
        if not evidence:
            return prompt
        before_task, task = prompt.user.rsplit("\n\n# Task\n", maxsplit=1)
        return replace(
            prompt,
            user=f"{before_task}\n\n{evidence}\n\n# Task\n{task}",
        )

    def _immediate_parent(self, current: Genome | None) -> Genome | None:
        """Resolve the selected program's direct ancestor while it remains in the population."""
        parent_id = getattr(current, "parent_id", None)
        if not parent_id:
            return None
        programs = getattr(self.population, "programs", None)
        if isinstance(programs, dict):
            return programs.get(parent_id)
        return next((genome for genome in self.population.all() if genome.id == parent_id), None)

    def _rendered_history(self, selection):
        """``H_{t-1}`` for the layer — OpenEvolve's *own* ``# Program Evolution History`` rendering.

        Reusing it means the gate judges the very text the mutation LLM is judging, instead of a
        second, differently-shaped description of the same search. Falls back to the recently
        admitted genomes for a builder that renders no history of its own.
        """
        render = getattr(self.prompt_builder, "_evolution_history", None)
        return render(selection) if callable(render) else list(self.ctx.recent)

    def after_step(self, child: Genome, result) -> None:
        """OpenEvolve's island bookkeeping, then ``UpdateDelta`` for this step."""
        super().after_step(child, result)
        if self.world_knowledge is None:
            return
        score = self._evaluated_score(child, result)
        self.world_knowledge.credit(
            self.ctx.iteration,
            delta=self._realized_delta(child, result),
            score=score,
            result_program_id=getattr(child, "id", "") or "",
            ctx=self.ctx,
        )

    def setup(self, task) -> None:
        # Before the seed evaluation, which can take minutes: a run configured to search must find
        # out that it cannot before it spends anything, not at the first retrieval.
        if self.world_knowledge is not None:
            self.world_knowledge.preflight()
        super().setup(task)
        # Pre-loop grounding runs off the evaluated seed, so it can ask about the problem the seed
        # actually solves. A no-op under the default `world_knowledge.grounding: none`.
        if self.world_knowledge is not None:
            self.world_knowledge.ground(parent=self.population.best() or self.ctx.best, ctx=self.ctx)

    def _resume(self, task, path: str) -> None:
        if self.world_knowledge is not None:
            self.world_knowledge.preflight()   # same fail-fast on the path that skips setup()
        super()._resume(task, path)

    # ---- checkpoint / resume ----------------------------------------------------------------------
    def _checkpoint_scaffold_state(self) -> dict:
        """``D_search`` is the layer's learned state, and the layer is not a component — so it rides
        here, with the rest of the search orchestration this scaffold owns. Without it a resumed run
        re-pays for every search and loses every ``Delta`` it had already bought."""
        state = super()._checkpoint_scaffold_state()
        if self.world_knowledge is not None:
            state = {**state, "world_knowledge": self.world_knowledge.state_dict()}
        return state

    def _load_checkpoint_scaffold_state(self, state: dict) -> None:
        super()._load_checkpoint_scaffold_state(state)
        payload = (state or {}).get("world_knowledge")
        if self.world_knowledge is not None and isinstance(payload, dict):
            self.world_knowledge.load_state_dict(payload)

    # ---- credit ------------------------------------------------------------------------------------
    @staticmethod
    def _evaluated_score(child: Genome | None, result) -> float | None:
        """Return the finite continuous evaluator score ``S`` for the generated child."""
        if child is None or result is None:
            return None
        score = child.fitness
        if score != score or abs(score) == float("inf"):
            return None
        return float(score)

    @classmethod
    def _realized_delta(cls, child: Genome | None, result) -> float | None:
        """``Delta_t = E(x_t) - E(x_{t-1})`` for the step just finished, or ``None``.

        ``None`` (rather than 0.0) whenever the step produced no verdict — a rejected diff, an
        over-length program, a parent that never scored. The entry still records that it was used;
        it just carries no opinion about what that use was worth, which is the honest reading and
        keeps an unevaluated step from dragging a query's mean ``Delta`` toward zero.
        """
        score = cls._evaluated_score(child, result)
        if score is None:
            return None
        parent_score = (child.metadata.get("parent_metrics") or {}).get("combined_score")
        if not isinstance(parent_score, (int, float)) or isinstance(parent_score, bool):
            return None
        return score - float(parent_score)