hms/crates/erp-ai/src/provider/claude.rs

use async_stream::stream;
use async_trait::async_trait;
use futures::{Stream, StreamExt};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::pin::Pin;

use super::AiProvider;
use crate::dto::GenerateRequest;
use crate::error::{AiError, AiResult};

#[derive(Debug, Clone)]
pub struct ClaudeProvider {
    client: Client,
    api_key: String,
    base_url: String,
}

impl ClaudeProvider {
    pub fn new(api_key: String) -> Self {
        Self {
            client: Client::new(),
            api_key,
            base_url: "https://api.anthropic.com".into(),
        }
    }

    pub fn with_base_url(mut self, url: String) -> Self {
        self.base_url = url;
        self
    }
}

#[derive(Serialize)]
struct ClaudeRequest {
    model: String,
    max_tokens: u32,
    temperature: f32,
    system: String,
    messages: Vec<ClaudeMessage>,
    stream: bool,
}

#[derive(Serialize)]
struct ClaudeMessage {
    role: String,
    content: String,
}

#[derive(Deserialize)]
struct ClaudeStreamEvent {
    #[serde(rename = "type")]
    event_type: String,
    delta: Option<ClaudeDelta>,
    #[allow(dead_code)] // serde 反序列化需要，但流式处理中不读取
    message: Option<ClaudeMessageResp>,
}

#[derive(Deserialize)]
struct ClaudeDelta {
    text: Option<String>,
}

#[derive(Deserialize)]
struct ClaudeMessageResp {
    #[allow(dead_code)] // serde 反序列化需要，暂未使用 usage 信息
    usage: Option<ClaudeUsage>,
}

#[derive(Deserialize)]
struct ClaudeUsage {
    #[allow(dead_code)] // serde 反序列化需要，暂未追踪 token 用量
    input_tokens: u32,
    #[allow(dead_code)]
    output_tokens: u32,
}

#[async_trait]
impl AiProvider for ClaudeProvider {
    async fn stream_generate(
        &self,
        req: GenerateRequest,
    ) -> AiResult<Pin<Box<dyn Stream<Item = AiResult<String>> + Send>>> {
        let claude_req = ClaudeRequest {
            model: req.model,
            max_tokens: req.max_tokens,
            temperature: req.temperature,
            system: req.system_prompt,
            messages: vec![ClaudeMessage {
                role: "user".into(),
                content: req.user_prompt,
            }],
            stream: true,
        };

        let response = self
            .client
            .post(format!("{}/v1/messages", self.base_url))
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("content-type", "application/json")
            .json(&claude_req)
            .send()
            .await
            .map_err(|e| AiError::ProviderError(format!("Claude API 请求失败: {e}")))?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(AiError::ProviderError(format!(
                "Claude API 错误 {status}: {body}"
            )));
        }

        let stream = Box::pin(stream! {
            let mut stream = response.bytes_stream();
            while let Some(chunk_result) = stream.next().await {
                let bytes = match chunk_result {
                    Ok(b) => b,
                    Err(e) => {
                        yield Err(AiError::ProviderError(format!("流读取错误: {e}")));
                        break;
                    }
                };

                let text = String::from_utf8_lossy(&bytes);
                for line in text.lines() {
                    if let Some(data) = line.strip_prefix("data: ") {
                        if data == "[DONE]" {
                            return;
                        }
                        if let Ok(event) = serde_json::from_str::<ClaudeStreamEvent>(data)
                            && event.event_type == "content_block_delta"
                                && let Some(delta) = event.delta
                                    && let Some(text) = delta.text {
                                        yield Ok(text);
                                    }
                    }
                }
            }
        });

        Ok(stream)
    }

    async fn generate(&self, req: GenerateRequest) -> AiResult<crate::dto::GenerateResponse> {
        let start = std::time::Instant::now();

        let claude_req = ClaudeRequest {
            model: req.model.clone(),
            max_tokens: req.max_tokens,
            temperature: req.temperature,
            system: req.system_prompt,
            messages: vec![ClaudeMessage {
                role: "user".into(),
                content: req.user_prompt,
            }],
            stream: false,
        };

        let resp = self
            .client
            .post(format!("{}/v1/messages", self.base_url))
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("content-type", "application/json")
            .json(&claude_req)
            .send()
            .await
            .map_err(|e| AiError::ProviderError(e.to_string()))?;

        let status = resp.status();
        let body = resp
            .text()
            .await
            .map_err(|e| AiError::ProviderError(e.to_string()))?;

        if !status.is_success() {
            return Err(AiError::ProviderError(format!("Claude {status}: {body}")));
        }

        let parsed: serde_json::Value = serde_json::from_str(&body)
            .map_err(|e| AiError::ProviderError(format!("解析响应失败: {e}")))?;

        let content = parsed["content"][0]["text"]
            .as_str()
            .unwrap_or("")
            .to_string();

        let input_tokens = parsed["usage"]["input_tokens"].as_u64().unwrap_or(0) as u32;
        let output_tokens = parsed["usage"]["output_tokens"].as_u64().unwrap_or(0) as u32;

        Ok(crate::dto::GenerateResponse {
            content,
            model: req.model,
            input_tokens,
            output_tokens,
            duration_ms: start.elapsed().as_millis() as u64,
        })
    }

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

    async fn health_check(&self) -> AiResult<bool> {
        let resp = self
            .client
            .post(format!("{}/v1/messages", self.base_url))
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("content-type", "application/json")
            .json(&serde_json::json!({
                "model": "claude-sonnet-4-6",
                "max_tokens": 1,
                "messages": [{"role": "user", "content": "hi"}]
            }))
            .send()
            .await;

        match resp {
            Ok(r) => Ok(r.status().is_success() || r.status().as_u16() == 400),
            Err(_) => Ok(false),
        }
    }
}