comprehensive_demo/

comprehensive_demo.rs

use std::time::Duration;

use chrono::Utc;
use llm_brain::{LLMBrain, Result};
use serde_json::json;
use tokio::time::sleep;

/// Demonstrate storing and retrieving different memory types
async fn demonstrate_memory_types(llm_brain: &LLMBrain) -> Result<()> {
    println!("\n=== Demonstrating Memory Types ===");

    // Create semantic memory - a Tesla car entity
    println!("\nAdding semantic memory (Tesla Model 3)...");
    let car_data = json!({
        "name": "Tesla_Model_3",
        "column": "Semantic",
        "properties": {
            "color": "red",
            "year": "2023",
            "mileage": "1000 miles",
            "features": ["autopilot capabilities", "glass roof"],
            "type": "electric vehicle"
        },
        "relationships": {
            "type_of": ["Vehicle"],
            "location": "garage"
        }
    });

    let car_content = "The Tesla Model 3 is a red electric vehicle from 2023 with 1000 miles. \
                      It features autopilot capabilities and a glass roof. It is parked in the garage.";

    let car_id = llm_brain
        .add_memory(car_content.to_owned(), car_data)
        .await?;
    println!("Added car memory with ID: {car_id}");

    // Create episodic memory - buying the car
    println!("\nAdding episodic memory (Buying the car)...");
    let purchase_data = json!({
        "name": "Tesla_Purchase",
        "column": "Episodic",
        "timestamp": Utc::now().to_rfc3339(),
        "properties": {
            "action": "Bought a new car",
            "location": "Tesla Dealership",
            "participants": ["Me", "Tesla Sales Representative"],
            "emotion": "excited"
        },
        "relationships": {
            "related_to": ["Tesla_Model_3"]
        }
    });

    let purchase_content = "I bought a new Tesla Model 3 today at the Tesla Dealership. \
                           The sales representative was very helpful. I was very excited about getting my new car.";

    let purchase_id = llm_brain
        .add_memory(purchase_content.to_owned(), purchase_data)
        .await?;
    println!("Added purchase memory with ID: {purchase_id}");

    // Create procedural memory - charging the car
    println!("\nAdding procedural memory (Charging the car)...");
    let charging_data = json!({
        "name": "Charge_Tesla",
        "column": "Procedural",
        "properties": {
            "steps": [
                "Park near charging station",
                "Open charging port on car",
                "Connect charging cable",
                "Initiate charging via app or car interface",
                "Wait for sufficient charge",
                "Disconnect charging cable",
                "Close charging port"
            ],
            "required_tools": ["Charging cable", "Tesla app"],
            "estimated_time": "30-60 minutes",
            "frequency": "weekly"
        },
        "relationships": {
            "applies_to": ["Tesla_Model_3"]
        }
    });

    let charging_content = "To charge a Tesla Model 3: First, park near a charging station. \
                          Open the charging port on the car and connect the charging cable. \
                          Then initiate charging via the app or car interface. Wait for 30-60 minutes \
                          until there is sufficient charge. Finally, disconnect the charging cable \
                          and close the charging port. This procedure should be performed weekly \
                          and requires a charging cable and the Tesla app.";

    let charging_id = llm_brain
        .add_memory(charging_content.to_owned(), charging_data)
        .await?;
    println!("Added charging memory with ID: {charging_id}");

    // Allow time for embeddings to be processed
    sleep(Duration::from_millis(200)).await;

    Ok(())
}

/// Demonstrate querying and memory recall capabilities
async fn demonstrate_memory_recall(llm_brain: &LLMBrain) -> Result<()> {
    println!("\n=== Demonstrating Memory Recall ===");

    // Query for Tesla car
    println!("\nQuerying for Tesla Model 3 information:");
    let tesla_query = "What do I know about my Tesla Model 3?";
    let tesla_results = llm_brain.recall(tesla_query, 3).await?;

    for (i, (memory, score)) in tesla_results.iter().enumerate() {
        println!("\nResult {}: Score {:.4}", i + 1, score);
        println!("Content: {}", memory.content);
        println!(
            "Metadata: {}",
            serde_json::to_string_pretty(&memory.metadata)?
        );
    }

    // Query for car purchase experience
    println!("\nQuerying for car purchase experience:");
    let purchase_query = "When did I buy my Tesla?";
    let purchase_results = llm_brain.recall(purchase_query, 2).await?;

    for (i, (memory, score)) in purchase_results.iter().enumerate() {
        println!("\nResult {}: Score {:.4}", i + 1, score);
        println!("Content: {}", memory.content);
        println!(
            "Metadata: {}",
            serde_json::to_string_pretty(&memory.metadata)?
        );
    }

    // Query for charging instructions
    println!("\nQuerying for charging instructions:");
    let charging_query = "How do I charge my car?";
    let charging_results = llm_brain.recall(charging_query, 2).await?;

    for (i, (memory, score)) in charging_results.iter().enumerate() {
        println!("\nResult {}: Score {:.4}", i + 1, score);
        println!("Content: {}", memory.content);
        println!(
            "Metadata: {}",
            serde_json::to_string_pretty(&memory.metadata)?
        );
    }

    Ok(())
}

/// Demonstrate adding a memory and then finding it by similarity
async fn demonstrate_memory_lookup(llm_brain: &LLMBrain) -> Result<()> {
    println!("\n=== Demonstrating Memory Lookup ===");

    // Add a new memory about a trip
    println!("\nAdding memory about a road trip...");
    let trip_data = json!({
        "name": "Road_Trip_Tahoe",
        "column": "Episodic",
        "timestamp": Utc::now().to_rfc3339(),
        "properties": {
            "action": "Road trip to Lake Tahoe",
            "location": "Lake Tahoe",
            "participants": ["Me", "Family"],
            "duration": "3 days",
            "transportation": "Tesla Model 3",
            "highlights": ["Beautiful scenery", "Hiking", "Swimming"]
        }
    });

    let trip_content = "Last weekend, we took our Tesla Model 3 on a road trip to Lake Tahoe. \
                       The drive was smooth and we enjoyed the beautiful scenery. \
                       We spent 3 days there with the family, hiking and swimming in the lake. \
                       The car performed excellently on the mountain roads.";

    let trip_id = llm_brain
        .add_memory(trip_content.to_owned(), trip_data)
        .await?;
    println!("Added trip memory with ID: {trip_id}");

    // Allow time for embeddings to be processed
    sleep(Duration::from_millis(200)).await;

    // Now try to find this memory with various related queries
    let queries = [
        "Did we go on any trips with our Tesla?",
        "What activities did we do at Lake Tahoe?",
        "How did the Tesla perform on our road trip?",
    ];

    for (i, query) in queries.iter().enumerate() {
        println!("\nQuery {}: \"{}\"", i + 1, query);
        let results = llm_brain.recall(query, 1).await?;

        if !results.is_empty() {
            let (memory, score) = &results[0];
            println!("Found memory (score: {score:.4}):");
            println!("Content: {}", memory.content);
        } else {
            println!("No relevant memories found.");
        }
    }

    // Direct lookup by ID
    println!("\nDirect lookup by ID:");
    if let Some(memory) = llm_brain
        .get_memory_by_id_string(&trip_id.to_string())
        .await?
    {
        println!("Successfully retrieved memory");
        println!("Content: {}", memory.content);
    } else {
        println!("Failed to retrieve memory by ID");
    }

    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize LLMBrain instance
    println!("Initializing LLMBrain...");
    let llm_brain = LLMBrain::launch().await?;

    // Demonstrate different memory types
    demonstrate_memory_types(&llm_brain).await?;

    // Demonstrate memory recall
    demonstrate_memory_recall(&llm_brain).await?;

    // Demonstrate memory lookup
    demonstrate_memory_lookup(&llm_brain).await?;

    println!("\nComprehensive demonstration completed successfully!");
    Ok(())
}