Agent Architecture Discovery

def solve_arc_task(task):
    """Naive agent: direct prompt to LLM."""
    prompt = f"""
    Given these input-output examples:
    {format_examples(task['train'])}
    
    Apply the pattern to: {task['test_input']}
    """
    
    response = llm.generate(prompt)
    return parse_output(response)

def solve_arc_task(task):
    """Sophisticated multi-stage agent with rule induction and verification."""
    
    # Stage 1: Rule Induction
    rules = induce_transformation_rules(task['train'])
    
    # Stage 2: Code Generation (Primary Attempt)
    code_solution = generate_transformation_code(
        examples=task['train'],
        rules=rules,
        test_input=task['test_input']
    )
    
    # Stage 3: Iterative Validation
    for iteration in range(3):
        validation_result = validate_on_training_examples(
            code=code_solution,
            examples=task['train']
        )
        
        if validation_result.all_correct:
            break
            
        # Refine code based on failures
        code_solution = refine_code(
            code=code_solution,
            failures=validation_result.failures,
            rules=rules
        )
    
    # Stage 4: Execute on Test Input
    try:
        prediction = execute_code(code_solution, task['test_input'])
    except Exception as e:
        # Fallback: Direct LLM prediction
        prediction = llm_direct_prediction(
            task=task,
            error_context=str(e)
        )
    
    return prediction


def induce_transformation_rules(examples):
    """Extract patterns from training examples."""
    prompt = f"""
    Analyze these input-output grid transformations:
    {format_examples(examples)}
    
    Extract the transformation rules:
    - What patterns appear?
    - What geometric operations are applied?
    - Are there color transformations?
    - What is the output size relative to input?
    
    List 3-5 specific, testable rules.
    """
    return llm.generate(prompt)


def generate_transformation_code(examples, rules, test_input):
    """Generate Python code implementing the transformation."""
    prompt = f"""
    Given transformation rules:
    {rules}
    
    And these examples:
    {format_examples(examples)}
    
    Write Python code to transform:
    {test_input}
    
    Requirements:
    - Function signature: transform(grid) -> grid
    - Use numpy for grid operations
    - Handle edge cases gracefully
    - Include comments explaining each step
    """
    return llm.generate(prompt)


def validate_on_training_examples(code, examples):
    """Execute code on training examples and check correctness."""
    failures = []
    
    for i, example in enumerate(examples):
        try:
            prediction = execute_code(code, example['input'])
            
            if not grids_match(prediction, example['output']):
                failures.append({
                    'example_id': i,
                    'expected': example['output'],
                    'got': prediction,
                    'diff': compute_diff(prediction, example['output'])
                })
        except Exception as e:
            failures.append({
                'example_id': i,
                'error': str(e),
                'traceback': traceback.format_exc()
            })
    
    return ValidationResult(
        all_correct=(len(failures) == 0),
        failures=failures
    )


def refine_code(code, failures, rules):
    """Improve code based on validation failures."""
    prompt = f"""
    Your code:
    {code}
    
    Failed on these cases:
    {format_failures(failures)}
    
    Original rules:
    {rules}
    
    Fix the code to handle these failures while maintaining correctness on 
    passing examples. Focus on:
    - Logic errors in transformation
    - Edge cases (boundary conditions, empty cells)
    - Off-by-one errors in grid indexing
    """
    return llm.generate(prompt)


def llm_direct_prediction(task, error_context=None):
    """Fallback: Direct LLM prediction without code execution."""
    prompt = f"""
    Task examples:
    {format_examples(task['train'])}
    
    {'Code generation failed: ' + error_context if error_context else ''}
    
    Predict the output grid for:
    {task['test_input']}
    
    Output the grid directly as a 2D array.
    """
    return llm.generate(prompt)

import gepa.optimize_anything as oa

def evaluate_arc_agent(candidate: dict, example: dict) -> tuple[float, dict]:
    """
    Run agent code on an ARC-AGI task.
    
    Args:
        candidate: Dict with 'agent_code' key containing full agent implementation
        example: ARC-AGI task with 'train' examples and 'test' input/output
    
    Returns:
        (score, side_info) where score is 1.0 for correct, 0.0 for incorrect
    """
    # Execute agent code
    exec_globals = {}
    exec(candidate['agent_code'], exec_globals)
    
    solve_fn = exec_globals['solve_arc_task']
    
    try:
        prediction = solve_fn(example)
        correct = grids_match(prediction, example['test_output'])
        
        return float(correct), {
            "Correct": correct,
            "Prediction": format_grid(prediction),
            "Expected": format_grid(example['test_output']),
            "VisualDiff": render_diff(prediction, example['test_output']),
        }
        
    except Exception as e:
        return 0.0, {
            "Error": str(e),
            "Traceback": traceback.format_exc(),
            "Expected": format_grid(example['test_output']),
        }

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

# Load ARC-AGI dataset
train_tasks, val_tasks, test_tasks = load_arc_agi_dataset()

# Seed: naive agent
seed_agent = {
    "agent_code": open("naive_arc_agent.py").read()
}

result = optimize_anything(
    seed_candidate=seed_agent,
    evaluator=evaluate_arc_agent,
    dataset=train_tasks,
    valset=val_tasks,
    objective="""
    Optimize the ARC-AGI agent to maximize accuracy on abstract reasoning tasks.
    The agent should induce rules from examples and generate/validate code.
    """,
    background="""
    ARC-AGI tasks involve grid transformations with abstract patterns.
    Strong agents typically:
    - Extract transformation rules explicitly
    - Generate executable code rather than direct predictions
    - Validate code on training examples before applying to test
    - Have fallback strategies when code fails
    """,
    config=GEPAConfig(
        engine=EngineConfig(max_metric_calls=100),
        reflection=ReflectionConfig(
            reflection_lm="vertex_ai/gemini-3-flash-preview",
            reflection_minibatch_size=3,  # Reflect on 3 tasks at a time
        ),
    ),
)

print("Final agent code:")
print(result.best_candidate['agent_code'])

print(f"\nTest accuracy: {evaluate_on_testset(result.best_candidate, test_tasks)}")

seed_agent = {
    "architecture": """
    # Simple RAG
    def answer_query(query, documents):
        relevant_docs = retrieve(query, documents, top_k=5)
        context = '\n'.join(relevant_docs)
        return llm.generate(f"Context: {context}\nQuestion: {query}")
    """
}

# After optimization → Multi-agent system:
optimized_agent = {
    "architecture": """
    # Specialized disease-expert sub-agents
    def answer_query(query, documents):
        # Route to specialist
        disease = classify_query(query)  # diabetes vs COPD
        
        if disease == 'diabetes':
            expert = diabetes_expert
        elif disease == 'copd':
            expert = copd_expert
        else:
            expert = general_expert
        
        # Expert does specialized retrieval + reasoning
        relevant_docs = expert.retrieve(query, documents)
        answer = expert.reason(query, relevant_docs)
        
        # Lead agent synthesizes and validates
        return lead_agent.synthesize(query, answer, disease_context=disease)
    """
}

from gepa.adapters.terminal_bench_adapter import TerminalBenchAdapter

adapter = TerminalBenchAdapter(
    task_lm="openai/gpt-4.5",
    terminal_bench_path="/path/to/terminalbench",
)

result = gepa.optimize(
    adapter=adapter,
    trainset=train_tasks,
    valset=val_tasks,
    max_metric_calls=100,
)

result = optimize_anything(
    seed_candidate=None,  # No seed!
    evaluator=evaluate_agent,
    dataset=train_tasks,
    valset=val_tasks,
    objective="""
    Create an agent to solve ARC-AGI abstract reasoning tasks.
    The agent should analyze input-output examples and predict test outputs.
    """,
    background="""
    ARC-AGI involves 2D grid transformations with abstract patterns.
    Successful approaches typically:
    - Extract transformation rules from examples
    - Use executable code rather than direct prediction
    - Validate generated code on training examples
    - Include fallback strategies
    
    You have access to: numpy, standard library, LLM API via llm.generate()
    """,
)