clawmates/crates/cm-llm/src/anthropic.rs

//! Anthropic Messages API provider (cloud target).

use async_stream::try_stream;
use eventsource_stream::Eventsource;
use futures::StreamExt;
use serde_json::{json, Value};

use crate::provider::{
    ChatRequest, ChatRole, ContentPart, EventStream, LlmError, LlmEvent, LlmProvider, StopReason,
};

pub struct AnthropicProvider {
    client: reqwest::Client,
    base_url: String,
    api_key: String,
}

impl AnthropicProvider {
    pub fn new(api_key: String) -> AnthropicProvider {
        AnthropicProvider::with_base_url(api_key, "https://api.anthropic.com".into())
    }

    /// Base URL override exists for self-hosted proxies and tests.
    pub fn with_base_url(api_key: String, base_url: String) -> AnthropicProvider {
        AnthropicProvider {
            client: reqwest::Client::new(),
            base_url,
            api_key,
        }
    }

    fn wire_messages(request: &ChatRequest) -> Vec<Value> {
        request
            .messages
            .iter()
            .map(|message| {
                let role = match message.role {
                    ChatRole::User => "user",
                    ChatRole::Assistant => "assistant",
                };
                let content: Vec<Value> = message
                    .parts
                    .iter()
                    .map(|part| match part {
                        ContentPart::Text { text } => json!({"type": "text", "text": text}),
                        ContentPart::ToolUse { id, name, input } => {
                            json!({"type": "tool_use", "id": id,
                                   "name": crate::wire_tool_name(name), "input": input})
                        }
                        ContentPart::ToolResult {
                            tool_use_id,
                            content,
                        } => json!({
                            "type": "tool_result",
                            "tool_use_id": tool_use_id,
                            "content": content.to_string(),
                        }),
                    })
                    .collect();
                json!({"role": role, "content": content})
            })
            .collect()
    }
}

fn stop_reason(wire: &str) -> StopReason {
    match wire {
        "tool_use" => StopReason::ToolUse,
        "max_tokens" => StopReason::MaxTokens,
        _ => StopReason::EndTurn,
    }
}

#[async_trait::async_trait]
impl LlmProvider for AnthropicProvider {
    async fn stream(&self, request: ChatRequest) -> Result<EventStream, LlmError> {
        let tools: Vec<Value> = request
            .tools
            .iter()
            .map(|t| {
                json!({
                    "name": crate::wire_tool_name(&t.name),
                    "description": t.description,
                    "input_schema": t.input_schema,
                })
            })
            .collect();
        let body = json!({
            "model": request.model,
            "max_tokens": request.max_tokens,
            "system": request.system,
            "messages": AnthropicProvider::wire_messages(&request),
            "tools": tools,
            "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")
            .json(&body)
            .send()
            .await
            .map_err(|e| LlmError::Transport(e.to_string()))?;
        if !response.status().is_success() {
            let status = response.status();
            let detail = response.text().await.unwrap_or_default();
            return Err(LlmError::Api(format!("{status}: {detail}")));
        }

        let mut sse = response.bytes_stream().eventsource();
        let stream = try_stream! {
            // tool_use input arrives as accumulated partial JSON between
            // content_block_start and content_block_stop.
            let mut pending_tool: Option<(String, String, String)> = None; // id, name, json
            let mut input_tokens: u32 = 0;
            while let Some(event) = sse.next().await {
                let event = event.map_err(|e| LlmError::Transport(e.to_string()))?;
                let data: Value = serde_json::from_str(&event.data)
                    .map_err(|e| LlmError::Wire(format!("{e}: {}", event.data)))?;
                match data["type"].as_str().unwrap_or_default() {
                    "content_block_start" => {
                        let block = &data["content_block"];
                        if block["type"] == "tool_use" {
                            pending_tool = Some((
                                block["id"].as_str().unwrap_or_default().to_owned(),
                                crate::internal_tool_name(
                                    block["name"].as_str().unwrap_or_default(),
                                ),
                                String::new(),
                            ));
                        }
                    }
                    "content_block_delta" => {
                        let delta = &data["delta"];
                        match delta["type"].as_str().unwrap_or_default() {
                            "text_delta" => {
                                if let Some(text) = delta["text"].as_str() {
                                    yield LlmEvent::TextDelta(text.to_owned());
                                }
                            }
                            "input_json_delta" => {
                                if let Some((_, _, buf)) = pending_tool.as_mut() {
                                    buf.push_str(delta["partial_json"].as_str().unwrap_or_default());
                                }
                            }
                            _ => {}
                        }
                    }
                    "content_block_stop" => {
                        if let Some((id, name, buf)) = pending_tool.take() {
                            let input: Value = if buf.is_empty() {
                                json!({})
                            } else {
                                serde_json::from_str(&buf)
                                    .map_err(|e| LlmError::Wire(format!("tool input: {e}")))?
                            };
                            yield LlmEvent::ToolUse { id, name, input };
                        }
                    }
                    "message_start" => {
                        input_tokens =
                            data["message"]["usage"]["input_tokens"].as_u64().unwrap_or(0) as u32;
                    }
                    "message_delta" => {
                        if let Some(out) = data["usage"]["output_tokens"].as_u64() {
                            yield LlmEvent::Usage {
                                input_tokens,
                                output_tokens: out as u32,
                            };
                        }
                        if let Some(reason) = data["delta"]["stop_reason"].as_str() {
                            yield LlmEvent::Stop(stop_reason(reason));
                        }
                    }
                    "error" => {
                        Err(LlmError::Api(data["error"]["message"]
                            .as_str()
                            .unwrap_or("unknown")
                            .to_owned()))?;
                    }
                    _ => {}
                }
            }
        };
        Ok(Box::pin(stream))
    }
}