Add alignment strategy (#11)

Author: guywaldman

README.md

@@ -112,7 +112,7 @@ async fn main() {
     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(&variants!(ResponseVariants))
+		.with_options(Box::new(variants!(ResponseVariants)))
     .try_build()
     .unwrap();
     let response = executor

core/Cargo.toml

@@ -13,7 +13,6 @@ orch_response = { path = "../response", version = "0.0.12" }
 orch_response_derive = { path = "../response_derive", version = "0.0.12" }
 async-gen = "0.2.3"
 dotenv = "0.15.0"
-dyn-clone = "1.0.17"
 openai-api-rs = "5.0.2"
 reqwest = { version = "0.12.5", features = ["blocking"] }
 serde = { version = "1.0.164", features = ["derive"] }
@@ -22,3 +21,5 @@ thiserror = "1.0.63"
 tokio = { version = "1.28.2", features = ["rt", "macros"] }
 tokio-stream = "0.1.15"
 async-trait = "0.1.81"
+dyn-clone = "1.0.17"
+async-recursion = "1.1.1"

core/examples/alignment.rs

@@ -0,0 +1,80 @@
+#![allow(dead_code)]
+
+use orch::alignment::AlignmentStrategyBuilder;
+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 answer",
+    description = "Result of the calculation"
+)]
+pub struct AnswerResponseVariant {
+    #[schema(description = "Result of the calculation", example = "42")]
+    pub result: String,
+}
+
+#[derive(Variant, serde::Deserialize)]
+#[variant(
+    variant = "Fail",
+    scenario = "You don't know the answer",
+    description = "Reason why the answer is not known"
+)]
+pub struct FailResponseVariant {
+    #[schema(
+        description = "Reason why the answer is not known",
+        example = "The phrase is not a mathematical related expression"
+    )]
+    pub reason: String,
+}
+
+#[tokio::main]
+async fn main() {
+    // We use a large foundational model for the main task.
+    let ollama_large = OllamaBuilder::new()
+        .with_model(ollama_model::LLAMA3_8B.to_string())
+        .try_build()
+        .unwrap();
+
+    // We use a smaller model for the correction.
+    let ollama_corrector = OllamaBuilder::new()
+        .with_model(ollama_model::LLAMA3_8B.to_string())
+        .try_build()
+        .unwrap();
+
+    // We define an alignment strategy that uses the correction model.
+    let alignment_strategy = AlignmentStrategyBuilder::new()
+        .with_retries(2)
+        .with_lm(Box::new(ollama_corrector))
+        .try_build()
+        .unwrap();
+
+    let executor = StructuredExecutorBuilder::new()
+	.with_lm(&ollama_large)
+	.with_preamble("
+		You are a mathematician who helps users understand the result of mathematical expressions.
+		You will receive a mathematical expression, and you will need to provide the result of that expression.
+	")
+	.with_options(Box::new(variants!(ResponseVariants)))
+    .with_alignment(alignment_strategy)
+	.try_build()
+	.unwrap();
+    let response = executor.execute("2 + 2").await.expect("Execution failed");
+
+    match response.content {
+        ResponseVariants::Answer(answer) => {
+            println!("Result: {}", answer.result);
+        }
+        ResponseVariants::Fail(fail) => {
+            println!("Model failed to generate a response: {}", fail.reason);
+        }
+    }
+}

core/examples/structured_data_generation_blog.rs

@@ -92,7 +92,7 @@ async fn main() {
 
             Be very specific and refer to specific sentences, paragraph and sections of the blog post.
         ")
-        .with_options(&variants!(ResponseVariants))
+        .with_options(Box::new(variants!(ResponseVariants)))
         .try_build()
         .unwrap();
     let response = executor.execute(&prompt).await.expect("Execution failed");

core/examples/structured_data_generation_capital.rs

@@ -77,7 +77,7 @@ async fn main() {
             You are a geography expert who helps users understand the capital city of countries around the world.
             You will receive a country name, and you will need to provide the capital city of that country.
             ")
-        .with_options(&variants!(ResponseVariants))
+        .with_options(Box::new(variants!(ResponseVariants)))
         .try_build()
         .unwrap();
     let response = executor.execute(prompt).await.expect("Execution failed");

core/src/alignment/mod.rs

@@ -0,0 +1,12 @@
+//! A module containing all logic related to alignment.
+//! Alignment, in this context, means "aligning" the model's output with the desired output.
+//! This takes the form of a so-called `[AlignmentStrategy]`, which is a trait that defines how to align the model's output.
+//!
+//! This concept has similarities to to traditional "resilience" techniques and libraries, such as .NET's [Polly](https://github.com/App-vNext/Polly),
+//! which I personally like a lot.
+
+mod strategy;
+mod strategy_builder;
+
+pub use strategy::*;
+pub use strategy_builder::*;