xz-provider 0.5.0

/// Anthropic Messages API protocol adapter.
///
/// Implements [`ProtocolAdapter`] for the Anthropic Messages API.
/// Handles request/response conversion, streaming via SSE, extended thinking,
/// tool use, prompt caching, and vision (image base64/url).
use serde_json::Value;
use tracing::warn;

use crate::error::ProviderError;
use crate::protocol::{AuthMethod, ProtocolAdapter};
use crate::types::{
    CacheControl, CacheInfo, Citation, CompletionRequest, CompletionResponse, ContentPart,
    EffortLevel, FinishReason, Message, MessageContent, StreamEvent, ThinkingType, TokenUsage,
    ToolCall, ToolChoice,
};

/// Adapter for the Anthropic Messages API protocol.
///
/// # Example
///
/// ```rust
/// use xz_provider::protocol::{ProtocolAdapter, AuthMethod};
/// use xz_provider::protocol::anthropic::AnthropicMessagesAdapter;
///
/// let adapter = AnthropicMessagesAdapter::new();
/// assert_eq!(adapter.endpoint_path(), "/messages");
/// assert_eq!(adapter.protocol_name(), "anthropic");
/// ```
#[derive(Debug)]
pub struct AnthropicMessagesAdapter {
    /// Anthropic beta feature flags (e.g. "prompt-caching-2025-02-19").
    anthropic_beta: Vec<String>,
    /// Default effort level for this adapter instance (from provider config).
    default_effort_level: Option<EffortLevel>,
    /// Override for the anthropic-version header (defaults to "2023-06-01").
    anthropic_version: Option<String>,
    /// Default thinking type for this adapter instance (from provider config).
    default_thinking_type: Option<ThinkingType>,
}

impl AnthropicMessagesAdapter {
    /// Creates a new `AnthropicMessagesAdapter` with no beta headers or default effort.
    pub fn new() -> Self {
        Self {
            anthropic_beta: Vec::new(),
            default_effort_level: None,
            anthropic_version: None,
            default_thinking_type: None,
        }
    }

    /// Configures anthropic beta feature headers.
    ///
    /// Each string in the list is emitted as a separate `anthropic-beta: <value>` header.
    pub fn with_beta_headers(mut self, beta: Vec<String>) -> Self {
        self.anthropic_beta = beta;
        self
    }

    /// Configures a default effort level used in `output_config`.
    pub fn with_effort_level(mut self, effort: EffortLevel) -> Self {
        self.default_effort_level = Some(effort);
        self
    }

    /// Overrides the default anthropic-version header value.
    ///
    /// If not set, defaults to `"2023-06-01"`.
    pub fn with_anthropic_version(mut self, version: String) -> Self {
        self.anthropic_version = Some(version);
        self
    }

    /// Configures a default thinking type used when no per-request thinking is set.
    pub fn with_default_thinking_type(mut self, thinking_type: ThinkingType) -> Self {
        self.default_thinking_type = Some(thinking_type);
        self
    }

    /// Converts a [`ContentPart`] into an Anthropic API content block [`Value`].
    fn convert_content_part(
        &self,
        part: &ContentPart,
        cache_control: Option<&CacheControl>,
    ) -> Value {
        match part {
            ContentPart::Text { text } => {
                let mut block = serde_json::json!({
                    "type": "text",
                    "text": text,
                });
                if cache_control.is_some() {
                    block["cache_control"] = serde_json::json!({"type": "ephemeral"});
                }
                block
            }
            ContentPart::ImageBase64 { media_type, data } => {
                serde_json::json!({
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": media_type,
                        "data": data,
                    }
                })
            }
            ContentPart::ImageUrl { url, .. } => {
                serde_json::json!({
                    "type": "image",
                    "source": {
                        "type": "url",
                        "url": url,
                    }
                })
            }
            ContentPart::Document { source, title, cache_control: part_cc } => {
                let source_val = serde_json::to_value(source).unwrap_or(
                    serde_json::json!({"type": "base64", "media_type": "text/plain", "data": ""}),
                );
                let mut block = serde_json::json!({
                    "type": "document",
                    "source": source_val,
                });
                if let Some(t) = title {
                    block["title"] = serde_json::Value::String(t.clone());
                }
                if part_cc.is_some() {
                    block["cache_control"] = serde_json::json!({"type": "ephemeral"});
                }
                block
            }
            ContentPart::ToolReference { tool_name, cache_control: part_cc } => {
                let mut block = serde_json::json!({
                    "type": "tool_reference",
                    "tool_name": tool_name,
                });
                if part_cc.is_some() {
                    block["cache_control"] = serde_json::json!({"type": "ephemeral"});
                }
                block
            }
            ContentPart::RedactedThinking { data } => {
                serde_json::json!({
                    "type": "redacted_thinking",
                    "data": data,
                })
            }
            _ => {
                // AudioBase64 and File content parts are not supported as image/text
                serde_json::json!({
                    "type": "text",
                    "text": "[unsupported content type]",
                })
            }
        }
    }

    /// Converts [`MessageContent`] into Anthropic API content blocks (array).
    fn content_to_blocks(
        &self,
        content: &MessageContent,
        cache_control: Option<&CacheControl>,
    ) -> Vec<Value> {
        match content {
            MessageContent::Text(text) => {
                let mut block = serde_json::json!({
                    "type": "text",
                    "text": text,
                });
                if cache_control.is_some() {
                    block["cache_control"] = serde_json::json!({"type": "ephemeral"});
                }
                vec![block]
            }
            MessageContent::MultiPart(parts) => {
                parts.iter().map(|part| self.convert_content_part(part, cache_control)).collect()
            }
            MessageContent::None => Vec::new(),
        }
    }

    /// Converts system messages and all other messages into Anthropic request format.
    ///
    /// Returns `(messages_array, system_value)` where `system_value` is:
    /// - `None` if no system messages
    /// - `Some(Value::String)` for a single text system message without cache control
    /// - `Some(Value::Array)` for multiple/blocks system content
    fn convert_request(&self, request: &CompletionRequest) -> (Vec<Value>, Option<Value>) {
        let mut system_blocks: Vec<Value> = Vec::new();
        let mut messages: Vec<Value> = Vec::new();

        for msg in &request.messages {
            match msg {
                Message::System { content, cache_control } => match content {
                    MessageContent::Text(text) => {
                        let mut block = serde_json::json!({
                            "type": "text",
                            "text": text,
                        });
                        if cache_control.is_some() {
                            block["cache_control"] = serde_json::json!({"type": "ephemeral"});
                        }
                        system_blocks.push(block);
                    }
                    MessageContent::MultiPart(parts) => {
                        for part in parts {
                            let block = self.convert_content_part(part, cache_control.as_ref());
                            system_blocks.push(block);
                        }
                    }
                    MessageContent::None => {}
                },
                Message::Developer { content, cache_control } => match content {
                    MessageContent::Text(text) => {
                        let mut block = serde_json::json!({
                            "type": "text",
                            "text": text,
                        });
                        if cache_control.is_some() {
                            block["cache_control"] = serde_json::json!({"type": "ephemeral"});
                        }
                        system_blocks.push(block);
                    }
                    MessageContent::MultiPart(parts) => {
                        for part in parts {
                            let block = self.convert_content_part(part, cache_control.as_ref());
                            system_blocks.push(block);
                        }
                    }
                    MessageContent::None => {}
                },
                Message::User { content } => {
                    let blocks = self.content_to_blocks(content, None);
                    messages.push(serde_json::json!({
                        "role": "user",
                        "content": blocks,
                    }));
                }
                Message::Assistant { content, tool_calls, cache_control, .. } => {
                    let mut blocks: Vec<Value> = Vec::new();

                    // tool_use blocks come first
                    if let Some(calls) = tool_calls {
                        for call in calls {
                            blocks.push(serde_json::json!({
                                "type": "tool_use",
                                "id": call.id,
                                "name": call.function_name,
                                "input": call.arguments,
                            }));
                        }
                    }

                    // Add text content blocks
                    if !matches!(content, MessageContent::None) {
                        let text_blocks = self.content_to_blocks(content, cache_control.as_ref());
                        blocks.extend(text_blocks);
                    }

                    messages.push(serde_json::json!({
                        "role": "assistant",
                        "content": blocks,
                    }));
                }
                Message::Tool { content, tool_call_id, is_error } => {
                    let text = match content {
                        MessageContent::Text(t) => t.clone(),
                        _ => String::new(),
                    };
                    messages.push(serde_json::json!({
                        "role": "user",
                        "content": [{
                            "type": "tool_result",
                            "tool_use_id": tool_call_id,
                            "content": text,
                            "is_error": is_error,
                        }],
                    }));
                }
            }
        }

        // Determine the system field format
        let system = if system_blocks.is_empty() {
            None
        } else if system_blocks.len() == 1 {
            let block = &system_blocks[0];
            // Use plain string for single text block without cache_control
            if block["type"] == "text" && block.get("cache_control").is_none() {
                Some(Value::String(block["text"].as_str().unwrap_or("").to_owned()))
            } else {
                Some(Value::Array(system_blocks))
            }
        } else {
            Some(Value::Array(system_blocks))
        };

        (messages, system)
    }

    /// Parses an Anthropic `stop_reason` string into a [`FinishReason`].
    ///
    /// Unknown reasons are mapped to [`FinishReason::Stop`] with a warning.
    fn parse_finish_reason(&self, reason: &str) -> FinishReason {
        match reason {
            "end_turn" => FinishReason::Stop,
            "tool_use" => FinishReason::ToolCall,
            "max_tokens" => FinishReason::MaxTokens,
            "stop_sequence" => FinishReason::Stop,
            "pause_turn" => FinishReason::PauseTurn,
            "refusal" => FinishReason::Refusal,
            other => {
                warn!(
                    reason = %other,
                    "Anthropic: unknown finish_reason, mapping to Stop"
                );
                FinishReason::Stop
            }
        }
    }
}

impl Default for AnthropicMessagesAdapter {
    fn default() -> Self {
        Self::new()
    }
}

impl ProtocolAdapter for AnthropicMessagesAdapter {
    fn endpoint_path(&self) -> &str {
        "/messages"
    }

    fn build_request_body(
        &self,
        request: &CompletionRequest,
        _stream: bool,
    ) -> Result<Value, ProviderError> {
        let (messages, system) = self.convert_request(request);

        let mut body = serde_json::json!({
            "model": request.model,
            "messages": messages,
            "max_tokens": request.max_tokens.unwrap_or(4096),
        });

        if let Some(s) = system {
            body["system"] = s;
        }

        if let Some(temp) = request.temperature {
            body["temperature"] = Value::from(temp);
        }
        if let Some(top_p) = request.top_p {
            body["top_p"] = Value::from(top_p);
        }
        if let Some(top_k) = request.top_k {
            body["top_k"] = Value::from(top_k);
        }
        if let Some(stop) = &request.stop {
            body["stop_sequences"] =
                Value::Array(stop.iter().map(|s| Value::String(s.clone())).collect());
        }

        if let Some(tools) = &request.tools {
            let tool_defs: Vec<Value> = tools
                .iter()
                .map(|t| {
                    serde_json::json!({
                        "name": t.name,
                        "description": t.description,
                        "input_schema": t.parameters,
                    })
                })
                .collect();
            body["tools"] = Value::Array(tool_defs);
        }

        if let Some(tc) = &request.tool_choice {
            body["tool_choice"] = match tc {
                ToolChoice::Auto => Value::String("auto".to_owned()),
                ToolChoice::Required => Value::String("any".to_owned()),
                ToolChoice::Disabled => Value::String("none".to_owned()),
                ToolChoice::Specific { name } => {
                    serde_json::json!({"type": "tool", "name": name})
                }
            };
        }

        // Thinking config mapping: request.thinking takes precedence over reasoning_effort.
        if let Some(thinking) = &request.thinking {
            match &thinking.thinking_type {
                ThinkingType::Enabled { budget_tokens } => {
                    body["thinking"] = serde_json::json!({
                        "type": "enabled",
                        "budget_tokens": budget_tokens.unwrap_or(2048),
                    });
                }
                ThinkingType::Adaptive => {
                    body["thinking"] = serde_json::json!({"type": "adaptive"});
                }
                ThinkingType::Disabled => {
                    // Omit the thinking field entirely for disabled.
                }
            }
        } else if let Some(reasoning) = &request.reasoning_effort {
            // Anthropic extended thinking: map reasoning_effort to thinking budget
            let budget_tokens = match reasoning {
                crate::types::ReasoningEffort::Low => 1024,
                crate::types::ReasoningEffort::Medium => 2048,
                crate::types::ReasoningEffort::High => 4096,
            };
            body["thinking"] = serde_json::json!({
                "type": "enabled",
                "budget_tokens": budget_tokens,
            });
        } else if let Some(default_thinking) = &self.default_thinking_type {
            // Default thinking type from provider config (fallback when no per-request thinking)
            match default_thinking {
                ThinkingType::Enabled { budget_tokens } => {
                    body["thinking"] = serde_json::json!({
                        "type": "enabled",
                        "budget_tokens": budget_tokens.unwrap_or(2048),
                    });
                }
                ThinkingType::Adaptive => {
                    body["thinking"] = serde_json::json!({"type": "adaptive"});
                }
                ThinkingType::Disabled => {
                    // Omit the thinking field entirely for disabled.
                }
            }
        }

        // Output config (effort level) from adapter default
        if let Some(effort) = &self.default_effort_level {
            let effort_val = serde_json::to_value(effort).ok();
            if let Some(val) = effort_val {
                body["output_config"] = serde_json::json!({"effort": val});
            }
        }

        // Metadata — map user field to Anthropic metadata.user_id
        if let Some(ref user_id) = request.user {
            body["metadata"] = serde_json::json!({"user_id": user_id});
        }

        // Anthropic does not include `stream` in the body — streaming is
        // controlled via the `Accept: text/event-stream` header by the caller.
        // The `_stream` parameter is unused intentionally.

        Ok(body)
    }

    fn build_auth_headers(&self, auth: &AuthMethod) -> Vec<(String, String)> {
        let mut headers = Vec::new();
        if let AuthMethod::ApiKey { header_name, key } = auth {
            headers.push((header_name.clone(), key.clone()));
        }
        headers.push((
            "anthropic-version".to_owned(),
            self.anthropic_version.clone().unwrap_or_else(|| "2023-06-01".to_owned()),
        ));
        for beta in &self.anthropic_beta {
            headers.push(("anthropic-beta".to_owned(), beta.clone()));
        }
        headers
    }

    fn parse_response(&self, body: &Value) -> Result<CompletionResponse, ProviderError> {
        let empty_vec = Vec::new();
        let content_blocks = body["content"].as_array().unwrap_or(&empty_vec);
        let mut text_content = String::new();
        let mut tool_calls: Vec<ToolCall> = Vec::new();
        let mut thinking: Option<String> = None;
        let mut signature: Option<String> = None;
        let mut redacted_thinking: Option<String> = None;
        let mut citations: Option<Vec<Citation>> = None;

        for block in content_blocks {
            if let Some(block_type) = block["type"].as_str() {
                match block_type {
                    "text" => {
                        if let Some(text) = block["text"].as_str() {
                            text_content.push_str(text);
                        }
                        // Parse citations from text block if present
                        if citations.is_none() {
                            if let Some(citation_array) = block["citations"].as_array() {
                                let parsed: Vec<Citation> = citation_array
                                    .iter()
                                    .filter_map(|c| serde_json::from_value(c.clone()).ok())
                                    .collect();
                                if !parsed.is_empty() {
                                    citations = Some(parsed);
                                }
                            }
                        }
                    }
                    "tool_use" => {
                        let id = block["id"].as_str().unwrap_or("").to_owned();
                        let name = block["name"].as_str().unwrap_or("").to_owned();
                        let input = block["input"].clone();
                        tool_calls.push(ToolCall { id, function_name: name, arguments: input });
                    }
                    "thinking" => {
                        if thinking.is_none() {
                            if let Some(t) = block["thinking"].as_str() {
                                thinking = Some(t.to_owned());
                            }
                        }
                        if signature.is_none() {
                            if let Some(sig) = block["signature"].as_str() {
                                signature = Some(sig.to_owned());
                            }
                        }
                    }
                    "redacted_thinking" => {
                        if let Some(data) = block["data"].as_str() {
                            redacted_thinking = Some(data.to_owned());
                        }
                    }
                    "document" | "tool_reference" => {
                        // Recognized block types — pass through without extracting text
                        warn!(
                            block_type = %block_type,
                            "Anthropic: encountered content block type in response, passing through"
                        );
                    }
                    "web_search_result" | "web_fetch" => {
                        warn!(
                            block_type = %block_type,
                            "Anthropic: encountered {} block in response, passing through",
                            block_type,
                        );
                    }
                    _ => {}
                }
            }
        }

        let model = body["model"].as_str().unwrap_or("").to_owned();
        let stop_reason_str = body["stop_reason"].as_str().unwrap_or("end_turn");
        let finish_reason = self.parse_finish_reason(stop_reason_str);

        let usage_data = &body["usage"];
        let cached_tokens = usage_data["cache_read_input_tokens"].as_u64().map(|v| v as u32);
        let cache_write_5m = usage_data
            .pointer("/cache_creation/ephemeral_5m_input_tokens")
            .and_then(|v| v.as_u64())
            .map(|v| v as u32);
        let cache_write_1h = usage_data
            .pointer("/cache_creation/ephemeral_1h_input_tokens")
            .and_then(|v| v.as_u64())
            .map(|v| v as u32);
        let usage = if usage_data.is_object() {
            let prompt = usage_data["input_tokens"].as_u64().unwrap_or(0) as u32;
            let completion = usage_data["output_tokens"].as_u64().unwrap_or(0) as u32;
            TokenUsage {
                prompt_tokens: prompt,
                completion_tokens: completion,
                total_tokens: prompt + completion,
                cached_tokens,
                cache_write_5m_input_tokens: cache_write_5m,
                cache_write_1h_input_tokens: cache_write_1h,
                ..Default::default()
            }
        } else {
            TokenUsage::new(0, 0)
        };

        let cache_info =
            cached_tokens.map(|ct| CacheInfo { cached_tokens: ct, cache_saved_cost: 0.0 });

        Ok(CompletionResponse {
            content: if text_content.is_empty() { None } else { Some(text_content) },
            thinking,
            tool_calls,
            usage,
            model,
            finish_reason,
            latency_ms: 0,
            cache_info,
            signature,
            redacted_thinking,
            citations,
            ..Default::default()
        })
    }

    fn parse_sse_event(&self, data: &str) -> Result<Option<StreamEvent>, ProviderError> {
        let parsed: Value = serde_json::from_str(data)?;

        let event_type = parsed["type"]
            .as_str()
            .ok_or_else(|| ProviderError::Format("SSE event missing type field".to_owned()))?;

        match event_type {
            "message_start" | "content_block_stop" | "message_stop" | "ping" => Ok(None),
            "content_block_start" => {
                let content_block = &parsed["content_block"];
                let block_type = content_block["type"].as_str();
                match block_type {
                    Some("tool_use") => {
                        let idx = parsed["index"].as_u64().unwrap_or(0) as usize;
                        Ok(Some(StreamEvent::ToolCallDelta {
                            index: idx,
                            id: content_block["id"].as_str().map(|s| s.to_owned()),
                            function_name: content_block["name"].as_str().map(|s| s.to_owned()),
                            arguments_delta: String::new(),
                        }))
                    }
                    Some("thinking") => {
                        let text = content_block["thinking"].as_str().unwrap_or("").to_owned();
                        if text.is_empty() {
                            Ok(None)
                        } else {
                            Ok(Some(StreamEvent::ThinkingDelta { delta: text }))
                        }
                    }
                    Some("redacted_thinking") => {
                        let data = content_block["data"].as_str().unwrap_or("").to_owned();
                        Ok(Some(StreamEvent::RedactedThinkingDelta { data }))
                    }
                    _ => Ok(None),
                }
            }
            "content_block_delta" => {
                let delta = &parsed["delta"];
                let delta_type = delta["type"].as_str();
                match delta_type {
                    Some("text_delta") => {
                        let text = delta["text"].as_str().unwrap_or("").to_owned();
                        Ok(Some(StreamEvent::ContentDelta { delta: text }))
                    }
                    Some("input_json_delta") => {
                        let idx = parsed["index"].as_u64().unwrap_or(0) as usize;
                        let partial = delta["partial_json"].as_str().unwrap_or("").to_owned();
                        Ok(Some(StreamEvent::ToolCallDelta {
                            index: idx,
                            id: None,
                            function_name: None,
                            arguments_delta: partial,
                        }))
                    }
                    Some("thinking_delta") => {
                        let text = delta["thinking"].as_str().unwrap_or("").to_owned();
                        Ok(Some(StreamEvent::ThinkingDelta { delta: text }))
                    }
                    Some("signature_delta") => {
                        let signature = delta["signature"].as_str().unwrap_or("").to_owned();
                        Ok(Some(StreamEvent::SignatureDelta { signature }))
                    }
                    Some("citations_delta") => {
                        let citations = delta["citations"].clone();
                        Ok(Some(StreamEvent::CitationsDelta { citations }))
                    }
                    _ => Ok(None),
                }
            }
            "message_delta" => {
                let delta = &parsed["delta"];
                let stop_reason_str = delta["stop_reason"].as_str().unwrap_or("end_turn");
                let finish_reason = self.parse_finish_reason(stop_reason_str);

                let usage_opt = parsed.get("usage").and_then(|u| {
                    let prompt = u["input_tokens"].as_u64()? as u32;
                    let completion = u["output_tokens"].as_u64()? as u32;
                    let cached = u["cache_read_input_tokens"].as_u64().map(|v| v as u32);
                    Some(TokenUsage {
                        prompt_tokens: prompt,
                        completion_tokens: completion,
                        total_tokens: prompt + completion,
                        cached_tokens: cached,
                        ..Default::default()
                    })
                });

                Ok(Some(StreamEvent::Done { finish_reason, usage: usage_opt }))
            }
            other => {
                warn!(
                    event_type = %other,
                    "Anthropic: unknown SSE event type, ignoring"
                );
                Ok(None)
            }
        }
    }

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

#[cfg(test)]
#[path = "anthropic_tests.rs"]
mod tests;