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
12pub 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 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}