ic-rig 0.2.0

//! Anthropic Messages API provider.
//!
//! # Usage
//!
//! ```rust,ignore
//! use irig::providers::anthropic::{Client, CLAUDE_OPUS_5};
//!
//! let client = Client::new(my_http, "sk-ant-...");
//! let model  = client.model(CLAUDE_OPUS_5);
//!
//! let agent = irig::Agent::builder(model)
//!     .preamble("You are a helpful assistant.")
//!     .max_tokens(1024)
//!     .build();
//!
//! let reply = agent.prompt("Hello!").await?;
//! ```

use crate::{
    completion::{CompletionError, CompletionModel, CompletionRequest, CompletionResponse, ModelChoice, Usage},
    http::{HttpClient, HttpRequest},
    message::{AssistantContent, Message, ToolCall, UserContent},
    tool::ToolDefinition,
};
use serde::{Deserialize, Serialize};

// ── Model constants ───────────────────────────────────────────────────────────

// Current generation (recommended). These IDs are bare aliases with no date
// suffix by design — Anthropic resolves them to the latest snapshot.
pub const CLAUDE_FABLE_5: &str = "claude-fable-5";
pub const CLAUDE_OPUS_5: &str = "claude-opus-5";
pub const CLAUDE_SONNET_5: &str = "claude-sonnet-5";
pub const CLAUDE_HAIKU_4_5: &str = "claude-haiku-4-5-20251001";

// Previous generation (still active).
pub const CLAUDE_OPUS_4_8: &str = "claude-opus-4-8";
pub const CLAUDE_OPUS_4_7: &str = "claude-opus-4-7";
pub const CLAUDE_OPUS_4_6: &str = "claude-opus-4-6";
pub const CLAUDE_SONNET_4_6: &str = "claude-sonnet-4-6";

// Legacy dated snapshots (still active, but superseded by the 5-series above).
pub const CLAUDE_OPUS_4_5: &str = "claude-opus-4-5-20251101";
/// Fixed 2026-08-27: this was previously (incorrectly) dated `-20251101`,
/// copy-pasted from `CLAUDE_OPUS_4_5`. Anthropic's actual Sonnet 4.5 snapshot
/// is dated `-20250929`.
pub const CLAUDE_SONNET_4_5: &str = "claude-sonnet-4-5-20250929";

/// Deprecated by Anthropic; retirement date TBD.
#[deprecated(note = "deprecated by Anthropic (retirement TBD); use CLAUDE_OPUS_5 instead")]
pub const CLAUDE_OPUS_4: &str = "claude-opus-4-20250514";
/// Deprecated by Anthropic; retirement date TBD.
#[deprecated(note = "deprecated by Anthropic (retirement TBD); use CLAUDE_SONNET_5 instead")]
pub const CLAUDE_SONNET_4: &str = "claude-sonnet-4-20250514";

/// Anthropic API version header value.
const API_VERSION: &str = "2023-06-01";
const BASE_URL: &str = "https://api.anthropic.com/v1";

// ── Client ────────────────────────────────────────────────────────────────────

/// Anthropic API client. Use [`model`](Client::model) to get a [`CompletionModel`].
pub struct Client<H> {
    http: H,
    api_key: String,
    base_url: String,
}

impl<H: HttpClient + Clone> Client<H> {
    pub fn new(http: H, api_key: impl Into<String>) -> Self {
        Self { http, api_key: api_key.into(), base_url: BASE_URL.to_owned() }
    }

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

    pub fn model(&self, model: impl Into<String>) -> Model<H> {
        Model {
            http: self.http.clone(),
            api_key: self.api_key.clone(),
            base_url: self.base_url.clone(),
            model: model.into(),
        }
    }
}

// ── Model ─────────────────────────────────────────────────────────────────────

pub struct Model<H> {
    http: H,
    api_key: String,
    base_url: String,
    model: String,
}

impl<H: HttpClient> CompletionModel for Model<H> {
    type Error = CompletionError;

    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, CompletionError> {
        let body = build_request(&self.model, request)?;
        let bytes = serde_json::to_vec(&body)?;

        let http_req = HttpRequest::new(format!("{}/messages", self.base_url))
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", API_VERSION)
            .json_body(bytes);

        let resp = self.http.post(http_req).await
            .map_err(|e| CompletionError::Http(e.to_string()))?;

        if !resp.is_success() {
            let message = String::from_utf8_lossy(&resp.body).into_owned();
            return Err(CompletionError::Provider { status: resp.status, message });
        }

        let api_resp: ApiResponse = resp.json()?;
        parse_response(api_resp)
    }
}

// ── Wire types (Anthropic JSON format) ───────────────────────────────────────

#[derive(Serialize)]
struct ApiRequest {
    model: String,
    /// `max_tokens` is required by the Anthropic API.
    max_tokens: u32,
    messages: Vec<ApiMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<ApiTool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    thinking: Option<ApiThinking>,
}

#[derive(Serialize)]
struct ApiThinking {
    #[serde(rename = "type")]
    kind: &'static str,
}

#[derive(Serialize)]
struct ApiMessage {
    role: &'static str,
    content: Vec<serde_json::Value>,
}

#[derive(Serialize)]
struct ApiTool {
    name: String,
    description: String,
    /// Anthropic calls this `input_schema`, not `parameters`.
    input_schema: serde_json::Value,
}

#[derive(Deserialize)]
struct ApiResponse {
    content: Vec<ApiContent>,
    usage: ApiUsage,
}

#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ApiContent {
    Text {
        text: String,
    },
    ToolUse {
        id: String,
        name: String,
        /// Anthropic delivers tool arguments as a JSON **object** (not a string).
        input: serde_json::Value,
    },
    /// Extended-thinking trace. Only appears when `thinking` is enabled on
    /// the request. Never treated as the answer — collected into
    /// [`CompletionResponse::reasoning`] instead.
    Thinking {
        thinking: String,
    },
    /// Thinking content redacted by Anthropic's safety systems. The `data`
    /// is encrypted and not human-readable, so there's nothing to surface.
    RedactedThinking {
        #[allow(dead_code)]
        data: String,
    },
}

#[derive(Deserialize)]
struct ApiUsage {
    input_tokens: u32,
    output_tokens: u32,
}

// ── Conversion helpers ────────────────────────────────────────────────────────

fn build_request(model: &str, req: CompletionRequest) -> Result<ApiRequest, CompletionError> {
    // Anthropic takes `system` as a top-level field; pull it out of messages.
    let mut system: Option<String> = None;
    let mut chat_messages: Vec<Message> = Vec::new();

    for msg in req.messages {
        match msg {
            Message::System { content } => {
                // Last system message wins (matches how most providers behave).
                system = Some(content);
            }
            other => chat_messages.push(other),
        }
    }

    let messages = convert_messages(chat_messages)?;
    let tools = req.tools.into_iter().map(convert_tool).collect();

    Ok(ApiRequest {
        model: model.to_owned(),
        // Anthropic requires max_tokens; default to 1024 if not specified.
        max_tokens: req.max_tokens.unwrap_or(1024),
        messages,
        system,
        tools,
        temperature: req.temperature,
        // "adaptive"/"disabled" match Claude 4.6 and newer; older dated
        // snapshots use a different (budget_tokens-based) scheme and may
        // reject this — see CompletionRequest::thinking.
        thinking: req.thinking.map(|enabled| ApiThinking {
            kind: if enabled { "adaptive" } else { "disabled" },
        }),
    })
}

/// Convert irig messages to Anthropic's content-array format.
///
/// Key differences from OpenAI:
/// - User tool results use `type: "tool_result"` with `tool_use_id` (not `tool_call_id`).
/// - Assistant tool calls use `type: "tool_use"` with an `input` object.
/// - Multiple content parts (text + tool results) can coexist in one message.
fn convert_messages(messages: Vec<Message>) -> Result<Vec<ApiMessage>, CompletionError> {
    let mut out = Vec::new();

    for msg in messages {
        match msg {
            Message::System { .. } => {
                // Already extracted above; skip any stragglers.
            }

            Message::User { content } => {
                let parts: Vec<serde_json::Value> = content
                    .into_iter()
                    .map(|part| match part {
                        UserContent::Text(t) => {
                            serde_json::json!({ "type": "text", "text": t.text })
                        }
                        UserContent::ToolResult(r) => {
                            serde_json::json!({
                                "type": "tool_result",
                                "tool_use_id": r.call_id,
                                "content": r.content,
                            })
                        }
                    })
                    .collect();

                out.push(ApiMessage { role: "user", content: parts });
            }

            Message::Assistant { content } => {
                let parts: Result<Vec<serde_json::Value>, CompletionError> = content
                    .into_iter()
                    .map(|part| match part {
                        AssistantContent::Text(t) => {
                            Ok(serde_json::json!({ "type": "text", "text": t.text }))
                        }
                        AssistantContent::ToolCall(c) => {
                            Ok(serde_json::json!({
                                "type": "tool_use",
                                "id": c.id,
                                "name": c.name,
                                // Anthropic expects the raw object, not a JSON string.
                                "input": c.arguments,
                            }))
                        }
                    })
                    .collect();

                out.push(ApiMessage { role: "assistant", content: parts? });
            }
        }
    }

    Ok(out)
}

fn convert_tool(def: ToolDefinition) -> ApiTool {
    ApiTool {
        name: def.name,
        description: def.description,
        input_schema: def.parameters,
    }
}

fn parse_response(resp: ApiResponse) -> Result<CompletionResponse, CompletionError> {
    let usage = Usage {
        prompt_tokens: resp.usage.input_tokens,
        completion_tokens: resp.usage.output_tokens,
    };

    // Collect tool calls, text, and thinking separately; prefer tool calls
    // if present. Thinking never becomes the answer — only `reasoning`.
    let mut text: Option<String> = None;
    let mut reasoning: Option<String> = None;
    let mut tool_calls: Vec<ToolCall> = Vec::new();

    for block in resp.content {
        match block {
            ApiContent::Text { text: t } => text = Some(t),
            ApiContent::ToolUse { id, name, input } => {
                tool_calls.push(ToolCall { id, name, arguments: input });
            }
            ApiContent::Thinking { thinking } => {
                reasoning = Some(match reasoning {
                    Some(existing) => format!("{existing}\n{thinking}"),
                    None => thinking,
                });
            }
            ApiContent::RedactedThinking { .. } => {
                // Encrypted trace; nothing human-readable to surface.
            }
        }
    }

    let choice = if !tool_calls.is_empty() {
        ModelChoice::ToolCall(tool_calls)
    } else {
        let t = text.ok_or_else(|| CompletionError::Response("empty content array".into()))?;
        ModelChoice::Message(t)
    };

    Ok(CompletionResponse { choice, reasoning, usage: Some(usage) })
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    #[test]
    fn thinking_toggle_serializes_expected_shape() {
        let mut on = CompletionRequest::new(vec![Message::user("hi")]);
        on.thinking = Some(true);
        let json = serde_json::to_value(build_request(CLAUDE_SONNET_5, on).unwrap()).unwrap();
        assert_eq!(json["thinking"]["type"], "adaptive");

        let mut off = CompletionRequest::new(vec![Message::user("hi")]);
        off.thinking = Some(false);
        let json = serde_json::to_value(build_request(CLAUDE_SONNET_5, off).unwrap()).unwrap();
        assert_eq!(json["thinking"]["type"], "disabled");

        let unset = CompletionRequest::new(vec![Message::user("hi")]);
        let json = serde_json::to_value(build_request(CLAUDE_SONNET_5, unset).unwrap()).unwrap();
        assert!(json.get("thinking").is_none());
    }

    fn make_response(json: &str) -> ApiResponse {
        serde_json::from_str(json).expect("test fixture must deserialise")
    }

    #[test]
    fn thinking_block_is_kept_out_of_the_answer() {
        // Regression test: before ApiContent had a Thinking variant, any
        // response with extended thinking enabled would fail to parse.
        let resp = make_response(
            r#"{
                "content": [
                    {"type": "thinking", "thinking": "Let me work this out..."},
                    {"type": "text", "text": "The answer is 4."}
                ],
                "usage": {"input_tokens": 10, "output_tokens": 5}
            }"#,
        );

        let result = parse_response(resp).unwrap();
        assert!(matches!(result.choice, ModelChoice::Message(ref s) if s == "The answer is 4."));
        assert_eq!(result.reasoning.as_deref(), Some("Let me work this out..."));
    }

    #[test]
    fn redacted_thinking_block_does_not_crash() {
        let resp = make_response(
            r#"{
                "content": [
                    {"type": "redacted_thinking", "data": "encrypted-blob"},
                    {"type": "text", "text": "ok"}
                ],
                "usage": {"input_tokens": 3, "output_tokens": 1}
            }"#,
        );

        let result = parse_response(resp).unwrap();
        assert!(matches!(result.choice, ModelChoice::Message(ref s) if s == "ok"));
        assert_eq!(result.reasoning, None);
    }

    #[test]
    fn plain_text_response_has_no_reasoning() {
        let resp = make_response(
            r#"{
                "content": [{"type": "text", "text": "Hello, world!"}],
                "usage": {"input_tokens": 10, "output_tokens": 5}
            }"#,
        );

        let result = parse_response(resp).unwrap();
        assert!(matches!(result.choice, ModelChoice::Message(ref s) if s == "Hello, world!"));
        assert_eq!(result.reasoning, None);
    }
}