use anyhow::Result;
use rag_module::RagModule;
use rag_module::services::search_service::{EstateSearchOptions, ChatSearchOptions};
use rag_module::StartSessionOptions;
#[tokio::main]
async fn main() -> Result<()> {
println!("🔐 TESTING: Decrypt-First Search Implementation");
println!("=============================================\n");
println!("🏢 TEST 1: Estate Search with Region Filter (Decrypt-First)");
let mut estate_rag = RagModule::new("./example-enhanced-estate-data").await?;
estate_rag.initialize().await?;
estate_rag.set_user_context("estate_user_example").await?;
let health = estate_rag.collection_manager.get_collections_health().await?;
println!("📊 Estate Collection: {} points", health.aws_estate.points_count);
let estate_search_options = EstateSearchOptions {
resource_types: None,
account_ids: None,
regions: Some(vec!["us-east-1".to_string()]), 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,
};
println!("\n🔍 Testing estate search WITH region filter (decrypt-first approach)...");
let estate_results = estate_rag.search_service
.search_estate_resources("AWS EC2 instances", estate_search_options, None)
.await?;
println!("✅ Estate Search Results (With Region Filter): {}", estate_results.len());
for (i, result) in estate_results.iter().take(3).enumerate() {
if let (Some(resource_type), Some(region), Some(account)) = (
result.get("resource_type").and_then(|v| v.as_str()),
result.get("region").and_then(|v| v.as_str()),
result.get("account_id").and_then(|v| v.as_str()),
) {
let score = result.get("score").and_then(|v| v.as_f64()).unwrap_or(0.0);
println!(" {}. {} in {} ({}) - Score: {:.4}",
i + 1, resource_type, region, account, score);
}
}
let multi_filter_options = EstateSearchOptions {
resource_types: Some(vec!["ec2-instance".to_string()]),
account_ids: None,
regions: Some(vec!["us-east-1".to_string()]),
services: Some(vec!["ec2".to_string()]),
states: None,
environment: None,
application: None,
synced_after: None,
limit: Some(5),
score_threshold: Some(0.1),
include_metadata: true,
use_anonymous_ids: false,
};
println!("\n🔍 Testing estate search WITH multiple filters...");
let multi_results = estate_rag.search_service
.search_estate_resources("EC2 instances", multi_filter_options, None)
.await?;
println!("✅ Estate Search Results (Multiple Filters): {}", multi_results.len());
println!("\n\n💬 TEST 2: Chat Search with Context Filter (Decrypt-First)");
let mut chat_rag = RagModule::new("./example-enhanced-chat-data").await?;
chat_rag.initialize().await?;
chat_rag.set_user_context("chat_user_example").await?;
let session_options = StartSessionOptions {
user_id: "test_user_decrypt".to_string(),
chat_title: Some("Decrypt Test Session".to_string()),
context_id: Some("decrypt_test_context".to_string()),
};
let session = chat_rag.start_session(session_options).await?;
println!("✅ Created session: {} (context: {})", session.id, session.context_id);
let conversations = vec![
("What is machine learning?", "Machine learning is AI that learns from data."),
("How does neural network work?", "Neural networks use interconnected nodes to process information."),
("Explain deep learning", "Deep learning uses multi-layer neural networks for complex pattern recognition."),
];
let conversation_count = conversations.len();
for (prompt, response) in conversations {
let _prompt_id = chat_rag.add_prompt(&session.id, prompt, "test_user_decrypt").await?;
let _response_id = chat_rag.add_response(&session.id, response, "test_user_decrypt").await?;
}
println!("✅ Added {} conversation pairs", conversation_count);
let chat_search_options = ChatSearchOptions {
context_id: Some(session.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,
};
println!("\n🔍 Testing chat search WITH context filter (decrypt-first approach)...");
let chat_results = chat_rag.search_service
.search_chat_history(chat_search_options)
.await?;
println!("✅ Chat Search Results (With Context Filter): {}", chat_results.len());
for (i, result) in chat_results.iter().take(3).enumerate() {
if let (Some(role), Some(content)) = (
result.get("role").and_then(|v| v.as_str()),
result.get("content").and_then(|v| v.as_str()),
) {
let preview = if content.len() > 50 {
format!("{}...", &content[..50])
} else {
content.to_string()
};
println!(" {}. {}: {}", i + 1, role, preview);
}
}
let role_filter_options = ChatSearchOptions {
context_id: Some(session.context_id.clone()),
role: Some("user".to_string()), from_timestamp: None,
to_timestamp: None,
from_message_index: None,
to_message_index: None,
limit: Some(10),
include_metadata: true,
};
println!("\n🔍 Testing chat search WITH role filter...");
let role_results = chat_rag.search_service
.search_chat_history(role_filter_options)
.await?;
println!("✅ Chat Search Results (User Messages Only): {}", role_results.len());
println!("\n\n🎉 DECRYPT-FIRST SEARCH VERIFICATION COMPLETED!");
println!("===============================================");
println!("✅ Estate Search with Region Filter: {} results", estate_results.len());
println!("✅ Estate Search with Multiple Filters: {} results", multi_results.len());
println!("✅ Chat Search with Context Filter: {} results", chat_results.len());
println!("✅ Chat Search with Role Filter: {} results", role_results.len());
if estate_results.len() > 0 && chat_results.len() > 0 {
println!("\n🚀 SUCCESS: Both search functions now work with encrypted metadata!");
println!(" - Decrypt-first approach successfully implemented");
println!(" - Filters work on decrypted data");
println!(" - Universal solution for all collections");
} else {
println!("\n⚠️ Some searches returned 0 results - check implementation");
}
Ok(())
}