rag-module 0.6.7

//! Test IAM Users Search

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

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    let home = std::env::var("HOME").expect("HOME environment variable not set");
    let base_path = format!("{}/.escher", home);
    let user_id = "a4d8c418-90e1-702a-64e8-604fc7b72f72";

    println!("🔍 Testing IAM Users Search\n");

    let rag = RagModule::new(&base_path).await?;
    rag.initialize().await?;
    rag.set_user_context(user_id).await?;

    let collection_name = "core_estate";
    
    // First, check what services are in the collection
    println!("📊 Checking services in collection...");
    let docs = rag.vector_store.list_documents(collection_name, Some(100), None).await?;
    
    let mut service_counts = std::collections::HashMap::new();
    for doc in &docs {
        if let Some(service) = doc.metadata.get("service") {
            if let Some(service_str) = service.as_str() {
                *service_counts.entry(service_str.to_string()).or_insert(0) += 1;
            }
        }
    }
    
    println!("Services found:");
    for (service, count) in service_counts.iter() {
        println!("  - {}: {} resources", service, count);
    }
    
    // Check for IAM-related resources
    let iam_docs: Vec<_> = docs.iter()
        .filter(|doc| {
            doc.metadata.get("service")
                .and_then(|s| s.as_str())
                .map(|s| s.to_lowercase().contains("iam"))
                .unwrap_or(false)
        })
        .collect();
    
    println!("\n📋 IAM resources found: {}", iam_docs.len());
    if iam_docs.len() > 0 {
        for (i, doc) in iam_docs.iter().take(5).enumerate() {
            println!("  {}. ID: {}", i + 1, doc.id);
            if let Some(resource_type) = doc.metadata.get("resourceType") {
                println!("     Type: {}", resource_type);
            }
        }
    }
    
    // Now try searching for IAM users
    println!("\n🔍 Searching for 'list my iam users'...");
    
    let estate_options = rag_module::services::search_service::EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: None,
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(10),
        score_threshold: Some(0.3), // Lower threshold
        include_metadata: true,
        use_anonymous_ids: false,
    };

    let results = rag.search_service.search_estate_resources(
        collection_name,
        "list my iam users",
        estate_options,
        None,
        user_id,
    ).await?;
    
    println!("\n📊 Search Results: {} found", results.len());
    for (i, result) in results.iter().enumerate() {
        if let Some(obj) = result.as_object() {
            println!("\n{}. Result:", i + 1);
            if let Some(score) = obj.get("score").and_then(|v| v.as_f64()) {
                println!("   Score: {:.4}", score);
            }
            if let Some(service) = obj.get("service").and_then(|v| v.as_str()) {
                println!("   Service: {}", service);
            }
            if let Some(resource_type) = obj.get("resourceType").and_then(|v| v.as_str()) {
                println!("   Resource Type: {}", resource_type);
            }
        }
    }

    Ok(())
}