AlgoTune/algotune_agent

"""AlgoTune Agent's one-command loop as a Galapagos Proposer."""
from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import tempfile
import time
import uuid
from dataclasses import dataclass

from ...components import Proposer
from ...components.proposer import (
    _emit_tool_events,
    _model_call_record,
    _tooling_summary,
)
from ...components.evaluator import ContainerEvaluator, _PROVIDER_SECRETS
from ...models import Prompt
from ...records import EvalResult, Genome
from .population import algotune_priority
from .validator import protected_filename_error
from .workspace import (
    delete_lines,
    numbered,
    pack_workspace,
    replace_lines,
    syntax_error,
    unpack_workspace,
    validate_filename,
)

_COMMANDS = {
    "edit", "ls", "view_file", "revert", "reference", "eval_input", "eval",
    "delete", "profile", "profile_lines",
}
_FENCE = re.compile(r"(?ms)^```([^\n`]*)[ \t]*\n(.*?)^```[ \t]*$")
_THINK = re.compile(r"(?is)<(?:think|thinking)>.*?</(?:think|thinking)>")


class CommandParseError(ValueError):
    pass


@dataclass(frozen=True)
class ParsedCommand:
    name: str
    args: dict
    raw: str


def _command_text(response: str) -> str:
    clean = _THINK.sub("", response or "").strip()
    blocks = list(_FENCE.finditer(clean))
    if len(blocks) > 1:
        raise CommandParseError("Multiple command blocks found; send one and only one command")
    if blocks:
        match = blocks[0]
        if clean[match.end() :].strip():
            raise CommandParseError("Text found after the command block")
        language, body = match.group(1).strip(), match.group(2).strip("\n")
        first = body.strip().split(maxsplit=1)[0] if body.strip() else ""
        if language in _COMMANDS:
            return f"{language}\n{body}".strip()
        if first in _COMMANDS:
            return body.strip()
        raise CommandParseError("The fenced block does not contain an AlgoTune command")
    first = clean.split(maxsplit=1)[0] if clean else ""
    if first in _COMMANDS:
        return clean
    raise CommandParseError("No command block found")


def parse_command(response: str) -> ParsedCommand:
    text = _command_text(response)
    name = text.split(maxsplit=1)[0]
    if name == "edit":
        match = re.fullmatch(
            r"edit\s*\nfile:\s*(\S+)\s*\nlines:\s*(\d+)\s*-\s*(\d+)\s*"
            r"\n---[ \t]*\n(.*?)\n---\s*",
            text,
            re.DOTALL,
        )
        if not match:
            raise CommandParseError("Invalid edit format")
        return ParsedCommand(name, {
            "file": match.group(1), "start": int(match.group(2)),
            "end": int(match.group(3)), "content": match.group(4),
        }, text)
    if name == "delete":
        match = re.fullmatch(
            r"delete\s*\nfile:\s*(\S+)\s*\nlines:\s*(\d+)\s*-\s*(\d+)\s*", text
        )
        if not match:
            raise CommandParseError("Invalid delete format")
        return ParsedCommand(name, {
            "file": match.group(1), "start": int(match.group(2)), "end": int(match.group(3)),
        }, text)
    if name == "view_file":
        match = re.fullmatch(r"view_file\s+(\S+)(?:\s+(\d+))?\s*", text)
        if not match:
            raise CommandParseError("Invalid view_file format")
        return ParsedCommand(name, {"file": match.group(1), "start": int(match.group(2) or 1)}, text)
    if name in {"ls", "revert", "eval"}:
        if not re.fullmatch(name + r"\s*", text):
            raise CommandParseError(f"Invalid {name} format")
        return ParsedCommand(name, {}, text)
    if name in {"reference", "eval_input"}:
        match = re.fullmatch(name + r"\s+([\s\S]+)", text)
        if not match:
            raise CommandParseError(f"Invalid {name} format")
        return ParsedCommand(name, {"input": match.group(1).strip()}, text)
    if name == "profile":
        match = re.fullmatch(r"profile\s+(\S+\.py)\s+([\s\S]+)", text)
        if not match:
            raise CommandParseError("Invalid profile format")
        return ParsedCommand(name, {"file": match.group(1), "input": match.group(2).strip()}, text)
    if name == "profile_lines":
        match = re.fullmatch(
            r"profile_lines\s+(\S+\.py)\s+((?:\d+(?:-\d+)?)(?:\s*,\s*\d+(?:-\d+)?)*)\s+([\s\S]+)",
            text,
        )
        if not match:
            raise CommandParseError("Invalid profile_lines format")
        return ParsedCommand(name, {
            "file": match.group(1), "lines": match.group(2), "input": match.group(3).strip(),
        }, text)
    raise CommandParseError(f"Unknown command {name!r}")


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


def format_evaluation(result: EvalResult) -> str:
    def render(value):
        return f"{value:.6g}" if isinstance(value, (int, float)) else str(value)

    metrics = ", ".join(f"{key}={render(value)}" for key, value in result.metrics.items()) or "none"
    feedback = result.text_feedback or (result.artifacts or {}).get("text_feedback") or ""
    return f"Evaluation complete. valid={bool(result.valid)}; metrics: {metrics}.\n{feedback}".strip()


_TOOL_RUNNER = r'''
import ast, contextlib, copy, cProfile, importlib.util, io, json, pathlib, pstats, re, sys, traceback

ev_path, candidate_path, mode, raw_input, focus, requested_file = sys.argv[1:7]

def load(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod

def parse_input(text):
    try:
        return ast.literal_eval(text)
    except Exception:
        return json.loads(text)

def plain(value):
    if hasattr(value, "tolist"):
        return plain(value.tolist())
    if isinstance(value, dict):
        return {str(key): plain(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [plain(item) for item in value]
    return value

def selected_lines(spec):
    chosen = set()
    for part in spec.split(","):
        if not part.strip():
            continue
        bounds = [int(item) for item in part.strip().split("-", 1)]
        chosen.update(range(bounds[0], bounds[-1] + 1))
    return chosen

def profile_target(candidate):
    if not requested_file or requested_file == "solver.py":
        target = candidate
    else:
        entrypoint = getattr(candidate, "run_solver", None)
        source = pathlib.Path(getattr(entrypoint, "__code__", None).co_filename)
        target_path = source.parent / requested_file
        if not target_path.is_file():
            raise FileNotFoundError("File not found: " + requested_file)
        target = load("_algotune_profile_target", str(target_path))
    function = getattr(target, "run_solver", None)
    if callable(function):
        return target, function
    solver_class = getattr(target, "Solver", None)
    if solver_class is not None:
        function = getattr(solver_class(), "solve", None)
    if not callable(function):
        raise AttributeError(
            "profile target missing run_solver(problem) or class Solver.solve(problem): "
            + (requested_file or "solver.py")
        )
    return target, function

try:
    evaluator = load("_algotune_tool_evaluator", ev_path)
    problem = parse_input(raw_input)
    if mode == "reference":
        out = evaluator.reference_solve(problem)
        result = {"reference": plain(out)}
    else:
        candidate = load("_algotune_tool_candidate", candidate_path)
        if not callable(getattr(candidate, "run_solver", None)):
            raise AttributeError("program missing run_solver(problem)")
        if mode == "eval_input":
            ref = evaluator.reference_solve(copy.deepcopy(problem))
            stream = io.StringIO()
            with contextlib.redirect_stdout(stream):
                out = candidate.run_solver(copy.deepcopy(problem))
            result = {"solver": plain(out), "reference": plain(ref), "stdout": stream.getvalue(),
                      "correct": bool(evaluator.is_correct(copy.deepcopy(problem), out))}
        else:
            stream = io.StringIO()
            target, target_function = profile_target(candidate)
            try:
                from line_profiler import LineProfiler
                profiler = LineProfiler()
                profiler.add_function(target_function)
                for value in vars(target).values():
                    if isinstance(value, type):
                        solve = getattr(value, "solve", None)
                        if callable(solve):
                            try:
                                profiler.add_function(solve)
                            except Exception:
                                pass
                profiler.runcall(target_function, problem)
                profiler.print_stats(stream=stream)
                report = stream.getvalue()
            except Exception:
                profiler = cProfile.Profile()
                profiler.runcall(target_function, problem)
                pstats.Stats(profiler, stream=stream).sort_stats("cumulative").print_stats(25)
                report = stream.getvalue()
            if mode == "profile_lines" and focus:
                wanted = selected_lines(focus)
                try:
                    source_lines = open(candidate_path, encoding="utf-8").read().splitlines()
                    marker = source_lines.index("# ALGOTUNE-SOLVER-PY-V1") + 1
                except (OSError, ValueError):
                    marker = 0
                actual = wanted | {line + marker for line in wanted}
                kept = []
                for line in report.splitlines():
                    match = re.match(r"^\s*(\d+)\s+", line)
                    if match is None or int(match.group(1)) in actual:
                        kept.append(line)
                report = "\n".join(kept)
            result = {"profile": report, "focus_lines": focus or None}
    sys.stdout.write("<<<ALGOTUNE_TOOL>>>" + json.dumps(result, default=str))
except Exception as exc:
    sys.stdout.write("<<<ALGOTUNE_TOOL>>>" + json.dumps(
        {"error": str(exc), "traceback": traceback.format_exc()}
    ))
'''


class AlgoTuneProposer(Proposer):
    """One LLM message, one parsed command, one workspace transition."""

    def _chat(self, prompt: Prompt, env):
        memory = env.memory
        memory.initialize(prompt.system, prompt.user)
        token_limit = (
            getattr(env.model, "context_length", None)
            or getattr(env.model, "context_window", None)
            or getattr(env.model, "max_tokens", None)
            or 4000
        )
        messages = memory.messages(token_limit=int(token_limit))
        started = time.monotonic()
        interface = "chat"
        try:
            try:
                generation = env.model.chat(messages)
            except (AttributeError, NotImplementedError):
                interface = "generate"
                flattened = "\n\n".join(
                    f"[{message['role'].upper()}]\n{message['content']}" for message in messages
                )
                generation = env.model.generate(Prompt(system="", user=flattened))
        except Exception as exc:
            if env.ctx is not None:
                env.ctx.emit_event(
                    event_type="model_call",
                    action="generate",
                    status="failed",
                    component_role="proposer",
                    component_name=type(self).__name__,
                    attempt_number=env.attempt_number,
                    evolution_track="solution",
                    caused_by_event_ids=([getattr(env.selection, "_etif_event_id", None)]
                                         if getattr(env.selection, "_etif_event_id", None) else []),
                    duration_seconds=time.monotonic() - started,
                    error=exc,
                    details={
                        "purpose": "candidate_proposal",
                        "call_id": None,
                        "model": str(getattr(env.model, "name", type(env.model).__name__)),
                        "prompt": {"system": None, "user": None, "messages": messages},
                        "response": {"text": None, "reasoning": None},
                        "usage": {
                            "prompt_tokens": None,
                            "cached_prompt_tokens": None,
                            "completion_tokens": None,
                            "cost_usd": None,
                        },
                        "parameters": {"interface": interface},
                    },
                )
            raise
        if not generation.latency_s:
            generation.latency_s = time.monotonic() - started
        if not generation.model:
            generation.model = str(getattr(env.model, "name", type(env.model).__name__))
        if env.ctx is not None:
            env.ctx.record_cost(
                generation.cost_usd, generation.prompt_tokens, generation.completion_tokens
            )
        call_record = _call_record(generation)
        model_call_event = (env.ctx.emit_event(
            event_type="model_call",
            action="generate",
            status="completed",
            component_role="proposer",
            component_name=type(self).__name__,
            attempt_number=env.attempt_number,
            evolution_track="solution",
            caused_by_event_ids=([getattr(env.selection, "_etif_event_id", None)]
                                 if getattr(env.selection, "_etif_event_id", None) else []),
            duration_seconds=call_record.get("latency_s"),
            details={
                "purpose": "candidate_proposal",
                "call_id": call_record.get("call_id"),
                "model": call_record.get("model"),
                "prompt": {"system": None, "user": None, "messages": messages},
                "response": {"text": generation.text, "reasoning": generation.reasoning},
                "usage": {
                    "prompt_tokens": call_record.get("prompt_tokens"),
                    "cached_prompt_tokens": call_record.get("cached_tokens"),
                    "completion_tokens": call_record.get("completion_tokens"),
                    "cost_usd": call_record.get("cost_usd"),
                },
                "parameters": {
                    "reasoning_effort": call_record.get("effort"),
                    "reasoning_summary": call_record.get("reasoning_summary"),
                    "interface": interface,
                },
                "accounting": {
                    "reasoning_tokens": call_record.get("reasoning_tokens"),
                    "reasoning_summary_present": call_record.get(
                        "reasoning_summary_present"
                    ),
                    "cache_write_tokens": call_record.get("cache_write_tokens"),
                    "cost_source": call_record.get("cost_source"),
                    "pricing_source": call_record.get("pricing_source"),
                },
                "tooling": _tooling_summary(call_record),
            },
        ) if env.ctx is not None else None)
        _emit_tool_events(
            env.ctx,
            model_call_event_id=model_call_event,
            call=call_record,
            attempt_number=env.attempt_number,
            evolution_track="solution",
        )
        return generation, messages, call_record, model_call_event

    @staticmethod
    def _best(pool: list[Genome]) -> Genome | None:
        return max(pool, key=algotune_priority) if pool else None

    @staticmethod
    def _tool(command: ParsedCommand, parent: Genome, env, files: dict[str, str]) -> str:
        if command.name == "eval":
            started = time.monotonic()
            try:
                evaluated = env.evaluator.evaluate(parent)
            except Exception as exc:
                if env.ctx is not None:
                    env.ctx.blackboard["_algotune_diagnostic_evaluation"] = {
                        "duration_seconds": time.monotonic() - started,
                        "data_split": "task_default",
                        "error": {
                            "error_type": type(exc).__name__,
                            "message": str(exc),
                        },
                    }
                return f"Diagnostic failed: {type(exc).__name__}: {exc}"
            if env.ctx is not None:
                env.ctx.blackboard["_algotune_diagnostic_evaluation"] = {
                    "duration_seconds": time.monotonic() - started,
                    "data_split": "task_default",
                    "result": {
                        "is_valid": bool(evaluated.valid),
                        "metrics": dict(evaluated.metrics or {}),
                        "text_feedback": evaluated.text_feedback,
                    },
                }
            return format_evaluation(evaluated)
        filename = command.args.get("file")
        if filename:
            filename = validate_filename(filename)
            if filename not in files:
                return f"File not found: {filename}"
        evaluator_path = getattr(env.evaluator, "evaluator_path", None)
        if not evaluator_path:
            return (
                "This diagnostic requires a local task evaluator. Run this scaffold through "
                "`galapagos run`; its whole-loop task-container mode provides that environment."
            )
        timeout = max(30, int(getattr(env.evaluator, "timeout", 120)))
        tool_args = [
            command.name, command.args["input"], command.args.get("lines", ""),
            command.args.get("file", ""),
        ]
        candidate_path = ""
        container = None
        try:
            if isinstance(env.evaluator, ContainerEvaluator):
                # Programmatic ``scaffold.run`` does not use the CLI's whole-loop orchestration.
                # Attach diagnostics to its task evaluator container so they still see the exact task
                # packages. The CLI path uses a local evaluator *inside* that image and takes the branch
                # below instead.
                env.evaluator._ensure_started()
                container = env.evaluator._container
                if container is None:
                    raise RuntimeError("task evaluator container did not start")
                candidate_path = f"/tmp/algotune_tool_{uuid.uuid4().hex}.py"
                container.put_bytes(
                    parent.content.encode(), candidate_path, user=env.evaluator.run_user
                )
                inside_evaluator = f"{env.evaluator._EVAL_DIR}/{env.evaluator.evaluator_name}"
                process = container.exec(
                    ["python3", "-c", _TOOL_RUNNER, inside_evaluator, candidate_path, *tool_args],
                    workdir=env.evaluator._EVAL_DIR,
                    timeout=timeout,
                    user=env.evaluator.run_user,
                )
            else:
                fd, candidate_path = tempfile.mkstemp(suffix=".py")
                with os.fdopen(fd, "w", encoding="utf-8") as handle:
                    handle.write(parent.content)
                environment = dict(os.environ)
                for key in _PROVIDER_SECRETS:
                    environment.pop(key, None)
                process = subprocess.run(
                    [
                        sys.executable, "-c", _TOOL_RUNNER, str(evaluator_path), candidate_path,
                        *tool_args,
                    ],
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                    cwd=os.path.dirname(os.path.abspath(str(evaluator_path))),
                    env=environment,
                )
            if "<<<ALGOTUNE_TOOL>>>" not in process.stdout:
                return "Diagnostic failed: " + (process.stderr or process.stdout or "no output")[-2000:]
            payload = json.loads(process.stdout.split("<<<ALGOTUNE_TOOL>>>", 1)[1])
            if payload.get("error"):
                return f"Diagnostic failed: {payload['error']}\n{payload.get('traceback', '')}".strip()
            return json.dumps(payload, indent=2, default=str)
        except subprocess.TimeoutExpired:
            return f"Diagnostic timed out after {timeout}s"
        except Exception as exc:  # noqa: BLE001 - a tool failure is feedback, not a run failure
            return f"Diagnostic failed: {type(exc).__name__}: {exc}"
        finally:
            if container is not None and candidate_path:
                try:
                    container.exec(["rm", "-f", candidate_path], user=env.evaluator.run_user)
                except Exception:  # noqa: BLE001 - diagnostic cleanup must not mask its result
                    pass
            elif candidate_path:
                try:
                    os.remove(candidate_path)
                except OSError:
                    pass

    def _execute(self, command: ParsedCommand, parent: Genome, env) -> tuple[str, bool, str, bool]:
        files = unpack_workspace(parent.content)
        name = command.name
        if name == "ls":
            listing = "\n".join(
                filename
                for filename in sorted(files)
                if "/" not in filename
                and not filename.startswith(".")
                and filename != "__init__.py"
            )
            return parent.content, False, f"File list:\n{listing}", False
        if name == "view_file":
            filename = validate_filename(command.args["file"])
            if filename not in files:
                return parent.content, False, f"File not found: {filename}", False
            view = numbered(files[filename], command.args["start"], 100)
            return parent.content, False, view or "(empty file)", False
        if name in {"reference", "eval_input", "profile", "profile_lines", "eval"}:
            return parent.content, False, self._tool(command, parent, env, files), False
        if name == "revert":
            best = self._best(env.selection.pool)
            if best is None:
                return parent.content, False, "No saved state to revert to", False
            detail = f"Restored best snapshot {best.id} (score={best.fitness:.6g})."
            # Upstream copies the saved snapshot and reloads it; it does not run dataset evaluation.
            # The existing measured Genome is already that snapshot, so after_step moves the mutable
            # current pointer to it without minting a duplicate scored candidate.
            return best.content, False, detail, False

        filename = validate_filename(command.args["file"])
        protected_error = protected_filename_error(filename)
        if protected_error:
            return parent.content, False, protected_error, False
        old = files.get(filename, "")
        try:
            if name == "edit":
                files[filename] = replace_lines(
                    old, command.args["start"], command.args["end"], command.args["content"]
                )
            else:
                files[filename] = delete_lines(old, command.args["start"], command.args["end"])
        except ValueError as exc:
            return parent.content, False, str(exc), False
        lint = syntax_error(files)
        if lint:
            return parent.content, False, f"Critical syntax error; edit was not applied: {lint}", False
        candidate = pack_workspace(files)
        changed = candidate != parent.content
        detail = (
            f"{name} applied to {filename}; trusted evaluation follows."
            if changed else "Command made no change."
        )
        return candidate, changed, detail, changed

    def propose(self, prompt: Prompt, env) -> Genome:
        parent = env.selection.parent
        if parent is None:
            raise RuntimeError("AlgoTune requires a current workspace parent")
        generation, messages, call_record, model_call_event = self._chat(prompt, env)
        env.memory.add("assistant", generation.text)
        if env.ctx is not None:
            env.ctx.blackboard.pop("_algotune_diagnostic_evaluation", None)
        command = None
        try:
            command = parse_command(generation.text)
            content, changed, result, needs_evaluation = self._execute(command, parent, env)
        except (CommandParseError, ValueError) as exc:
            content, changed, result, needs_evaluation = parent.content, False, str(exc), False
        child = parent.child(content, generation=parent.metadata.get("generation", 0) + 1)
        child.metadata.update({
            "changed": bool(changed),
            "parent_metrics": dict(parent.scores),
            "changes": command.raw[:500] if command else "invalid command",
            "algotune_command": command.name if command else "parse_error",
            "algotune_command_raw": command.raw if command else "",
            "algotune_result": result,
            "algotune_needs_evaluation": bool(needs_evaluation),
            "algotune_edited_file": command.args.get("file") if command else None,
            "algotune_revert_target_id": (
                self._best(env.selection.pool).id
                if command is not None and command.name == "revert" and env.selection.pool
                else None
            ),
            "model_call": call_record,
        })
        child.artifacts["response"] = generation.text
        child.artifacts["reasoning"] = generation.reasoning
        child.trace["algotune_chat_messages"] = messages
        child.trace["proposal_model_call_event_id"] = model_call_event
        child.trace["proposal_memory_writes"] = [{
            "memory_name": "conversation_history",
            "entry_type": "assistant_message",
            "content": generation.text,
        }]
        child.trace["proposal_prompt"] = {
            "system": "",
            "user": "\n\n".join(
                f"[{message['role'].upper()}]\n{message['content']}" for message in messages
            ),
        }
        diagnostic = (
            env.ctx.blackboard.pop("_algotune_diagnostic_evaluation", None)
            if env.ctx is not None else None
        )
        if diagnostic:
            child.trace["diagnostic_evaluation"] = diagnostic
        if not changed:
            child.trace["etif_no_candidate"] = True
            child.trace["no_candidate_reason"] = (
                "command_parse_error" if command is None else
                "control_action" if command.name in {
                    "ls", "view_file", "reference", "eval_input", "eval",
                    "profile", "profile_lines", "revert",
                } else "no_change"
            )
        return child