rag_in_memory/
rag_in_memory.rs

1//! # Example: Agent with In-Memory RAG
2//!
3//! This example demonstrates how to create an agent with an in-memory RAG system.
4//! The in-memory vector store is perfect for development, testing, or when you
5//! don't need persistence across restarts.
6//!
7//! ## Prerequisites
8//!
9//! - **OpenAI API Key**: You need an OpenAI API key for generating embeddings.
10//!   Set it as an environment variable:
11//!   ```sh
12//!   export OPENAI_API_KEY=your-key
13//!   ```
14
15use helios_engine::{Agent, Config, RAGTool};
16
17#[tokio::main]
18async fn main() -> helios_engine::Result<()> {
19    println!("šŸš€ Helios Engine - Agent with In-Memory RAG Example");
20    println!("===================================================\n");
21
22    // Check for the required OpenAI API key
23    let embedding_api_key = std::env::var("OPENAI_API_KEY").unwrap_or_else(|_| {
24        println!("⚠ Warning: OPENAI_API_KEY not set. Using placeholder.");
25        "your-api-key-here".to_string()
26    });
27
28    // Load configuration
29    let config = Config::from_file("config.toml").unwrap_or_else(|_| {
30        println!("⚠ No config.toml found, using default configuration");
31        Config::new_default()
32    });
33
34    // Create a new RAG tool with in-memory vector store
35    let rag_tool =
36        RAGTool::new_in_memory("https://api.openai.com/v1/embeddings", embedding_api_key);
37
38    // Create an agent with RAG capabilities
39    let mut agent = Agent::builder("KnowledgeAgent")
40        .config(config)
41        .system_prompt(
42            "You are a helpful assistant with access to an in-memory RAG (Retrieval-Augmented Generation) system. \
43             You can store documents and retrieve relevant information to answer questions. \
44             When answering questions, first search for relevant documents, then provide informed answers based on the retrieved context."
45        )
46        .tool(Box::new(rag_tool))
47        .max_iterations(10)
48        .build()
49        .await?;
50
51    println!("āœ“ Agent created with in-memory RAG capabilities\n");
52
53    // --- Example 1: Add knowledge about programming languages ---
54    println!("Example 1: Building a Knowledge Base");
55    println!("=====================================\n");
56
57    let documents = [
58        "Rust is a systems programming language that runs blazingly fast, prevents segfaults, \
59         and guarantees thread safety. It was created by Mozilla Research and first released in 2010.",
60        "Python is a high-level, interpreted programming language known for its clear syntax \
61         and readability. It was created by Guido van Rossum and first released in 1991.",
62        "JavaScript is a programming language commonly used for web development. It allows \
63         developers to create interactive web pages and runs in web browsers. It was created in 1995.",
64        "Go is a statically typed, compiled programming language designed at Google. It is \
65         syntactically similar to C, but with memory safety and garbage collection.",
66        "TypeScript is a strongly typed programming language that builds on JavaScript. It was \
67         developed and maintained by Microsoft and first released in 2012.",
68    ];
69
70    for (i, doc) in documents.iter().enumerate() {
71        println!("Adding document {}...", i + 1);
72        let response = agent
73            .chat(&format!("Store this information: {}", doc))
74            .await?;
75        println!("Agent: {}\n", response);
76    }
77
78    // --- Example 2: Semantic search with different queries ---
79    println!("\nExample 2: Semantic Search");
80    println!("==========================\n");
81
82    let queries = vec![
83        "What programming language prevents segfaults?",
84        "Tell me about the language created by Guido van Rossum",
85        "Which language is used for web development in browsers?",
86        "What language was developed by Google?",
87    ];
88
89    for query in queries {
90        println!("Query: {}", query);
91        let response = agent.chat(query).await?;
92        println!("Agent: {}\n", response);
93    }
94
95    // --- Example 3: Check document count ---
96    println!("\nExample 3: Document Count");
97    println!("=========================\n");
98
99    let response = agent.chat("How many documents are stored?").await?;
100    println!("Agent: {}\n", response);
101
102    // --- Example 4: Complex query requiring multiple documents ---
103    println!("\nExample 4: Complex Query");
104    println!("========================\n");
105
106    let response = agent
107        .chat("Compare the programming languages that were created in the 1990s")
108        .await?;
109    println!("Agent: {}\n", response);
110
111    println!("\nāœ… Example completed successfully!");
112    println!("\nšŸ’” Key Features Demonstrated:");
113    println!("  • In-memory vector storage (no external dependencies)");
114    println!("  • Document embedding with OpenAI embeddings");
115    println!("  • Semantic search with cosine similarity");
116    println!("  • RAG workflow for context-aware answers");
117    println!("  • Fast performance for development and testing");
118
119    println!("\nšŸ“ Use Cases for In-Memory RAG:");
120    println!("  • Development and testing");
121    println!("  • Short-lived sessions");
122    println!("  • When persistence is not required");
123    println!("  • Rapid prototyping");
124
125    println!("\nšŸ”§ Advantages:");
126    println!("  • No external dependencies (no database needed)");
127    println!("  • Fast setup and execution");
128    println!("  • Simple deployment");
129    println!("  • Perfect for demos and examples");
130
131    Ok(())
132}