rag-module 0.6.7

use rag_module::RagModule;
use uuid::Uuid;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    println!("=== Request ID Pairing Example ===\n");

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

    let context_id = "demo-conversation";
    let user_id = "user123";
    let chat_title = "AWS Infrastructure Questions";

    println!("Starting conversation with context_id: {}\n", context_id);

    // ==========================================
    // Conversation Turn 1
    // ==========================================

    // Client generates a unique request_id for this prompt-response pair
    let request_id_1 = Uuid::new_v4().to_string();
    println!("🔹 Turn 1 - Request ID: {}", request_id_1);

    // Add user prompt with request_id
    let prompt_1 = "How many EC2 instances do we have?";
    rag.add_prompt(
        context_id,
        prompt_1,
        user_id,
        Some(chat_title),
        &request_id_1  // Client sends request_id
    ).await?;
    println!("   User: {}", prompt_1);

    // Add assistant response with the SAME request_id
    let response_1 = "You have 16 EC2 instances running in ap-south-1 region.";
    rag.add_response(
        context_id,
        response_1,
        user_id,
        Some(chat_title),
        &request_id_1  // Same request_id links them together
    ).await?;
    println!("   Assistant: {}\n", response_1);

    // ==========================================
    // Conversation Turn 2
    // ==========================================

    // Client generates a NEW request_id for the next pair
    let request_id_2 = Uuid::new_v4().to_string();
    println!("🔹 Turn 2 - Request ID: {}", request_id_2);

    let prompt_2 = "Which ones are running PostgreSQL?";
    rag.add_prompt(
        context_id,
        prompt_2,
        user_id,
        Some(chat_title),
        &request_id_2
    ).await?;
    println!("   User: {}", prompt_2);

    let response_2 = "13 instances are running PostgreSQL DR configurations.";
    rag.add_response(
        context_id,
        response_2,
        user_id,
        Some(chat_title),
        &request_id_2
    ).await?;
    println!("   Assistant: {}\n", response_2);

    // ==========================================
    // Conversation Turn 3
    // ==========================================

    let request_id_3 = Uuid::new_v4().to_string();
    println!("🔹 Turn 3 - Request ID: {}", request_id_3);

    let prompt_3 = "What about unattached IAM policies?";
    rag.add_prompt(
        context_id,
        prompt_3,
        user_id,
        Some(chat_title),
        &request_id_3
    ).await?;
    println!("   User: {}", prompt_3);

    let response_3 = "You have 16 IAM policies not attached to any user or role.";
    rag.add_response(
        context_id,
        response_3,
        user_id,
        Some(chat_title),
        &request_id_3
    ).await?;
    println!("   Assistant: {}\n", response_3);

    println!("✅ All messages stored with request_id linking!\n");

    // Now retrieve and show the stored format
    println!("=== Stored Format ===");
    println!("Each prompt-response pair shares the same 'ri' (request_id):");
    println!("Message 0 (User):      ri = {}", request_id_1);
    println!("Message 1 (Assistant): ri = {} (linked!)", request_id_1);
    println!("Message 2 (User):      ri = {}", request_id_2);
    println!("Message 3 (Assistant): ri = {} (linked!)", request_id_2);
    println!("Message 4 (User):      ri = {}", request_id_3);
    println!("Message 5 (Assistant): ri = {} (linked!)", request_id_3);

    Ok(())
}