prvng_platform/crates/rag/examples/rag_agent_with_tools.rs

//! Example: RAG Agent with Function Calling / Tool Use
//!
//! Demonstrates how to use the ToolRegistry and ToolEnabledRagAgent to:
//! - Register tools for Claude to invoke
//! - Execute tools with security validation
//! - Track tool usage statistics
//! - Integrate tool results into RAG responses

#![allow(unused_imports)]

use std::sync::Arc;

use provisioning_rag::{
    CreateServerTool, RagTool, TaskservManagementTool, ToolCall, ToolDefinition, ToolInput,
    ToolOutput, ToolRegistry, WorkspaceStatusTool,
};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("=== RAG Agent with Function Calling ===\n");

    // 1. Create tool registry
    println!("1. Setting up tool registry...");
    let registry = Arc::new(ToolRegistry::new());
    println!("   ✓ Tool registry created\n");

    // 2. Register core tools
    println!("2. Registering tools...");
    registry.register(Arc::new(CreateServerTool::new())).await?;
    registry
        .register(Arc::new(WorkspaceStatusTool::new("production".to_string())))
        .await?;
    registry
        .register(Arc::new(TaskservManagementTool::new()))
        .await?;
    println!("   ✓ 3 tools registered\n");

    // 3. Display available tools
    println!("3. Available tools:");
    let definitions = registry.get_definitions().await;
    for def in &definitions {
        println!("   - {} ({})", def.name, def.description);
    }
    println!();

    // 4. Execute tools with validation
    println!("=== Tool Execution Examples ===\n");

    // Example 1: Workspace Status (no auth required)
    println!("Example 1: Get workspace status");
    let input = ToolInput {
        params: serde_json::json!({"include_metrics": true}),
    };
    let output = registry
        .call_tool("get_workspace_status", input, "user@example.com")
        .await?;
    println!("Result: {}", output.result);
    println!("Success: {}\n", output.success);

    // Example 2: Create Server (requires admin)
    println!("Example 2: Create server (non-admin user)");
    let input = ToolInput {
        params: serde_json::json!({
            "hostname": "web-01",
            "cores": 4,
            "memory_gb": 8,
            "region": "us-east-1"
        }),
    };
    let output = registry
        .call_tool("create_server", input, "user@example.com")
        .await?;
    println!("Result: {}", output.result);
    if let Some(error) = &output.error {
        println!("Error (expected): {}", error);
    }
    println!();

    // Example 3: Create Server (admin user)
    println!("Example 3: Create server (admin user)");
    let input = ToolInput {
        params: serde_json::json!({
            "hostname": "web-02",
            "cores": 8,
            "memory_gb": 16,
            "region": "eu-west-1"
        }),
    };
    let output = registry
        .call_tool("create_server", input, "admin@example.com")
        .await?;
    println!("Result: {}", output.result);
    println!("Success: {}\n", output.success);

    // Example 4: Manage Taskserv
    println!("Example 4: Manage taskserv");
    let input = ToolInput {
        params: serde_json::json!({
            "operation": "create",
            "service": "kubernetes",
            "version": "1.28.0"
        }),
    };
    let output = registry
        .call_tool("manage_taskserv", input, "user@example.com")
        .await?;
    println!("Result: {}", output.result);
    println!("Success: {}\n", output.success);

    // 5. Display tool statistics
    println!("=== Tool Call Statistics ===\n");
    let stats = registry.get_stats().await;
    println!("Total calls: {}", stats.total_calls);
    println!("Successful calls: {}", stats.successful_calls);
    println!("Failed calls: {}", stats.failed_calls);
    println!("Blocked calls (auth/rate limit): {}", stats.blocked_calls);
    println!();

    println!("Calls by tool:");
    for (tool, count) in stats.calls_by_tool {
        println!("  - {}: {} calls", tool, count);
    }
    println!();

    // 6. Demonstrate security validation
    println!("=== Security Validation ===\n");
    println!("Tool registration ensures:");
    println!("✓ Access control (admin vs user)");
    println!("✓ Rate limiting (prevents abuse)");
    println!("✓ Audit logging (tracks all invocations)");
    println!("✓ Input validation (rejects invalid params)");
    println!();

    // 7. Tool integration workflow
    println!("=== Tool Integration Workflow ===\n");
    println!("1. Claude generates response with tool_use tags");
    println!("2. Agent parses tool calls from Claude response");
    println!("3. Registry validates user access and rate limits");
    println!("4. Tool executes with validated parameters");
    println!("5. Results logged to audit trail");
    println!("6. Results integrated into final response");
    println!("7. Statistics updated for monitoring\n");

    // 8. Production integration
    println!("=== Production Integration ===\n");
    println!("In production RAG system:");
    println!("- Tools are registered at agent initialization");
    println!("- Claude response is parsed for tool_use tags");
    println!("- Tools execute in controlled environment");
    println!("- Results enhance RAG responses with actions");
    println!("- Audit logs enable compliance and debugging");
    println!("- Statistics drive monitoring and optimization\n");

    println!("=== Tool System Benefits ===\n");
    println!("✓ Extensible: Add new tools by implementing RagTool trait");
    println!("✓ Secure: Access control, rate limiting, input validation");
    println!("✓ Observable: Complete audit logging and statistics");
    println!("✓ Flexible: Custom tools for any system operation");
    println!("✓ Production-ready: Error handling, async support, testing\n");

    Ok(())
}