use dsrs::prelude::*;
use std::collections::HashMap;
use indexmap::IndexMap;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Define a signature for Q&A
    let signature = Signature::builder()
        .name("QASignature".to_string())
        .instruction("Answer the question concisely and accurately.".to_string())
        .input_fields(IndexMap::from([(
            "question".to_string(),
            Field::InputField {
                prefix: "Question:".to_string(),
                desc: "The question to answer".to_string(),
                format: None,
                output_type: "String".to_string(),
            },
        )]))
        .output_fields(IndexMap::from([(
            "answer".to_string(),
            Field::OutputField {
                prefix: "Answer:".to_string(),
                desc: "The answer to the question".to_string(),
                format: None,
                output_type: "String".to_string(),
            },
        )]))
        .build()?;

    // Create a predictor
    let predictor = Predict { 
        signature: &signature 
    };

    // Prepare input
    let inputs = HashMap::from([
        ("question".to_string(), "What is the capital of France?".to_string())
    ]);

    // Execute prediction
    let result = predictor.forward(inputs, None, None).await?;
    
    println!("Answer: {}", result.get("answer", None));

    Ok(())
}