SakanaAI/ale_agent

"""ALE-Agent multi-turn program proposer."""
from __future__ import annotations

import time

from ...components.proposer import (
    Env,
    LLMProposer,
    WholeLineDiff,
    _emit_tool_events,
    _model_call_record,
    _tooling_summary,
)
from ...models import Generation, Prompt


def _call_record(generation: Generation) -> dict:
    return _model_call_record(generation)


def _aggregate_calls(calls: list[Generation]) -> dict:
    records = [_call_record(call) for call in calls]
    last = records[-1]
    return {
        "model": last["model"],
        "prompt_tokens": sum(r["prompt_tokens"] for r in records),
        "completion_tokens": sum(r["completion_tokens"] for r in records),
        "cached_tokens": sum(r["cached_tokens"] for r in records),
        "cost_usd": sum(r["cost_usd"] for r in records),
        "latency_s": sum(r["latency_s"] for r in records),
        "effort": last["effort"],
        "reasoning_summary": last["reasoning_summary"],
        "reasoning_tokens": sum(
            int(r.get("reasoning_tokens") or 0) for r in records
        ),
        "reasoning_summary_present": any(
            r["reasoning_summary_present"] for r in records
        ),
        "cache_write_tokens": sum(
            int(r.get("cache_write_tokens") or 0) for r in records
        ),
        "cost_source": last["cost_source"],
        "pricing_source": last["pricing_source"],
        "tool_call_on": any(r.get("tool_call_on") for r in records),
        "tool_names": list(dict.fromkeys(
            name for record in records for name in record.get("tool_names", [])
        )),
        "tool_rounds": sum(int(r.get("tool_rounds") or 0) for r in records),
        "max_tool_rounds": last.get("max_tool_rounds"),
        "tool_calls": [
            call for record in records for call in record.get("tool_calls", [])
        ],
        "tool_errors": [
            error for record in records for error in record.get("tool_errors", [])
        ],
        "llm_rounds": sum(int(r.get("llm_rounds") or 1) for r in records),
        "llm_latency_s": sum(float(r.get("llm_latency_s") or 0.0) for r in records),
        "tool_latency_s": sum(float(r.get("tool_latency_s") or 0.0) for r in records),
        "calls": len(records),
    }


class ALEAgentProposer(LLMProposer):
    """Strategy call on turn one, then one complete implementation per refinement turn.

    The domain prompts in the paper explicitly ask the model to reason before implementation.
    Therefore the first refinement turn makes two calls: a strategy-only response followed by an
    implementation response that receives that strategy.  Later turns make one implementation call
    using the previous code and its fresh evaluator feedback.
    """

    edit_strategy = WholeLineDiff()

    @staticmethod
    def _emit_model_call(ctx, prompt: Prompt, *, generation: Generation | None,
                         purpose: str, attempt_number: int | None,
                         caused_by_event_id: str | None = None,
                         error: Exception | None = None,
                         duration_seconds: float | None = None) -> str | None:
        if ctx is None:
            return None
        call = _call_record(generation) if generation is not None else {}
        model_event = ctx.emit_event(
            event_type="model_call",
            action="generate",
            status="failed" if error else "completed",
            component_role="proposer",
            component_name="ALEAgentProposer",
            attempt_number=attempt_number,
            evolution_track="solution",
            caused_by_event_ids=[caused_by_event_id] if caused_by_event_id else [],
            duration_seconds=(
                duration_seconds if duration_seconds is not None else call.get("latency_s")
            ),
            error=error,
            details={
                "purpose": purpose,
                "call_id": call.get("call_id"),
                "model": call.get("model"),
                "prompt": {"system": prompt.system, "user": prompt.user},
                "response": {
                    "text": generation.text if generation is not None else None,
                    "reasoning": generation.reasoning if generation is not None else None,
                },
                "usage": {
                    "prompt_tokens": call.get("prompt_tokens"),
                    "cached_prompt_tokens": call.get("cached_tokens"),
                    "completion_tokens": call.get("completion_tokens"),
                    "cost_usd": call.get("cost_usd"),
                },
                "parameters": {
                    "reasoning_effort": call.get("effort"),
                    "reasoning_summary": call.get("reasoning_summary"),
                },
                "accounting": {
                    "reasoning_tokens": call.get("reasoning_tokens"),
                    "reasoning_summary_present": call.get("reasoning_summary_present"),
                    "cache_write_tokens": call.get("cache_write_tokens"),
                    "cost_source": call.get("cost_source"),
                    "pricing_source": call.get("pricing_source"),
                },
                "tooling": _tooling_summary(call),
            },
        )
        _emit_tool_events(
            ctx,
            model_call_event_id=model_event,
            call=call,
            attempt_number=attempt_number,
            evolution_track="solution",
        )
        return model_event

    @staticmethod
    def _generate(model, prompt: Prompt, ctx) -> Generation:
        started = time.monotonic()
        generation = model.generate(prompt)
        if not generation.latency_s:
            generation.latency_s = time.monotonic() - started
        if not generation.model:
            generation.model = str(getattr(model, "name", type(model).__name__))
        if ctx is not None:
            ctx.record_cost(
                generation.cost_usd,
                generation.prompt_tokens,
                generation.completion_tokens,
            )
        return generation

    def propose(self, prompt: Prompt, env: Env):
        parent = env.selection.parent
        if parent is None:
            raise RuntimeError("ALE-Agent cannot propose without a parent")
        state = (env.ctx.blackboard.get("ale_agent") if env.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)))
        calls: list[Generation] = []
        strategy_prompt: Prompt | None = None
        strategy_generation: Generation | None = None
        selection_event = getattr(env.selection, "_etif_event_id", None)
        strategy_event = None

        if turn == 1:
            strategy_prompt = prompt
            started = time.monotonic()
            try:
                strategy_generation = self._generate(env.model, strategy_prompt, env.ctx)
            except Exception as exc:
                self._emit_model_call(
                    env.ctx,
                    strategy_prompt,
                    generation=None,
                    purpose="strategy_guidance",
                    attempt_number=env.attempt_number,
                    caused_by_event_id=selection_event,
                    error=exc,
                    duration_seconds=time.monotonic() - started,
                )
                raise
            strategy_event = self._emit_model_call(
                env.ctx,
                strategy_prompt,
                generation=strategy_generation,
                purpose="strategy_guidance",
                attempt_number=env.attempt_number,
                caused_by_event_id=selection_event,
            )
            calls.append(strategy_generation)
            strategy_text = strategy_generation.text.strip()
            implementation_prompt = Prompt(
                system=prompt.system,
                user=(
                    f"{prompt.user}\n\n# Chosen branch strategy\n{strategy_text}\n\n"
                    f"# Implementation refinement turn 1/{turns}\n"
                    "Now implement that strategy as a complete, compilable program. Preserve the required "
                    "input/output protocol and use the evaluator feedback above.\n\n"
                    f"{self.mutation_approach.add_prompt()}"
                ),
            )
        else:
            strategy_text = str(parent.metadata.get("ale_strategy") or "")
            implementation_prompt = prompt

        started = time.monotonic()
        try:
            implementation = self._generate(env.model, implementation_prompt, env.ctx)
        except Exception as exc:
            self._emit_model_call(
                env.ctx,
                implementation_prompt,
                generation=None,
                purpose="candidate_proposal",
                attempt_number=env.attempt_number,
                caused_by_event_id=strategy_event or selection_event,
                error=exc,
                duration_seconds=time.monotonic() - started,
            )
            raise
        implementation_event = self._emit_model_call(
            env.ctx,
            implementation_prompt,
            generation=implementation,
            purpose="candidate_proposal",
            attempt_number=env.attempt_number,
            caused_by_event_id=strategy_event or selection_event,
        )
        calls.append(implementation)
        parser = self.mutation_approach.parser() or self.edit_strategy
        new_code, changed, changes = parser.apply(parent.content, implementation.text)
        child = parent.child(new_code, generation=parent.metadata.get("generation", 0) + 1)
        # A child inherits parent metadata by design; reset the outer-search flags so an expanded
        # parent cannot make every descendant accidentally tabu.
        child.metadata.update({
            "changed": changed,
            "parent_metrics": dict(parent.scores),
            "changes": changes,
            # Live logs retain the turn aggregate. trajectory.proposal uses trace["proposal_call"]
            # below for the implementation call itself; the strategy call lives in extensions.
            "model_call": _aggregate_calls(calls),
            "ale_model_call_count": len(calls),
            "ale_outer_parent_id": state.get("outer_parent_id", parent.id),
            "ale_branch_index": int(state.get("branch_index", 1)),
            "ale_refinement_turn": turn,
            "ale_refinement_turns": turns,
            "ale_guidance_index": int(state.get("guidance_index", 0)),
            "ale_strategy": strategy_text,
            "ale_frontier": False,
            "ale_expanded": False,
            "ale_branch_winner": False,
        })
        child.artifacts["response"] = implementation.text
        child.artifacts["reasoning"] = implementation.reasoning
        child.artifacts["implementation_prompt"] = {
            "system": implementation_prompt.system,
            "user": implementation_prompt.user,
        }
        child.artifacts["model_calls"] = [_call_record(call) for call in calls]
        if strategy_prompt is not None and strategy_generation is not None:
            child.artifacts["strategy_prompt"] = {
                "system": strategy_prompt.system,
                "user": strategy_prompt.user,
            }
            child.artifacts["strategy_response"] = strategy_generation.text
            child.artifacts["strategy_reasoning"] = strategy_generation.reasoning
        child.trace.update({
            "proposal_prompt": {
                "system": implementation_prompt.system,
                "user": implementation_prompt.user,
            },
            "proposal_call": _call_record(implementation),
            "proposal_response_text": implementation.text,
            "proposal_reasoning": implementation.reasoning,
            "proposal_model_call_event_id": implementation_event,
        })
        return child