use anyhow::Result;
use rag_module::{RagModule, services::search_service::EstateSearchOptions};
use serde_json::json;
use std::fs;
use std::path::Path;
#[tokio::main]
async fn main() -> Result<()> {
println!("🚀 AWS Estate Data Storage Test with Encryption\n");
let aws_estate_data = json!([
{
"account_id": "288761761556",
"role_arn": null,
"account_name": "",
"scan_timestamp": null,
"services": {
"ec2": {
"instances": [
{
"instance_id": "i-08619fa3385e63bab",
"instance_type": "c4.large",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-08-26T12:37:51+00:00",
"name": "escher-scan-stream-v2"
},
{
"instance_id": "i-0f6ce804a1a464e1d",
"instance_type": "c4.large",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-08-23T06:45:12+00:00",
"name": "slack-ec2-stream"
},
{
"instance_id": "i-01697357246074533",
"instance_type": "t2.micro",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-09-05T07:06:47+00:00",
"name": "space-argo-dev-inst-1"
},
{
"instance_id": "i-0b1e8b61d9b846a3e",
"instance_type": "t2.micro",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-07-21T04:22:19+00:00",
"name": "escher-ec2"
},
{
"instance_id": "i-0a5ff673a52c7a1c7",
"instance_type": "c4.large",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-07-28T05:05:24+00:00",
"name": "escher-stream-ec2"
},
{
"instance_id": "i-06a446d8aaec15d27",
"instance_type": "c4.large",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-08-29T04:11:05+00:00",
"name": "finops-ec2-stream"
},
{
"instance_id": "i-07901ab76d1aca487",
"instance_type": "c1.xlarge",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-09-29T10:42:10+00:00",
"name": "dev-benchmark"
}
],
"volumes_count": 13,
"elastic_ips_total": 7,
"elastic_ips_attached": 3,
"elastic_ips_unattached": 4,
"elastic_ips_instance_attached": 0,
"elastic_ips_eni_attached": 3,
"vpcs_count": 7
},
"rds": {
"instances": [
{
"instance_id": "db-VXDRDM3CQ6OIKGF7HLGVPPO454",
"db_instance_identifier": "dev-eshop-mysql-rds",
"db_instance_class": "db.t3.micro",
"engine": "mysql",
"db_instance_status": "available",
"region": "us-east-1",
"name": "dev-eshop-mysql-rds"
},
{
"instance_id": "db-2XYHZFX3PDXT3CD3LGIK435ALY",
"db_instance_identifier": "dev-tesserade-pg-rds",
"db_instance_class": "db.m5.xlarge",
"engine": "postgres",
"db_instance_status": "available",
"region": "us-east-1",
"name": "dev-tesserade-pg-rds"
},
{
"instance_id": "db-QTKHKSX6HCJBUAIZSSTOLAHQMY",
"db_instance_identifier": "mysql-test-instance1",
"db_instance_class": "db.t3.micro",
"engine": "mysql",
"db_instance_status": "available",
"region": "us-east-1",
"name": "mysql-test-instance1"
}
]
},
"s3": {
"total_buckets": 23,
"public_buckets": 0,
"latest_buckets_info": [
{
"name": "escher-bucket-test-dev",
"s3_identifier": "escher-bucket-test-dev",
"creation_date": "2025-10-03T05:05:37+00:00",
"region": "us-east-1",
"is_public": false
},
{
"name": "test-escher-prod",
"s3_identifier": "test-escher-prod",
"creation_date": "2025-10-01T08:44:14+00:00",
"region": "us-east-1",
"is_public": false
}
]
},
"lambda": {
"functions": [
{
"name": "eShopping_User_Simulator",
"lambda_identifier": "eShopping_User_Simulator",
"runtime": "python3.13",
"memory_size": 128,
"timeout": 3,
"region": "us-east-1"
},
{
"name": "prod-escher-storage-service",
"lambda_identifier": "prod-escher-storage-service",
"runtime": "provided.al2",
"memory_size": 256,
"timeout": 30,
"region": "us-east-1"
}
]
}
}
},
{
"account_id": "547639878546",
"role_arn": "arn:aws:iam::547639878546:role/escher_cross_role",
"account_name": "",
"scan_timestamp": null,
"services": {
"ec2": {
"instances": [
{
"instance_id": "i-0f9f85890fc371f78",
"instance_type": "t2.micro",
"state": "running",
"region": "us-east-1",
"launch_time": "2025-08-22T18:19:11+00:00",
"name": "i-0f9f85890fc371f78"
}
]
},
"s3": {
"latest_buckets_info": [
{
"name": "aparna-dev-test",
"s3_identifier": "aparna-dev-test",
"creation_date": "2025-10-03T07:16:28+00:00",
"region": "us-east-1",
"is_public": false
},
{
"name": "tex-instru-test-bckt",
"s3_identifier": "tex-instru-test-bckt",
"creation_date": "2025-07-24T04:50:07+00:00",
"region": "us-east-1",
"is_public": false
}
]
}
}
}
]);
store_aws_estate_data(aws_estate_data).await?;
Ok(())
}
async fn store_aws_estate_data(aws_estate_data: serde_json::Value) -> Result<()> {
let mut rag_module = RagModule::new("./example-enhanced-chat-data").await?;
println!("🔧 Initializing RAG module with encryption...");
rag_module.initialize().await?;
println!("✅ RAG module initialized successfully!\n");
println!("📦 Processing AWS Estate Data...");
let mut total_resources_stored = 0;
let user_id = "cudewvsbxakj1n";
println!("Processing complete AWS estate data structure for user: {}...", user_id);
let document_ids = rag_module.process_aws_estate(aws_estate_data.clone(), user_id).await?;
total_resources_stored = document_ids.len();
if let Some(accounts) = aws_estate_data.as_array() {
for account in accounts {
if let Some(account_obj) = account.as_object() {
let account_id = account_obj.get("account_id")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
println!("\n🏢 Processing Account: {}", account_id);
if let Some(services) = account_obj.get("services").and_then(|s| s.as_object()) {
if let Some(ec2) = services.get("ec2").and_then(|e| e.as_object()) {
if let Some(instances) = ec2.get("instances").and_then(|i| i.as_array()) {
println!(" 📱 Processing {} EC2 instances...", instances.len());
for instance in instances {
if let Some(instance_obj) = instance.as_object() {
let instance_id = instance_obj.get("instance_id").and_then(|v| v.as_str()).unwrap_or("unknown");
println!(" ✅ Processed EC2: {}", instance_id);
}
}
}
}
if let Some(rds) = services.get("rds").and_then(|r| r.as_object()) {
if let Some(instances) = rds.get("instances").and_then(|i| i.as_array()) {
println!(" 🗄️ Processing {} RDS instances...", instances.len());
for db in instances {
if let Some(db_obj) = db.as_object() {
let db_instance_identifier = db_obj.get("db_instance_identifier").and_then(|v| v.as_str()).unwrap_or("unknown");
println!(" ✅ Processed RDS: {}", db_instance_identifier);
}
}
}
}
if let Some(s3) = services.get("s3").and_then(|s| s.as_object()) {
if let Some(buckets) = s3.get("latest_buckets_info").and_then(|b| b.as_array()) {
println!(" 📁 Processing {} S3 buckets...", buckets.len());
for bucket in buckets {
if let Some(bucket_obj) = bucket.as_object() {
let name = bucket_obj.get("name").and_then(|v| v.as_str()).unwrap_or("unknown");
println!(" ✅ Processed S3: {}", name);
}
}
}
}
if let Some(lambda) = services.get("lambda").and_then(|l| l.as_object()) {
if let Some(functions) = lambda.get("functions").and_then(|f| f.as_array()) {
println!(" ⚡ Processing {} Lambda functions...", functions.len());
for function in functions {
if let Some(lambda_obj) = function.as_object() {
let name = lambda_obj.get("name").and_then(|v| v.as_str()).unwrap_or("unknown");
println!(" ✅ Processed Lambda: {}", name);
}
}
}
}
}
}
}
}
println!("\n🎉 Successfully stored {} AWS resources with encryption!", total_resources_stored);
println!("\n🔍 Testing Search Functionality...");
println!("Test 1: Basic search with low threshold...");
let estate_search_options = 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.0), include_metadata: false,
use_anonymous_ids: true,
};
println!(" Generating embedding for query 'EC2'...");
let query_embedding = rag_module.embedding_service.generate_embedding("EC2 running instances").await?;
println!(" Generated embedding with {} dimensions", query_embedding.len());
let search_results = rag_module.search_service.search_estate_resources("EC2", estate_search_options, None).await?;
println!("✅ Found {} results for \"EC2\" with zero threshold", search_results.len());
for (i, result) in search_results.iter().take(3).enumerate() {
println!("\n📋 Result {}: ID={}", i + 1, result.get("id").unwrap_or(&serde_json::Value::Null));
if let Some(service) = result.get("service") {
println!(" 🔧 Service: {}", service);
}
if let Some(account_id) = result.get("account_id") {
println!(" 🏢 Account ID: {}", account_id);
}
if let Some(account_name) = result.get("account_name") {
println!(" 📛 Account Name: {}", account_name);
}
if let Some(content) = result.get("content") {
if let Some(content_str) = content.as_str() {
println!(" 📄 Decrypted Content:");
if let Ok(json_content) = serde_json::from_str::<serde_json::Value>(content_str) {
let pretty_json = serde_json::to_string_pretty(&json_content).unwrap_or_else(|_| content_str.to_string());
let display_content = if pretty_json.len() > 500 {
format!("{}...\n [Content truncated - {} total chars]", &pretty_json[..500], pretty_json.len())
} else {
pretty_json
};
println!(" {}", display_content.replace('\n', "\n "));
} else {
let display_content = if content_str.len() > 200 {
format!("{}...", &content_str[..200])
} else {
content_str.to_string()
};
println!(" {}", display_content);
}
}
}
println!(" ─────────────────────────────────────────────────────");
}
println!("Test 2: Search with EC2 filter...");
let estate_search_options_filtered = EstateSearchOptions {
resource_types: Some(vec!["instance".to_string()]),
account_ids: None,
regions: None,
services: Some(vec!["ec2".to_string()]),
states: Some(vec!["running".to_string()]),
environment: None,
application: None,
synced_after: None,
limit: Some(5),
score_threshold: Some(0.01), include_metadata: false,
use_anonymous_ids: true,
};
let filtered_results = rag_module.search_service.search_estate_resources("instance", estate_search_options_filtered, None).await?;
println!("✅ Found {} results for filtered EC2 instance search", filtered_results.len());
println!("Test 3: Check encrypted documents and vector database...");
let all_encrypted_docs = rag_module.get_encrypted_aws_estate_documents().await?;
println!("✅ Total encrypted estate documents: {}", all_encrypted_docs.len());
println!("Test 4: Direct vector database search...");
let search_options = rag_module::types::SearchOptions {
limit: Some(10),
score_threshold: Some(0.0), with_payload: Some(true),
filter: None,
..Default::default()
};
let direct_results = rag_module.vector_store.search("aws_estate", query_embedding, search_options).await?;
println!("✅ Direct vector search found {} results", direct_results.len());
for (i, result) in direct_results.iter().take(2).enumerate() {
println!(" Direct Result {}: ID={}, Score={:.4}", i + 1, result.id, result.score);
if let Some(payload) = &result.payload {
if let Some(service) = payload.get("service") {
println!(" Service: {}", service);
}
if let Some(account_id) = payload.get("account_id") {
println!(" Account: {}", account_id);
}
}
}
println!("\n📊 RAG Module Statistics:");
let encrypted_docs = rag_module.get_encrypted_aws_estate_documents().await?;
println!(" 📋 Total Estate Documents: {}", encrypted_docs.len());
println!(" 🔒 Encryption Enabled: true");
println!(" 📁 Base Path: ./test-estate-data");
println!("\n📁 Checking stored file structure...");
let data_path = "./test-estate-data/qdrant-data";
if Path::new(data_path).exists() {
println!("✅ Qdrant data directory created");
let estate_docs_path = format!("{}/aws_estate-metadata.json", data_path);
if Path::new(&estate_docs_path).exists() {
println!("✅ AWS estate metadata file created");
if let Ok(content) = fs::read_to_string(&estate_docs_path) {
let preview = if content.len() > 100 {
format!("{}...", &content[..100])
} else {
content
};
println!(" 🔒 Estate metadata preview (encrypted): {}", preview);
}
}
let segments_path = format!("{}/segments", data_path);
if Path::new(&segments_path).exists() {
println!("✅ Vector segments directory created");
}
}
println!("\n🏁 Test completed");
Ok(())
}