llm_brain_rag/
llm_brain_rag.rs

1use std::fs;
2use std::path::PathBuf;
3use std::time::Instant;
4
5use anyhow::Result;
6use chrono::Utc;
7use llm_brain::LLMBrain;
8use serde_json::{Value as JsonValue, json};
9
10// --- Config Helper (Copied and adapted) ---
11fn ensure_config_exists() -> Result<()> {
12    let config_dir = PathBuf::from("./config");
13    let default_config_path = config_dir.join("default.toml");
14    if default_config_path.exists() {
15        return Ok(());
16    }
17    if !config_dir.exists() {
18        fs::create_dir_all(&config_dir)?;
19    }
20    // Use a different DB for this specific RAG example
21    let default_toml_content = r#"
22[database]
23path = "./llm_brain_rag_db"
24namespace = "rag_ns"
25database = "rag_db"
26
27[llm]
28# Provide necessary LLM config
29"#;
30    fs::write(&default_config_path, default_toml_content)?;
31    println!("Created default config: {}", default_config_path.display());
32    Ok(())
33}
34// --- End Config Helper ---
35
36// Structure to hold content and metadata together for test data
37struct TestMemoryData {
38    content: String,
39    metadata: JsonValue,
40}
41
42async fn run_rag_queries(llm_brain: &LLMBrain, queries: &[&str]) -> Result<()> {
43    println!("\n--- Running Simulated RAG Queries (Retrieval Only) ---");
44    for query in queries {
45        println!("\nQuery: {query}");
46        let start_time = Instant::now();
47        match llm_brain.recall(query, 3).await {
48            // Retrieve top 3 relevant fragments
49            Ok(results) => {
50                println!(
51                    "Retrieved {} relevant fragments in {:?}:",
52                    results.len(),
53                    start_time.elapsed()
54                );
55                if results.is_empty() {
56                    println!("  (No relevant fragments found for RAG context)");
57                } else {
58                    println!("  --- Potential RAG Context: ---");
59                    for (fragment, score) in results {
60                        println!(
61                            "  - Score: {:.4}, Entity: {:?}\n    Content: {:.150}...",
62                            score,
63                            fragment.metadata.get("entity_name"), // Display name if available
64                            fragment.content.replace('\n', " ")
65                        );
66                    }
67                    println!("  --- End Context ---");
68                    println!(
69                        "  (Note: In a full RAG system, these fragments would be passed to an LLM with the query to generate a final answer.)"
70                    );
71                }
72            }
73            Err(e) => {
74                eprintln!("  Error during recall: {e}");
75            }
76        }
77    }
78    Ok(())
79}
80
81#[tokio::main]
82async fn main() -> Result<()> {
83    println!("--- Starting LLMBrain RAG Example (Rust/LLMBrain Version) ---");
84    ensure_config_exists()?;
85    let llm_brain = LLMBrain::launch().await?;
86
87    // --- Add Memories Step ---
88    println!("\nAdding example memories...");
89    let examples = vec![
90        TestMemoryData {
91            content: "Python is a programming language created by Guido van Rossum in 1991. It supports object-oriented, imperative, and functional programming. Commonly used for web development, data science, automation.".to_owned(),
92            metadata: json!({
93                "entity_name": "Python_Language",
94                "memory_type": "Semantic",
95                "properties": {"creator": "Guido van Rossum", "year": 1991, "paradigms": ["OOP", "Imperative", "Functional"], "uses": ["web dev", "data science", "automation"]},
96            }),
97        },
98        TestMemoryData {
99            content: "Completed first Python project today in home office. Took 2 hours, successful. Reviewed code.".to_owned(),
100            metadata: json!({
101                "entity_name": "First_Python_Project",
102                "memory_type": "Episodic",
103                "properties": {"timestamp": Utc::now().timestamp(), "location": "home office", "duration_hours": 2, "outcome": "successful"},
104            }),
105        },
106        TestMemoryData {
107            content: "Tesla Model 3 is red, made in 2023, parked in the garage. Range 358 miles, 0-60 mph in 3.1 seconds.".to_owned(),
108            metadata: json!({
109                "entity_name": "Tesla_Model_3_Red",
110                "memory_type": "Semantic",
111                "properties": {"color": "red", "year": 2023, "location": "garage", "range_miles": 358, "acceleration_0_60_sec": 3.1},
112            }),
113        },
114    ];
115
116    for (i, data) in examples.into_iter().enumerate() {
117        println!("\nAdding Example {}", i + 1);
118        match llm_brain
119            .add_memory(data.content.clone(), data.metadata.clone())
120            .await
121        {
122            Ok(id) => println!("  Added successfully. ID: {id}"),
123            Err(e) => eprintln!("  Failed to add: {e}"),
124        }
125    }
126
127    // --- Query Step (Simulating RAG Retrieval) ---
128    let test_queries = [
129        "What programming language was created by Guido van Rossum?",
130        "Tell me about the Tesla Model 3's specifications.",
131        "What happened during the first Python project?",
132    ];
133
134    if let Err(e) = run_rag_queries(&llm_brain, &test_queries).await {
135        eprintln!("Error during querying: {e}");
136    }
137
138    println!("\n--- LLMBrain RAG Example Finished ---");
139    println!("Note: Database is at ./llm_brain_rag_db");
140
141    Ok(())
142}
llm_brain_rag/llm_brain_rag.rs

llm_brain_rag/
llm_brain_rag.rs
