rag-module 0.6.7

Enterprise RAG module with chat context storage, vector search, session management, and model downloading. Rust implementation with Node.js compatibility.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Quick Search Demo - Proves search works with encrypted data
//! Run with: cargo run --example quick_search_demo

use anyhow::Result;
use serde_json::json;
use uuid::Uuid;
use rag_module::*;
use rag_module::services::search_service::{ChatSearchOptions, EstateSearchOptions};

#[tokio::main]
async fn main() -> Result<()> {
    println!("🔍 Quick Search Demo: Encrypted Data + Search Functionality");
    println!("============================================================\n");

    let base_path = std::env::current_dir()?.join("quick-demo-data");
    let rag_module = create_rag_module(base_path).await?;
    rag_module.initialize().await?;

    let user_id = "demo_user";
    let uuid_str = Uuid::new_v4().to_string();
    let context_id = format!("chat_{}", &uuid_str[..8]);

    // Test 1: Chat History Search (Decrypt-First)
    println!("📝 Test 1: Chat History with Context ID Filtering");
    
    let session = rag_module.start_session(StartSessionOptions {
        user_id: user_id.to_string(),
        chat_title: Some("Search Test".to_string()),
        context_id: Some(context_id.clone()),
    }).await?;

    // Add messages
    rag_module.add_prompt(&session.id, "How to create Lambda functions?", user_id).await?;
    rag_module.add_response(&session.id, "Lambda functions can be created using AWS console or CLI", user_id).await?;
    
    println!("✅ Added encrypted chat messages");

    // Search by context_id (requires decrypting metadata first)
    let chat_options = ChatSearchOptions {
        context_id: Some(context_id.clone()),
        role: None,
        from_timestamp: None,
        to_timestamp: None,
        from_message_index: None,
        to_message_index: None,
        limit: Some(10),
        include_metadata: true,
    };

    let chat_results = rag_module.search_service.search_chat_history(chat_options).await?;
    println!("✅ Found {} messages for context_id: {}", chat_results.len(), context_id);
    
    for result in &chat_results {
        if let Some(role) = result.get("role").and_then(|r| r.as_str()) {
            if let Some(content) = result.get("content").and_then(|c| c.as_str()) {
                let preview = if content.len() > 50 { 
                    format!("{}...", &content[..50])
                } else { 
                    content.to_string() 
                };
                println!("  [{}] {}", role.to_uppercase(), preview);
            }
        }
    }

    // Test 2: Estate Search (Vector Search)
    println!("\n🏢 Test 2: Estate Resources with Vector Search");
    
    let estate_data = json!([{
        "account_id": "123456789",
        "services": {
            "lambda": {
                "functions": [{
                    "function_name": "DataProcessor", 
                    "description": "Lambda function for processing data files"
                }]
            }
        }
    }]);

    rag_module.process_aws_estate(estate_data, user_id).await?;
    println!("✅ Added estate data with embeddings from plain text");

    tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

    let estate_options = EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: None,
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(5),
        score_threshold: Some(0.1),
        include_metadata: true,
        use_anonymous_ids: false,
        parameters: None, // Return all fields (default behavior)
    };

    let estate_results = rag_module.search_service
        .search_estate_resources("lambda functions", estate_options, None)
        .await?;

    println!("✅ Found {} estate resources for 'lambda functions'", estate_results.len());
    
    for result in &estate_results {
        if let Some(service) = result.get("service").and_then(|s| s.as_str()) {
            if let Some(score) = result.get("score").and_then(|s| s.as_f64()) {
                println!("  Service: {} | Relevance: {:.3}", service, score);
            }
        }
    }

    // Verify encryption
    println!("\n🔐 Test 3: Verify Data is Actually Encrypted");
    let encrypted_docs = rag_module.get_collection_documents("chat_history", user_id).await?;
    
    if let Some(doc) = encrypted_docs.first() {
        let content_sample = if doc.content.len() > 80 {
            format!("{}...", &doc.content[..80])
        } else {
            doc.content.clone()
        };
        println!("✅ Encrypted content sample: {}", content_sample);
        println!("✅ Content length: {} chars (encrypted base64)", doc.content.len());
    }

    println!("\n🎉 DEMO RESULTS:");
    println!("✅ Chat search worked despite encrypted metadata (decrypt-first approach)");
    println!("✅ Estate search worked despite encrypted content (embeddings from plain text)"); 
    println!("✅ Both methods successfully handle encryption while preserving search functionality");
    
    rag_module.end_session(&session.id).await?;
    Ok(())
}