Self-Evolving Agents - A Cookbook for Autonomous Agent Retraining

# gepa and litellm are only required for the Section 4.b (prompt optimization with GEPA)
%pip install --upgrade openai openai-agents pydantic pandas gepa litellm python-dotenv -qqq 
%load_ext dotenv
%dotenv

# Place your API key in a file called .env
# OPENAI_API_KEY=sk-...

import os
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

data_source_config = {
    "type": "custom",
    "item_schema": {
        "type": "object",
        "properties": {"section": {"type": "string"}, "summary": {"type": "string"}},
        "required": ["section", "summary"],
    },
    "include_sample_schema": False,
}

testing_criteria = [
    {
        "type": "python",
        "name": "chemical_name_grader",
        "image_tag": "2025-05-08",
        "pass_threshold": 0.8,
        "source": r"""def grade(sample: dict, item: dict) -> float:
    section = item["section"]
    summary = item["summary"]
    CHEMICALS_MASTER = ["[1-¹³C]Pyruvic acid","[1-¹³C]Pyruvate","¹²C Pyruvic acid","Sodium [1-¹³C]pyruvate","Sodium pyruvate (¹²C)","AH111501 (Trityl radical)","Tris{8-carboxyl-2,2,6,6-tetra[2-(1-methoxyethyl)]-benzo(1,2-d:4,5-d’)bis(1,3)dithiole-4-yl}methyl acid","AH111501 sodium salt","Methyl, tris[8-carboxy-2,2,6,6-tetrakis(2-methoxyethyl)benzo[1,2-d:4,5-d’]bis[1,3]dithiol-4-yl]-, trisodium salt","AH111501 trisodium salt","AH111576","2,2′,2″,2‴-(4,8-Dibromobenzo[1,2-d:4,5-d′]bis([1,3]dithiole)-2,2,6,6-tetrayl)tetraethanol","AH111586","4,8-Dibromo-2,2,6,6-tetrakis(2-methoxyethyl)benzo[1,2-d:4,5-d′]bis([1,3]dithiole)","AH111709","AH111743","AH112615","4,4-Bis-hydroxymethyl-2-methyl-oxazolidine-2-carboxylic acid","AH112623","Parapyruvate","2-Hydroxy-2-methyl-4-oxo-pentanedioic acid","AH113127","(4-Hydroxymethyl-oxazolidin-4-yl)-methanol","AH113462/E","Enol lactone","AH113462/K","Keto lactone","Acetyl bromide","Methanol","Dimethyl sulfoxide","DMSO","Tetrahydrofuran","THF","Acetonitrile","ACN","Diethyl ether","Et₂O","N,N-Dimethylacetamide","DMA","1,3-Dimethyl-2-imidazolidinone","DMI","Hydrochloric acid","HCl","Sodium hydroxide","NaOH","Disodium ethylenediaminetetraacetate","Na₂EDTA","Ethylenediaminetetraacetic acid","EDTA","Tris(hydroxymethyl)aminomethane","TRIS","Trometamol","Trifluoroacetic acid","TFA","Toluene","Heptane","Ethyl acetate","Ethanol","Water","H₂O","Sodium chloride","NaCl","Cuprous [1-¹³C]cyanide","Cu¹³CN","Gadolinium","Gd","Tin","Sn","Phosphorus","P","Carbon dioxide","CO₂","Sodium [1-13C]pyruvate","[1-13C]Pyruvic acid","1-13C pyruvate"]

    # Identify the chemicals present in the section
    present = [chem for chem in CHEMICALS_MASTER if chem in section]

    # If no chemicals present, consider it satisfied
    if not present:
        return 1.0

    correct = 0
    for chem in present:
        # Only count as correct if the exact chemical string appears in the summary
        if chem in summary:
            correct += 1

    return correct / len(present)""",
    },
    {
    "type": "python",
    "name": "word_length_deviation_grader",
    "image_tag": "2025-05-08",
    "pass_threshold": 0.85,
    "source": r"""
def grade(sample: dict, item: dict) -> float:
    summary = item["summary"]
    word_count = len(summary.split())
    
    expected_summary_length = 100
    tolerance = 0.2  # 20% band around target
    
    # relative deviation
    deviation = abs(word_count - expected_summary_length) / expected_summary_length
    
    # If within tolerance band → full score
    if deviation <= tolerance:
        return 1.0
    
    # Outside band → score decays linearly, capped at 0
    # e.g., deviation 0.3 → score 0.8, deviation 1.0+ → 0.0
    score = 1.0 - (deviation - tolerance)
    return max(0.0, score)
""",
},
    {
        "name": "cosine_similarity",
        "type": "text_similarity",
        "input": "{{ item.summary }}",
        "reference": "{{ item.section }}",
        "evaluation_metric": "cosine",
        "pass_threshold": 0.85,
    },
    {
        "name": "llm_as_judge",
        "type": "score_model",
        "model": "gpt-4.1",
        "input": [
            {
                "role": "system",
                "content": (
                    "You are an expert technical summarization evaluator. "
                    "Evaluate whether the summary captures and preserves the important technical facts and specific details from the section, allowing for occasional minor rewording or omissions of less important points, but not major technical inaccuracies or information loss.\n\n"
                    "Scoring Guidelines:\n"
                    "- Return a numerical score between 0 and 1 (with up to two decimal places).\n"
                    "- A score of 1 means the summary is almost flawless: it is comprehensive, highly faithful, and technically accurate, with virtually no important or meaningful details missing, and no significant misstatements or distortions.\n"
                    "- 0.75-0.99 indicates excellent work: all main facts are represented, but there may be trivial omissions or very minor rewording that do not materially affect understanding.\n"
                    "- 0.5-0.75 indicates good but imperfect: most technical information is retained and correctly presented, some less critical details might be missing or slightly rephrased, but overall fidelity is preserved.\n"
                    "- 0.3-0.5 means significant information is missing, or some technical inaccuracies are present, but the summary retains a reasonable portion of key facts.\n"
                    "- 0.0-0.3 means there are major omissions, misunderstandings, or a failure to capture the most important technical content.\n\n"
                    "Respond only with a single number between 0 and 1 indicating summary quality by these criteria."
                ),
            },
            {
                "role": "user",
                "content": (
                    "Section:\n{{item.section}}\n"
                    "Summary:\n{{sample.output_text}}"
                ),
            },
        ],
        "range": [0, 1],
        "pass_threshold": 0.85,
    },
]

eval = client.evals.create(
    name="self_evolving_eval",
    data_source_config=data_source_config,
    testing_criteria=testing_criteria,
)
print(f"Created Eval: {eval.id}")

import time
import json

def run_eval(eval_id: str, section: str, summary: str):
  """Creates a run of the eval with the input section and output summary."""
  return client.evals.runs.create(
    eval_id=eval_id,
    name="self-evolving-eval",
    data_source={
      "type": "jsonl",
      "source": {
        "type": "file_content",
        "content": [
          {
            "item": {
              "section": section,
              "summary": summary,
            }
          }
        ],
      },
    },
  )


def poll_eval_run(eval_id: str, run_id: str, max_polls = 10):
    """
    Polls the evaluation run until completion or timeout.

    This function exists to handle asynchronous behavior in the eval service by
    periodically checking run status. It balances responsiveness and resource use by
    polling at fixed intervals rather than blocking indefinitely. The retry limit
    prevents runaway loops in cases where the service never returns a completed status.
    """
    run = None
    for attempt in range(1, max_polls + 1):
        run = client.evals.runs.retrieve(eval_id=eval_id, run_id=run_id)
        if run.status == "completed":
            break
        if attempt == max_polls:
            print("Exceeded retries, aborting")
            break

        time.sleep(5)

    run_output_items = client.evals.runs.output_items.list(
        eval_id=eval_id, run_id=run_id
    )
    return run_output_items


def parse_eval_run_output(items):
    """Extract all grader scores and any available conclusion outputs."""
    all_results = []

    for item in items.data:
        for result in item.results:
            grader_name_full = result.name
            score = result.score
            passed = result.passed
            reasoning = None
            try:
                sample = result.sample
                if sample:
                    content = result.sample["output"][0]["content"]
                    content_json = json.loads(content)
                    steps = content_json["steps"]
                    reasoning = " ".join([step["conclusion"] for step in steps])
            except Exception:
                pass

            all_results.append(
                {
                    "grader_name": grader_name_full,
                    "score": score,
                    "passed": passed,
                    "reasoning": reasoning,
                }
            )

    return all_results

EVAL_ID = eval.id #Created eval ID from above cell
SECTION = "3.2.S.1 General Information ([1-13C]pyruvic acid) The active ingredient in Hyperpolarized Pyruvate (13C) Injection is hyperpolarized [1-13C]pyruvate. The drug substance is defined as [13C]pyruvic acid, which is neutralized to [1-13C]pyruvate during the compounding process. In several pre-clinical and clinical studies and during evaluation of stability, pyruvic acid has been used instead of [1-13C]pyruvic acid (see Sections 3.2.P.2.2.1 Formulation Development for Hyperpolarized Pyruvate (13C) Injection and Section 8.1 Introduction for Item 8 Pharmacology and Toxicology Info). In the Section 3.2.S Drug Substance, data are presented for both pyruvic acid and for [1-13C]pyruvic acid. For simplicity, the terminology used in headings and captions is [1-13C]pyruvic acid. Batches containing pyruvic acid are specified by footnotes. 3.2.S.1.1 Nomenclature ([1-13C]pyruvic acid) The drug substance used for compounding of Hyperpolarized Pyruvate (13C) Injection is [1-13C]pyruvic acid. Company code: W6578 Chemical name: [1-13C]pyruvic acid CAS registry number: 127-17-3 3.2.S.1.2 Structure ([1-13C]pyruvic acid) Figure 1 Structure of [1-13C]pyruvic acid Molecular formula: C H O 3 4 3 Molecular weight: 89.06 3.2.S.1.3 General Properties ([1-13C]pyruvic acid) Appearance: Colorless to yellow, clear, viscous liquid pKa:Ka:aranWater solubility: Complete The structure of [1-13C]pyruvic acid has been confirmed by spectroscopic analysis (see Section 3.2.S.3.1 Elucidation of Structure and other Characteristics)."
SUMMARY = "The active ingredient in Hyperpolarized Pyruvate (13C) Injection is hyperpolarized [1-13C]pyruvate, derived from [1-13C]pyruvic acid (neutralized during compounding). Both pyruvic acid and [1-13C]pyruvic acid were used in studies and stability evaluations, but the documentation refers to [1-13C]pyruvic acid unless otherwise noted. The drug substance ([1-13C]pyruvic acid, CAS 127-17-3) is a colorless to yellow, clear, viscous liquid with a molecular formula C3H4O3 and molecular weight 89.06. Its structure has been confirmed by spectroscopic analysis, and it is completely soluble in water."

eval_run = run_eval(EVAL_ID, section=SECTION, summary=SUMMARY)
run_output = poll_eval_run(eval_id=EVAL_ID, run_id=eval_run.id)

grader_scores = parse_eval_run_output(run_output)
print(grader_scores)

from datetime import datetime
from typing import Any, Optional

from pydantic import BaseModel, Field, ConfigDict, field_validator

class PromptVersionEntry(BaseModel):
    """Data model for a prompt and associated data for observability"""
    version: int = Field(
        ..., ge=0, description="Version number of the prompt (increments)"
    )
    model: str = Field(
        "gpt-5",
        min_length=1,
        description="The model version to use for this version of the prompt, defaults to gpt-5",
    )
    prompt: str = Field(
        ..., min_length=1, description="The prompt text for this version"
    )
    timestamp: datetime = Field(
        default_factory=datetime.utcnow,
        description="UTC timestamp when this version was created",
    )
    eval_id: Optional[str] = Field(
        None, description="ID of the evaluation associated with this prompt version"
    )
    run_id: Optional[str] = Field(
        None, description="ID of the run associated with this prompt version"
    )
    metadata: Optional[dict[str, Any]] = Field(
        None, description="Free-form metadata dict (e.g., section, summary)"
    )

    model_config = ConfigDict(
        str_strip_whitespace=True, validate_assignment=True, extra="forbid"
    )

    @field_validator("prompt")
    @classmethod
    def prompt_not_blank(cls, v: str) -> str:
        if not v.strip():
            raise ValueError("prompt must not be blank or only whitespace")
        return v


class VersionedPrompt:
    """Manages a collection of prompt versions and provides controlled updates and rollbacks."""
    def __init__(
        self,
        initial_prompt: str,
        model: Optional[str] = "gpt-5",
        eval_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ):
        if not initial_prompt or not initial_prompt.strip():
            raise ValueError("initial_prompt must be non-empty")
        self._versions: list[PromptVersionEntry] = []
        first_entry = PromptVersionEntry(
            version=0,
            prompt=initial_prompt,
            model=model,
            eval_id=eval_id,
            run_id=run_id,
            metadata=metadata,
        )
        self._versions.append(first_entry)

    def update(
        self,
        new_prompt: str,
        model: Optional[str] = "gpt-5",
        eval_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> PromptVersionEntry:
        if not new_prompt or not new_prompt.strip():
            raise ValueError("new_prompt must be non-empty")

        version = self.current().version + 1
        entry = PromptVersionEntry(
            version=version,
            prompt=new_prompt,
            model=model,
            eval_id=eval_id,
            run_id=run_id,
            metadata=metadata,
        )
        self._versions.append(entry)
        return entry

    def current(self) -> PromptVersionEntry:
        return self._versions[-1]

    def revert_to_version(self, version: int) -> PromptVersionEntry:
        idx = None
        for i, entry in enumerate(self._versions):
            if entry.version == version:
                idx = i
                break

        if idx is None:
            raise ValueError(f"No version found with version={version}")

        self._versions = self._versions[: idx + 1]
        return self._versions[-1]

from agents import Agent

METAPROMPT_TEMPLATE = """
# Context:
## Original prompt:
{original_prompt}

## Section:
{section}

## Summary:
{summary}

## Reason to improve the prompt:
{reasoning}

# Task:
Write a new summarization prompt that is significantly improved and more specific than the original.  
The new prompt should instruct the model to produce concise yet comprehensive technical summaries that precisely preserve all explicit information from the source text. It should emphasize the inclusion of all named entities, quantities, compounds, and technical terminology without paraphrasing or omission. The resulting prompt should read like a clear, directive system message for a technical summarization assistant—structured, unambiguous, and generalizable across scientific or regulatory document sections.
"""

metaprompt_agent = Agent(
    name="MetapromptAgent", instructions="You are a prompt optimizer."
)

summarization_prompt = VersionedPrompt(
    initial_prompt="""You are a summarization assistant.
Given a section of text, produce a summary."""
)

def make_summarization_agent(prompt_entry: PromptVersionEntry) -> Agent:
    return Agent(
        name="SummarizationAgent",
        instructions=prompt_entry.prompt,
        model=prompt_entry.model,
    )

summarization_agent = make_summarization_agent(summarization_prompt.current())

# Cache eval results by section + summary so repeated attempts do not trigger redundant grader runs.
eval_cache: dict[tuple[str, str], list[dict[str, Any]]] = {}

# Track the highest-scoring candidate that also passes the lenient score threshold.
best_candidate: dict[str, Any] = {
    "score": float("-inf"),
    "prompt": summarization_prompt.current().prompt,
    "model": summarization_prompt.current().model,
    "summary": None,
    "metadata": None,
    "version": summarization_prompt.current().version,
    "passed_lenient": False,
    "total_score": float("-inf"),
}

# Aggregate per-version performance so we can pick the strongest total scorer at the end.
aggregate_prompt_stats: dict[int, dict[str, Any]] = {}

import asyncio
from typing import Any, Optional
from agents import Runner

LENIENT_PASS_RATIO = 0.75 # 75% of graders must pass (binary) 
LENIENT_AVERAGE_THRESHOLD = 0.85 # 85% average score across graders 

def reset_best_candidate() -> None:
    """Reset the best candidate tracker for a new optimization run."""
    global best_candidate

    current = summarization_prompt.current()
    best_candidate = {
        "score": float("-inf"),
        "prompt": current.prompt,
        "model": current.model,
        "summary": None,
        "metadata": None,
        "version": current.version,
    }

def reset_best_trackers() -> None:
    """Reset both the best-candidate tracker and aggregate stats."""
    reset_best_candidate()
    aggregate_prompt_stats.clear()


def update_best_candidate(
    *,
    average_score: Optional[float] = None,
    prompt_text: str,
    model_name: str,
    summary_text: str = None,
    metadata: dict[str, Any] = None,
    lenient_passed: bool = False,
    prompt_version: int = None,
    total_score: Optional[float] = None,
    score: Optional[float] = None,
) -> None:
    """Persist the best lenient-passing candidate."""
    global best_candidate

    if prompt_version is None:
        prompt_version = summarization_prompt.current().version

    if average_score is None:
        average_score = score

    if average_score is None:
        return

    if lenient_passed:
        best_candidate.update(
            {
                "score": average_score,
                "prompt": prompt_text,
                "model": model_name,
                "summary": summary_text,
                "metadata": metadata,
                "version": prompt_version,
                "total_score": total_score if total_score is not None else average_score,
            }
        )


def apply_best_candidate_if_needed() -> Agent:
    """Ensure summarization_prompt reflects the best prompt candidate."""
    if best_candidate["score"] > float("-inf"):
        current = summarization_prompt.current()
        target = best_candidate
        # Only update if different
        if (
            current.prompt != target["prompt"]
            or current.model != target["model"]
            or current.version != target.get("version")
        ):
            summarization_prompt.update(
                new_prompt=target["prompt"],
                model=target["model"],
                metadata=target.get("metadata"),
            )
            target["version"] = summarization_prompt.current().version
        return make_summarization_agent(summarization_prompt.current())

    return make_summarization_agent(summarization_prompt.current())


def record_aggregate_prompt_score(
    *,
    prompt_version: int,
    prompt_text: str,
    model_name: str,
    average_score: float,
    total_score: Optional[float] = None,
) -> None:
    """Accumulate per-version grader scores for aggregate selection."""
    stats = aggregate_prompt_stats.setdefault(
        prompt_version,
        {
            "version": prompt_version,
            "prompt": prompt_text,
            "model": model_name,
            "total_score": 0.0,
            "total_average": 0.0,
            "count": 0,
        },
    )
    stats["total_score"] += total_score if total_score is not None else average_score
    stats["total_average"] += average_score
    stats["count"] += 1
    stats["prompt"] = prompt_text
    stats["model"] = model_name


def select_best_aggregate_prompt() -> Optional[dict[str, Any]]:
    """Return the prompt version with the highest cumulative score."""
    if not aggregate_prompt_stats:
        return None
    return max(
        aggregate_prompt_stats.values(),
        key=lambda entry: (
            entry.get("total_score", float("-inf")),
            entry.get("version", -1),
        ),
    )


async def get_eval_grader_score(eval_id: str, section: str, summary: str):
    """Retrieve grader scores for a section-summary pair with caching."""
    cache_key = (section, summary)
    if cache_key in eval_cache:
        return eval_cache[cache_key]

    eval_run = run_eval(eval_id=eval_id, section=section, summary=summary)
    run_output = poll_eval_run(eval_id=eval_id, run_id=eval_run.id)
    results = parse_eval_run_output(run_output)
    eval_cache[cache_key] = results
    return results


def calculate_grader_score(grader_scores):
    """Simple average score of all graders from the eval."""
    if not grader_scores:
        return 0.0

    score_sum = 0.0
    for entry in grader_scores:
        score_sum += entry.get("score", 0.0)

    return score_sum / len(grader_scores)



def calculate_total_grader_score(grader_scores):
    """Sum of all grader scores for aggregate tracking."""
    if not grader_scores:
        return 0.0

    return sum(entry.get("score", 0.0) for entry in grader_scores)


DEFAULT_PASSING_FEEDBACK = (
    "All graders passed; tighten factual coverage, chemical completeness, and conciseness."
)


def is_lenient_pass(grader_scores, average_score: float) -> bool:
    if not grader_scores:
        return False

    passed_count = sum(1 for entry in grader_scores if entry.get("passed"))
    total_graders = len(grader_scores)

    if total_graders and (passed_count / total_graders) >= LENIENT_PASS_RATIO:
        return True
    return average_score >= LENIENT_AVERAGE_THRESHOLD


def collect_grader_feedback(grader_scores):
    """Consolidate grader reasoning into actionable feedback for the metaprompt agent."""
    feedback_lines = []

    for entry in grader_scores:
        grader = entry.get("grader_name", "")
        passed = entry.get("passed", False)
        reasoning = entry.get("reasoning")

        if not passed:
            if grader.startswith("chemical_name_grader"):
                feedback_lines.append(
                    "Not all chemical names in the input section were included in the summary."
                )
            elif grader.startswith("word_length_deviation_grader"):
                feedback_lines.append(
                    "The summary length deviates too much from the expected length."
                )
            elif grader.startswith("cosine_similarity"):
                feedback_lines.append(
                    "The summary is not sufficiently similar to the source section (cosine similarity too low)."
                )
            elif grader.startswith("llm_as_judge") and reasoning:
                feedback_lines.append(reasoning)

    if not feedback_lines:
        feedback_lines.append(DEFAULT_PASSING_FEEDBACK)

    return "".join(feedback_lines)

import pandas as pd

from agents import Agent, trace

EVAL_ID = eval.id #Created eval ID from above cell
MAX_OPTIMIZATION_RETRIES = 3

async def self_evolving_loop(summarization_agent: Agent) -> Agent:
    print(f"Starting self-evolving loop | Initial prompt v{summarization_prompt.current().version}")
    print(f"Prompt:{summarization_prompt.current().prompt}")
    print("-" * 80)

    reset_best_trackers()
    df = pd.read_csv("data/dataset.csv")

    with trace("Self-evolving Optimization Workflow"):
        for _, row in df.head().iterrows():
            content = row.get("content")
            if pd.isna(content) or (isinstance(content, str) and not content.strip()):
                continue

            section_number = str(row["section_number"])
            section = str(content)
            current_version = summarization_prompt.current().version

            print(f"[Section {section_number}] Using prompt v{current_version}")

            optimization_success = False

            for attempt in range(1, MAX_OPTIMIZATION_RETRIES + 1):
                print(f"  Attempt {attempt}: evaluating summary...")

                summary_result = await Runner.run(summarization_agent, section)
                summary = summary_result.final_output

                grader_scores = await get_eval_grader_score(eval_id=EVAL_ID, summary=summary, section=section)
                average_score = calculate_grader_score(grader_scores)
                total_score = calculate_total_grader_score(grader_scores)
                lenient_passed = is_lenient_pass(grader_scores, average_score)
                print(
                    f"	Scores — avg={average_score:.3f}, total={total_score:.3f}, lenient_passed={lenient_passed}"
                )

                record_aggregate_prompt_score(
                    prompt_version=summarization_prompt.current().version,
                    prompt_text=summarization_prompt.current().prompt,
                    model_name=summarization_prompt.current().model,
                    average_score=average_score,
                    total_score=total_score,
                )

                update_best_candidate(
                    average_score=average_score,
                    prompt_text=summarization_prompt.current().prompt,
                    model_name=summarization_prompt.current().model,
                    summary_text=summary,
                    metadata={
                        "section": section_number,
                        "average_score": average_score,
                        "grader_results": grader_scores,
                        "prompt_version": summarization_prompt.current().version,
                    },
                    lenient_passed=lenient_passed,
                    prompt_version=summarization_prompt.current().version,
                )

                if lenient_passed:
                    optimization_success = True
                    print(f"	Passed with prompt v{summarization_prompt.current().version}")
                    break

                print("	Failed eval. Improving prompt...")
                eval_feedback = collect_grader_feedback(grader_scores)

                metaprompt_result = await Runner.run(
                    metaprompt_agent,
                    input=METAPROMPT_TEMPLATE.format(
                        original_prompt=summarization_prompt.current().prompt,
                        section=section,
                        summary=summary,
                        reasoning=eval_feedback,
                    ),
                )
                improved_prompt = metaprompt_result.final_output
                summarization_prompt.update(
                    new_prompt=improved_prompt,
                    metadata={"section": section, "summary": summary},
                )
                summarization_agent = make_summarization_agent(summarization_prompt.current())

                print(f"	Prompt improved → v{summarization_prompt.current().version}")

            if not optimization_success:
                print(
                    "	All attempts failed; keeping latest prompt version "
                    f"v{summarization_prompt.current().version} for the next section."
                )

    summarization_agent = apply_best_candidate_if_needed()

    print("" + "-" * 80)
    print("Completed optimization loop.")
    print(f"Final prompt version: v{summarization_prompt.current().version}")
    if best_candidate["score"] > float("-inf"):
        print(
            f"Best lenient prompt: v{best_candidate.get('version')} (avg={best_candidate['score']:.3f})"
        )

    aggregate_best = select_best_aggregate_prompt()
    if aggregate_best:
        per_section = (
            aggregate_best.get("total_average", 0.0) / aggregate_best.get("count", 1)
            if aggregate_best.get("count")
            else 0.0
        )
        print(
            f"Aggregate best prompt: v{aggregate_best.get('version')} "
            f"(total={aggregate_best.get('total_score', 0.0):.3f}, avg/section={per_section:.3f}, model={aggregate_best.get('model', 'unknown')})"
        )

    print(f"Final prompt:{summarization_prompt.current().prompt}")
    return summarization_agent

summarization_agent = await self_evolving_loop(summarization_agent)

# this cell is pseudo-code and not meant to be run as-is

import time

def continuous_monitoring(interval_hours=24):
    """Periodically check for new data and trigger the evaluation loop."""
    while True:
        print("Checking for new data...")
        if new_data_detected():
            print("New data found — running evaluation and optimization loop.")
            self_evolving_loop()
        else:
            print("No new data. Sleeping until next cycle.")
        time.sleep(interval_hours * 3600)

continuous_monitoring(interval_hours=24)

from agents import Agent, Runner

async def eval_agent_candidate(agent: Agent, section: str, prompt_text: str, model_name: str):
    summary_result = await Runner.run(agent, section)
    summary = summary_result.final_output

    scores = await get_eval_grader_score(
        eval_id=EVAL_ID, summary=summary, section=section
    )
    average = calculate_grader_score(scores)
    lenient_passed = is_lenient_pass(scores, average)
    passed = all(entry.get("passed") is True for entry in scores)

    update_best_candidate(
        average_score=average,
        prompt_text=prompt_text,
        model_name=model_name,
        summary_text=summary,
        metadata={
            "section": section,
            "average_score": average,
            "grader_results": scores,
        },
        lenient_passed=lenient_passed,
    )

    return {"summary": summary, "scores": scores, "average": average, "passed": passed}

async def compare_model_candidates(
    summarization_prompt,
    eval_feedback: str,
    section: str,
    summary: str,
    model_candidates=None,
):
    """Improve the prompt, evaluate it across candidate models, and adopt the top performer."""
    if model_candidates is None:
        model_candidates = ["gpt-5", "gpt-5-mini"]

    metaprompt_result = await Runner.run(
        metaprompt_agent,
        input=METAPROMPT_TEMPLATE.format(
            original_prompt=summarization_prompt.current().prompt,
            section=section,
            summary=summary,
            reasoning=eval_feedback,
        ),
    )
    improved_prompt = metaprompt_result.final_output

    async def evaluate_model(model_name: str):
        candidate_agent = Agent(
            name=f"SummarizationAgent:{model_name}",
            instructions=improved_prompt,
            model=model_name,
        )
        result = await eval_agent_candidate(candidate_agent, section, improved_prompt, model_name)
        return model_name, candidate_agent, result

    best = {
        "average": float("-inf"),
        "passed": False,
        "agent": None,
        "model": None,
        "summary": None,
    }

    tasks = [asyncio.create_task(evaluate_model(model_name)) for model_name in model_candidates]
    for task in asyncio.as_completed(tasks):
        model_name, candidate_agent, result = await task
        print(
            f"Candidate average — {model_name}: {result['average']:.4f} "
            f"(passed={result.get('passed', False)})"
        )
        if result["average"] > best["average"]:
            best.update(
                {
                    "average": result["average"],
                    "model": model_name,
                    "summary": result.get("summary"),
                    "agent": candidate_agent,
                    "passed": result.get("passed", False),
                }
            )

    for task in tasks:
        if not task.done():
            task.cancel()

    if best["passed"] and best["model"]:
        summarization_prompt.update(
            new_prompt=improved_prompt,
            model=best["model"],
            metadata={"section": section, "summary": best["summary"]},
        )
        print(f"Updated summarization_prompt with passing model: {best['model']}")
        return make_summarization_agent(summarization_prompt.current())

    print(
        f"No passing models. Best candidate (model={best['model']}, "
        f"avg={best['average']:.4f}) did not pass. Prompt not updated."
    )
    return None

async def self_evolving_loop_with_model_comparison(summarization_agent: Agent) -> Agent:
    print(
        f"Starting self-evolving loop | Initial prompt v{summarization_prompt.current().version}"
    )
    print(f"Prompt: {summarization_prompt.current().prompt}")
    print(f"Model: {summarization_prompt.current().model}")
    print("-" * 80)

    reset_best_trackers()
    df = pd.read_csv("data/dataset.csv")

    with trace("Self-evolving Optimization Workflow: model comparison"):
        for _, row in df.head(5).iterrows():
            content = row.get("content")
            if pd.isna(content) or (isinstance(content, str) and not content.strip()):
                continue

            section_number = str(row["section_number"])
            section = str(content)
            current_version = summarization_prompt.current().version

            print(f"[Section {section_number}] Using prompt v{current_version}")

            summary_passed = False

            for attempt in range(1, MAX_OPTIMIZATION_RETRIES + 1):
                print(f"\tAttempt {attempt}: evaluating summary...")

                summary_result = await Runner.run(summarization_agent, section)
                summary = summary_result.final_output

                grader_scores = await get_eval_grader_score(
                    eval_id=EVAL_ID, summary=summary, section=section
                )
                average_score = calculate_grader_score(grader_scores)
                total_score = calculate_total_grader_score(grader_scores)
                lenient_passed = is_lenient_pass(grader_scores, average_score)
                print(
                    f"\tScores — avg={average_score:.3f}, total={total_score:.3f}, lenient_passed={lenient_passed}"
                )

                record_aggregate_prompt_score(
                    prompt_version=summarization_prompt.current().version,
                    prompt_text=summarization_prompt.current().prompt,
                    model_name=summarization_prompt.current().model,
                    average_score=average_score,
                    total_score=total_score,
                )

                update_best_candidate(
                    average_score=average_score,
                    total_score=total_score,
                    prompt_text=summarization_prompt.current().prompt,
                    model_name=summarization_prompt.current().model,
                    summary_text=summary,
                    metadata={
                        "section": section_number,
                        "average_score": average_score,
                        "grader_results": grader_scores,
                        "prompt_version": summarization_prompt.current().version,
                    },
                    lenient_passed=lenient_passed,
                    prompt_version=summarization_prompt.current().version,
                )

                if lenient_passed:
                    summary_passed = True
                    print(
                        f"\tPassed with prompt v{summarization_prompt.current().version} (model={summarization_prompt.current().model})"
                    )
                    break

                print("\tFailed eval. Improving prompt...")
                eval_feedback = collect_grader_feedback(grader_scores)

                new_agent = await compare_model_candidates(
                    summarization_prompt=summarization_prompt,
                    eval_feedback=eval_feedback,
                    section=section,
                    summary=summary,
                    # model_candidates could be given as an argument if you want to expand options.
                )

                if new_agent is None:
                    print(
                        "\tNo passing model found. Optimization failed for this section."
                    )
                    summary_passed = False
                else:
                    summarization_agent = new_agent
                    summary_passed = True
                    print(
                        f"\tPrompt improved → v{summarization_prompt.current().version} "
                        f"(model={summarization_prompt.current().model})"
                    )
                    break

            if not summary_passed:
                print(
                    "\tAll attempts failed; keeping latest prompt version "
                    f"v{summarization_prompt.current().version} (model={summarization_prompt.current().model}) for the next section."
                )

    summarization_agent = apply_best_candidate_if_needed()

    print("" + "-" * 80)
    print("Completed optimization loop.")
    print(f"Final prompt version: v{summarization_prompt.current().version}")
    print(f"Final model: {summarization_prompt.current().model}")
    aggregate_best = select_best_aggregate_prompt()
    if best_candidate["score"] > float("-inf"):
        print(
            f"Best lenient prompt: v{best_candidate.get('version')} (avg={best_candidate['score']:.3f}, model={best_candidate.get('model', 'unknown')})"
        )
    if aggregate_best:
        per_section = (
            aggregate_best.get("total_average", 0.0) / aggregate_best.get("count", 1)
            if aggregate_best.get("count")
            else 0.0
        )
        print(
            f"Aggregate best prompt: v{aggregate_best.get('version')} "
            f"(total={aggregate_best.get('total_score', 0.0):.3f}, avg/section={per_section:.3f}, model={aggregate_best.get('model', 'unknown')})"
        )
    print(f"Final prompt: {summarization_prompt.current().prompt}")
    print(f"Final model: {summarization_prompt.current().model}")
    return summarization_agent

summarization_agent = await self_evolving_loop_with_model_comparison(summarization_agent)

import pandas as pd
import gepa
from gepa import EvaluationBatch

# Extract sections from dataset
def read_csv_content(file_path: str) -> list[dict]:
    """Read csv and return section to summarize."""
    df = pd.read_csv(file_path)
    return [{'content': content} for content in df['content'].tolist()]

# Split dataset into training and validation sets
trainset = read_csv_content("data/dataset.csv")
val_cut = max(1, int(0.1 * len(trainset)))
valset = trainset[:val_cut] if len(trainset) > 1 else trainset

class EvalsBackedSummarizationAdapter:
    """
    Minimal adapter for GEPA:
      - evaluate(...) -> EvaluationBatch (scores + outputs + feedback-rich trajectories)
      - get_components_to_update(...) returns the prompt to update
      - make_reflective_dataset(...) packages examples for reflection
    """
    propose_new_texts = None  # use GEPA's default reflection flow

    def __init__(self, client, eval_id: str, gen_model: str = "gpt-5", user_prefix: str | None = None):
        self.client = client
        self.eval_id = eval_id
        self.gen_model = gen_model
        self.user_prefix = user_prefix or "Summarize:\n\n"

    # Same summarization agent as in the previous section
    def _summarize(self, system_prompt: str, section: str) -> str:
        resp = self.client.chat.completions.create(
            model=self.gen_model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"{self.user_prefix}{section}"},
            ],
        )
        return resp.choices[0].message.content.strip()

    # Required by GEPA: run eval minibatch
    def evaluate(self, inputs: list[dict], candidate: dict, capture_traces: bool = True) -> EvaluationBatch:
        system_prompt = candidate["system_prompt"]

        scores: list[float] = []
        outputs: list[str] = []
        trajectories: list[dict] = []

        for item in inputs:
            section = item["content"]

            # 1) Generate with the candidate prompt
            summary = self._summarize(system_prompt, section)
            outputs.append(summary)

            # 2) Grade using previous evals pipeline
            run = run_eval(eval_id=self.eval_id, section=section, summary=summary)
            out_items = poll_eval_run(eval_id=self.eval_id, run_id=run.id)
            grader_scores = parse_eval_run_output(out_items)

            # 3) Score + actionable feedback
            scalar = calculate_grader_score(grader_scores)
            feedback = collect_grader_feedback(grader_scores) or "All graders passed; keep precision and coverage."

            scores.append(float(scalar))
            trajectories.append(
                {
                    "inputs": {"section": section},
                    "generated_output": summary,
                    "metrics": {
                        "combined": float(scalar),
                        "by_grader": grader_scores,  # keeping for analysis if needed
                    },
                    "feedback": feedback,
                }
            )

        return EvaluationBatch(scores=scores, outputs=outputs, trajectories=trajectories)

    # Required by GEPA: text field to evolve
    def get_components_to_update(self, candidate: dict) -> list[str]:
        return ["system_prompt"]

    # Required by GEPA: build the reflective dataset the reflection LM will read
    def make_reflective_dataset(self, candidate: dict, eval_batch: EvaluationBatch, components_to_update: list[str]) -> dict:
        examples = []
        for traj in (eval_batch.trajectories or []):
            examples.append(
                {
                    "Inputs": {"section": traj["inputs"]["section"]},
                    "Generated Outputs": traj["generated_output"],
                    "Feedback": traj["feedback"],
                }
            )
        return {"system_prompt": examples}

seed_candidate = {"system_prompt": "You are a summarization assistant. Given a section of text, produce a summary."}

adapter = EvalsBackedSummarizationAdapter(
    client=client,
    eval_id=EVAL_ID,
    gen_model=summarization_prompt.current().model,  
)

# Keeping max_metric_calls small for the cookbook. 
# In practice, use a larger value to allow more optimization iterations.
result = gepa.optimize(
    seed_candidate=seed_candidate,
    trainset=trainset,
    valset=valset,
    adapter=adapter,
    reflection_lm="gpt-5",
    max_metric_calls=10,
    track_best_outputs=True,
    display_progress_bar=True
)

best_prompt = result.best_candidate["system_prompt"]
print("\n=== Best evolved instruction ===\n")
print(best_prompt)