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

use anyhow::Result;
use rag_module::RagModule;
use rag_module::services::search_service::{ChatSearchOptions, EstateSearchOptions};
use rag_module::types::Document;
use serde_json::json;
use indexmap::IndexMap;

#[tokio::main]
async fn main() -> Result<()> {
    println!("🔍 Debug: Search and Indexing Test");
    
    // Initialize RAG module
    let mut rag_module = RagModule::new("./debug-search-test").await?;
    rag_module.initialize().await?;
    
    // Set user context
    let user_id = "test_search_user";
    rag_module.set_user_context(user_id).await?;
    
    println!("✅ RAG Module initialized with user context: {}", user_id);

    // Create a test estate document
    let mut metadata = IndexMap::new();
    metadata.insert("resource_type".to_string(), json!("ec2-instance"));
    metadata.insert("account_id".to_string(), json!("123456789012"));
    metadata.insert("region".to_string(), json!("us-east-1"));
    metadata.insert("service".to_string(), json!("ec2"));
    metadata.insert("state".to_string(), json!("running"));
    
    let estate_doc = Document::new(
        "arn:aws:ec2:us-east-1:123456789012:instance/i-search-test".to_string(), 
        "Test EC2 Instance for search testing - running in us-east-1".to_string()
    ).with_metadata(metadata);
    
    // Create document
    println!("🔧 Creating estate document...");
    let create_result = rag_module.create(vec![estate_doc]).await?;
    println!("✅ Document creation: {} created, {} failed", 
        create_result.created, create_result.failed.len());
    
    // Check collection health immediately
    println!("\n📊 Collection health after document creation:");
    let health = rag_module.collection_manager.get_collections_health().await?;
    println!("   Estate Collection: {} points", health.aws_estate.points_count);
    
    // Test direct vector store search (bypassing search service)
    println!("\n🔧 Testing direct vector store search...");
    let query_embedding = rag_module.embedding_service.generate_embedding("EC2 instance running").await?;
    
    use rag_module::types::SearchOptions;
    let search_options = SearchOptions {
        limit: Some(10),
        score_threshold: Some(0.0), // Very low threshold
        filter: None,
        with_payload: Some(true),
        ..Default::default()
    };
    
    let direct_results = rag_module.vector_store.search("aws_estate", query_embedding, search_options).await?;
    println!("✅ Direct vector store search returned {} results", direct_results.len());
    
    for (i, result) in direct_results.iter().take(3).enumerate() {
        println!("   {}. ID: {}, Score: {:.4}", i + 1, result.id, result.score);
        if let Some(doc) = &result.document {
            println!("      Content: {}...", &doc.content[..50.min(doc.content.len())]);
        }
    }
    
    // Test enhanced search service
    println!("\n🔧 Testing enhanced search service...");
    let estate_search_options = EstateSearchOptions {
        resource_types: Some(vec!["ec2-instance".to_string()]),
        account_ids: Some(vec!["123456789012".to_string()]),
        regions: Some(vec!["us-east-1".to_string()]),
        services: Some(vec!["ec2".to_string()]),
        states: Some(vec!["running".to_string()]),
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(10),
        score_threshold: Some(0.0), // Very low threshold
        include_metadata: true,
        use_anonymous_ids: false,
    };
    
    let search_results = rag_module.search_service
        .search_estate_resources("EC2 instance running", estate_search_options, None)
        .await?;
        
    println!("✅ Enhanced search service returned {} results", search_results.len());
    for (i, result) in search_results.iter().take(3).enumerate() {
        if let Some(resource_type) = result.get("resource_type") {
            if let Some(account_id) = result.get("account_id") {
                if let Some(state) = result.get("state") {
                    println!("   {}. {} in account {} - state: {}", 
                        i + 1, resource_type, account_id, state);
                }
            }
        }
    }
    
    Ok(())
}