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                                match block.get("type").and_then(|v| v.as_str()) {
320                                    Some("tool_use") => {
321                                        if let (Some(id), Some(name)) = (
322                                            block.get("id").and_then(|v| v.as_str()),
323                                            block.get("name").and_then(|v| v.as_str()),
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                                    Some("thinking") => {
333                                        if let Some(sig) =
334                                            block.get("signature").and_then(|v| v.as_str())
335                                        {
336                                            acc_signature = Some(sig.to_string());
337                                        }
338                                    }
339                                    _ => {}
340                                }
341                            }
342                        }
343                        "content_block_delta" => {
344                            if let Some(delta) = parsed.get("delta") {
345                                let delta_ty =
346                                    delta.get("type").and_then(|v| v.as_str()).unwrap_or("");
347                                if delta_ty == "text_delta" {
348                                    if let Some(text) = delta.get("text").and_then(|v| v.as_str()) {
349                                        acc_text.push_str(text);
350                                        cumulative += estimate_tokens(text);
351                                        let _ = tx.send(NodeEvent::LlmChunk {
352                                            text: text.to_string(),
353                                            cumulative_tokens: cumulative,
354                                        });
355                                    }
356                                } else if delta_ty == "thinking_delta" {
357                                    if let Some(text) =
358                                        delta.get("thinking").and_then(|v| v.as_str())
359                                    {
360                                        acc_thinking.push_str(text);
361                                        let _ = tx.send(NodeEvent::ThinkingChunk {
362                                            text: text.to_string(),
363                                        });
364                                    }
365                                } else if let Some(text) =
366                                    delta.get("reasoning_content").and_then(|v| v.as_str())
367                                {
368                                    acc_thinking.push_str(text);
369                                    let _ = tx.send(NodeEvent::ThinkingChunk {
370                                        text: text.to_string(),
371                                    });
372                                } else if delta_ty == "signature_delta" {
373                                    if let Some(sig) =
374                                        delta.get("signature").and_then(|v| v.as_str())
375                                    {
376                                        acc_signature = Some(sig.to_string());
377                                    }
378                                } else if delta_ty == "input_json_delta"
379                                    && let Some(partial) =
380                                        delta.get("partial_json").and_then(|v| v.as_str())
381                                    && let Some(last) = tool_use_partial.last_mut()
382                                {
383                                    last.input_json.push_str(partial);
384                                }
385                            }
386                        }
387                        "message_delta" => {
388                            if let Some(out) = parsed
389                                .pointer("/usage/output_tokens")
390                                .and_then(|v| v.as_u64())
391                            {
392                                cumulative = out;
393                            }
394                            if let Some(inp) = parsed
395                                .pointer("/usage/input_tokens")
396                                .and_then(|v| v.as_u64())
397                            {
398                                input_tokens = inp;
399                            }
400                            if let Some(cr) = parsed
401                                .pointer("/usage/cache_read_input_tokens")
402                                .and_then(|v| v.as_u64())
403                            {
404                                cache_read_tokens = cr;
405                            }
406                            if let Some(cw) = parsed
407                                .pointer("/usage/cache_creation_input_tokens")
408                                .and_then(|v| v.as_u64())
409                            {
410                                cache_write_tokens = cw;
411                            }
412                            if let Some(reason) = parsed
413                                .pointer("/delta/stop_reason")
414                                .and_then(|v| v.as_str())
415                            {
416                                stop_reason = parse_stop_reason(reason);
417                            }
418                        }
419                        "message_stop" => break,
420                        _ => {}
421                    }
422                }
423                if cancel_for_task.is_cancelled() {
424                    let _ = tx.send(NodeEvent::LlmDone {
425                        total_tokens: cumulative,
426                    });
427                    return Err(RuntimeError::Cancelled(
428                        "anthropic cancelled mid-stream".into(),
429                    ));
430                }
431                let _ = tx.send(NodeEvent::LlmDone {
432                    total_tokens: cumulative,
433                });
434
435                let mut parts: Vec<MessagePart> = Vec::new();
436                if !acc_thinking.is_empty() {
437                    if req.thinking_enabled && acc_signature.is_none() {
438                        return Err(RuntimeError::ThinkingSignatureMissing);
439                    }
440                    parts.push(MessagePart::Thinking {
441                        thinking: acc_thinking,
442                        signature: acc_signature,
443                    });
444                }
445                if !acc_text.is_empty() {
446                    parts.push(MessagePart::Text { text: acc_text });
447                }
448                for pu in tool_use_partial {
449                    let input: serde_json::Value = if pu.input_json.is_empty() {
450                        serde_json::Value::Object(Default::default())
451                    } else {
452                        serde_json::from_str(&pu.input_json).unwrap_or(serde_json::Value::Null)
453                    };
454                    parts.push(MessagePart::ToolUse {
455                        id: pu.id,
456                        name: name_from_provider(&pu.name, &tools),
457                        input,
458                    });
459                }
460                Ok(AssistantMessage {
461                    message: Message {
462                        role: MessageRole::Assistant,
463                        parts,
464                        turn_id,
465                    },
466                    stop_reason,
467                    token_usage: TokenUsage {
468                        input: input_tokens,
469                        cached_input: cache_read_tokens,
470                        output: cumulative,
471                        cache_write: cache_write_tokens,
472                        ..Default::default()
473                    },
474                    timing: CallTiming::default(),
475                    model: String::new(),
476                    response_id: None,
477                })
478            },
479        );
480        Observable {
481            output,
482            events,
483            cancel,
484        }
485    }
486}
487
488struct PartialToolUse {
489    id: String,
490    name: String,
491    input_json: String,
492}
493
494fn response_to_assistant(
495    body: MessagesResponse,
496    turn_id: crate::event::TurnId,
497    tools: &[crate::tool::ToolSpec],
498) -> AssistantMessage {
499    let mut parts: Vec<MessagePart> = Vec::new();
500    for block in body.content {
501        match block {
502            ContentBlock::Text { text } => parts.push(MessagePart::Text { text }),
503            ContentBlock::Thinking {
504                thinking,
505                signature,
506            } => parts.push(MessagePart::Thinking {
507                thinking,
508                signature,
509            }),
510            ContentBlock::ToolUse { id, name, input } => parts.push(MessagePart::ToolUse {
511                id,
512                name: name_from_provider(&name, tools),
513                input,
514            }),
515            ContentBlock::Other => {}
516        }
517    }
518    let stop_reason = body
519        .stop_reason
520        .as_deref()
521        .map(parse_stop_reason)
522        .unwrap_or(StopReason::End);
523    let usage = body
524        .usage
525        .map(|u| TokenUsage {
526            input: u.input_tokens.unwrap_or(0),
527            cached_input: u.cache_read_input_tokens.unwrap_or(0),
528            output: u.output_tokens.unwrap_or(0),
529            cache_write: u.cache_creation_input_tokens.unwrap_or(0),
530            ..Default::default()
531        })
532        .unwrap_or_default();
533    AssistantMessage {
534        message: Message {
535            role: MessageRole::Assistant,
536            parts,
537            turn_id,
538        },
539        stop_reason,
540        token_usage: usage,
541        timing: CallTiming::default(),
542        model: body.model.unwrap_or_default(),
543        response_id: body.id,
544    }
545}
546
547fn parse_stop_reason(s: &str) -> StopReason {
548    match s {
549        "tool_use" => StopReason::ToolUse,
550        "max_tokens" => StopReason::Length,
551        _ => StopReason::End,
552    }
553}
554
555fn next_turn_id_from_req(req: &LlmRequest) -> crate::event::TurnId {
556    req.messages
557        .first()
558        .map(|m| m.turn_id.clone())
559        .unwrap_or_else(crate::event::TurnId::now)
560}
561
562fn net_err(e: reqwest::Error) -> RuntimeError {
563    RuntimeError::ToolFailed(format!("anthropic net: {e}"))
564}
565
566#[derive(Serialize, Clone)]
567struct MessagesRequest {
568    model: String,
569    max_tokens: u32,
570    stream: bool,
571    #[serde(skip_serializing_if = "Option::is_none")]
572    system: Option<String>,
573    messages: Vec<WireMessage>,
574    #[serde(skip_serializing_if = "Vec::is_empty")]
575    tools: Vec<WireTool>,
576    #[serde(skip_serializing_if = "Option::is_none")]
577    thinking: Option<ThinkingConfig>,
578    #[serde(skip_serializing_if = "Option::is_none")]
579    cache_control: Option<CacheControl>,
580}
581
582#[derive(Serialize, Clone)]
583struct ThinkingConfig {
584    #[serde(rename = "type")]
585    kind: &'static str,
586    #[serde(skip_serializing_if = "Option::is_none")]
587    budget_tokens: Option<u32>,
588}
589
590#[derive(Serialize, Clone)]
591struct WireTool {
592    name: String,
593    #[serde(skip_serializing_if = "Option::is_none")]
594    description: Option<String>,
595    input_schema: serde_json::Value,
596}
597
598#[derive(Serialize, Clone)]
599struct WireMessage {
600    role: &'static str,
601    content: MessageContent,
602}
603
604#[derive(Serialize, Clone)]
605#[serde(untagged)]
606enum MessageContent {
607    Blocks(Vec<ContentPart>),
608}
609
610#[derive(Serialize, Clone)]
611#[serde(tag = "type", rename_all = "snake_case")]
612enum ContentPart {
613    Text {
614        text: String,
615        #[serde(skip_serializing_if = "Option::is_none")]
616        cache_control: Option<CacheControl>,
617    },
618    Thinking {
619        thinking: String,
620        #[serde(skip_serializing_if = "Option::is_none")]
621        signature: Option<String>,
622    },
623    Image {
624        source: ImageSourceWire,
625    },
626    ToolUse {
627        id: String,
628        name: String,
629        input: serde_json::Value,
630    },
631    ToolResult {
632        tool_use_id: String,
633        content: String,
634        #[serde(skip_serializing_if = "core::ops::Not::not")]
635        is_error: bool,
636    },
637}
638
639#[derive(Serialize, Clone)]
640struct ImageSourceWire {
641    #[serde(rename = "type")]
642    kind: &'static str,
643    media_type: String,
644    data: String,
645}
646
647#[derive(Serialize, Clone)]
648struct CacheControl {
649    #[serde(rename = "type")]
650    kind: &'static str,
651}
652
653#[derive(Deserialize)]
654struct MessagesResponse {
655    content: Vec<ContentBlock>,
656    #[serde(default)]
657    stop_reason: Option<String>,
658    #[serde(default)]
659    usage: Option<AnthropicUsage>,
660    #[serde(default)]
661    model: Option<String>,
662    #[serde(default)]
663    id: Option<String>,
664}
665
666#[derive(Deserialize, Default)]
667struct AnthropicUsage {
668    #[serde(default)]
669    input_tokens: Option<u64>,
670    #[serde(default)]
671    output_tokens: Option<u64>,
672    #[serde(default)]
673    cache_read_input_tokens: Option<u64>,
674    #[serde(default)]
675    cache_creation_input_tokens: Option<u64>,
676}
677
678#[derive(Deserialize)]
679#[serde(tag = "type", rename_all = "snake_case")]
680enum ContentBlock {
681    Text {
682        text: String,
683    },
684    Thinking {
685        thinking: String,
686        #[serde(default)]
687        signature: Option<String>,
688    },
689    ToolUse {
690        id: String,
691        name: String,
692        input: serde_json::Value,
693    },
694    #[serde(other)]
695    Other,
696}