Skip to main content

cairn_models/
anthropic.rs

1use async_trait::async_trait;
2use cairn_runtime::{
3    Content, Message, Model, ModelError, ModelRequest, ModelResponse, ModelStream, Role, StopReason, StreamEvent, Usage,
4};
5use serde_json::{Value, json};
6
7use crate::sse::{self, SseEvent};
8use crate::{Pricing, merge_options, result_string, send_json};
9
10const PROVIDER: &str = "anthropic";
11
12/// Native Anthropic Messages API adapter.
13///
14/// Signed thinking blocks are kept and replayed to Anthropic on later turns,
15/// and dropped automatically when the conversation is sent to another provider.
16pub struct AnthropicModel {
17    model: String,
18    api_key: String,
19    base_url: String,
20    version: String,
21    prompt_caching: bool,
22    pricing: Option<Pricing>,
23    client: reqwest::Client,
24}
25
26impl AnthropicModel {
27    pub fn new(model: impl Into<String>, api_key: impl Into<String>) -> Self {
28        Self {
29            model: model.into(),
30            api_key: api_key.into(),
31            base_url: "https://api.anthropic.com".into(),
32            version: "2023-06-01".into(),
33            prompt_caching: true,
34            pricing: None,
35            client: reqwest::Client::new(),
36        }
37    }
38
39    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
40        self.base_url = url.into();
41        self
42    }
43
44    /// Cache the system prompt and tool definitions (on by default). Agent
45    /// loops resend them every step, so this cuts cost and latency sharply.
46    pub fn with_prompt_caching(mut self, on: bool) -> Self {
47        self.prompt_caching = on;
48        self
49    }
50
51    pub fn with_pricing(mut self, pricing: Pricing) -> Self {
52        self.pricing = Some(pricing);
53        self
54    }
55
56    pub(crate) fn build_body(&self, req: &ModelRequest) -> Value {
57        let cache = json!({ "type": "ephemeral" });
58        let messages: Vec<Value> = req
59            .messages
60            .iter()
61            .filter_map(|m| {
62                let content: Vec<Value> = m.content.iter().filter_map(block_to_json).collect();
63                (!content.is_empty()).then(|| json!({ "role": role(m), "content": content }))
64            })
65            .collect();
66
67        let mut body = json!({
68            "model": self.model,
69            "max_tokens": req.max_tokens,
70            "messages": messages,
71        });
72
73        if let Some(system) = &req.system {
74            body["system"] = if self.prompt_caching {
75                json!([{ "type": "text", "text": system, "cache_control": cache }])
76            } else {
77                json!(system)
78            };
79        }
80        if !req.tools.is_empty() {
81            let mut tools: Vec<Value> = req
82                .tools
83                .iter()
84                .map(|t| json!({ "name": t.name, "description": t.description, "input_schema": t.input_schema }))
85                .collect();
86            if self.prompt_caching {
87                if let Some(last) = tools.last_mut() {
88                    last["cache_control"] = cache.clone();
89                }
90            }
91            body["tools"] = json!(tools);
92        }
93        if let Some(t) = req.temperature {
94            body["temperature"] = json!(t);
95        }
96        merge_options(&mut body, req.provider_options.get(PROVIDER));
97        body
98    }
99
100    pub(crate) fn parse_response(v: &Value) -> Result<ModelResponse, ModelError> {
101        let blocks = v["content"].as_array().ok_or_else(|| ModelError::Decode("missing `content`".into()))?;
102        let content = blocks
103            .iter()
104            .filter_map(|b| match b["type"].as_str()? {
105                "text" => Some(Content::Text { text: b["text"].as_str()?.to_owned() }),
106                "tool_use" => Some(Content::ToolUse {
107                    id: b["id"].as_str()?.to_owned(),
108                    name: b["name"].as_str()?.to_owned(),
109                    input: b["input"].clone(),
110                }),
111                "thinking" => Some(Content::Thinking {
112                    text: b["thinking"].as_str().unwrap_or_default().to_owned(),
113                    provider_data: Some(json!({ "provider": PROVIDER, "signature": b["signature"] })),
114                }),
115                "redacted_thinking" => Some(Content::Thinking {
116                    text: String::new(),
117                    provider_data: Some(json!({ "provider": PROVIDER, "redacted_data": b["data"] })),
118                }),
119                _ => None,
120            })
121            .collect();
122
123        let stop_reason = match v["stop_reason"].as_str() {
124            Some("end_turn") | Some("stop_sequence") => StopReason::EndTurn,
125            Some("tool_use") => StopReason::ToolUse,
126            Some("max_tokens") => StopReason::MaxTokens,
127            Some(other) => StopReason::Other(other.to_owned()),
128            None => StopReason::Other("unknown".into()),
129        };
130        let u = &v["usage"];
131        let n = |k: &str| u[k].as_u64().unwrap_or(0);
132        Ok(ModelResponse {
133            model: v["model"].as_str().unwrap_or_default().to_owned(),
134            content,
135            stop_reason,
136            usage: Usage {
137                input_tokens: n("input_tokens"),
138                output_tokens: n("output_tokens"),
139                cache_read_tokens: n("cache_read_input_tokens"),
140                cache_write_tokens: n("cache_creation_input_tokens"),
141            },
142            ..Default::default()
143        })
144    }
145}
146
147fn role(m: &Message) -> &'static str {
148    match m.role {
149        Role::User => "user",
150        Role::Assistant => "assistant",
151    }
152}
153
154fn block_to_json(c: &Content) -> Option<Value> {
155    Some(match c {
156        Content::Text { text } => json!({ "type": "text", "text": text }),
157        Content::ToolUse { id, name, input } => json!({ "type": "tool_use", "id": id, "name": name, "input": input }),
158        Content::ToolResult { tool_use_id, content, is_error } => json!({
159            "type": "tool_result",
160            "tool_use_id": tool_use_id,
161            "content": result_string(content),
162            "is_error": is_error,
163        }),
164        Content::Thinking { text, provider_data } => {
165            let data = provider_data.as_ref()?;
166            if data["provider"] != PROVIDER {
167                return None;
168            }
169            if !data["redacted_data"].is_null() {
170                json!({ "type": "redacted_thinking", "data": data["redacted_data"] })
171            } else {
172                json!({ "type": "thinking", "thinking": text, "signature": data["signature"] })
173            }
174        }
175    })
176}
177
178#[async_trait]
179impl Model for AnthropicModel {
180    fn id(&self) -> String {
181        format!("{PROVIDER}:{}", self.model)
182    }
183
184    async fn complete(&self, request: &ModelRequest) -> Result<ModelResponse, ModelError> {
185        let body = self.build_body(request);
186        let http = self
187            .client
188            .post(format!("{}/v1/messages", self.base_url))
189            .header("x-api-key", &self.api_key)
190            .header("anthropic-version", &self.version);
191        let v = send_json(http, &body).await?;
192        Self::parse_response(&v)
193    }
194
195    async fn stream(&self, request: &ModelRequest) -> Result<ModelStream, ModelError> {
196        let mut body = self.build_body(request);
197        body["stream"] = json!(true);
198        let resp = self
199            .client
200            .post(format!("{}/v1/messages", self.base_url))
201            .header("x-api-key", &self.api_key)
202            .header("anthropic-version", &self.version)
203            .json(&body)
204            .send()
205            .await
206            .map_err(|e| ModelError::Transport(e.to_string()))?;
207        let resp = sse::check_status(resp).await?;
208        Ok(sse::translate(sse::events(resp), StreamAcc::default(), StreamAcc::step, StreamAcc::end))
209    }
210
211    fn pricing(&self) -> Option<Pricing> {
212        self.pricing
213    }
214}
215
216enum Block {
217    Text(String),
218    Thinking { text: String, signature: String },
219    Redacted(Value),
220    Tool { id: String, name: String, json: String, index: usize },
221}
222
223#[derive(Default)]
224pub(crate) struct StreamAcc {
225    model: String,
226    blocks: Vec<Block>,
227    tools: usize,
228    usage: Usage,
229    stop_reason: Option<String>,
230}
231
232impl StreamAcc {
233    pub(crate) fn step(&mut self, ev: SseEvent) -> Result<Vec<StreamEvent>, ModelError> {
234        let v: Value = serde_json::from_str(&ev.data).map_err(|e| ModelError::Decode(e.to_string()))?;
235        let mut out = Vec::new();
236        match v["type"].as_str().unwrap_or_default() {
237            "message_start" => {
238                let m = &v["message"];
239                self.model = m["model"].as_str().unwrap_or_default().to_owned();
240                self.read_usage(&m["usage"]);
241            }
242            "content_block_start" => {
243                let b = &v["content_block"];
244                self.blocks.push(match b["type"].as_str().unwrap_or_default() {
245                    "tool_use" => {
246                        let (id, name) = (b["id"].as_str().unwrap_or_default(), b["name"].as_str().unwrap_or_default());
247                        out.push(StreamEvent::ToolCallStart { index: self.tools, id: id.into(), name: name.into() });
248                        self.tools += 1;
249                        Block::Tool { id: id.into(), name: name.into(), json: String::new(), index: self.tools - 1 }
250                    }
251                    "thinking" => Block::Thinking { text: String::new(), signature: String::new() },
252                    "redacted_thinking" => Block::Redacted(b["data"].clone()),
253                    _ => Block::Text(String::new()),
254                });
255            }
256            "content_block_delta" => {
257                let d = &v["delta"];
258                let piece = |k: &str| d[k].as_str().unwrap_or_default().to_owned();
259                match (self.blocks.last_mut(), d["type"].as_str().unwrap_or_default()) {
260                    (Some(Block::Text(t)), "text_delta") => {
261                        t.push_str(&piece("text"));
262                        out.push(StreamEvent::TextDelta(piece("text")));
263                    }
264                    (Some(Block::Thinking { text, .. }), "thinking_delta") => {
265                        text.push_str(&piece("thinking"));
266                        out.push(StreamEvent::ThinkingDelta(piece("thinking")));
267                    }
268                    (Some(Block::Thinking { signature, .. }), "signature_delta") => {
269                        signature.push_str(&piece("signature"));
270                        out.push(StreamEvent::ThinkingSignature(piece("signature")));
271                    }
272                    (Some(Block::Tool { json, index, .. }), "input_json_delta") => {
273                        json.push_str(&piece("partial_json"));
274                        out.push(StreamEvent::ToolInputDelta { index: *index, partial_json: piece("partial_json") });
275                    }
276                    _ => {}
277                }
278            }
279            "message_delta" => {
280                if let Some(r) = v["delta"]["stop_reason"].as_str() {
281                    self.stop_reason = Some(r.to_owned());
282                }
283                self.read_usage(&v["usage"]);
284            }
285            "message_stop" => out.push(StreamEvent::Completed(self.finish()?)),
286            "error" => {
287                let kind = v["error"]["type"].as_str().unwrap_or_default();
288                let status = match kind {
289                    "overloaded_error" => 529,
290                    "rate_limit_error" => 429,
291                    _ => 500,
292                };
293                return Err(ModelError::Api { status, body: v["error"].to_string() });
294            }
295            _ => {}
296        }
297        Ok(out)
298    }
299
300    pub(crate) fn end(&mut self) -> Result<Vec<StreamEvent>, ModelError> {
301        Err(ModelError::Transport("stream closed before message_stop".into()))
302    }
303
304    fn read_usage(&mut self, u: &Value) {
305        let set = |field: &mut u64, k: &str| {
306            if let Some(n) = u[k].as_u64() {
307                *field = n;
308            }
309        };
310        set(&mut self.usage.input_tokens, "input_tokens");
311        set(&mut self.usage.output_tokens, "output_tokens");
312        set(&mut self.usage.cache_read_tokens, "cache_read_input_tokens");
313        set(&mut self.usage.cache_write_tokens, "cache_creation_input_tokens");
314    }
315
316    fn finish(&mut self) -> Result<ModelResponse, ModelError> {
317        let content = std::mem::take(&mut self.blocks)
318            .into_iter()
319            .map(|b| match b {
320                Block::Text(text) => Content::Text { text },
321                Block::Thinking { text, signature } => Content::Thinking {
322                    text,
323                    provider_data: Some(json!({ "provider": PROVIDER, "signature": signature })),
324                },
325                Block::Redacted(data) => Content::Thinking {
326                    text: String::new(),
327                    provider_data: Some(json!({ "provider": PROVIDER, "redacted_data": data })),
328                },
329                Block::Tool { id, name, json, .. } => Content::ToolUse {
330                    id,
331                    name,
332                    input: if json.trim().is_empty() {
333                        json!({})
334                    } else {
335                        serde_json::from_str(&json).unwrap_or(json!({}))
336                    },
337                },
338            })
339            .collect();
340        let response = Self::parse_stop(self.stop_reason.as_deref());
341        Ok(ModelResponse {
342            model: self.model.clone(),
343            content,
344            stop_reason: response,
345            usage: self.usage,
346            ..Default::default()
347        })
348    }
349
350    fn parse_stop(r: Option<&str>) -> StopReason {
351        match r {
352            Some("end_turn") | Some("stop_sequence") => StopReason::EndTurn,
353            Some("tool_use") => StopReason::ToolUse,
354            Some("max_tokens") => StopReason::MaxTokens,
355            Some(other) => StopReason::Other(other.to_owned()),
356            None => StopReason::Other("unknown".into()),
357        }
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use cairn_runtime::{Effect, ToolSpec};
365
366    fn request() -> ModelRequest {
367        ModelRequest {
368            system: Some("be terse".into()),
369            messages: vec![
370                Message::user("refund order 7"),
371                Message::assistant(vec![
372                    Content::Thinking {
373                        text: "need order".into(),
374                        provider_data: Some(json!({"provider": "anthropic", "signature": "sig"})),
375                    },
376                    Content::Thinking { text: "foreign".into(), provider_data: Some(json!({"provider": "openai"})) },
377                    Content::ToolUse { id: "t1".into(), name: "get_order".into(), input: json!({"id": 7}) },
378                ]),
379                Message {
380                    role: Role::User,
381                    content: vec![Content::ToolResult {
382                        tool_use_id: "t1".into(),
383                        content: json!({"total": 10}),
384                        is_error: false,
385                    }],
386                },
387            ],
388            tools: vec![ToolSpec {
389                name: "get_order".into(),
390                description: "d".into(),
391                input_schema: json!({"type": "object"}),
392                effect: Effect::Read,
393            }],
394            max_tokens: 100,
395            provider_options: [(
396                "anthropic".to_owned(),
397                json!({"thinking": {"type": "enabled", "budget_tokens": 1024}}),
398            )]
399            .into(),
400            ..Default::default()
401        }
402    }
403
404    #[test]
405    fn builds_body_with_caching_thinking_and_options() {
406        let body = AnthropicModel::new("claude-x", "k").build_body(&request());
407        assert_eq!(body["system"][0]["cache_control"]["type"], "ephemeral");
408        assert_eq!(body["tools"][0]["cache_control"]["type"], "ephemeral");
409        assert_eq!(body["thinking"]["budget_tokens"], 1024);
410        let assistant = &body["messages"][1]["content"];
411        assert_eq!(assistant.as_array().unwrap().len(), 2, "foreign thinking dropped");
412        assert_eq!(assistant[0]["signature"], "sig");
413        assert_eq!(body["messages"][2]["content"][0]["content"], r#"{"total":10}"#);
414    }
415
416    #[test]
417    fn parses_response() {
418        let v = json!({
419            "model": "claude-x",
420            "stop_reason": "tool_use",
421            "content": [
422                {"type": "thinking", "thinking": "hmm", "signature": "s"},
423                {"type": "text", "text": "checking"},
424                {"type": "tool_use", "id": "t9", "name": "get_order", "input": {"id": 1}}
425            ],
426            "usage": {"input_tokens": 10, "output_tokens": 5, "cache_read_input_tokens": 90}
427        });
428        let r = AnthropicModel::parse_response(&v).unwrap();
429        assert_eq!(r.stop_reason, StopReason::ToolUse);
430        assert_eq!(r.content.len(), 3);
431        assert_eq!(r.usage.cache_read_tokens, 90);
432    }
433
434    #[test]
435    fn accumulates_stream() {
436        let raw = r#"event: message_start
437data: {"type":"message_start","message":{"model":"claude-x","usage":{"input_tokens":12}}}
438
439event: content_block_start
440data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}
441
442event: content_block_delta
443data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"hmm"}}
444
445event: content_block_delta
446data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig"}}
447
448event: content_block_start
449data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}
450
451event: ping
452data: {"type":"ping"}
453
454event: content_block_delta
455data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"Look"}}
456
457event: content_block_delta
458data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"ing"}}
459
460event: content_block_start
461data: {"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"tu1","name":"get_order","input":{}}}
462
463event: content_block_delta
464data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"id\":"}}
465
466event: content_block_delta
467data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"7}"}}
468
469event: message_delta
470data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":30}}
471
472event: message_stop
473data: {"type":"message_stop"}
474"#;
475        let mut acc = StreamAcc::default();
476        let events: Vec<StreamEvent> = sse::parse_text(raw).into_iter().flat_map(|e| acc.step(e).unwrap()).collect();
477        assert_eq!(events[0], StreamEvent::ThinkingDelta("hmm".into()));
478        assert!(events.contains(&StreamEvent::TextDelta("Look".into())));
479        assert!(events.contains(&StreamEvent::ToolCallStart { index: 0, id: "tu1".into(), name: "get_order".into() }));
480        let StreamEvent::Completed(r) = events.last().unwrap() else { panic!("no completion") };
481        assert_eq!(r.stop_reason, StopReason::ToolUse);
482        assert_eq!(r.usage.input_tokens, 12);
483        assert_eq!(r.usage.output_tokens, 30);
484        assert_eq!(r.content[1], Content::Text { text: "Looking".into() });
485        assert_eq!(
486            r.content[2],
487            Content::ToolUse { id: "tu1".into(), name: "get_order".into(), input: json!({"id": 7}) }
488        );
489        assert!(matches!(&r.content[0], Content::Thinking { provider_data: Some(d), .. } if d["signature"] == "sig"));
490    }
491
492    #[test]
493    fn stream_error_event_is_retryable_when_overloaded() {
494        let mut acc = StreamAcc::default();
495        let ev = SseEvent {
496            event: Some("error".into()),
497            data: r#"{"type":"error","error":{"type":"overloaded_error","message":"busy"}}"#.into(),
498        };
499        assert!(acc.step(ev).unwrap_err().retryable());
500    }
501}