chain_example/
chain_example.rs

1//! This example demonstrates how to use the PromptStore with multi-step chains
2//! that execute different prompts using different LLM providers.
3//!
4//! The example shows:
5//! - Setting up multiple LLM providers (OpenAI and Anthropic)
6//! - Creating a chain that uses stored prompts by ID
7//! - Passing outputs between chain steps using template variables
8//! - Using different providers for different steps
9
10use llm::builder::{LLMBackend, LLMBuilder};
11use llm::chain::LLMRegistry;
12use llm::chain::MultiChainStepMode;
13use prompt_store::{PromptStore, RunError, RunOutput};
14
15#[tokio::main]
16async fn main() -> Result<(), RunError> {
17    // 1. Initialize the store once. This loads keys and configuration.
18    // Or use with_password("password") to use a password to decrypt the vault of the prompts.
19    let password = std::env::var("PROMPT_STORE_PASSWORD")
20        .expect("PROMPT_STORE_PASSWORD must be set for this example.");
21    let store = PromptStore::with_password(&password)?;
22
23    // 2. Set up the LLM providers and a registry to hold them.
24    let openai_llm = LLMBuilder::new()
25        .backend(LLMBackend::OpenAI)
26        .api_key(std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY must be set"))
27        .model("gpt-4o-mini")
28        .max_tokens(1000)
29        .build()
30        .unwrap();
31
32    let anthropic_llm = LLMBuilder::new()
33        .backend(LLMBackend::Anthropic)
34        .api_key(std::env::var("ANTHROPIC_API_KEY").expect("ANTHROPIC_API_KEY must be set"))
35        .model("claude-3-5-sonnet-20240620")
36        .max_tokens(1000)
37        .build()
38        .unwrap();
39
40    let mut registry = LLMRegistry::new();
41    registry.insert("openai_fast", openai_llm);
42    registry.insert("anthropic_strong", anthropic_llm);
43
44    // 3. Define and run the chain fluently, loading prompts from the store.
45    let user_question = "How does photosynthesis work at the molecular level?";
46
47    println!("Executing prompt chain for: \"{}\"", user_question);
48
49    let outputs = store
50        .chain(&registry) // Start a chain with the provider registry.
51        // Step 1: uses the prompt with id "9k6zezem".
52        // Its output will be available as the `{{analyse}}` variable.
53        .step("analyse", "9k6zezem")
54        .with_mode(MultiChainStepMode::Chat)
55        .with_provider("openai_fast")
56        // Step 2: uses the prompt with id "uetgwnq1".
57        // It implicitly uses the `{{analyse}}` output from the previous step.
58        .step("suggestions", "uetgwnq1")
59        .with_mode(MultiChainStepMode::Chat)
60        .with_provider("anthropic_strong")
61        // Step 3: uses the prompt with id "dkeodfyp".
62        // It can use both the initial `{{query}}` and `{{suggestions}}`.
63        .step("final_response", "dkeodfyp")
64        .with_mode(MultiChainStepMode::Chat)
65        .with_provider("anthropic_strong")
66        .step_raw(
67            "raw",
68            "Synthesize the following: {{final_response}} in 2 sentences.",
69        )
70        .with_mode(MultiChainStepMode::Chat)
71        .with_provider("anthropic_strong")
72        // Provide the initial variable for the first step.
73        .vars([("query", user_question)])
74        .run()
75        .await?;
76
77    // 4. Process the results.
78    if let RunOutput::Chain(map) = outputs {
79        println!("\n--- Chain Execution Complete ---");
80        println!(
81            "\n[✅] Final Answer (from 'final_response' step):\n{}",
82            map.get("final_response").unwrap_or(&"N/A".to_string())
83        );
84        println!("\n--- Intermediate Steps ---");
85        println!(
86            "\n[1] Analysis ('analyse'):\n{}",
87            map.get("analyse").unwrap_or(&"N/A".to_string())
88        );
89        println!(
90            "\n[2] Suggestions ('suggestions'):\n{}",
91            map.get("suggestions").unwrap_or(&"N/A".to_string())
92        );
93        println!(
94            "\n[3] Raw ('raw'):\n{}",
95            map.get("raw").unwrap_or(&"N/A".to_string())
96        );
97    }
98
99    Ok(())
100}