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