SkyDiscover/gepa_native

"""Reflective prompt builder for SkyDiscover's GEPA Native search strategy."""
from __future__ import annotations

import re

from ...components.prompt import DefaultPromptBuilder
from ...models.base import Prompt
from ...records import LoopContext, Selection
from .population import GepaNativePopulation, genome_score

_TEXT_LANGUAGES = {"text", "prompt", "text/plain"}


def _prompt_optimization_task(language: str, timeout_warning: str) -> str:
    """SkyDiscover ``full_rewrite_prompt_opt_user_message.txt`` from its ``# Task`` onward."""
    return (
        "# Task\n"
        "Suggest improvements to the prompt that will improve its COMBINED_SCORE.\n"
        "The system maintains diversity across these dimensions: score, complexity.\n"
        "Different prompts with similar combined_score but different strategies are valuable.\n"
        "\n"
        "Provide the complete rewritten prompt.\n"
        "\n"
        "IMPORTANT: Your rewritten prompt must serve the same task as the original prompt,\n"
        "but with improved instructions, structure, or reasoning guidance.\n"
        "Do NOT change the task the prompt is designed for - only improve HOW it instructs the LLM.\n"
        "\n"
        f"```{language}\n"
        "# Your rewritten prompt here\n"
        "```\n"
        "\n"
        "Be thoughtful about your changes and explain your reasoning thoroughly.\n"
        "Consider these prompt engineering strategies:\n"
        "- Clearer step-by-step reasoning instructions\n"
        "- Better output format constraints\n"
        "- More precise task framing and role definitions\n"
        "- Adding or refining examples\n"
        "- Reducing ambiguity in instructions\n"
        "\n"
        'IMPORTANT: If an instruction header of "## IMPORTANT: ..." is given below the '
        '"# Current Solution", you MUST follow it. Otherwise,\n'
        "focus on targeted improvements of the prompt.\n"
        "\n"
        f"{timeout_warning}"
    )


class GepaNativePromptBuilder(DefaultPromptBuilder):
    """Insert recent stored-parent rejections immediately before the normal mutation task."""

    def __init__(
        self,
        population: GepaNativePopulation,
        max_recent_failures: int = 5,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.population = population
        self.max_recent_failures = int(max_recent_failures)

    def build(
        self,
        selection: Selection,
        memory=None,
        ctx: LoopContext | None = None,
    ) -> Prompt:
        prompt = super().build(selection, memory, ctx)
        guidance = self._build_search_guidance()
        user = prompt.user
        if (
            self.mutation_approach.name == "full_rewrite"
            and self.language.lower() in _TEXT_LANGUAGES
        ):
            head, separator, _tail = user.rpartition("\n# Task\n")
            if separator:
                head = head.replace(
                    "# Current Solution Information",
                    "# Current Prompt Information",
                    1,
                ).replace(
                    "# Program Generation History",
                    "# Prompt Generation History",
                    1,
                )
                user = (
                    f"{head.rstrip()}\n\n"
                    f"{_prompt_optimization_task(self.language, self._timeout_warning())}"
                )
        if guidance:
            head, separator, tail = user.rpartition("\n# Task\n")
            if separator:
                user = f"{head.rstrip()}\n\n{guidance}\n\n# Task\n{tail}"
        # Upstream collapses the blank placeholder left by an empty guidance section too.
        user = re.sub(r"\n{3,}", "\n\n", user)
        return Prompt(system=prompt.system, user=user)

    def _build_search_guidance(self) -> str:
        rejected = self.population.get_rejection_history(self.max_recent_failures)
        if not rejected:
            return ""
        rejection_section = self._format_rejection_history(rejected)
        if not rejection_section:
            return ""
        header = (
            "## Reflective Analysis\n"
            "Review the evaluation results and diagnostics in the program "
            "information above. Identify root causes and domain-specific "
            "insights. Address these failure modes in your solution."
        )
        return f"{header}\n\n{rejection_section}"

    def _format_rejection_history(self, rejected: list) -> str | None:
        if not rejected:
            return None
        entries: list[str] = []
        for index, genome in enumerate(rejected, 1):
            parent_score = ""
            parent = self.population.get(genome.parent_id)
            if parent is not None:
                parent_score = f", parent_score: {genome_score(parent):.4f}"

            error_message = ""
            if genome.scores:
                error_message = (
                    genome.scores.get("error", "")
                    or genome.scores.get("error_message", "")
                )
            changes = genome.metadata.get("changes", "") if genome.metadata else ""

            lines = [
                f"#### Attempt {index} "
                f"(score: {genome_score(genome):.4f}{parent_score})"
            ]
            if changes:
                lines.append(f"Changes: {changes}")
            if error_message:
                lines.append(f"Error: {error_message}")
            if genome.content:
                code_lines = genome.content.splitlines()
                snippet = (
                    "\n".join(code_lines[:30]) + "\n... (truncated)"
                    if len(code_lines) > 30
                    else genome.content
                )
                lines.append(f"Code tried:\n```\n{snippet}\n```")
            entries.append("\n".join(lines))

        return "\n\n".join([
            "### Recent Rejected Attempts",
            "The following mutations were rejected because they did not "
            "improve on the parent. Avoid repeating these approaches:",
            *entries,
        ])