advanced_chain_example/
advanced_chain_example.rs

1//! This example demonstrates advanced chain logic: conditional steps.
2//!
3//! Pre-requisites:
4//! 1. Run `prompt-store new` to create the following prompts:
5//!    - Title: "Sentiment Check", Content: "Is the following user feedback positive or negative? Answer with only the word 'positive' or 'negative'. Feedback: {{feedback}}"
6//!    - Title: "Positive Reply", Content: "Write a short, cheerful thank you message for this positive feedback: {{feedback}}"
7//!    - Title: "Negative Reply", Content: "Write a short, apologetic response and offer to help with this negative feedback: {{feedback}}"
8//! 2. Ensure your key is password protected: `prompt-store rotate-key --password`
9//! 3. Set `PROMPT_STORE_PASSWORD` and `OPENAI_API_KEY` environment variables.
10
11use llm::builder::{LLMBackend, LLMBuilder};
12use llm::chain::LLMRegistry;
13use prompt_store::{PromptStore, RunError, RunOutput};
14
15#[tokio::main]
16async fn main() -> Result<(), RunError> {
17    let password = std::env::var("PROMPT_STORE_PASSWORD")
18        .expect("PROMPT_STORE_PASSWORD must be set for this example.");
19    let store = PromptStore::with_password(&password)?;
20
21    let openai_llm = LLMBuilder::new()
22        .backend(LLMBackend::OpenAI)
23        .api_key(std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY must be set"))
24        .model("gpt-4o-mini")
25        .build()
26        .unwrap();
27
28    let mut registry = LLMRegistry::new();
29    registry.insert("openai", openai_llm);
30
31    // --- Test with negative feedback ---
32    println!("--- Testing with NEGATIVE feedback ---");
33    let user_feedback_negative = "The app keeps crashing, it's unusable!";
34    let outputs_neg = run_chain(&store, &registry, user_feedback_negative).await?;
35    if let RunOutput::Chain(map) = outputs_neg {
36        // We expect `negative_reply` to exist, but `positive_reply` should not.
37        assert!(map.contains_key("negative_reply"));
38        assert!(!map.contains_key("positive_reply"));
39        println!("\nFinal Response:\n{}", map.get("negative_reply").unwrap());
40    }
41
42    // --- Test with positive feedback ---
43    println!("\n--- Testing with POSITIVE feedback ---");
44    let user_feedback_positive = "I love the new update, it's so fast!";
45    let outputs_pos = run_chain(&store, &registry, user_feedback_positive).await?;
46    if let RunOutput::Chain(map) = outputs_pos {
47        // We expect `positive_reply` to exist, but `negative_reply` should not.
48        assert!(map.contains_key("positive_reply"));
49        assert!(!map.contains_key("negative_reply"));
50        println!("\nFinal Response:\n{}", map.get("positive_reply").unwrap());
51    }
52
53    Ok(())
54}
55
56async fn run_chain(
57    store: &PromptStore,
58    registry: &LLMRegistry,
59    feedback: &str,
60) -> Result<RunOutput, RunError> {
61    store
62        .chain(registry)
63        // Step 1: Always run sentiment analysis.
64        .step("sentiment", "Sentiment Check")
65            .with_provider("openai")
66        // Step 2 (Conditional): Only run if the sentiment is "positive".
67        .step_if("positive_reply", "Positive Reply", |prev_outputs| {
68            matches!(prev_outputs.get("sentiment"), Some(s) if s.trim().eq_ignore_ascii_case("positive"))
69        })
70            .with_provider("openai")
71
72        // Step 3 (Conditional): Only run if the sentiment is "negative".
73        .step_if("negative_reply", "Negative Reply", |prev_outputs| {
74            matches!(prev_outputs.get("sentiment"), Some(s) if s.trim().eq_ignore_ascii_case("negative"))
75        })
76            .with_provider("openai")
77
78        .vars([("feedback", feedback)])
79        .run()
80        .await
81}