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.
#![recursion_limit = "1"]

use rag_module::services::search_service::{SearchService, EstateSearchOptions};
use rag_module::{create_rag_module};
use serde_json::json;
use anyhow::Result;

#[tokio::main]
async fn main() -> Result<()> {
    println!("🎯 Parameter Filtering Demo");
    println!("This demo shows how to use the new 'parameters' field to filter returned metadata");

    let rag_module = create_rag_module("./test_parameter_filtering").await?;
    rag_module.initialize().await?;

    let user_id = "demo_user";
    let collection_name = "aws_estate";
    
    // Add some test estate data
    let test_estate_data = json!({
        "type": "rds",
        "keywords": [
          "rds",
          "database",
          "relational"
        ],
        "profile": "default",
        "accountId": "288761761556",
        "region": "us-east-1",
        "permissions": {
          "accessLevel": "FullAccess",
          "hasRead": true,
          "hasWrite": true,
          "hasFullAccess": true
        },
        "content": "RDS instance dev-eshop-mysql-rds, id dev-eshop-mysql-rds, profile default, region us-east-1, engine mysql 8.0.42, class db.t3.micro, state stopped, storage 20GB, created 2025-07-07T10:02:58.453000+00:00",
        "dbInstanceIdentifier": "dev-eshop-mysql-rds",
        "dbInstanceClass": "db.t3.micro",
        "engine": "mysql",
        "dbInstanceStatus": "stopped",
        "automaticRestartTime": "2025-12-08T11:37:25.225000+00:00",
        "masterUsername": "mysql_admin",
        "endpoint": {
          "address": "dev-eshop-mysql-rds.cuxmiwm0ulok.us-east-1.rds.amazonaws.com",
          "port": 3306,
          "hostedZoneId": "Z2R2ITUGPM61AM"
        },
        "allocatedStorage": 20,
        "instanceCreateTime": "2025-07-07T10:02:58.453000+00:00",
        "preferredBackupWindow": "10:27-10:57",
        "backupRetentionPeriod": 1,
        "dbSecurityGroups": [],
        "vpcSecurityGroups": [
          {
            "vpcSecurityGroupId": "sg-0a017fbb383b24395",
            "status": "active"
          }
        ],
        "dbParameterGroups": [
          {
            "dbParameterGroupName": "default.mysql8.0",
            "parameterApplyStatus": "in-sync"
          }
        ],
        "availabilityZone": "us-east-1f",
        "dbSubnetGroup": {
          "dbSubnetGroupName": "default",
          "dbSubnetGroupDescription": "default",
          "vpcId": "vpc-0a82df6f3f43f638f",
          "subnetGroupStatus": "Complete",
          "subnets": [
            {
              "subnetIdentifier": "subnet-031538e931e215314",
              "subnetAvailabilityZone": {
                "name": "us-east-1b"
              },
              "subnetOutpost": {},
              "subnetStatus": "Active"
            },
            {
              "subnetIdentifier": "subnet-0a3f5c2eaf6cf3628",
              "subnetAvailabilityZone": {
                "name": "us-east-1f"
              },
              "subnetOutpost": {},
              "subnetStatus": "Active"
            },
            {
              "subnetIdentifier": "subnet-07a5e332fba1aec62",
              "subnetAvailabilityZone": {
                "name": "us-east-1a"
              },
              "subnetOutpost": {},
              "subnetStatus": "Active"
            },
            {
              "subnetIdentifier": "subnet-077a94e2e77fbd626",
              "subnetAvailabilityZone": {
                "name": "us-east-1c"
              },
              "subnetOutpost": {},
              "subnetStatus": "Active"
            },
            {
              "subnetIdentifier": "subnet-095aa065d3815957a",
              "subnetAvailabilityZone": {
                "name": "us-east-1e"
              },
              "subnetOutpost": {},
              "subnetStatus": "Active"
            },
            {
              "subnetIdentifier": "subnet-03344c284627f58a8",
              "subnetAvailabilityZone": {
                "name": "us-east-1d"
              },
              "subnetOutpost": {},
              "subnetStatus": "Active"
            }
          ]
        },
        "preferredMaintenanceWindow": "thu:04:36-thu:05:06",
        "pendingModifiedValues": {},
        "latestRestorableTime": "2025-12-01T11:25:00+00:00",
        "multiAZ": false,
        "engineVersion": "8.0.42",
        "autoMinorVersionUpgrade": true,
        "readReplicaDBInstanceIdentifiers": [],
        "licenseModel": "general-public-license",
        "optionGroupMemberships": [
          {
            "optionGroupName": "default:mysql-8-0",
            "status": "in-sync"
          }
        ],
        "publiclyAccessible": true,
        "storageType": "gp2",
        "dbInstancePort": 0,
        "storageEncrypted": true,
        "kmsKeyId": "arn:aws:kms:us-east-1:288761761556:key/3469cdf7-b2dc-4d99-8dba-816a95e9465a",
        "dbiResourceId": "db-VXDRDM3CQ6OIKGF7HLGVPPO454",
        "caCertificateIdentifier": "rds-ca-rsa2048-g1",
        "domainMemberships": [],
        "copyTagsToSnapshot": true,
        "monitoringInterval": 0,
        "dbInstanceArn": "arn:aws:rds:us-east-1:288761761556:db:dev-eshop-mysql-rds",
        "iamDatabaseAuthenticationEnabled": false,
        "databaseInsightsMode": "standard",
        "performanceInsightsEnabled": false,
        "deletionProtection": false,
        "associatedRoles": [],
        "tagList": [
          {
            "key": "app",
            "value": "e-shopping"
          },
          {
            "key": "environment",
            "value": "dev"
          }
        ],
        "customerOwnedIpEnabled": false,
        "activityStreamStatus": "stopped",
        "backupTarget": "region",
        "networkType": "IPV4",
        "storageThroughput": 0,
        "certificateDetails": {
          "caIdentifier": "rds-ca-rsa2048-g1",
          "validTill": "2026-07-09T06:40:17+00:00"
        },
        "dedicatedLogVolume": false,
        "isStorageConfigUpgradeAvailable": false,
        "engineLifecycleSupport": "open-source-rds-extended-support-disabled"
    });

    println!("\n📊 Ingesting test estate data...");
    rag_module.ingest_aws_estate(test_estate_data, user_id, collection_name).await?;
    
    // Wait a moment for processing
    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;

    println!("\n🔍 Test 1: Search with ALL fields (parameters = None)");
    let options_all = EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: None,
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(1),
        score_threshold: Some(0.0),
        include_metadata: true,
        use_anonymous_ids: false,
        parameters: None, // Return all fields
    };

    let results_all = rag_module.search_service
        .search_estate_resources(collection_name, "EC2 instance", options_all, None, user_id)
        .await?;

    println!("✅ Found {} results with ALL fields", results_all.len());
    if let Some(result) = results_all.first() {
        println!("📋 All fields returned:");
        for (key, _value) in result.as_object().unwrap() {
            if key != "score" && key != "id" && key != "created_at" && key != "updated_at" {
                println!("  - {}", key);
            }
        }
    }

    println!("\n🎯 Test 2: Search with SPECIFIC parameters only");
    let specific_params = vec![
        "instanceId".to_string(),
        "name".to_string(),
        "instanceClass".to_string(),
        "region".to_string(),
        "status".to_string(),
        "launchTime".to_string(),
        "dnsName".to_string(),
        "tags".to_string(),
        "storage".to_string(),
        "storageType".to_string(),
    ];

    let options_filtered = EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: None,
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(1),
        score_threshold: Some(0.0),
        include_metadata: true,
        use_anonymous_ids: false,
        parameters: Some(specific_params.clone()), // Only return specified fields
    };

    let results_filtered = rag_module.search_service
        .search_estate_resources(collection_name, "EC2 instance", options_filtered, None, user_id)
        .await?;

    println!("✅ Found {} results with FILTERED fields", results_filtered.len());
    if let Some(result) = results_filtered.first() {
        println!("📋 Filtered fields returned:");
        for (key, value) in result.as_object().unwrap() {
            if key != "score" && key != "id" && key != "created_at" && key != "updated_at" {
                println!("  - {}: {:?}", key, value);
            }
        }
        
        // Verify that extra fields are NOT included
        let extra_fields = vec!["extraField1", "extraField2", "extraField3"];
        let found_extra = extra_fields.iter().any(|field| result.get(field).is_some());
        
        if found_extra {
            println!("❌ ERROR: Extra fields found when they should be filtered out!");
        } else {
            println!("✅ SUCCESS: Extra fields properly filtered out");
        }
    }

    println!("\n🎯 Test 3: Search with MINIMAL parameters");
    let minimal_params = vec![
        "instanceId".to_string(),
        "name".to_string(),
        "endpoint.port".to_string(),
    ];

    let options_minimal = EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: None,
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(1),
        score_threshold: Some(0.0),
        include_metadata: true,
        use_anonymous_ids: false,
        parameters: Some(minimal_params.clone()), // Only return 3 fields
    };

    let results_minimal = rag_module.search_service
        .search_estate_resources(collection_name, "EC2 instance", options_minimal, None, user_id)
        .await?;

    println!("✅ Found {} results with MINIMAL fields", results_minimal.len());
    if let Some(result) = results_minimal.first() {
        println!("📋 Minimal fields returned:");
        for (key, value) in result.as_object().unwrap() {
            if key != "score" && key != "id" && key != "created_at" && key != "updated_at" {
                println!("  - {}: {:?}", key, value);
            }
        }
        
        // Count non-system fields
        let field_count = result.as_object().unwrap().iter()
            .filter(|(key, _)| !["score", "id", "created_at", "updated_at"].contains(&key.as_str()))
            .count();
        
        if field_count == 3 {
            println!("✅ SUCCESS: Exactly 3 fields returned as expected");
        } else {
            println!("❌ ERROR: Expected 3 fields, got {}", field_count);
        }
    }

    println!("\n🎉 Parameter filtering demo completed!");
    println!("💡 Key takeaways:");
    println!("   - Set 'parameters: None' to get all fields (default behavior)");
    println!("   - Set 'parameters: Some(vec![...])' to get only specific fields");
    println!("   - Field filtering happens after decryption, from encrypted metadata");
    println!("   - System fields (id, score, timestamps) are always included");

    Ok(())
}