GEPA/gepa

"""GEPA PromptBuilder component — the reflection prompt (Algorithm 3 in the paper).

Two upstream pieces, both rendered verbatim:

* the user message is ``optimize_anything_reflection_prompt_template``
  (``gepa/optimize_anything.py``), the template the single-task ``optimize_anything`` API uses, with
  ``<curr_param>`` filled by the parent program and ``<side_info>`` by its evaluation feedback;
* ``<side_info>`` itself is the **reflective dataset**, rendered by
  ``InstructionProposalSignature.prompt_renderer``'s ``format_samples``
  (``gepa/strategies/instruction_proposal.py``) — the ``# Example N`` / ``## key`` / ``### subkey``
  markdown walk, nesting one heading level per dict depth and capping at ``######``.

The record fed into it is what ``OptimizeAnythingAdapter.make_reflective_dataset`` builds from a
``SideInfo`` dict: the ``"scores"`` entry renamed to ``"Scores (Higher is Better)"`` followed by every
other diagnostic key. galapagos's ``SideInfo`` is the Evaluator's verdict on the parent — its
``metrics`` and its ``artifacts`` (``text_feedback``, tracebacks, stdout, whatever the task returns).
This is the ASI ("Actionable Side Information") the paper's sample-efficiency claim rests on: the LM
is told *why* the current program scored what it scored, not merely *that* it scored it.

Nothing is truncated. The whole point of reflective mutation is that the proposer reads the full
diagnostic text, and upstream caps none of it — a task that emits enormous feedback should cap it in
its evaluator, where the same cap also reaches every other scaffold.
"""
from __future__ import annotations

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

# ``gepa/optimize_anything.py::optimize_anything_reflection_prompt_template``, VERBATIM — the long
# lines are upstream's own and are deliberately not rewrapped.
_REFLECTION_TEMPLATE = """I am optimizing a parameter in my system. The current parameter value is:
```
<curr_param>
```

Below is evaluation data showing how this parameter value performed across multiple test cases. The data contains performance metrics, diagnostic information, and other relevant details from the evaluation:
```
<side_info>
```

Your task is to propose a new, improved parameter value that can be used as a drop-in replacement for the current one.

Carefully analyze all the evaluation data provided above. Look for patterns that indicate what works and what doesn't. Pay special attention to:
- Performance metrics and how they correlate with parameter behavior
- Recurring issues, errors, or failure patterns across multiple test cases
- Successful patterns or behaviors that should be preserved or enhanced
- Any domain-specific requirements, constraints, or factual information revealed in the evaluation data
- Specific technical details that are crucial for understanding the parameter's role

Based on your analysis, propose a new parameter value that addresses the identified issues while maintaining or improving upon what works well. Your proposal should be directly informed by the patterns and insights from the evaluation data.

Provide the new parameter value within ``` blocks."""

# The "system context" line of ``_build_reflection_prompt_template`` (same module), used only when the
# task card ships no ``system_message`` of its own. GEPA sends the reflection prompt as a bare user
# message; galapagos routes the task's persona through the system slot, so this is just the fallback.
_SYSTEM = ("You are an expert optimization assistant. Your task is to analyze evaluation "
           "feedback and propose an improved version of a system component.")

#: Raw model output the Proposer stashes on the genome — never evaluator feedback, so it is no more
#: part of the reflective dataset than it is of the default builder's ``## Evaluator Feedback``.
_NOT_FEEDBACK = ("response", "reasoning")


def render_value(value, level: int = 3) -> str:
    """``format_samples.render_value``: dicts and lists become nested markdown headings, scalars
    become a stripped paragraph. Heading depth saturates at ``######``."""
    if isinstance(value, dict):
        out = ""
        for key, item in value.items():
            out += f"{'#' * level} {key}\n"
            out += render_value(item, min(level + 1, 6))
        if not value:
            out += "\n"
        return out
    if isinstance(value, (list, tuple)):
        out = ""
        for index, item in enumerate(value):
            out += f"{'#' * level} Item {index + 1}\n"
            out += render_value(item, min(level + 1, 6))
        if not value:
            out += "\n"
        return out
    return f"{str(value).strip()}\n\n"


def render_reflective_dataset(records: list[dict]) -> str:
    """``format_samples``: one ``# Example N`` block per record, blank-line separated."""
    blocks = []
    for index, record in enumerate(records, start=1):
        block = f"# Example {index}\n"
        for key, value in record.items():
            block += f"## {key}\n"
            block += render_value(value, level=3)
        blocks.append(block)
    return "\n\n".join(blocks)


class GepaReflectionPromptBuilder(PromptBuilder):
    """Render the parent program plus its evaluation feedback into GEPA's reflection prompt."""

    def build(self, selection: Selection, memory=None, ctx: LoopContext | None = None) -> Prompt:
        # The task card's system_message REPLACES the generic fallback: it already opens with its own
        # persona and problem statement, and appending would stack two.
        system = ctx.task_context if (ctx and ctx.task_context) else _SYSTEM
        parent = selection.parent
        if parent is None:   # delegated selection (GEPA always selects a parent; kept for safety)
            return Prompt(system=system, user="")
        side_info = render_reflective_dataset([self.reflective_record(parent)])
        user = _REFLECTION_TEMPLATE.replace("<curr_param>", parent.content)
        user = user.replace("<side_info>", side_info)
        return Prompt(system=system, user=user)

    def reflective_record(self, parent: Genome) -> dict:
        """One ``SideInfo`` record for the candidate under reflection, in
        ``OptimizeAnythingAdapter.make_reflective_dataset`` order: the renamed score block first, then
        every remaining diagnostic key."""
        record: dict = {}
        if parent.scores:
            record["Scores (Higher is Better)"] = dict(parent.scores)
        for key, value in parent.artifacts.items():
            if value is None or key in _NOT_FEEDBACK:
                continue
            record[key] = value
        return record