1use async_trait::async_trait;
7use reqwest::Client;
8use serde::Deserialize;
9use serde_json::Value;
10use std::time::Duration;
11
12use super::{Completion, LlmProvider, StreamSender, TokenUsage, ToolCall, ToolSpec};
13use crate::error::{Error, Result};
14use crate::message::{Message, Role};
15
16#[derive(Debug, Clone)]
18pub struct RetryPolicy {
19 pub max_retries: usize,
20 pub initial_backoff: Duration,
21 pub max_backoff: Duration,
22}
23
24impl Default for RetryPolicy {
25 fn default() -> Self {
26 Self {
27 max_retries: 2,
28 initial_backoff: Duration::from_secs(1),
29 max_backoff: Duration::from_secs(8),
30 }
31 }
32}
33
34impl RetryPolicy {
35 pub fn backoff_for(
37 &self,
38 attempt: usize,
39 status: Option<u16>,
40 is_network_error: bool,
41 ) -> Option<Duration> {
42 if attempt >= self.max_retries {
43 return None;
44 }
45
46 let is_transient = is_network_error || status.is_some_and(|s| (500..600).contains(&s));
47
48 if !is_transient {
49 return None;
50 }
51
52 let backoff = self.initial_backoff * 2u32.pow(attempt as u32);
53 Some(backoff.min(self.max_backoff))
54 }
55}
56
57#[derive(Debug, Clone)]
58pub struct AnthropicProvider {
59 base_url: String,
60 api_key: String,
61 model: String,
62 client: Client,
63 temperature: f64,
64 max_tokens: u32,
65 retry: RetryPolicy,
66}
67
68impl AnthropicProvider {
69 pub fn new(
70 base_url: impl Into<String>,
71 api_key: impl Into<String>,
72 model: impl Into<String>,
73 ) -> Self {
74 Self {
75 base_url: base_url.into().trim_end_matches('/').to_string(),
76 api_key: api_key.into(),
77 model: model.into(),
78 client: Client::builder()
79 .timeout(Duration::from_secs(180))
80 .build()
81 .expect("reqwest client build"),
82 temperature: 0.2,
83 max_tokens: 4096,
84 retry: RetryPolicy::default(),
85 }
86 }
87
88 fn make_err(&self, ctx: impl Into<String>) -> Error {
90 Error::Llm {
91 provider: self.model.clone(),
92 message: ctx.into(),
93 }
94 }
95
96 pub fn with_temperature(mut self, t: f64) -> Self {
97 self.temperature = t;
98 self
99 }
100
101 pub fn with_max_tokens(mut self, n: u32) -> Self {
102 self.max_tokens = n;
103 self
104 }
105
106 pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
107 self.retry = policy;
108 self
109 }
110}
111
112#[async_trait]
113impl LlmProvider for AnthropicProvider {
114 async fn complete(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Completion> {
115 let (system, messages) = extract_system_message(messages);
116
117 let messages = filter_leading_assistant(&messages);
119
120 let body = build_request(
121 &self.model,
122 self.temperature,
123 self.max_tokens,
124 system.as_deref(),
125 &messages,
126 tools,
127 );
128 let url = format!("{}/v1/messages", self.base_url);
129
130 let mut attempt = 0;
131 loop {
132 tracing::debug!(target: "recursive::llm", request = %body, "POST {}", url);
133 let result = self
134 .client
135 .post(&url)
136 .header("x-api-key", &self.api_key)
137 .header("anthropic-version", "2023-06-01")
138 .header("content-type", "application/json")
139 .json(&body)
140 .send()
141 .await;
142
143 match result {
144 Ok(resp) => {
145 let status = resp.status();
146 let is_network_error = false;
147
148 if status.is_success() {
149 let text = resp.text().await?;
150 let parsed: AnthropicResponse =
151 serde_json::from_str(&text).map_err(|e| {
152 self.make_err(format!(
153 "failed to parse response: {e}; body: {text}"
154 ))
155 })?;
156 return Ok(parse_completion(parsed));
157 }
158
159 let text = resp.text().await?;
161 tracing::debug!(target: "recursive::llm", body = %text, "error response");
162
163 if let Some(backoff) =
164 self.retry
165 .backoff_for(attempt, Some(status.as_u16()), is_network_error)
166 {
167 tracing::warn!(
168 target: "recursive::llm",
169 attempt,
170 backoff_ms = backoff.as_millis(),
171 status = status.as_u16(),
172 "transient HTTP error, retrying"
173 );
174 tokio::time::sleep(backoff).await;
175 attempt += 1;
176 continue;
177 }
178
179 return Err(self.make_err(format!("HTTP {}: {}", status, text)));
181 }
182 Err(e) => {
183 if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
185 tracing::warn!(
186 target: "recursive::llm",
187 attempt,
188 backoff_ms = backoff.as_millis(),
189 error = %e,
190 "network error, retrying"
191 );
192 tokio::time::sleep(backoff).await;
193 attempt += 1;
194 continue;
195 }
196
197 return Err(self.make_err(format!("request failed: {e}")));
198 }
199 }
200 }
201 }
202
203 async fn stream(
204 &self,
205 messages: &[Message],
206 tools: &[ToolSpec],
207 stream_tx: Option<StreamSender>,
208 ) -> Result<Completion> {
209 self.stream_inner(messages, tools, stream_tx).await
210 }
211}
212
213impl AnthropicProvider {
214 async fn stream_inner(
216 &self,
217 messages: &[Message],
218 tools: &[ToolSpec],
219 stream_tx: Option<StreamSender>,
220 ) -> Result<Completion> {
221 let (system, messages) = extract_system_message(messages);
222 let messages = filter_leading_assistant(&messages);
223
224 let mut body = build_request(
225 &self.model,
226 self.temperature,
227 self.max_tokens,
228 system.as_deref(),
229 &messages,
230 tools,
231 );
232 body["stream"] = Value::Bool(true);
233
234 let url = format!("{}/v1/messages", self.base_url);
235
236 let mut attempt = 0;
237 loop {
238 tracing::debug!(target: "recursive::llm", request = %body, "POST {} (stream)", url);
239 let result = self
240 .client
241 .post(&url)
242 .header("x-api-key", &self.api_key)
243 .header("anthropic-version", "2023-06-01")
244 .header("content-type", "application/json")
245 .json(&body)
246 .send()
247 .await;
248
249 match result {
250 Ok(resp) => {
251 let status = resp.status();
252 if status.is_success() {
253 return self.parse_sse_stream(resp, stream_tx.clone()).await;
254 }
255
256 let text = resp.text().await?;
258 tracing::debug!(target: "recursive::llm", body = %text, "error response (stream)");
259
260 if let Some(backoff) =
261 self.retry
262 .backoff_for(attempt, Some(status.as_u16()), false)
263 {
264 tracing::warn!(
265 target: "recursive::llm",
266 attempt,
267 backoff_ms = backoff.as_millis(),
268 status = status.as_u16(),
269 "transient HTTP error, retrying (stream)"
270 );
271 tokio::time::sleep(backoff).await;
272 attempt += 1;
273 continue;
274 }
275
276 return Err(self.make_err(format!("HTTP {}: {}", status, text)));
277 }
278 Err(e) => {
279 if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
280 tracing::warn!(
281 target: "recursive::llm",
282 attempt,
283 backoff_ms = backoff.as_millis(),
284 error = %e,
285 "network error, retrying (stream)"
286 );
287 tokio::time::sleep(backoff).await;
288 attempt += 1;
289 continue;
290 }
291
292 return Err(self.make_err(format!("request failed: {e}")));
293 }
294 }
295 }
296 }
297
298 async fn parse_sse_stream(
312 &self,
313 resp: reqwest::Response,
314 stream_tx: Option<StreamSender>,
315 ) -> Result<Completion> {
316 let mut content = String::new();
317 let mut tool_calls: Vec<StreamToolCall> = Vec::new();
318 let mut finish_reason: Option<String> = None;
319 let mut usage: Option<TokenUsage> = None;
320 let mut input_tokens: Option<u32> = None;
321 let mut output_tokens: Option<u32> = None;
322 let mut cache_creation: Option<u32> = None;
323 let mut cache_read: Option<u32> = None;
324
325 let reader = resp.text().await?;
327 let mut current_event: Option<String> = None;
328
329 for line in reader.lines() {
330 if let Some(event_name) = line.strip_prefix("event: ") {
331 current_event = Some(event_name.to_string());
332 continue;
333 }
334
335 if let Some(data) = line.strip_prefix("data: ") {
336 let event_type = current_event.as_deref().unwrap_or("unknown");
337
338 match event_type {
339 "message_start" => {
340 let parsed: Value = serde_json::from_str(data).map_err(|e| {
342 self.make_err(format!(
343 "SSE parse error (message_start): {e}; data: {data}"
344 ))
345 })?;
346 if let Some(msg) = parsed.get("message") {
347 if let Some(u) = msg.get("usage") {
348 input_tokens = u
349 .get("input_tokens")
350 .and_then(|v| v.as_u64())
351 .map(|v| v as u32);
352 cache_creation = u
353 .get("cache_creation_input_tokens")
354 .and_then(|v| v.as_u64())
355 .map(|v| v as u32);
356 cache_read = u
357 .get("cache_read_input_tokens")
358 .and_then(|v| v.as_u64())
359 .map(|v| v as u32);
360 }
361 }
362 }
363 "content_block_start" => {
364 let parsed: Value = serde_json::from_str(data).map_err(|e| {
365 self.make_err(format!(
366 "SSE parse error (content_block_start): {e}; data: {data}"
367 ))
368 })?;
369 if let Some(block) = parsed.get("content_block") {
370 let block_type =
371 block.get("type").and_then(|v| v.as_str()).unwrap_or("");
372 let index =
373 parsed.get("index").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
374 match block_type {
375 "tool_use" => {
376 let id = block
377 .get("id")
378 .and_then(|v| v.as_str())
379 .unwrap_or("")
380 .to_string();
381 let name = block
382 .get("name")
383 .and_then(|v| v.as_str())
384 .unwrap_or("")
385 .to_string();
386 while tool_calls.len() <= index {
388 tool_calls.push(StreamToolCall::default());
389 }
390 tool_calls[index].id = id;
391 tool_calls[index].name = name;
392 }
393 "text" => {
394 if let Some(text) = block.get("text").and_then(|v| v.as_str()) {
396 if !text.is_empty() {
397 content.push_str(text);
398 if let Some(ref tx) = stream_tx {
399 let _ = tx.send(text.to_string());
400 }
401 }
402 }
403 }
404 _ => {}
405 }
406 }
407 }
408 "content_block_delta" => {
409 let parsed: Value = serde_json::from_str(data).map_err(|e| {
410 self.make_err(format!(
411 "SSE parse error (content_block_delta): {e}; data: {data}"
412 ))
413 })?;
414 if let Some(delta) = parsed.get("delta") {
415 let delta_type =
416 delta.get("type").and_then(|v| v.as_str()).unwrap_or("");
417 match delta_type {
418 "text_delta" => {
419 if let Some(text) = delta.get("text").and_then(|v| v.as_str()) {
420 if !text.is_empty() {
421 content.push_str(text);
422 if let Some(ref tx) = stream_tx {
423 let _ = tx.send(text.to_string());
424 }
425 }
426 }
427 }
428 "input_json_delta" => {
429 if let Some(partial) =
430 delta.get("partial_json").and_then(|v| v.as_str())
431 {
432 let index = parsed
433 .get("index")
434 .and_then(|v| v.as_u64())
435 .unwrap_or(0)
436 as usize;
437 while tool_calls.len() <= index {
438 tool_calls.push(StreamToolCall::default());
439 }
440 tool_calls[index].partial_json.push_str(partial);
441 }
442 }
443 _ => {}
444 }
445 }
446 }
447 "message_delta" => {
448 let parsed: Value = serde_json::from_str(data).map_err(|e| {
449 self.make_err(format!(
450 "SSE parse error (message_delta): {e}; data: {data}"
451 ))
452 })?;
453 if let Some(delta) = parsed.get("delta") {
454 if let Some(reason) = delta.get("stop_reason").and_then(|v| v.as_str())
455 {
456 finish_reason = Some(reason.to_string());
457 }
458 }
459 if let Some(u) = parsed.get("usage") {
460 output_tokens = u
461 .get("output_tokens")
462 .and_then(|v| v.as_u64())
463 .map(|v| v as u32);
464 }
465 }
466 "message_stop" => {
467 }
469 "ping" => {
470 }
472 _ => {
473 tracing::debug!(target: "recursive::llm", event = %event_type, data = %data, "unhandled SSE event");
474 }
475 }
476
477 current_event = None;
478 }
479 }
480
481 if input_tokens.is_some() || output_tokens.is_some() {
483 let prompt = input_tokens.unwrap_or(0);
484 let completion = output_tokens.unwrap_or(0);
485 usage = Some(TokenUsage {
486 prompt_tokens: prompt,
487 completion_tokens: completion,
488 total_tokens: prompt.saturating_add(completion),
489 cache_hit_tokens: cache_read.unwrap_or(0),
490 cache_miss_tokens: cache_creation.unwrap_or(0),
491 });
492 }
493
494 let final_tool_calls: Vec<ToolCall> = tool_calls
496 .into_iter()
497 .filter(|tc| !tc.id.is_empty())
498 .map(|tc| {
499 let args: Value = if tc.partial_json.trim().is_empty() {
500 Value::Object(Default::default())
501 } else {
502 serde_json::from_str(&tc.partial_json)
503 .unwrap_or_else(|_| Value::String(tc.partial_json.clone()))
504 };
505 ToolCall {
506 id: tc.id,
507 name: tc.name,
508 arguments: args,
509 }
510 })
511 .collect();
512
513 Ok(Completion {
514 content,
515 tool_calls: final_tool_calls,
516 finish_reason: finish_reason.map(|r| match r.as_str() {
517 "end_turn" => "stop".to_string(),
518 "max_tokens" => "length".to_string(),
519 "tool_use" => "tool_calls".to_string(),
520 other => other.to_string(),
521 }),
522 usage,
523 })
524 }
525}
526
527fn extract_system_message(messages: &[Message]) -> (Option<String>, Vec<Message>) {
529 if messages.first().is_some_and(|m| m.role == Role::System) {
530 let system = messages[0].content.clone();
531 let rest = messages[1..].to_vec();
532 (Some(system), rest)
533 } else {
534 (None, messages.to_vec())
535 }
536}
537
538fn filter_leading_assistant(messages: &[Message]) -> Vec<Message> {
541 let mut result = messages.to_vec();
542 while result
543 .first()
544 .is_some_and(|m| m.role == Role::Assistant && m.tool_calls.is_empty())
545 {
546 result.remove(0);
548 }
549 result
550}
551
552fn build_request(
553 model: &str,
554 temperature: f64,
555 max_tokens: u32,
556 system: Option<&str>,
557 messages: &[Message],
558 tools: &[ToolSpec],
559) -> Value {
560 let mut req = serde_json::json!({
561 "model": model,
562 "max_tokens": max_tokens,
563 "temperature": temperature,
564 });
565
566 if let Some(sys) = system {
567 req["system"] = Value::String(sys.to_string());
568 }
569
570 let msgs: Vec<Value> = messages.iter().map(serialize_message).collect();
571 req["messages"] = Value::Array(msgs);
572
573 if !tools.is_empty() {
574 let tools_json: Vec<Value> = tools
575 .iter()
576 .map(|t| {
577 serde_json::json!({
578 "name": t.name,
579 "description": t.description,
580 "input_schema": t.parameters,
581 })
582 })
583 .collect();
584 req["tools"] = Value::Array(tools_json);
585 }
586
587 req
588}
589
590fn serialize_message(m: &Message) -> Value {
591 let role = match m.role {
592 Role::System => "system",
593 Role::User => "user",
594 Role::Assistant => "assistant",
595 Role::Tool => "user", };
597
598 if !m.tool_calls.is_empty() {
600 let tool_uses: Vec<Value> = m
601 .tool_calls
602 .iter()
603 .map(|c| {
604 serde_json::json!({
605 "type": "tool_use",
606 "id": c.id,
607 "name": c.name,
608 "input": c.arguments,
609 })
610 })
611 .collect();
612 serde_json::json!({
613 "role": role,
614 "content": tool_uses,
615 })
616 } else if let Some(id) = &m.tool_call_id {
617 serde_json::json!({
619 "role": role,
620 "content": [{
621 "type": "tool_result",
622 "tool_use_id": id,
623 "content": m.content,
624 }],
625 })
626 } else {
627 serde_json::json!({
629 "role": role,
630 "content": m.content,
631 })
632 }
633}
634
635#[derive(Debug, Deserialize)]
636struct AnthropicResponse {
637 #[serde(default)]
638 content: Vec<ContentBlock>,
639 #[serde(default)]
640 usage: Option<AnthropicUsage>,
641 #[serde(default)]
642 stop_reason: Option<String>,
643}
644
645#[derive(Debug, Deserialize)]
646#[serde(tag = "type")]
647enum ContentBlock {
648 #[serde(rename = "text")]
649 Text {
650 #[serde(default)]
651 text: String,
652 },
653 #[serde(rename = "tool_use")]
654 ToolUse {
655 id: String,
656 name: String,
657 #[serde(default)]
658 input: Value,
659 },
660}
661
662#[derive(Debug, Deserialize)]
663struct AnthropicUsage {
664 #[serde(default)]
665 input_tokens: Option<u32>,
666 #[serde(default)]
667 output_tokens: Option<u32>,
668 #[serde(default)]
669 cache_creation_input_tokens: Option<u32>,
670 #[serde(default)]
671 cache_read_input_tokens: Option<u32>,
672}
673
674impl AnthropicUsage {
675 fn to_token_usage(&self) -> TokenUsage {
676 TokenUsage {
677 prompt_tokens: self.input_tokens.unwrap_or(0),
678 completion_tokens: self.output_tokens.unwrap_or(0),
679 total_tokens: self
680 .input_tokens
681 .unwrap_or(0)
682 .saturating_add(self.output_tokens.unwrap_or(0)),
683 cache_hit_tokens: self.cache_read_input_tokens.unwrap_or(0),
684 cache_miss_tokens: self.cache_creation_input_tokens.unwrap_or(0),
685 }
686 }
687}
688
689#[derive(Default)]
691struct StreamToolCall {
692 id: String,
693 name: String,
694 partial_json: String,
695}
696
697fn parse_completion(response: AnthropicResponse) -> Completion {
698 let mut content = String::new();
699 let mut tool_calls = Vec::new();
700
701 for block in response.content {
702 match block {
703 ContentBlock::Text { text } => {
704 content.push_str(&text);
705 }
706 ContentBlock::ToolUse { id, name, input } => {
707 tool_calls.push(ToolCall {
708 id,
709 name,
710 arguments: input,
711 });
712 }
713 }
714 }
715
716 let finish_reason = response.stop_reason.map(|r| match r.as_str() {
717 "end_turn" => "stop".to_string(),
718 "max_tokens" => "length".to_string(),
719 "tool_use" => "tool_calls".to_string(),
720 other => other.to_string(),
721 });
722
723 Completion {
724 content,
725 tool_calls,
726 finish_reason,
727 usage: response.usage.map(|u| u.to_token_usage()),
728 }
729}
730
731#[cfg(test)]
732mod tests {
733 use super::*;
734
735 #[test]
736 fn extract_system_message_when_present() {
737 let msgs = vec![
738 Message::system("You are helpful.".to_string()),
739 Message::user("Hello".to_string()),
740 ];
741 let (system, rest) = extract_system_message(&msgs);
742 assert_eq!(system, Some("You are helpful.".to_string()));
743 assert_eq!(rest.len(), 1);
744 assert_eq!(rest[0].role, Role::User);
745 }
746
747 #[test]
748 fn extract_system_message_when_not_present() {
749 let msgs = vec![Message::user("Hello".to_string())];
750 let (system, rest) = extract_system_message(&msgs);
751 assert_eq!(system, None);
752 assert_eq!(rest.len(), 1);
753 }
754
755 #[test]
756 fn filter_leading_assistant_removes_empty_assistant() {
757 let msgs = vec![
758 Message::assistant("Hi".to_string()),
759 Message::user("Hello".to_string()),
760 ];
761 let filtered = filter_leading_assistant(&msgs);
762 assert_eq!(filtered.len(), 1);
763 assert_eq!(filtered[0].role, Role::User);
764 }
765
766 #[test]
767 fn filter_leading_assistant_keeps_assistant_with_tool_calls() {
768 let msgs = vec![
769 Message::assistant_with_tool_calls(
770 "I'll use a tool".to_string(),
771 vec![ToolCall {
772 id: "call_1".to_string(),
773 name: "test".to_string(),
774 arguments: Value::Object(Default::default()),
775 }],
776 ),
777 Message::user("Hello".to_string()),
778 ];
779 let filtered = filter_leading_assistant(&msgs);
780 assert_eq!(filtered.len(), 2);
782 }
783
784 #[test]
785 fn parses_text_response() {
786 let raw = r#"{
787 "type": "message",
788 "content": [{"type": "text", "text": "Hello"}],
789 "stop_reason": "end_turn",
790 "usage": {"input_tokens": 10, "output_tokens": 5}
791 }"#;
792 let parsed: AnthropicResponse = serde_json::from_str(raw).unwrap();
793 let c = parse_completion(parsed);
794 assert_eq!(c.content, "Hello");
795 assert!(c.tool_calls.is_empty());
796 assert_eq!(c.finish_reason, Some("stop".to_string()));
797 assert!(c.usage.is_some());
798 }
799
800 #[test]
801 fn parses_tool_use_response() {
802 let raw = r#"{
803 "type": "message",
804 "content": [
805 {"type": "tool_use", "id": "call_1", "name": "read_file", "input": {"path": "src/lib.rs"}}
806 ],
807 "stop_reason": "tool_use",
808 "usage": {"input_tokens": 100, "output_tokens": 50}
809 }"#;
810 let parsed: AnthropicResponse = serde_json::from_str(raw).unwrap();
811 let c = parse_completion(parsed);
812 assert!(c.content.is_empty());
813 assert_eq!(c.tool_calls.len(), 1);
814 assert_eq!(c.tool_calls[0].id, "call_1");
815 assert_eq!(c.tool_calls[0].name, "read_file");
816 assert_eq!(c.tool_calls[0].arguments["path"], "src/lib.rs");
817 }
818
819 #[test]
820 fn parses_mixed_content_response() {
821 let raw = r#"{
822 "type": "message",
823 "content": [
824 {"type": "text", "text": "I'll read that file. "},
825 {"type": "tool_use", "id": "call_1", "name": "read_file", "input": {"path": "src/lib.rs"}}
826 ],
827 "stop_reason": "tool_use",
828 "usage": {"input_tokens": 50, "output_tokens": 25}
829 }"#;
830 let parsed: AnthropicResponse = serde_json::from_str(raw).unwrap();
831 let c = parse_completion(parsed);
832 assert_eq!(c.content, "I'll read that file. ");
833 assert_eq!(c.tool_calls.len(), 1);
834 }
835
836 #[test]
837 fn parses_usage_correctly() {
838 let raw = r#"{
839 "type": "message",
840 "content": [{"type": "text", "text": "hi"}],
841 "stop_reason": "end_turn",
842 "usage": {
843 "input_tokens": 100,
844 "output_tokens": 50,
845 "cache_creation_input_tokens": 20,
846 "cache_read_input_tokens": 30
847 }
848 }"#;
849 let parsed: AnthropicResponse = serde_json::from_str(raw).unwrap();
850 let c = parse_completion(parsed);
851 let u = c.usage.unwrap();
852 assert_eq!(u.prompt_tokens, 100);
853 assert_eq!(u.completion_tokens, 50);
854 assert_eq!(u.cache_hit_tokens, 30);
855 assert_eq!(u.cache_miss_tokens, 20);
856 }
857
858 #[test]
859 fn serializes_tool_result_message() {
860 let msg = Message::tool_result("call_123".to_string(), "result content".to_string());
861 let v = serialize_message(&msg);
862 assert_eq!(v["role"], "user");
863 let content = &v["content"][0];
864 assert_eq!(content["type"], "tool_result");
865 assert_eq!(content["tool_use_id"], "call_123");
866 assert_eq!(content["content"], "result content");
867 }
868
869 #[test]
870 fn serializes_assistant_with_tool_calls() {
871 let msg = Message::assistant_with_tool_calls(
872 "I'll use a tool".to_string(),
873 vec![ToolCall {
874 id: "abc".to_string(),
875 name: "write_file".to_string(),
876 arguments: serde_json::json!({"path": "a", "contents": "b"}),
877 }],
878 );
879 let v = serialize_message(&msg);
880 assert_eq!(v["role"], "assistant");
881 let content = &v["content"][0];
882 assert_eq!(content["type"], "tool_use");
883 assert_eq!(content["name"], "write_file");
884 }
885
886 #[test]
887 fn builds_request_with_system() {
888 let req = build_request(
889 "claude-3",
890 0.2,
891 4096,
892 Some("You are helpful."),
893 &[Message::user("Hi".to_string())],
894 &[],
895 );
896 assert_eq!(req["system"], "You are helpful.");
897 assert_eq!(req["model"], "claude-3");
898 }
899
900 #[test]
901 fn builds_request_without_system() {
902 let req = build_request(
903 "claude-3",
904 0.2,
905 4096,
906 None,
907 &[Message::user("Hi".to_string())],
908 &[],
909 );
910 assert!(req.get("system").is_none());
911 }
912
913 #[test]
914 fn builds_request_with_tools() {
915 let tools = vec![ToolSpec {
916 name: "test".to_string(),
917 description: "A test tool".to_string(),
918 parameters: serde_json::json!({"type": "object"}),
919 }];
920 let req = build_request(
921 "claude-3",
922 0.2,
923 4096,
924 None,
925 &[Message::user("Hi".to_string())],
926 &tools,
927 );
928 assert!(req.get("tools").is_some());
929 }
930
931 #[tokio::test]
932 async fn test_a_mock_server_returns_canned_response() {
933 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
935 let addr = listener.local_addr().unwrap();
936
937 let handle = std::thread::spawn(move || {
938 let (mut stream, _) = listener.accept().unwrap();
939 use std::io::{Read, Write};
940 let mut buf = [0u8; 4096];
941 let _ = stream.read(&mut buf);
942
943 let body = r#"{"type":"message","content":[{"type":"text","text":"Hello from Anthropic"}],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}"#;
944 let response = format!(
945 "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
946 body.len(),
947 body
948 );
949 write!(stream, "{}", response).unwrap();
950 stream.flush().unwrap();
951 });
952
953 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
955
956 let provider =
957 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
958 let result = provider
959 .complete(&[Message::user("hi".to_string())], &[])
960 .await;
961
962 let _ = handle.join();
963
964 let completion = result.expect("should succeed");
965 assert!(completion.content.contains("Hello from Anthropic"));
966 assert!(completion.usage.is_some());
967 let u = completion.usage.unwrap();
968 assert_eq!(u.prompt_tokens, 10);
969 assert_eq!(u.completion_tokens, 5);
970 }
971
972 #[tokio::test]
973 async fn test_b_error_response_includes_model_name() {
974 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
976 let addr = listener.local_addr().unwrap();
977
978 let handle = std::thread::spawn(move || {
979 let (mut stream, _) = listener.accept().unwrap();
980 use std::io::{Read, Write};
981 let mut buf = [0u8; 4096];
982 let _ = stream.read(&mut buf);
983
984 let body = r#"{"error":{"type":"authentication_error","message":"Invalid API Key"}}"#;
985 let response = format!(
986 "HTTP/1.1 401 Unauthorized\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
987 body.len(),
988 body
989 );
990 write!(stream, "{}", response).unwrap();
991 stream.flush().unwrap();
992 });
993
994 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
995
996 let provider =
997 AnthropicProvider::new(format!("http://{addr}"), "sk-invalid", "claude-3-opus");
998 let err = provider
999 .complete(&[Message::user("hi".to_string())], &[])
1000 .await
1001 .unwrap_err();
1002
1003 let _ = handle.join();
1004
1005 let msg = err.to_string();
1006 assert!(
1007 msg.contains("claude-3-opus"),
1008 "error should contain model name: {msg}"
1009 );
1010 assert!(
1011 msg.contains("401"),
1012 "error should contain HTTP status: {msg}"
1013 );
1014 }
1015
1016 #[tokio::test]
1017 async fn test_c_tool_use_extracts_correct_tool_call() {
1018 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1020 let addr = listener.local_addr().unwrap();
1021
1022 let handle = std::thread::spawn(move || {
1023 let (mut stream, _) = listener.accept().unwrap();
1024 use std::io::{Read, Write};
1025 let mut buf = [0u8; 4096];
1026 let _ = stream.read(&mut buf);
1027
1028 let body = r#"{"type":"message","content":[{"type":"tool_use","id":"call_abc123","name":"read_file","input":{"path":"src/lib.rs"}}],"stop_reason":"tool_use","usage":{"input_tokens":50,"output_tokens":30}}"#;
1029 let response = format!(
1030 "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1031 body.len(),
1032 body
1033 );
1034 write!(stream, "{}", response).unwrap();
1035 stream.flush().unwrap();
1036 });
1037
1038 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1039
1040 let provider =
1041 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
1042 let result = provider
1043 .complete(&[Message::user("read the file".to_string())], &[])
1044 .await;
1045
1046 let _ = handle.join();
1047
1048 let completion = result.expect("should succeed");
1049 assert_eq!(completion.tool_calls.len(), 1);
1050 let tc = &completion.tool_calls[0];
1051 assert_eq!(tc.id, "call_abc123");
1052 assert_eq!(tc.name, "read_file");
1053 assert_eq!(tc.arguments["path"], "src/lib.rs");
1054 }
1055
1056 #[tokio::test]
1057 async fn test_d_stream_request_includes_stream_true() {
1058 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1061 let addr = listener.local_addr().unwrap();
1062 let captured = std::sync::Arc::new(std::sync::Mutex::new(String::new()));
1063 let captured_clone = captured.clone();
1064
1065 std::thread::spawn(move || {
1066 let (mut stream, _) = listener.accept().unwrap();
1067 use std::io::{Read, Write};
1068 let mut buf = [0u8; 8192];
1069 let n = stream.read(&mut buf).unwrap();
1070 let request = String::from_utf8_lossy(&buf[..n]).to_string();
1071 if let Some(body_start) = request.find("\r\n\r\n") {
1072 let body = request[body_start + 4..].trim().to_string();
1073 *captured_clone.lock().unwrap() = body.clone();
1074 }
1075 let body = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
1077 write!(
1078 stream,
1079 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1080 body.len(),
1081 body,
1082 ).unwrap();
1083 stream.flush().unwrap();
1084 });
1085
1086 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1087
1088 let provider =
1089 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
1090 let _ = provider
1091 .stream(&[Message::user("hi".to_string())], &[], None)
1092 .await;
1093
1094 let body_str = captured.lock().unwrap().clone();
1095 let body: serde_json::Value = serde_json::from_str(&body_str).unwrap();
1096 assert_eq!(body["stream"], true);
1097 }
1098
1099 #[tokio::test]
1100 async fn test_e_stream_text_deltas_accumulate() {
1101 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1104 let addr = listener.local_addr().unwrap();
1105
1106 std::thread::spawn(move || {
1107 let (mut stream, _) = listener.accept().unwrap();
1108 use std::io::{Read, Write};
1109 let mut buf = [0u8; 4096];
1110 let _ = stream.read(&mut buf);
1111 let body = "\
1112event: message_start
1113data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}
1114
1115event: content_block_start
1116data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"Hello\"}}
1117
1118event: content_block_delta
1119data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" \"}}
1120
1121event: content_block_delta
1122data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"World\"}}
1123
1124event: message_delta
1125data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}
1126
1127event: message_stop
1128data: {\"type\":\"message_stop\"}
1129
1130";
1131 write!(
1132 stream,
1133 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1134 body.len(),
1135 body,
1136 ).unwrap();
1137 stream.flush().unwrap();
1138 });
1139
1140 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1141
1142 let provider =
1143 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
1144 let completion = provider
1145 .stream(&[Message::user("hi".to_string())], &[], None)
1146 .await
1147 .unwrap();
1148 assert_eq!(completion.content, "Hello World");
1149 assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
1150 assert!(completion.usage.is_some());
1151 let u = completion.usage.unwrap();
1152 assert_eq!(u.prompt_tokens, 10);
1153 assert_eq!(u.completion_tokens, 5);
1154 }
1155
1156 #[tokio::test]
1157 async fn test_f_stream_tool_use_assembles_tool_calls() {
1158 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1161 let addr = listener.local_addr().unwrap();
1162
1163 std::thread::spawn(move || {
1164 let (mut stream, _) = listener.accept().unwrap();
1165 use std::io::{Read, Write};
1166 let mut buf = [0u8; 4096];
1167 let _ = stream.read(&mut buf);
1168 let body = "\
1169event: message_start
1170data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":50,\"output_tokens\":0}}}
1171
1172event: content_block_start
1173data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_abc123\",\"name\":\"read_file\",\"input\":{}}}
1174
1175event: content_block_delta
1176data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"path\\\":\\\"src/\"}}
1177
1178event: content_block_delta
1179data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"lib.rs\\\"}\"}}
1180
1181event: message_delta
1182data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":30}}
1183
1184event: message_stop
1185data: {\"type\":\"message_stop\"}
1186
1187";
1188 write!(
1189 stream,
1190 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1191 body.len(),
1192 body,
1193 ).unwrap();
1194 stream.flush().unwrap();
1195 });
1196
1197 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1198
1199 let provider =
1200 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
1201 let completion = provider
1202 .stream(&[Message::user("read the file".to_string())], &[], None)
1203 .await
1204 .unwrap();
1205 assert_eq!(completion.tool_calls.len(), 1);
1206 let tc = &completion.tool_calls[0];
1207 assert_eq!(tc.id, "call_abc123");
1208 assert_eq!(tc.name, "read_file");
1209 assert_eq!(tc.arguments["path"], "src/lib.rs");
1210 assert_eq!(completion.finish_reason.as_deref(), Some("tool_calls"));
1211 }
1212
1213 #[tokio::test]
1214 async fn test_g_stream_tx_receives_text_chunks() {
1215 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1218 let addr = listener.local_addr().unwrap();
1219
1220 std::thread::spawn(move || {
1221 let (mut stream, _) = listener.accept().unwrap();
1222 use std::io::{Read, Write};
1223 let mut buf = [0u8; 4096];
1224 let _ = stream.read(&mut buf);
1225 let body = "\
1226event: message_start
1227data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":5,\"output_tokens\":0}}}
1228
1229event: content_block_start
1230data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"A\"}}
1231
1232event: content_block_delta
1233data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"B\"}}
1234
1235event: content_block_delta
1236data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"C\"}}
1237
1238event: message_delta
1239data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":3}}
1240
1241event: message_stop
1242data: {\"type\":\"message_stop\"}
1243
1244";
1245 write!(
1246 stream,
1247 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1248 body.len(),
1249 body,
1250 ).unwrap();
1251 stream.flush().unwrap();
1252 });
1253
1254 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1255
1256 let provider =
1257 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
1258 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1259 let completion = provider
1260 .stream(&[Message::user("hi".to_string())], &[], Some(tx))
1261 .await
1262 .unwrap();
1263 assert_eq!(completion.content, "ABC");
1264
1265 let mut deltas = Vec::new();
1267 while let Some(d) = rx.recv().await {
1268 deltas.push(d);
1269 if deltas.len() >= 3 {
1270 break;
1271 }
1272 }
1273 assert_eq!(deltas, vec!["A", "B", "C"]);
1274 }
1275
1276 #[tokio::test]
1277 async fn test_h_stream_with_end_turn_produces_stop_reason() {
1278 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1281 let addr = listener.local_addr().unwrap();
1282
1283 std::thread::spawn(move || {
1284 let (mut stream, _) = listener.accept().unwrap();
1285 use std::io::{Read, Write};
1286 let mut buf = [0u8; 4096];
1287 let _ = stream.read(&mut buf);
1288 let body = "\
1289event: message_start
1290data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}
1291
1292event: content_block_start
1293data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"Hello\"}}
1294
1295event: message_delta
1296data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}
1297
1298event: message_stop
1299data: {\"type\":\"message_stop\"}
1300
1301";
1302 write!(
1303 stream,
1304 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1305 body.len(),
1306 body,
1307 ).unwrap();
1308 stream.flush().unwrap();
1309 });
1310
1311 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1312
1313 let provider =
1314 AnthropicProvider::new(format!("http://{addr}"), "sk-noop", "claude-3-sonnet");
1315 let completion = provider
1316 .stream(&[Message::user("hi".to_string())], &[], None)
1317 .await
1318 .unwrap();
1319 assert_eq!(completion.content, "Hello");
1320 assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
1321 assert!(completion.usage.is_some());
1322 }
1323
1324 #[test]
1325 fn policy_retries_5xx_with_exponential_backoff() {
1326 let policy = RetryPolicy::default();
1327 assert_eq!(
1328 policy.backoff_for(0, Some(503), false),
1329 Some(Duration::from_secs(1))
1330 );
1331 assert_eq!(
1332 policy.backoff_for(1, Some(500), false),
1333 Some(Duration::from_secs(2))
1334 );
1335 assert_eq!(policy.backoff_for(2, Some(500), false), None);
1336 }
1337
1338 #[test]
1339 fn policy_retries_network_errors() {
1340 let policy = RetryPolicy::default();
1341 assert_eq!(
1342 policy.backoff_for(0, None, true),
1343 Some(Duration::from_secs(1))
1344 );
1345 }
1346
1347 #[test]
1348 fn policy_does_not_retry_4xx() {
1349 let policy = RetryPolicy::default();
1350 assert_eq!(policy.backoff_for(0, Some(400), false), None);
1351 assert_eq!(policy.backoff_for(0, Some(401), false), None);
1352 assert_eq!(policy.backoff_for(0, Some(404), false), None);
1353 }
1354}