prvng_platform/crates/rag/examples/rag_agent.rs

//! Complete RAG Agent Example
//!
//! Demonstrates the full RAG pipeline:
//! 1. Ingest documents (chunking + embedding + storage)
//! 2. Create RAG agent with workspace context
//! 3. Answer questions using semantic search
//!
//! Run with: `cargo run --example rag_agent`

#![allow(clippy::field_reassign_with_default)]

use provisioning_rag::{
    config::{EmbeddingConfig, IngestionConfig, RetrievalConfig, VectorDbConfig},
    context::WorkspaceContext,
    db::DbConnection,
    embeddings::EmbeddingEngine,
    ingestion::DocumentIngester,
    retrieval::RetrieverEngine,
    RagAgent,
};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    tracing_subscriber::fmt::init();

    println!("🚀 RAG Agent System - Complete Example\n");

    // ==== PHASE 1: Setup Database ====
    println!("Phase 1: Setting up SurrealDB connection...");
    let mut db_config = VectorDbConfig::default();
    db_config.url = "memory".to_string();
    db_config.database = "rag_agent_demo".to_string();

    let db = DbConnection::new(db_config).await?;
    db.initialize_schema().await?;
    println!("✓ Database ready\n");

    // ==== PHASE 2: Ingest Documents ====
    println!("Phase 2: Ingesting sample documentation...");
    let mut embedding_config = EmbeddingConfig::default();
    embedding_config.provider = "local".to_string();
    let embedding_engine = EmbeddingEngine::new(embedding_config.clone())?;

    let ingestion_config = IngestionConfig::default();
    let ingester = DocumentIngester::new(ingestion_config, embedding_engine);

    // Sample documentation
    let architecture_doc = r#"
# Provisioning Platform Architecture

## Overview
The provisioning platform is a unified infrastructure automation system. It provides:
- Multi-cloud support (AWS, UpCloud, local)
- Kubernetes cluster management
- Automatic resource provisioning
- Configuration management via KCL

## Core Components

### Orchestrator
Central task coordination service. Manages all infrastructure operations.
Features:
- REST API for task submission
- Workflow management
- Task scheduling and execution
- Progress tracking

### Control Center
Web-based management interface for monitoring deployments and managing infrastructure.

### MCP Server
Model Context Protocol integration for AI-powered operations and context-aware assistance.
"#;

    let deployment_doc = r#"
# Deployment Guide

## Prerequisites
- Workspace configured with infrastructure target
- Cloud provider credentials set up
- Required taskservs installed (kubernetes, containerd, etc.)

## Deployment Steps

### 1. Create Infrastructure
```bash
provisioning server create --count 3 --size medium
```

### 2. Install Taskservs
```bash
provisioning taskserv create kubernetes
provisioning taskserv create prometheus
provisioning taskserv create cilium
```

### 3. Create Cluster
```bash
provisioning cluster create --name production
```

## Troubleshooting
Check logs with: `provisioning logs view`
"#;

    let security_doc = r#"
# Security System

## Authentication
The system uses JWT (RS256) for authentication with automatic token rotation.
- Login with: `provisioning login`
- Session tokens expire after 24 hours
- Refresh tokens enable automatic renewal

## Authorization
Cedar-based authorization system with role-based access control (RBAC).
Roles:
- admin: Full access to all resources
- operator: Can create/delete resources, modify configurations
- viewer: Read-only access

## Secrets Management
Dynamic secrets for temporary access:
- AWS STS for temporary credentials
- SSH key generation with TTL
- KMS encryption for sensitive data
"#;

    // Chunk and embed documents
    let chunking_engine = ingester.chunking_engine();
    let mut all_chunks = Vec::new();

    let arch_chunks = chunking_engine.chunk_markdown(architecture_doc, "docs/architecture.md")?;
    let deploy_chunks = chunking_engine.chunk_markdown(deployment_doc, "docs/deployment.md")?;
    let security_chunks = chunking_engine.chunk_markdown(security_doc, "docs/security.md")?;

    all_chunks.extend(arch_chunks);
    all_chunks.extend(deploy_chunks);
    all_chunks.extend(security_chunks);

    // Embed and store documents
    let embedded_docs = ingester.embedding_engine().embed_batch(&all_chunks).await?;
    let stored = db.store_documents(&embedded_docs).await?;
    println!("✓ Stored {} document chunks\n", stored);

    // ==== PHASE 3: Setup Retriever & Agent ====
    println!("Phase 3: Creating RAG agent with workspace context...");

    // Create a new embedding engine for retrieval
    let retrieval_embedding = EmbeddingEngine::new(embedding_config)?;

    let retrieval_config = RetrievalConfig::default();
    let retriever = RetrieverEngine::new(retrieval_config, db, retrieval_embedding).await?;

    // Create workspace context
    let mut workspace = WorkspaceContext::new(
        "librecloud".to_string(),
        "/workspace/librecloud".to_string(),
    );
    workspace.add_taskserv("kubernetes".to_string());
    workspace.add_taskserv("prometheus".to_string());
    workspace.add_provider("aws".to_string());

    let agent = RagAgent::new(retriever, workspace, "claude-opus-4-1".to_string())?;

    println!("✓ RAG Agent ready\n");

    // Print agent metadata
    let metadata = agent.metadata();
    println!("Agent Metadata:");
    println!("  Model: {}", metadata.model);
    println!("  Workspace: {}", metadata.workspace);
    println!("  Version: {}\n", metadata.version);

    // ==== PHASE 4: Ask Questions ====
    println!("Phase 4: Testing semantic search and Q&A...\n");

    let questions = vec![
        "How do I deploy the platform?",
        "What are the core components of the system?",
        "How does authentication work?",
        "What cloud providers are supported?",
    ];

    let num_questions = questions.len();
    for question in questions {
        println!("❓ Question: {}", question);
        match agent.ask(question).await {
            Ok(response) => {
                println!("📝 Answer: {}\n", response.answer);
                println!("   Confidence: {:.0}%", response.confidence * 100.0);
                println!("   Sources used: {}", response.sources.len());
                for (i, source) in response.sources.iter().enumerate() {
                    println!(
                        "     {}. {} ({})",
                        i + 1,
                        source.source_path,
                        source.doc_type
                    );
                }
            }
            Err(e) => {
                println!("❌ Error: {}\n", e);
            }
        }
        println!();
    }

    // ==== PHASE 5: Statistics ====
    println!("Phase 5: Deployment Complete\n");
    println!("✅ RAG Agent successfully deployed and tested");
    println!("   Documents ingested: {}", stored);
    println!("   Questions answered: {}", num_questions);

    println!("\n✅ RAG Agent example completed successfully!");
    println!("\nKey features demonstrated:");
    println!("  ✓ Document ingestion (markdown chunking)");
    println!("  ✓ Vector embedding (local)");
    println!("  ✓ SurrealDB storage");
    println!("  ✓ Semantic search");
    println!("  ✓ Context-aware Q&A");
    println!("  ✓ Workspace integration");

    Ok(())
}