COPRO

let copro = COPRO::builder()
    .breadth(10)              // Candidates per round (must be > 1)
    .depth(3)                 // Number of refinement rounds
    .build();

use anyhow::Result;
use bon::Builder;
use dspy_rs::{
    COPRO, Eval, LM, Module, Predict, PredictError,
    Predicted, Signature, Trace, TypedMetric, configure, init_tracing,
};

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

    #[output]
    answer: String,
}

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

dspy_rs::predictors!(MyModule { predictor });

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<(QAInput, QAOutput), MyModule> for ExactMatchMetric {
    async fn evaluate(
        &self,
        example: &(QAInput, QAOutput),
        prediction: &Predicted<QAOutput>,
        _trace: Option<&Trace>,
    ) -> Result<Eval> {
        let expected = example.1.answer.trim().to_lowercase();
        let actual = prediction.answer.trim().to_lowercase();
        Ok(Eval::score((expected == actual) as u8 as f64))
    }
}

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

    // API key automatically read from OPENAI_API_KEY env var
    configure(
        LM::builder()
            .model("openai:gpt-4o-mini".to_string())
            .build()
            .await?,
    );

    let mut module = MyModule::builder().build();
    let trainset = vec![
        (
            QAInput {
                question: "What is 2+2?".to_string(),
            },
            QAOutput {
                answer: "4".to_string(),
            },
        ),
        (
            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_module(&mut module, &trainset, &metric).await?;

    Ok(())
}