QuiCommit/src/llm/anthropic.rs

use super::{create_http_client, LlmProvider};
use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Anthropic Claude API client
pub struct AnthropicClient {
    api_key: String,
    model: String,
    client: reqwest::Client,
}

#[derive(Debug, Serialize)]
struct MessagesRequest {
    model: String,
    max_tokens: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    messages: Vec<AnthropicMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
struct AnthropicMessage {
    role: String,
    content: String,
}

#[derive(Debug, Deserialize)]
struct MessagesResponse {
    content: Vec<ContentBlock>,
}

#[derive(Debug, Deserialize)]
struct ContentBlock {
    #[serde(rename = "type")]
    content_type: String,
    text: String,
}

#[derive(Debug, Deserialize)]
struct ErrorResponse {
    error: AnthropicError,
}

#[derive(Debug, Deserialize)]
struct AnthropicError {
    #[serde(rename = "type")]
    error_type: String,
    message: String,
}

impl AnthropicClient {
    /// Create new Anthropic client
    pub fn new(api_key: &str, model: &str) -> Result<Self> {
        let client = create_http_client(Duration::from_secs(60))?;
        
        Ok(Self {
            api_key: api_key.to_string(),
            model: model.to_string(),
            client,
        })
    }

    /// Set timeout
    pub fn with_timeout(mut self, timeout: Duration) -> Result<Self> {
        self.client = create_http_client(timeout)?;
        Ok(self)
    }

    /// Validate API key
    pub async fn validate_key(&self) -> Result<bool> {
        let url = "https://api.anthropic.com/v1/messages";
        
        let request = MessagesRequest {
            model: self.model.clone(),
            max_tokens: 5,
            temperature: Some(0.0),
            messages: vec![AnthropicMessage {
                role: "user".to_string(),
                content: "Hi".to_string(),
            }],
            system: None,
        };
        
        let response = self.client
            .post(url)
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await;
        
        match response {
            Ok(resp) => {
                if resp.status().is_success() {
                    Ok(true)
                } else {
                    let status = resp.status();
                    if status.as_u16() == 401 {
                        Ok(false)
                    } else {
                        let text = resp.text().await.unwrap_or_default();
                        bail!("Anthropic API error: {} - {}", status, text)
                    }
                }
            }
            Err(e) => Err(e.into()),
        }
    }
}

#[async_trait]
impl LlmProvider for AnthropicClient {
    async fn generate(&self, prompt: &str) -> Result<String> {
        let messages = vec![AnthropicMessage {
            role: "user".to_string(),
            content: prompt.to_string(),
        }];
        
        self.messages_request(messages, None).await
    }

    async fn generate_with_system(&self, system: &str, user: &str) -> Result<String> {
        let messages = vec![AnthropicMessage {
            role: "user".to_string(),
            content: user.to_string(),
        }];
        
        let system = if system.is_empty() {
            None
        } else {
            Some(system.to_string())
        };
        
        self.messages_request(messages, system).await
    }

    async fn is_available(&self) -> bool {
        self.validate_key().await.unwrap_or(false)
    }

    fn name(&self) -> &str {
        "anthropic"
    }
}

impl AnthropicClient {
    async fn messages_request(
        &self,
        messages: Vec<AnthropicMessage>,
        system: Option<String>,
    ) -> Result<String> {
        let url = "https://api.anthropic.com/v1/messages";
        
        let request = MessagesRequest {
            model: self.model.clone(),
            max_tokens: 500,
            temperature: Some(0.7),
            messages,
            system,
        };
        
        let response = self.client
            .post(url)
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send request to Anthropic")?;
        
        let status = response.status();
        
        if !status.is_success() {
            let text = response.text().await.unwrap_or_default();
            
            // Try to parse error
            if let Ok(error) = serde_json::from_str::<ErrorResponse>(&text) {
                bail!("Anthropic API error: {} ({})", error.error.message, error.error.error_type);
            }
            
            bail!("Anthropic API error: {} - {}", status, text);
        }
        
        let result: MessagesResponse = response
            .json()
            .await
            .context("Failed to parse Anthropic response")?;
        
        result.content
            .into_iter()
            .find(|c| c.content_type == "text")
            .map(|c| c.text.trim().to_string())
            .ok_or_else(|| anyhow::anyhow!("No text response from Anthropic"))
    }
}

/// Available Anthropic models
pub const ANTHROPIC_MODELS: &[&str] = &[
    "claude-3-opus-20240229",
    "claude-3-sonnet-20240229",
    "claude-3-haiku-20240307",
    "claude-2.1",
    "claude-2.0",
    "claude-instant-1.2",
];

/// Check if a model name is valid
pub fn is_valid_model(model: &str) -> bool {
    ANTHROPIC_MODELS.contains(&model)
}

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

    #[test]
    fn test_model_validation() {
        assert!(is_valid_model("claude-3-sonnet-20240229"));
        assert!(!is_valid_model("invalid-model"));
    }
}