SakanaAI/ale_agent

"""Domain-guided ALE-Agent prompts from the ALE-Bench paper appendix."""
from __future__ import annotations

from ...components.prompt import PromptBuilder
from ...models import Prompt
from ...records import Genome, LoopContext, Selection


DOMAIN_GUIDANCE = (
    (
        "Speed and complexity",
        "Based on the code and feedback: What are the key algorithms and data structures used? "
        "What are its computational complexity bottlenecks (including Time Limit Exceeded feedback)? "
        "How might the time or space complexity be improved (including Memory Limit Exceeded feedback)? "
        "Consider both small optimizations and completely different approaches. Think deeply and "
        "broadly about possible improvements before implementing anything.",
    ),
    (
        "Simulated annealing state representation",
        "If this solution uses simulated annealing, analyze the problem, solution properties, and "
        "feedback to suggest a better state representation. Consider how the current encoding may "
        "limit the search space or convergence speed, and alternatives that could reach better local "
        "optima or converge faster. Think deeply and broadly before implementing anything.",
    ),
    (
        "Simulated annealing neighborhood",
        "If this solution uses simulated annealing, use the problem and feedback to design a better "
        "neighborhood. Consider (1) balancing small and large moves, (2) whether every valid solution "
        "is reachable, and (3) moves that preserve feasibility while exploring new regions. Think "
        "deeply and broadly before implementing anything.",
    ),
    (
        "Beam search",
        "Consider implementing or improving a beam-search approach. Analyze a useful beam width and "
        "evaluation function, and how to balance diversity against solution quality. Think deeply and "
        "broadly before implementing anything.",
    ),
)


def _metrics(genome: Genome) -> str:
    if not genome.scores:
        return "- No evaluation metrics are available."
    return "\n".join(f"- {key}: {value}" for key, value in genome.scores.items())


def _feedback(genome: Genome) -> str:
    artifacts = genome.artifacts or {}
    value = (
        artifacts.get("text_feedback")
        or artifacts.get("message")
        or artifacts.get("standard_error")
        or artifacts.get("traceback")
        or artifacts.get("error")
    )
    return str(value or "No textual evaluator feedback was returned.")


def _state_block(title: str, genome: Genome, language: str) -> str:
    return (
        f"# {title}\n"
        f"Genome: {genome.id}\n\n"
        f"## Evaluation metrics\n{_metrics(genome)}\n\n"
        f"## Evaluator feedback\n{_feedback(genome)}\n\n"
        f"## Complete program\n```{language}\n{genome.content}\n```"
    )


class ALEAgentPromptBuilder(PromptBuilder):
    """Render current state, historical best, optional history, and one domain directive."""

    def __init__(self, *, domain_guidance: bool = True, include_history: bool = False):
        self.domain_guidance = bool(domain_guidance)
        self.include_history = bool(include_history)
        self.language = "cpp"

    @staticmethod
    def _system(ctx: LoopContext | None) -> str:
        if ctx is not None and ctx.task_context:
            return ctx.task_context
        return (
            "You are ALE-Agent, an expert algorithm engineer improving a scored heuristic-programming "
            "solution through measured experiments."
        )

    def build(self, selection: Selection, memory=None, ctx: LoopContext | None = None) -> Prompt:
        parent = selection.parent
        if parent is None:
            raise RuntimeError("ALE-Agent requires an explicit parent state")
        state = (ctx.blackboard.get("ale_agent") if ctx is not None else None) or {}
        turn = max(1, int(state.get("refinement_turn", 1)))
        turns = max(1, int(state.get("refinement_turns", 1)))
        branch = max(1, int(state.get("branch_index", 1)))
        branches = max(1, int(state.get("children_per_parent", 1)))
        best = (
            selection.inspirations[0]
            if selection.inspirations
            else (ctx.best if ctx is not None and ctx.best is not None else parent)
        )

        sections = [
            f"# ALE-Agent expansion\nSibling branch {branch}/{branches}; refinement turn {turn}/{turns}.",
            _state_block("Current state and performance feedback", parent, self.language),
        ]
        if best.id != parent.id:
            sections.append(_state_block("Historically best state and feedback", best, self.language))
        else:
            sections.append("# Historically best state and feedback\nThe current state is also the best so far.")

        if self.include_history and memory is not None:
            history = memory.read()
            if history:
                sections.append(f"# Summary of the recent search trajectory\n{history}")

        if self.domain_guidance:
            index = int(state.get("guidance_index", 0)) % len(DOMAIN_GUIDANCE)
            title, guidance = DOMAIN_GUIDANCE[index]
            sections.append(f"# Targeted improvement guidance: {title}\n{guidance}")

        inherited_strategy = parent.metadata.get("ale_strategy")
        if turn == 1:
            sections.append(
                "# Strategy turn\n"
                "Formulate one concrete, high-level improvement strategy for this branch. Analyze the "
                "current code, its measured feedback, and the historical best. Do not output code yet; "
                "explain the proposed algorithmic or implementation change precisely enough to implement "
                "in the next message."
            )
        else:
            if inherited_strategy:
                sections.append(f"# Branch strategy\n{inherited_strategy}")
            sections.append(
                f"# Implementation refinement turn {turn}/{turns}\n"
                "Use the latest evaluation feedback to correct or improve this branch. Implement the "
                "strategy as a complete, compilable program. Preserve the required input/output protocol.\n\n"
                f"{self.mutation_approach.add_prompt()}"
            )
        return Prompt(system=self._system(ctx), user="\n\n".join(sections))