rag-module 0.6.7

//! Example to retrieve and display documents from our batch test data
//! 
//! This demonstrates how to access the documents we created with batch ingestion (up to 1).

use anyhow::Result;
use rag_module::RagModule;

#[tokio::main]
async fn main() -> Result<()> {
    println!("🔍 Retrieving documents from batch test data...\n");

    // Initialize RAG module with the batch test data directory
    let rag = RagModule::new("./batch-test-data").await?;
    rag.initialize().await?;

    let user_id = "batch_test_user";
    let collection_name = "aws_estate";

    // Set user context before retrieving documents
    rag.set_user_context(user_id).await?;

    // Method 1: Get all AWS estate documents
    println!("📦 Fetching AWS estate documents from batch test data...");
    let aws_docs = rag.list_documents(Some(collection_name), None).await?;
    println!("✅ Retrieved {} AWS estate documents\n", aws_docs.len());

    // Display summary of what we have
    if !aws_docs.is_empty() {
        println!("📊 Document Summary:");
        println!("═══════════════════════");
        
        let mut ec2_count = 0;
        let mut s3_count = 0; 
        let mut rds_count = 0;
        let mut lambda_count = 0;
        let mut other_count = 0;

        for (i, doc) in aws_docs.iter().enumerate() {
            // Categorize based on content since metadata is encrypted
            let content = &doc.content;
            if content.contains("EC2 instance") {
                ec2_count += 1;
            } else if content.contains("S3 bucket") {
                s3_count += 1;
            } else if content.contains("RDS") || content.contains("database") {
                rds_count += 1;
            } else if content.contains("Lambda function") {
                lambda_count += 1;
            } else {
                other_count += 1;
            }

            // Show first few documents in detail
            if i < 3 {
                println!("\n📄 Document {} (ID: {}):", i + 1, doc.id);
                println!("   Content: {}", 
                    if doc.content.len() > 100 { 
                        format!("{}...", &doc.content[..100])
                    } else {
                        doc.content.clone()
                    }
                );
                println!("   Created: {}", doc.created_at);
            }
        }

        println!("\n🏷️  Resource Type Summary:");
        println!("   🖥️  EC2 Instances: {}", ec2_count);
        println!("   🗄️  S3 Buckets: {}", s3_count);
        println!("   🗃️  RDS Databases: {}", rds_count);
        println!("   ⚡ Lambda Functions: {}", lambda_count);
        if other_count > 0 {
            println!("   📋 Other Resources: {}", other_count);
        }
    }

    // Method 2: Test search functionality on the batch data
    println!("\n🔍 Testing search functionality on batch data...");
    
    let search_options = rag_module::SearchOptions {
        limit: Some(5),
        score_threshold: Some(0.1),
        ..Default::default()
    };

    // Search for EC2 instances
    let ec2_results = rag.search(collection_name, "EC2 instance production", user_id, search_options.clone()).await?;
    println!("🖥️  EC2 search results: {} matches", ec2_results.len());
    
    for (i, result) in ec2_results.iter().enumerate() {
        if let Some(doc) = &result.document {
            let content = &doc.content;
            println!("   {}. Score: {:.3} - {}", 
                     i + 1, 
                     result.score, 
                     if content.len() > 80 { 
                         format!("{}...", &content[..80])
                     } else {
                         content.clone()
                     });
        }
    }

    // Search for databases
    let db_results = rag.search(collection_name, "database MySQL PostgreSQL", user_id, search_options.clone()).await?;
    println!("\n🗃️  Database search results: {} matches", db_results.len());
    
    for (i, result) in db_results.iter().enumerate() {
        if let Some(doc) = &result.document {
            let content = &doc.content;
            println!("   {}. Score: {:.3} - {}", 
                     i + 1, 
                     result.score,
                     if content.len() > 80 { 
                         format!("{}...", &content[..80])
                     } else {
                         content.clone()
                     });
        }
    }

    // Search for Lambda functions  
    let lambda_results = rag.search(collection_name, "Lambda function python", user_id, search_options).await?;
    println!("\n⚡ Lambda search results: {} matches", lambda_results.len());
    
    for (i, result) in lambda_results.iter().enumerate() {
        if let Some(doc) = &result.document {
            let content = &doc.content;
            println!("   {}. Score: {:.3} - {}", 
                     i + 1, 
                     result.score,
                     if content.len() > 80 { 
                         format!("{}...", &content[..80])
                     } else {
                         content.clone()
                     });
        }
    }

    println!("\n✅ Batch test document retrieval completed!");
    println!("💡 Your batch ingestion created {} searchable documents", aws_docs.len());
    
    Ok(())
}