pip install gepa

oa.optimize_anything(seed_candidate=..., evaluator=...)

oa.optimize_anything(seed_candidate=..., evaluator=..., dataset=tasks)

oa.optimize_anything(seed_candidate=..., evaluator=..., dataset=train, valset=val)

seed_agent = """
import json
from typing import List, Dict, Any

def solve_puzzle(train_examples: List[Dict], test_input: Any) -> Any:
    """
    Solve an ARC-AGI puzzle given training examples.
    
    Args:
        train_examples: List of {"input": grid, "output": grid} pairs
        test_input: Input grid to solve
    
    Returns:
        Predicted output grid
    """
    # Naive baseline: return first training output
    if train_examples:
        return train_examples[0]["output"]
    return test_input
"""

import gepa.optimize_anything as oa

def evaluate_agent(candidate: str, example: Dict) -> tuple[float, dict]:
    """
    Evaluate agent on a single puzzle.
    
    Returns:
        (score, diagnostics): Score in [0, 1] and diagnostic information
    """
    try:
        # Execute the agent code
        namespace = {}
        exec(candidate, namespace)
        solve_fn = namespace["solve_puzzle"]
        
        # Run on test input
        prediction = solve_fn(
            train_examples=example["train"],
            test_input=example["test_input"]
        )
        
        # Score accuracy
        ground_truth = example["test_output"]
        correct = (prediction == ground_truth)
        score = 1.0 if correct else 0.0
        
        # Capture diagnostics as ASI (Actionable Side Information)
        diagnostics = {
            "Correct": correct,
            "Prediction": str(prediction)[:200],
            "GroundTruth": str(ground_truth)[:200],
            "TrainExamples": len(example["train"]),
        }
        
        # Log for reflection
        if not correct:
            oa.log(f"Failed on puzzle {example['id']}: predicted {prediction}, expected {ground_truth}")
        
        return score, diagnostics
        
    except Exception as e:
        # Execution errors are valuable feedback
        oa.log(f"Execution error: {type(e).__name__}: {e}")
        return 0.0, {"Error": str(e)}

# Load or create ARC-AGI-style puzzles
train_puzzles = [
    {
        "id": "train_001",
        "train": [
            {"input": [[0, 1], [1, 0]], "output": [[1, 0], [0, 1]]},
            {"input": [[1, 1], [0, 0]], "output": [[0, 0], [1, 1]]},
        ],
        "test_input": [[0, 0], [1, 1]],
        "test_output": [[1, 1], [0, 0]],
    },
    # Add 20-50 training puzzles for good results
]

val_puzzles = [
    {
        "id": "val_001",
        "train": [...],
        "test_input": ...,
        "test_output": ...,
    },
    # Add 10-20 validation puzzles
]

print(f"Training puzzles: {len(train_puzzles)}")
print(f"Validation puzzles: {len(val_puzzles)}")

from gepa.optimize_anything import (
    optimize_anything,
    GEPAConfig,
    EngineConfig,
    ReflectionConfig,
)

result = optimize_anything(
    seed_candidate=seed_agent,
    evaluator=evaluate_agent,
    dataset=train_puzzles,
    valset=val_puzzles,
    objective="""Optimize a Python agent that solves ARC-AGI puzzles.
    The agent should:
    - Analyze training examples to identify patterns
    - Apply learned rules to test inputs
    - Handle various transformation types (rotation, mirroring, color mapping, etc.)
    - Verify outputs when possible
    """,
    background="""ARC-AGI puzzles test abstract reasoning through grid transformations.
    Each puzzle has 2-4 training input/output pairs and one test input.
    Grids are 2D lists of integers (0-9 representing colors).
    Successful agents identify the transformation rule from examples.
    """,
    config=GEPAConfig(
        engine=EngineConfig(
            max_metric_calls=100,  # Budget: 100 evaluations
        ),
        reflection=ReflectionConfig(
            reflection_lm="openai/gpt-4o",  # Strong model for reflection
            reflection_minibatch_size=3,     # Focus on 3 puzzles per iteration
        ),
    ),
)

print(f"\nOptimization complete!")
print(f"Best validation score: {result.val_aggregate_scores[result.best_idx]:.2%}")
print(f"Total iterations: {result.total_metric_calls}")

print("\nEvolved Agent Architecture:")
print("=" * 60)
print(result.best_candidate[:1000])  # First 1000 chars
print("...")
print(f"\nTotal code size: {len(result.best_candidate)} characters")

# Save the optimized agent
with open("optimized_agent.py", "w") as f:
    f.write(result.best_candidate)
    
print("\nSaved to optimized_agent.py")

test_puzzles = [...]  # Load unseen test puzzles

test_scores = []
for puzzle in test_puzzles:
    score, _ = evaluate_agent(result.best_candidate, puzzle)
    test_scores.append(score)

test_accuracy = sum(test_scores) / len(test_scores)
print(f"\nTest Accuracy: {test_accuracy:.2%}")

# Compare to baseline
baseline_scores = []
for puzzle in test_puzzles:
    score, _ = evaluate_agent(seed_agent, puzzle)
    baseline_scores.append(score)

baseline_accuracy = sum(baseline_scores) / len(baseline_scores)
print(f"Baseline Accuracy: {baseline_accuracy:.2%}")
print(f"Improvement: {test_accuracy - baseline_accuracy:+.2%}")