Skip to main content

atman_runtime/providers/
anthropic.rs

1use serde::{Deserialize, Serialize};
2use tokio::sync::broadcast;
3use tokio_util::sync::CancellationToken;
4
5use crate::error::RuntimeError;
6use crate::event::{NodeEvent, Observable};
7use crate::message::{ImageData, Message, MessagePart, MessageRole};
8use crate::provider::{
9    AssistantMessage, CallTiming, DEFAULT_STREAM_BUFFER, LlmRequest, Provider, StopReason,
10    TokenUsage, estimate_tokens,
11};
12use crate::providers::classify_attachment_error;
13use crate::tool::BoxFut;
14
15pub struct AnthropicProvider {
16    name: String,
17    api_key: String,
18    base_url: String,
19    client: reqwest::Client,
20    max_tokens: u32,
21    anthropic_version: String,
22}
23
24impl AnthropicProvider {
25    pub fn new(name: impl Into<String>, api_key: impl Into<String>) -> Self {
26        Self {
27            name: name.into(),
28            api_key: api_key.into(),
29            base_url: "https://api.anthropic.com".into(),
30            client: reqwest::Client::new(),
31            max_tokens: 16384,
32            anthropic_version: "2023-06-01".into(),
33        }
34    }
35
36    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
37        self.base_url = url.into();
38        self
39    }
40
41    pub fn with_max_tokens(mut self, n: u32) -> Self {
42        self.max_tokens = n;
43        self
44    }
45
46    pub fn with_anthropic_version(mut self, v: impl Into<String>) -> Self {
47        self.anthropic_version = v.into();
48        self
49    }
50
51    fn build_body(&self, req: &LlmRequest, stream: bool) -> MessagesRequest {
52        let raw_wire: Vec<WireMessage> = req
53            .messages
54            .iter()
55            .map(|m| build_wire_message(m, false))
56            .collect();
57        let wire_messages = merge_consecutive_same_role(raw_wire);
58        let tools: Vec<WireTool> = req
59            .tools
60            .iter()
61            .map(|t| WireTool {
62                name: name_to_provider(&t.name),
63                description: t.description.clone(),
64                input_schema: t.input_schema.clone(),
65            })
66            .collect();
67        MessagesRequest {
68            model: req.model.clone(),
69            max_tokens: self.max_tokens,
70            stream,
71            system: req.system.clone(),
72            messages: wire_messages,
73            tools,
74            thinking: if req.thinking_enabled {
75                Some(ThinkingConfig {
76                    kind: "enabled",
77                    budget_tokens: self.max_tokens.saturating_sub(4096).max(1024),
78                })
79            } else {
80                None
81            },
82            cache_control: if req.cache_prompt {
83                Some(CacheControl { kind: "ephemeral" })
84            } else {
85                None
86            },
87        }
88    }
89
90    fn build_request(&self, req: &LlmRequest, stream: bool) -> reqwest::RequestBuilder {
91        let body = self.build_body(req, stream);
92        self.client
93            .post(format!("{}/v1/messages", self.base_url))
94            .header("x-api-key", &self.api_key)
95            .header("anthropic-version", &self.anthropic_version)
96            .json(&body)
97    }
98
99    #[doc(hidden)]
100    pub fn wire_body_bytes(&self, req: &LlmRequest, stream: bool) -> Vec<u8> {
101        serde_json::to_vec(&self.build_body(req, stream)).expect("serialize wire body")
102    }
103}
104
105fn name_to_provider(flow_name: &str) -> String {
106    flow_name.replace('.', "_")
107}
108
109fn name_from_provider(native: &str, tools: &[crate::tool::ToolSpec]) -> String {
110    for t in tools {
111        if name_to_provider(&t.name) == native {
112            return t.name.clone();
113        }
114    }
115    native.to_string()
116}
117
118fn build_wire_message(m: &Message, apply_cache_control: bool) -> WireMessage {
119    let role = match m.role {
120        MessageRole::User => "user",
121        MessageRole::Assistant => "assistant",
122        MessageRole::System => "user",
123        MessageRole::Tool => "user",
124    };
125    let mut blocks: Vec<ContentPart> = Vec::with_capacity(m.parts.len());
126    let last_idx = m.parts.len().saturating_sub(1);
127    for (i, part) in m.parts.iter().enumerate() {
128        blocks.push(match part {
129            MessagePart::CompactSummary { summary, .. } => ContentPart::Text {
130                text: summary.clone(),
131                cache_control: if apply_cache_control && i == last_idx {
132                    Some(CacheControl { kind: "ephemeral" })
133                } else {
134                    None
135                },
136            },
137            MessagePart::Text { text } => ContentPart::Text {
138                text: text.clone(),
139                cache_control: if apply_cache_control && i == last_idx {
140                    Some(CacheControl { kind: "ephemeral" })
141                } else {
142                    None
143                },
144            },
145            MessagePart::Image { source } => {
146                let data = match &source.data {
147                    ImageData::Base64 { data } => data.clone(),
148                    ImageData::Path { path } => {
149                        let bytes = std::fs::read(path).unwrap_or_default();
150                        use base64::Engine;
151                        base64::engine::general_purpose::STANDARD.encode(&bytes)
152                    }
153                };
154                ContentPart::Image {
155                    source: ImageSourceWire {
156                        kind: "base64",
157                        media_type: source.media_type.clone(),
158                        data,
159                    },
160                }
161            }
162            MessagePart::ToolUse { id, name, input } => ContentPart::ToolUse {
163                id: id.clone(),
164                name: name_to_provider(name),
165                input: input.clone(),
166            },
167            MessagePart::Thinking {
168                thinking,
169                signature,
170            } => ContentPart::Thinking {
171                thinking: thinking.clone(),
172                signature: signature.clone(),
173            },
174            MessagePart::ToolResult {
175                tool_use_id,
176                content,
177                is_error,
178            } => ContentPart::ToolResult {
179                tool_use_id: tool_use_id.clone(),
180                content: content.clone(),
181                is_error: *is_error,
182            },
183        });
184    }
185    WireMessage {
186        role,
187        content: MessageContent::Blocks(blocks),
188    }
189}
190
191fn merge_consecutive_same_role(wire: Vec<WireMessage>) -> Vec<WireMessage> {
192    let mut out: Vec<WireMessage> = Vec::with_capacity(wire.len());
193    for msg in wire {
194        let WireMessage { role, content } = msg;
195        let mut content = Some(content);
196        if let Some(last) = out.last_mut()
197            && last.role == role
198            && let Some(msg_content) = content.take()
199        {
200            let MessageContent::Blocks(last_blocks) = &mut last.content;
201            let MessageContent::Blocks(msg_blocks) = msg_content;
202            last_blocks.extend(msg_blocks);
203        }
204        if let Some(content) = content {
205            out.push(WireMessage { role, content });
206        }
207    }
208    out
209}
210
211impl Provider for AnthropicProvider {
212    fn name(&self) -> &str {
213        &self.name
214    }
215
216    fn call<'a>(&'a self, req: LlmRequest) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>> {
217        let request = self.build_request(&req, false);
218        Box::pin(async move {
219            let resp = request.send().await.map_err(net_err)?;
220            let status = resp.status();
221            let body: MessagesResponse = if status.is_success() {
222                resp.json().await.map_err(net_err)?
223            } else {
224                let body_text = resp.text().await.unwrap_or_default();
225                if let Some(reason) = classify_attachment_error(status.as_u16(), &body_text) {
226                    return Err(RuntimeError::AttachmentError { reason });
227                }
228                return Err(RuntimeError::ToolFailed(format!(
229                    "anthropic http {status}: {body_text}"
230                )));
231            };
232            Ok(response_to_assistant(
233                body,
234                next_turn_id_from_req(&req),
235                &req.tools,
236            ))
237        })
238    }
239
240    fn call_streaming(&self, req: LlmRequest) -> Observable<AssistantMessage> {
241        let request = self.build_request(&req, true);
242        let turn_id = next_turn_id_from_req(&req);
243        let tools: Vec<crate::tool::ToolSpec> = req.tools.clone();
244        let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
245        let cancel = CancellationToken::new();
246        let cancel_for_task = cancel.clone();
247        let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> = Box::pin(
248            async move {
249                use eventsource_stream::Eventsource;
250                use futures::StreamExt;
251
252                let resp = tokio::select! {
253                    biased;
254                    _ = cancel_for_task.cancelled() => return Err(RuntimeError::Cancelled("anthropic cancelled before send".into())),
255                    r = request.send() => r.map_err(net_err)?,
256                };
257                let status = resp.status();
258                if !status.is_success() {
259                    let body = resp.text().await.unwrap_or_default();
260                    if let Some(reason) = classify_attachment_error(status.as_u16(), &body) {
261                        return Err(RuntimeError::AttachmentError { reason });
262                    }
263                    return Err(RuntimeError::ToolFailed(format!(
264                        "anthropic http {status}: {body}"
265                    )));
266                }
267
268                let mut stream = resp.bytes_stream().eventsource();
269                let mut acc_text = String::new();
270                let mut acc_thinking = String::new();
271                let mut acc_signature: Option<String> = None;
272                let mut cumulative = 0u64;
273                let mut input_tokens: u64 = 0;
274                let mut cache_read_tokens: u64 = 0;
275                let mut cache_write_tokens: u64 = 0;
276                let mut tool_use_partial: Vec<PartialToolUse> = Vec::new();
277                let mut stop_reason = StopReason::End;
278                while let Some(event) = tokio::select! {
279                    biased;
280                    _ = cancel_for_task.cancelled() => None,
281                    next = stream.next() => next,
282                } {
283                    let event = event.map_err(|e| RuntimeError::ToolFailed(format!("sse: {e}")))?;
284                    if event.data.is_empty() {
285                        continue;
286                    }
287                    let parsed: serde_json::Value = match serde_json::from_str(&event.data) {
288                        Ok(v) => v,
289                        Err(_) => continue,
290                    };
291                    let ty = parsed.get("type").and_then(|v| v.as_str()).unwrap_or("");
292                    match ty {
293                        "message_start" => {
294                            if let Some(usage) = parsed.pointer("/message/usage") {
295                                input_tokens = usage
296                                    .get("input_tokens")
297                                    .and_then(|v| v.as_u64())
298                                    .unwrap_or(0);
299                                cache_read_tokens = usage
300                                    .get("cache_read_input_tokens")
301                                    .and_then(|v| v.as_u64())
302                                    .unwrap_or(0);
303                                cache_write_tokens = usage
304                                    .get("cache_creation_input_tokens")
305                                    .and_then(|v| v.as_u64())
306                                    .unwrap_or(0);
307                            }
308                        }
309                        "content_block_start" => {
310                            if let Some(block) = parsed.get("content_block")
311                                && block.get("type").and_then(|v| v.as_str()) == Some("tool_use")
312                                && let (Some(id), Some(name)) = (
313                                    block.get("id").and_then(|v| v.as_str()),
314                                    block.get("name").and_then(|v| v.as_str()),
315                                )
316                            {
317                                tool_use_partial.push(PartialToolUse {
318                                    id: id.to_string(),
319                                    name: name.to_string(),
320                                    input_json: String::new(),
321                                });
322                            }
323                        }
324                        "content_block_delta" => {
325                            if let Some(delta) = parsed.get("delta") {
326                                let delta_ty =
327                                    delta.get("type").and_then(|v| v.as_str()).unwrap_or("");
328                                if delta_ty == "text_delta" {
329                                    if let Some(text) = delta.get("text").and_then(|v| v.as_str()) {
330                                        acc_text.push_str(text);
331                                        cumulative += estimate_tokens(text);
332                                        let _ = tx.send(NodeEvent::LlmChunk {
333                                            text: text.to_string(),
334                                            cumulative_tokens: cumulative,
335                                        });
336                                    }
337                                } else if delta_ty == "thinking_delta" {
338                                    if let Some(text) =
339                                        delta.get("thinking").and_then(|v| v.as_str())
340                                    {
341                                        acc_thinking.push_str(text);
342                                        let _ = tx.send(NodeEvent::ThinkingChunk {
343                                            text: text.to_string(),
344                                        });
345                                    }
346                                } else if let Some(text) =
347                                    delta.get("reasoning_content").and_then(|v| v.as_str())
348                                {
349                                    acc_thinking.push_str(text);
350                                    let _ = tx.send(NodeEvent::ThinkingChunk {
351                                        text: text.to_string(),
352                                    });
353                                } else if delta_ty == "signature_delta" {
354                                    if let Some(sig) =
355                                        delta.get("signature").and_then(|v| v.as_str())
356                                    {
357                                        acc_signature = Some(sig.to_string());
358                                    }
359                                } else if delta_ty == "input_json_delta"
360                                    && let Some(partial) =
361                                        delta.get("partial_json").and_then(|v| v.as_str())
362                                    && let Some(last) = tool_use_partial.last_mut()
363                                {
364                                    last.input_json.push_str(partial);
365                                }
366                            }
367                        }
368                        "message_delta" => {
369                            if let Some(out) = parsed
370                                .pointer("/usage/output_tokens")
371                                .and_then(|v| v.as_u64())
372                            {
373                                cumulative = out;
374                            }
375                            if let Some(inp) = parsed
376                                .pointer("/usage/input_tokens")
377                                .and_then(|v| v.as_u64())
378                            {
379                                input_tokens = inp;
380                            }
381                            if let Some(cr) = parsed
382                                .pointer("/usage/cache_read_input_tokens")
383                                .and_then(|v| v.as_u64())
384                            {
385                                cache_read_tokens = cr;
386                            }
387                            if let Some(cw) = parsed
388                                .pointer("/usage/cache_creation_input_tokens")
389                                .and_then(|v| v.as_u64())
390                            {
391                                cache_write_tokens = cw;
392                            }
393                            if let Some(reason) = parsed
394                                .pointer("/delta/stop_reason")
395                                .and_then(|v| v.as_str())
396                            {
397                                stop_reason = parse_stop_reason(reason);
398                            }
399                        }
400                        "message_stop" => break,
401                        _ => {}
402                    }
403                }
404                if cancel_for_task.is_cancelled() {
405                    let _ = tx.send(NodeEvent::LlmDone {
406                        total_tokens: cumulative,
407                    });
408                    return Err(RuntimeError::Cancelled(
409                        "anthropic cancelled mid-stream".into(),
410                    ));
411                }
412                let _ = tx.send(NodeEvent::LlmDone {
413                    total_tokens: cumulative,
414                });
415
416                let mut parts: Vec<MessagePart> = Vec::new();
417                if !acc_thinking.is_empty() {
418                    parts.push(MessagePart::Thinking {
419                        thinking: acc_thinking,
420                        signature: acc_signature,
421                    });
422                }
423                if !acc_text.is_empty() {
424                    parts.push(MessagePart::Text { text: acc_text });
425                }
426                for pu in tool_use_partial {
427                    let input: serde_json::Value = if pu.input_json.is_empty() {
428                        serde_json::Value::Object(Default::default())
429                    } else {
430                        serde_json::from_str(&pu.input_json).unwrap_or(serde_json::Value::Null)
431                    };
432                    parts.push(MessagePart::ToolUse {
433                        id: pu.id,
434                        name: name_from_provider(&pu.name, &tools),
435                        input,
436                    });
437                }
438                Ok(AssistantMessage {
439                    message: Message {
440                        role: MessageRole::Assistant,
441                        parts,
442                        turn_id,
443                    },
444                    stop_reason,
445                    token_usage: TokenUsage {
446                        input: input_tokens,
447                        cached_input: cache_read_tokens,
448                        output: cumulative,
449                        cache_write: cache_write_tokens,
450                        ..Default::default()
451                    },
452                    timing: CallTiming::default(),
453                    model: String::new(),
454                    response_id: None,
455                })
456            },
457        );
458        Observable {
459            output,
460            events,
461            cancel,
462        }
463    }
464}
465
466struct PartialToolUse {
467    id: String,
468    name: String,
469    input_json: String,
470}
471
472fn response_to_assistant(
473    body: MessagesResponse,
474    turn_id: crate::event::TurnId,
475    tools: &[crate::tool::ToolSpec],
476) -> AssistantMessage {
477    let mut parts: Vec<MessagePart> = Vec::new();
478    for block in body.content {
479        match block {
480            ContentBlock::Text { text } => parts.push(MessagePart::Text { text }),
481            ContentBlock::Thinking {
482                thinking,
483                signature,
484            } => parts.push(MessagePart::Thinking {
485                thinking,
486                signature,
487            }),
488            ContentBlock::ToolUse { id, name, input } => parts.push(MessagePart::ToolUse {
489                id,
490                name: name_from_provider(&name, tools),
491                input,
492            }),
493            ContentBlock::Other => {}
494        }
495    }
496    let stop_reason = body
497        .stop_reason
498        .as_deref()
499        .map(parse_stop_reason)
500        .unwrap_or(StopReason::End);
501    let usage = body
502        .usage
503        .map(|u| TokenUsage {
504            input: u.input_tokens.unwrap_or(0),
505            cached_input: u.cache_read_input_tokens.unwrap_or(0),
506            output: u.output_tokens.unwrap_or(0),
507            cache_write: u.cache_creation_input_tokens.unwrap_or(0),
508            ..Default::default()
509        })
510        .unwrap_or_default();
511    AssistantMessage {
512        message: Message {
513            role: MessageRole::Assistant,
514            parts,
515            turn_id,
516        },
517        stop_reason,
518        token_usage: usage,
519        timing: CallTiming::default(),
520        model: body.model.unwrap_or_default(),
521        response_id: body.id,
522    }
523}
524
525fn parse_stop_reason(s: &str) -> StopReason {
526    match s {
527        "tool_use" => StopReason::ToolUse,
528        "max_tokens" => StopReason::Length,
529        _ => StopReason::End,
530    }
531}
532
533fn next_turn_id_from_req(req: &LlmRequest) -> crate::event::TurnId {
534    req.messages
535        .first()
536        .map(|m| m.turn_id.clone())
537        .unwrap_or_else(crate::event::TurnId::now)
538}
539
540fn net_err(e: reqwest::Error) -> RuntimeError {
541    RuntimeError::ToolFailed(format!("anthropic net: {e}"))
542}
543
544#[derive(Serialize, Clone)]
545struct MessagesRequest {
546    model: String,
547    max_tokens: u32,
548    stream: bool,
549    #[serde(skip_serializing_if = "Option::is_none")]
550    system: Option<String>,
551    messages: Vec<WireMessage>,
552    #[serde(skip_serializing_if = "Vec::is_empty")]
553    tools: Vec<WireTool>,
554    #[serde(skip_serializing_if = "Option::is_none")]
555    thinking: Option<ThinkingConfig>,
556    #[serde(skip_serializing_if = "Option::is_none")]
557    cache_control: Option<CacheControl>,
558}
559
560#[derive(Serialize, Clone)]
561struct ThinkingConfig {
562    #[serde(rename = "type")]
563    kind: &'static str,
564    budget_tokens: u32,
565}
566
567#[derive(Serialize, Clone)]
568struct WireTool {
569    name: String,
570    #[serde(skip_serializing_if = "Option::is_none")]
571    description: Option<String>,
572    input_schema: serde_json::Value,
573}
574
575#[derive(Serialize, Clone)]
576struct WireMessage {
577    role: &'static str,
578    content: MessageContent,
579}
580
581#[derive(Serialize, Clone)]
582#[serde(untagged)]
583enum MessageContent {
584    Blocks(Vec<ContentPart>),
585}
586
587#[derive(Serialize, Clone)]
588#[serde(tag = "type", rename_all = "snake_case")]
589enum ContentPart {
590    Text {
591        text: String,
592        #[serde(skip_serializing_if = "Option::is_none")]
593        cache_control: Option<CacheControl>,
594    },
595    Thinking {
596        thinking: String,
597        #[serde(skip_serializing_if = "Option::is_none")]
598        signature: Option<String>,
599    },
600    Image {
601        source: ImageSourceWire,
602    },
603    ToolUse {
604        id: String,
605        name: String,
606        input: serde_json::Value,
607    },
608    ToolResult {
609        tool_use_id: String,
610        content: String,
611        #[serde(skip_serializing_if = "core::ops::Not::not")]
612        is_error: bool,
613    },
614}
615
616#[derive(Serialize, Clone)]
617struct ImageSourceWire {
618    #[serde(rename = "type")]
619    kind: &'static str,
620    media_type: String,
621    data: String,
622}
623
624#[derive(Serialize, Clone)]
625struct CacheControl {
626    #[serde(rename = "type")]
627    kind: &'static str,
628}
629
630#[derive(Deserialize)]
631struct MessagesResponse {
632    content: Vec<ContentBlock>,
633    #[serde(default)]
634    stop_reason: Option<String>,
635    #[serde(default)]
636    usage: Option<AnthropicUsage>,
637    #[serde(default)]
638    model: Option<String>,
639    #[serde(default)]
640    id: Option<String>,
641}
642
643#[derive(Deserialize, Default)]
644struct AnthropicUsage {
645    #[serde(default)]
646    input_tokens: Option<u64>,
647    #[serde(default)]
648    output_tokens: Option<u64>,
649    #[serde(default)]
650    cache_read_input_tokens: Option<u64>,
651    #[serde(default)]
652    cache_creation_input_tokens: Option<u64>,
653}
654
655#[derive(Deserialize)]
656#[serde(tag = "type", rename_all = "snake_case")]
657enum ContentBlock {
658    Text {
659        text: String,
660    },
661    Thinking {
662        thinking: String,
663        #[serde(default)]
664        signature: Option<String>,
665    },
666    ToolUse {
667        id: String,
668        name: String,
669        input: serde_json::Value,
670    },
671    #[serde(other)]
672    Other,
673}