llm-universal 0.2.0

use async_trait::async_trait;
use reqwest::header::HeaderValue;
use serde_json::{json, Value};

use crate::{
    send_json, ChatRequest, ChatResponse, LlmError, LlmProvider, Message, Part, Result, Role,
    StopReason, Usage, DEFAULT_MAX_TOKENS,
};

/// Value of the `anthropic-version` header.
const API_VERSION: &str = "2023-06-01";

/// Client for the Anthropic Messages API (`POST {base_url}/v1/messages`).
#[derive(Clone)]
pub struct AnthropicProvider {
    client: reqwest::Client,
    api_key: String,
    base_url: String,
}

impl std::fmt::Debug for AnthropicProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AnthropicProvider")
            .field("api_key", &"<redacted>")
            .field("base_url", &self.base_url)
            .finish()
    }
}

impl AnthropicProvider {
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            client: crate::http_client(),
            api_key: api_key.into(),
            base_url: "https://api.anthropic.com".to_string(),
        }
    }

    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
        self.base_url = base_url.into().trim_end_matches('/').to_string();
        self
    }

    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
    pub fn with_client(mut self, client: reqwest::Client) -> Self {
        self.client = client;
        self
    }

    fn build_body(req: &ChatRequest) -> Value {
        // Anthropic has no `system` role in `messages`: system messages are merged into `system`.
        let mut system_parts: Vec<String> = req.system.iter().cloned().collect();
        let mut messages = Vec::new();
        for message in &req.messages {
            if message.role == Role::System {
                system_parts.push(message.text_content());
            } else {
                messages.push(to_anthropic_message(message));
            }
        }

        let mut body = json!({
            "model": req.model,
            "max_tokens": req.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS),
            "messages": messages,
        });
        if !system_parts.is_empty() {
            body["system"] = json!(system_parts.join("\n\n"));
        }
        if !req.tools.is_empty() {
            let tools: Vec<Value> = req
                .tools
                .iter()
                .map(|t| {
                    json!({
                        "name": t.name,
                        "description": t.description,
                        "input_schema": t.parameters,
                    })
                })
                .collect();
            body["tools"] = json!(tools);
        }
        if let Some(temperature) = req.temperature {
            body["temperature"] = json!(temperature);
        }
        body
    }
}

fn to_anthropic_message(message: &Message) -> Value {
    // Tool results are `tool_result` blocks inside a `user` turn.
    let role = if message.role == Role::Assistant {
        "assistant"
    } else {
        "user"
    };
    // The API rejects empty text blocks, so they are skipped.
    let content: Vec<Value> = message
        .parts
        .iter()
        .filter(|p| !matches!(p, Part::Text { text } if text.is_empty()))
        .map(|p| match p {
            Part::Text { text } => json!({ "type": "text", "text": text }),
            Part::ToolCall {
                id,
                name,
                arguments,
            } => {
                json!({ "type": "tool_use", "id": id, "name": name, "input": arguments })
            }
            Part::ToolResult {
                call_id,
                content,
                is_error,
            } => json!({
                "type": "tool_result",
                "tool_use_id": call_id,
                "content": content,
                "is_error": is_error,
            }),
        })
        .collect();
    json!({ "role": role, "content": content })
}

fn parse_response(v: &Value) -> Result<ChatResponse> {
    let blocks = v["content"]
        .as_array()
        .ok_or_else(|| LlmError::InvalidResponse("missing field `content`".into()))?;

    let mut parts = Vec::new();
    for block in blocks {
        match block["type"].as_str() {
            Some("text") => {
                let text = block["text"].as_str().unwrap_or_default();
                if !text.is_empty() {
                    parts.push(Part::Text {
                        text: text.to_string(),
                    });
                }
            }
            Some("tool_use") => parts.push(Part::ToolCall {
                id: block["id"].as_str().unwrap_or_default().to_string(),
                name: block["name"].as_str().unwrap_or_default().to_string(),
                arguments: block["input"].clone(),
            }),
            // Thinking and other blocks are not part of the normalized response.
            _ => {}
        }
    }

    let stop_reason = match v["stop_reason"].as_str() {
        Some("end_turn") | Some("stop_sequence") => StopReason::EndTurn,
        Some("tool_use") => StopReason::ToolUse,
        Some("max_tokens") | Some("model_context_window_exceeded") => StopReason::MaxTokens,
        Some("pause_turn") => StopReason::Paused,
        Some("refusal") => StopReason::Refusal,
        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
    };

    let tokens =
        |field: &str| u32::try_from(v["usage"][field].as_u64().unwrap_or(0)).unwrap_or(u32::MAX);
    let usage = Usage {
        input_tokens: tokens("input_tokens"),
        output_tokens: tokens("output_tokens"),
    };

    Ok(ChatResponse {
        model: v["model"].as_str().unwrap_or_default().to_string(),
        message: Message::new(Role::Assistant, parts),
        stop_reason,
        usage,
    })
}

#[async_trait]
impl LlmProvider for AnthropicProvider {
    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
        // Anthropic allows tool names up to 128 characters.
        crate::validate(
            "AnthropicProvider",
            &self.api_key,
            true,
            &self.base_url,
            128,
            req,
        )?;
        let mut api_key = HeaderValue::from_str(&self.api_key).map_err(|_| {
            LlmError::InvalidArgument(
                "AnthropicProvider: api_key contains characters not allowed in an HTTP header"
                    .into(),
            )
        })?;
        api_key.set_sensitive(true);

        let url = format!("{}/v1/messages", self.base_url);
        let builder = self
            .client
            .post(url)
            .header("x-api-key", api_key)
            .header("anthropic-version", API_VERSION)
            .json(&Self::build_body(req));
        let body = send_json(builder).await?;
        parse_response(&body)
    }
}

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

    #[test]
    fn moves_system_and_maps_tool_results() {
        let req = ChatRequest::new("claude-opus-5-5")
            .system("sys")
            .tool(Tool::new("f", "d", json!({"type": "object"})))
            .message(Message::user("hi"))
            .message(Message::tool_result("toolu_1", "42", false));
        let body = AnthropicProvider::build_body(&req);
        assert_eq!(body["system"], "sys");
        assert_eq!(body["max_tokens"], DEFAULT_MAX_TOKENS);
        assert_eq!(body["tools"][0]["input_schema"]["type"], "object");
        assert_eq!(body["messages"][1]["role"], "user");
        assert_eq!(body["messages"][1]["content"][0]["type"], "tool_result");
        assert_eq!(body["messages"][1]["content"][0]["tool_use_id"], "toolu_1");
    }

    #[test]
    fn parses_tool_use_response() {
        let v = json!({
            "model": "claude-opus-5-5",
            "content": [
                { "type": "text", "text": "Let me check." },
                { "type": "tool_use", "id": "toolu_1", "name": "f", "input": { "x": 1 } }
            ],
            "stop_reason": "tool_use",
            "usage": { "input_tokens": 3, "output_tokens": 9 }
        });
        let resp = parse_response(&v).unwrap();
        assert_eq!(resp.stop_reason, StopReason::ToolUse);
        assert_eq!(resp.text(), "Let me check.");
        assert_eq!(resp.tool_calls()[0].2["x"], 1);
        assert_eq!(resp.usage.input_tokens, 3);
    }

    #[test]
    fn skips_empty_text_blocks() {
        let msg = Message::new(
            Role::Assistant,
            vec![Part::Text {
                text: String::new(),
            }],
        );
        assert_eq!(to_anthropic_message(&msg)["content"], json!([]));
    }

    #[test]
    fn debug_redacts_api_key() {
        let debug = format!("{:?}", AnthropicProvider::new("sk-ant-secret"));
        assert!(!debug.contains("sk-ant-secret"));
    }
}