GEPA: Reflective Optimizer

Incorrect classification
  Expected: "positive"
  Predicted: "negative"
  Input text: "Great product but shipping was slow"
  May have misunderstood mixed sentiment

"The current instruction doesn't handle mixed sentiments well. 
Suggest modifying to explicitly consider both positive and negative aspects..."

use dspy_rs::*;

#[derive(Builder, facet::Facet)]
#[facet(crate = facet)]
struct MyModule {
    predictor: Predict<MySignature>,
}

impl Module for MyModule {
    type Input = MySignatureInput;
    type Output = MySignatureOutput;

    async fn forward(&self, inputs: MySignatureInput) -> Result<Predicted<MySignatureOutput>, PredictError> {
        self.predictor.call(inputs).await
    }
}

struct MyMetric;

impl TypedMetric<MySignature, MyModule> for MyMetric {
    async fn evaluate(
        &self,
        example: &Example<MySignature>,
        prediction: &Predicted<<MyModule as Module>::Output>,
    )
        -> Result<MetricOutcome>
    {
        let predicted = prediction.answer.as_str();
        let expected = example.output.answer.as_str();
        
        let correct = predicted == expected;
        let score = if correct { 1.0 } else { 0.0 };
        
        let feedback = if correct {
            format!("Correct answer: {}", predicted)
        } else {
            format!("Incorrect\n  Expected: {}\n  Predicted: {}", expected, predicted)
        };
        
        Ok(MetricOutcome::with_feedback(score, FeedbackMetric::new(score, feedback)))
    }
}

let gepa = GEPA::builder()
    .num_iterations(20)
    .minibatch_size(25)
    .num_trials(10)
    .temperature(0.9)
    .track_stats(true)
    .maybe_max_rollouts(Some(500))  // Budget control
    .build();

let result = gepa.compile(&mut module, trainset, &metric).await?;

println!("Best score: {:.3}", result.best_candidate.average_score());
println!("Best instruction: {}", result.best_candidate.instruction);

use dspy_rs::feedback_helpers::retrieval_feedback;

let feedback = retrieval_feedback(
    &retrieved_docs,
    &expected_docs,
    Some(&all_available_docs)
);

// Output:
// Retrieved 3/5 correct documents (Precision: 0.6, Recall: 0.6, F1: 0.6)
// Correctly retrieved: doc1, doc2, doc3
// Missed: doc4, doc5

use dspy_rs::feedback_helpers::{code_pipeline_feedback, CodeStage, StageResult};

let stages = vec![
    (CodeStage::Parse, StageResult::Success),
    (CodeStage::Compile, StageResult::Success),
    (CodeStage::Execute, StageResult::Failure {
        error: "Division by zero on line 42".to_string(),
    }),
];

let feedback = code_pipeline_feedback(&stages, 0.6);

// Output:
// Parse: Success
// Compile: Success
// Execute: Division by zero on line 42

use dspy_rs::feedback_helpers::multi_objective_feedback;

let mut objectives = HashMap::new();
objectives.insert("accuracy".to_string(), (0.9, "High accuracy".to_string()));
objectives.insert("latency".to_string(), (0.7, "Acceptable latency".to_string()));
objectives.insert("cost".to_string(), (0.8, "Within budget".to_string()));

let mut weights = HashMap::new();
weights.insert("accuracy".to_string(), 2.0);
weights.insert("latency".to_string(), 1.0);
weights.insert("cost".to_string(), 1.0);

let feedback = multi_objective_feedback(&objectives, Some(&weights));

GEPA::builder()
    .num_iterations(20)          // Number of evolutionary iterations
    .minibatch_size(25)          // Examples per rollout
    .num_trials(10)              // Trials per candidate evaluation
    .temperature(0.9)            // LLM temperature for mutations
    .track_stats(true)           // Track detailed statistics
    .track_best_outputs(false)   // Store best outputs per example
    .maybe_max_rollouts(Some(500))   // Budget: max rollouts
    .maybe_max_lm_calls(Some(1000))  // Budget: max LM calls
    .maybe_prompt_model(Some(lm))    // Separate LM for meta-prompting
    .build()

let result = gepa
    .compile_with_valset(&mut module, trainset, Some(valset), &metric)
    .await?;

let result = gepa.compile(&mut module, trainset, &metric).await?;

// Best candidate found
println!("Best instruction: {}", result.best_candidate.instruction);
println!("Average score: {:.3}", result.best_candidate.average_score());
println!("Generation: {}", result.best_candidate.generation);

// Resource usage
println!("Total rollouts: {}", result.total_rollouts);
println!("Total LM calls: {}", result.total_lm_calls);

// Evolution over time
for (generation, score) in &result.evolution_history {
    println!("Gen {}: {:.3}", generation, score);
}

pub struct FeedbackMetric {
    pub score: f32,
    pub feedback: String,
    pub metadata: HashMap<String, serde_json::Value>,
}

pub struct ExecutionTrace {
    pub inputs: Example,
    pub outputs: Option<BamlValue>,
    pub feedback: Option<FeedbackMetric>,
    pub intermediate_steps: Vec<(String, serde_json::Value)>,
    pub errors: Vec<String>,
    pub metadata: HashMap<String, serde_json::Value>,
}

pub struct GEPACandidate {
    pub id: usize,
    pub instruction: String,
    pub module_name: String,
    pub example_scores: Vec<f32>,
    pub parent_id: Option<usize>,
    pub generation: usize,
}

// BAD: Vague feedback
FeedbackMetric::new(0.5, "Wrong answer")

// GOOD: Specific, actionable feedback
FeedbackMetric::new(0.5, 
    "Incorrect answer\n\
     Expected: 'Paris'\n\
     Predicted: 'France'\n\
     Issue: Returned country instead of city")

// For development/testing
GEPA::builder()
    .num_iterations(5)
    .maybe_max_rollouts(Some(100))
    .build()

// For production optimization
GEPA::builder()
    .num_iterations(20)
    .maybe_max_rollouts(Some(1000))
    .build()

// Solution: Return MetricOutcome::with_feedback(...) from TypedMetric::evaluate
impl TypedMetric<MySignature, MyModule> for MyMetric {
    async fn evaluate(
        &self,
        example: &Example<MySignature>,
        prediction: &Predicted<<MyModule as Module>::Output>,
    ) -> Result<MetricOutcome> {
        Ok(MetricOutcome::with_feedback(
            1.0,
            FeedbackMetric::new(1.0, "detailed textual feedback"),
        ))
    }
}

// Increase minibatch size for better gradient
GEPA::builder().minibatch_size(50).build()

// Increase temperature for more exploration
GEPA::builder().temperature(1.2).build()

// Reduce iterations or increase budget
GEPA::builder()
    .num_iterations(10)
    .maybe_max_rollouts(Some(2000))
    .build()

let gepa = GEPA::builder()
    .track_stats(true)
    .track_best_outputs(true)
    .build();

let result = gepa
    .compile_with_valset(&mut module, my_tasks.clone(), Some(my_tasks), &metric)
    .await?;

// Access per-task best scores and outputs
let best_scores = result.highest_score_achieved_per_val_task;
let best_outputs = result.best_outputs_valset;