use orch::execution::*;
use orch::lm::*;

#[tokio::main]
async fn main() {
  let lm = OllamaBuilder::new().try_build().unwrap();
  let executor = TextExecutorBuilder::new().with_lm(&lm).try_build().unwrap();
  let response = executor.execute("What is 2 2?").await.expect("Execution failed");
  println!("{}", response.content);
}

use orch::execution::*;
use orch::lm::*;
use tokio_stream::StreamExt;

#[tokio::main]
async fn main() {
  let lm = OllamaBuilder::new().try_build().unwrap();
  let executor = TextExecutorBuilder::new().with_lm(&lm).try_build().unwrap();
  let mut response = executor.execute_stream("What is 2 2?").await.expect("Execution failed");
  while let Some(chunk) = response.stream.next().await {
    match chunk {
      Ok(chunk) => print!("{chunk}"),
      Err(e) => {
        println!("Error: {e}");
        break;
      }
    }
  }
  println!();
}

use orch::execution::*;
use orch::lm::*;
use orch::response::*;

#[derive(Variants, serde::Deserialize)]
pub enum ResponseVariants {
    Answer(AnswerResponseVariant),
    Fail(FailResponseVariant),
}

#[derive(Variant, serde::Deserialize)]
#[variant(
    variant = "Answer",
    scenario = "You know the capital city of the country",
    description = "Capital city of the country"
)]
pub struct AnswerResponseVariant {
    #[schema(
        description = "Capital city of the received country",
        example = "London"
    )]
    pub capital: String,
}

#[derive(Variant, serde::Deserialize)]
#[variant(
    variant = "Fail",
    scenario = "You don't know the capital city of the country",
    description = "Reason why the capital city is not known"
)]
pub struct FailResponseVariant {
    #[schema(
        description = "Reason why the capital city is not known",
        example = "Country 'foobar' does not exist"
    )]
    pub reason: String,
}

#[tokio::main]
async fn main() {
    let lm = OllamaBuilder::new().try_build().unwrap();
    let executor = StructuredExecutorBuilder::new()
    .with_lm(&lm)
    .with_preamble("You are a geography expert who helps users understand the capital city of countries around the world.")
		.with_options(Box::new(variants!(ResponseVariants)))
    .try_build()
    .unwrap();
    let response = executor
        .execute("What is the capital of Fooland?")
        .await
        .expect("Execution failed");

    println!("Response:");
    match response.content {
        ResponseVariants::Answer(answer) => {
            println!("Capital city: {}", answer.capital);
        }
        ResponseVariants::Fail(fail) => {
            println!("Model failed to generate a response: {}", fail.reason);
        }
    }
}

use orch::execution::*;
use orch::lm::*;

#[tokio::main]
async fn main() {
  let lm = OllamaBuilder::new().try_build().unwrap();
  let executor = TextExecutorBuilder::new()
    .with_lm(&lm)
    .try_build()
    .unwrap();
  let embedding = executor
    .generate_embedding("Phrase to generate an embedding for")
    .await
    .expect("Execution failed");

  println!("Embedding:");
  println!("{:?}", embedding);
}