COPRO

let copro = COPRO::builder()
    .breadth(10)              // Number of candidates per iteration
    .depth(3)                 // Number of refinement iterations
    .init_temperature(1.4)    // Temperature for generation
    .track_stats(false)       // Track optimization statistics
    .build();

use anyhow::Result;
use bon::Builder;
use facet;
use dspy_rs::{
    COPRO, ChatAdapter, Example, LM, MetricOutcome, Module, Optimizer, Predict, PredictError,
    Predicted, Signature, TypedMetric, configure, init_tracing,
};

#[derive(Signature, Clone, Debug)]
struct QA {
    #[input]
    question: String,

    #[output]
    answer: String,
}

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

impl Module for MyModule {
    type Input = QAInput;
    type Output = QAOutput;

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

struct ExactMatchMetric;

impl TypedMetric<QA, MyModule> for ExactMatchMetric {
    async fn evaluate(&self, example: &Example<QA>, prediction: &Predicted<QAOutput>) -> Result<MetricOutcome> {
        let expected = example.output.answer.trim().to_lowercase();
        let actual = prediction.answer.trim().to_lowercase();
        Ok(MetricOutcome::score((expected == actual) as u8 as f32))
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    init_tracing()?;

    // API key automatically read from OPENAI_API_KEY env var
    configure(
        LM::builder()
            .model("gpt-4o-mini".to_string())
            .build()
            .await?,
        ChatAdapter,
    );
    
    let mut module = MyModule::builder().build();
    let trainset = vec![
        Example::new(
            QAInput {
                question: "What is 2+2?".to_string(),
            },
            QAOutput {
                answer: "4".to_string(),
            },
        ),
        Example::new(
            QAInput {
                question: "Capital of France?".to_string(),
            },
            QAOutput {
                answer: "Paris".to_string(),
            },
        ),
    ];
    
    let copro = COPRO::builder()
        .breadth(10)
        .depth(3)
        .build();
    let metric = ExactMatchMetric;
    
copro.compile::<QA, _, _>(&mut module, trainset, &metric).await?;
    
    Ok(())
}