prvng_platform/crates/rag/src/llm.rs

//! LLM (Large Language Model) integration module
//! Provides Claude API integration for RAG-based answer generation

use serde::{Deserialize, Serialize};
use tracing::info;

use crate::error::Result;

/// Claude API request message
#[derive(Debug, Clone, Serialize)]
pub struct ClaudeMessage {
    pub role: String,
    pub content: String,
}

/// Claude API response
#[derive(Debug, Clone, Deserialize)]
pub struct ClaudeResponse {
    pub content: Vec<ClaudeContent>,
    pub stop_reason: String,
}

/// Claude response content
#[derive(Debug, Clone, Deserialize)]
pub struct ClaudeContent {
    #[serde(rename = "type")]
    pub content_type: String,
    pub text: Option<String>,
}

/// LLM Client for Claude API
pub struct LlmClient {
    api_key: String,
    pub model: String,
    base_url: String,
}

impl LlmClient {
    /// Create a new Claude LLM client
    pub fn new(model: String) -> Result<Self> {
        // Get API key from environment
        let api_key = std::env::var("ANTHROPIC_API_KEY").unwrap_or_else(|_| {
            tracing::warn!("ANTHROPIC_API_KEY not set - Claude API calls will fail");
            String::new()
        });

        Ok(Self {
            api_key,
            model,
            base_url: "https://api.anthropic.com/v1".to_string(),
        })
    }

    /// Generate an answer using Claude
    pub async fn generate_answer(&self, query: &str, context: &str) -> Result<String> {
        // If no API key, return placeholder
        if self.api_key.is_empty() {
            return Ok(self.generate_placeholder(query, context));
        }

        // Build the system prompt
        let system_prompt = format!(
            r#"You are a helpful assistant answering questions about a provisioning platform.
You have been provided with relevant documentation context below.
Answer the user's question based on this context.
Be concise and accurate.

# Retrieved Context
{}
"#,
            context
        );

        // Build the user message
        let user_message = query.to_string();

        // Call Claude API
        self.call_claude_api(&system_prompt, &user_message).await
    }

    /// Call Claude API with messages
    async fn call_claude_api(&self, system: &str, user_message: &str) -> Result<String> {
        let client = reqwest::Client::new();

        // Build request payload
        let payload = serde_json::json!({
            "model": self.model,
            "max_tokens": 1024,
            "system": system,
            "messages": [
                {
                    "role": "user",
                    "content": user_message
                }
            ]
        });

        // Make the API request
        let response = client
            .post(format!("{}/messages", self.base_url))
            .header("anthropic-version", "2023-06-01")
            .header("x-api-key", &self.api_key)
            .json(&payload)
            .send()
            .await
            .map_err(|e| {
                crate::error::RagError::LlmError(format!("Claude API request failed: {}", e))
            })?;

        // Check status
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(crate::error::RagError::LlmError(format!(
                "Claude API error {}: {}",
                status, error_text
            )));
        }

        // Parse response
        let claude_response: ClaudeResponse = response.json().await.map_err(|e| {
            crate::error::RagError::LlmError(format!("Failed to parse Claude response: {}", e))
        })?;

        // Extract text from response
        let answer = claude_response
            .content
            .first()
            .and_then(|c| c.text.clone())
            .ok_or_else(|| {
                crate::error::RagError::LlmError("No text in Claude response".to_string())
            })?;

        info!(
            "Claude API call successful, generated {} characters",
            answer.len()
        );
        Ok(answer)
    }

    /// Generate placeholder answer when API key is missing
    fn generate_placeholder(&self, query: &str, context: &str) -> String {
        format!(
            "Based on the provided context about the provisioning platform:\n\n{}\n\n(Note: This \
             is a placeholder response. Set ANTHROPIC_API_KEY environment variable for full \
             Claude integration.)",
            self.format_context_summary(query, context)
        )
    }

    /// Format a summary of the context
    fn format_context_summary(&self, query: &str, context: &str) -> String {
        let context_lines = context.lines().count();
        let query_lower = query.to_lowercase();

        if query_lower.contains("deploy") || query_lower.contains("create") {
            format!(
                "Your question about deployment is addressed in {} lines of documentation. The \
                 system supports multi-cloud deployment across AWS, UpCloud, and local \
                 environments.",
                context_lines
            )
        } else if query_lower.contains("architecture") || query_lower.contains("design") {
            format!(
                "The provisioning platform uses a modular architecture as described in {} lines \
                 of documentation. Core components include Orchestrator, Control Center, and MCP \
                 Server integration.",
                context_lines
            )
        } else if query_lower.contains("security") || query_lower.contains("auth") {
            format!(
                "Security features are documented in {} lines. The system implements JWT-based \
                 authentication, Cedar-based authorization, and dynamic secrets management.",
                context_lines
            )
        } else {
            format!(
                "Your question is addressed in the provided {} lines of documentation. Please \
                 review the context above for details.",
                context_lines
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_llm_client_creation() {
        let client = LlmClient::new("claude-opus-4-1".to_string());
        assert!(client.is_ok());
    }

    #[test]
    fn test_placeholder_generation() {
        let client = LlmClient {
            api_key: String::new(),
            model: "claude-opus-4-1".to_string(),
            base_url: "https://api.anthropic.com/v1".to_string(),
        };

        let query = "How do I deploy the platform?";
        let context = "Deployment is done using provisioning commands";

        let answer = client.generate_placeholder(query, context);
        assert!(answer.contains("deployment"));
        assert!(answer.contains("placeholder"));
    }

    #[test]
    fn test_context_summary_formatting() {
        let client = LlmClient {
            api_key: String::new(),
            model: "claude-opus-4-1".to_string(),
            base_url: "https://api.anthropic.com/v1".to_string(),
        };

        let deployment_query = "How do I deploy?";
        let context = "Line 1\nLine 2\nLine 3";
        let summary = client.format_context_summary(deployment_query, context);
        assert!(summary.contains("deployment"));
        assert!(summary.contains("3"));
    }
}