Skip to main content

agent_base/llm/
anthropic.rs

1use async_trait::async_trait;
2use eventsource_stream::Eventsource;
3use futures_core::Stream;
4use futures_util::StreamExt;
5use reqwest::Client;
6use serde_json::{json, Value};
7use std::pin::Pin;
8
9use crate::types::{AgentResult, AgentError, ChatMessage, ImageAttachment, ResponseFormat};
10use super::{LlmCapabilities, LlmClient, ReasoningConfig, StreamChunk, UsageInfo};
11
12pub struct AnthropicClient {
13    api_key: String,
14    model: String,
15    base_url: String,
16    client: Client,
17}
18
19impl AnthropicClient {
20    pub fn new(api_key: String, model: String, base_url: Option<String>) -> Self {
21        Self::new_with_config(api_key, model, base_url, crate::llm::LlmClientConfig::default())
22    }
23
24    pub fn new_with_config(api_key: String, model: String, base_url: Option<String>, config: crate::llm::LlmClientConfig) -> Self {
25        let client = Client::builder()
26            .connect_timeout(config.connect_timeout)
27            .timeout(config.request_timeout)
28            .pool_max_idle_per_host(config.pool_max_idle_per_host)
29            .pool_idle_timeout(config.pool_idle_timeout)
30            .build()
31            .unwrap_or_else(|e| {
32                tracing::warn!(error = %e, "Failed to build reqwest client with custom config, falling back to default");
33                Client::new()
34            });
35        Self {
36            api_key,
37            model,
38            base_url: base_url
39                .unwrap_or_else(|| "https://api.anthropic.com".to_string()),
40            client,
41        }
42    }
43
44    fn convert_messages(messages: &[ChatMessage]) -> (Option<String>, Vec<Value>) {
45        let mut system_prompt: Option<String> = None;
46        let mut result: Vec<Value> = Vec::new();
47
48        for msg in messages {
49            match msg {
50                ChatMessage::System { content } => {
51                    system_prompt = Some(content.clone());
52                }
53                ChatMessage::User { content, images } => {
54                    let mut content_parts: Vec<Value> = Vec::new();
55                    content_parts.push(json!({"type": "text", "text": content}));
56                    for img in images {
57                        match img {
58                            ImageAttachment::Url { url, detail: _ } => {
59                                content_parts.push(json!({
60                                    "type": "image",
61                                    "source": {
62                                        "type": "url",
63                                        "url": url,
64                                    }
65                                }));
66                            }
67                            ImageAttachment::Base64 { data, media_type, detail: _ } => {
68                                let mime = media_type.as_deref().unwrap_or("image/jpeg");
69                                content_parts.push(json!({
70                                    "type": "image",
71                                    "source": {
72                                        "type": "base64",
73                                        "media_type": mime,
74                                        "data": data,
75                                    }
76                                }));
77                            }
78                        }
79                    }
80                    result.push(json!({
81                        "role": "user",
82                        "content": content_parts,
83                    }));
84                }
85                ChatMessage::Assistant { content, reasoning_content: _, tool_calls } => {
86                    let mut parts: Vec<Value> = Vec::new();
87                    if let Some(text) = content {
88                        if !text.is_empty() {
89                            parts.push(json!({"type": "text", "text": text}));
90                        }
91                    }
92                    if let Some(tc) = tool_calls {
93                        for t in tc {
94                            let input: Value = serde_json::from_str(&t.arguments)
95                                .unwrap_or(Value::Null);
96                            parts.push(json!({
97                                "type": "tool_use",
98                                "id": t.id,
99                                "name": t.name,
100                                "input": input,
101                            }));
102                        }
103                    }
104                    if !parts.is_empty() {
105                        result.push(json!({"role": "assistant", "content": parts}));
106                    }
107                }
108                ChatMessage::Tool { tool_call_id, content } => {
109                    result.push(json!({
110                        "role": "user",
111                        "content": [{
112                            "type": "tool_result",
113                            "tool_use_id": tool_call_id,
114                            "content": content,
115                        }]
116                    }));
117                }
118            }
119        }
120
121        (system_prompt, result)
122    }
123
124    fn convert_tools(tools: &[Value]) -> Vec<Value> {
125        tools
126            .iter()
127            .filter_map(|tool| {
128                let func = tool.get("function")?;
129                let name = func.get("name")?.as_str()?;
130                let description = func.get("description")
131                    .and_then(Value::as_str)
132                    .unwrap_or("");
133                let input_schema = func.get("parameters")
134                    .cloned()
135                    .unwrap_or_else(|| json!({"type": "object"}));
136                Some(json!({
137                    "name": name,
138                    "description": description,
139                    "input_schema": input_schema,
140                }))
141            })
142            .collect()
143    }
144
145    fn build_body(
146        messages: &[ChatMessage],
147        tools: &[Value],
148        model: &str,
149        reasoning: Option<&ReasoningConfig>,
150    ) -> Value {
151        let (system_prompt, anthropic_messages) = Self::convert_messages(messages);
152        let anthropic_tools = Self::convert_tools(tools);
153
154        let mut body = json!({
155            "model": model,
156            "max_tokens": 8192,
157            "messages": anthropic_messages,
158        });
159
160        if !anthropic_tools.is_empty() {
161            if let Some(obj) = body.as_object_mut() {
162                obj.insert("tools".to_string(), json!(anthropic_tools));
163            }
164        }
165
166        if let Some(system) = system_prompt {
167            if let Some(obj) = body.as_object_mut() {
168                obj.insert("system".to_string(), json!(system));
169            }
170        }
171
172        if let Some(config) = reasoning {
173            if config.enabled == Some(true) || config.budget_tokens.is_some() {
174                let mut thinking = serde_json::Map::new();
175                thinking.insert("type".to_string(), json!("enabled"));
176                if let Some(budget) = config.budget_tokens {
177                    thinking.insert("budget_tokens".to_string(), json!(budget));
178                }
179                if let Some(obj) = body.as_object_mut() {
180                    obj.insert("thinking".to_string(), Value::Object(thinking));
181                }
182            } else if config.enabled == Some(false) {
183                let mut thinking = serde_json::Map::new();
184                thinking.insert("type".to_string(), json!("disabled"));
185                if let Some(obj) = body.as_object_mut() {
186                    obj.insert("thinking".to_string(), Value::Object(thinking));
187                }
188            }
189        }
190
191        body
192    }
193
194    fn parse_sse(data_str: &str, event_type: &str) -> AgentResult<StreamChunk> {
195        if data_str.is_empty() {
196            return Ok(StreamChunk::Text(String::new()));
197        }
198
199        let data: Value = serde_json::from_str(data_str)
200            .map_err(|e| AgentError::json(format!("Anthropic SSE JSON: {e}")))?;
201
202        match event_type {
203            "message_start" => {
204                let input_tokens = data
205                    .get("message")
206                    .and_then(|m| m.get("usage"))
207                    .and_then(|u| u.get("input_tokens"))
208                    .and_then(Value::as_u64)
209                    .map(|v| v as u32);
210                let output_tokens = data
211                    .get("message")
212                    .and_then(|m| m.get("usage"))
213                    .and_then(|u| u.get("output_tokens"))
214                    .and_then(Value::as_u64)
215                    .map(|v| v as u32);
216                Ok(StreamChunk::Usage(UsageInfo {
217                    prompt_tokens: input_tokens,
218                    completion_tokens: output_tokens,
219                    total_tokens: None,
220                }))
221            }
222            "content_block_start" => {
223                let cb = data.get("content_block");
224                let idx = data.get("index").and_then(Value::as_u64).unwrap_or(0);
225                if let Some(cb) = cb {
226                    if cb.get("type").and_then(Value::as_str) == Some("tool_use") {
227                        let id = cb.get("id").and_then(Value::as_str).unwrap_or("").to_string();
228                        let name = cb.get("name").and_then(Value::as_str).unwrap_or("").to_string();
229                        return Ok(StreamChunk::ToolCall(json!({
230                            "delta": {
231                                "tool_calls": [{
232                                    "index": idx,
233                                    "id": if id.is_empty() { Value::Null } else { json!(id) },
234                                    "function": {
235                                        "name": name,
236                                        "arguments": "",
237                                    }
238                                }]
239                            }
240                        })));
241                    }
242                }
243                Ok(StreamChunk::Text(String::new()))
244            }
245            "content_block_delta" => {
246                let delta = data.get("delta");
247                let idx = data.get("index").and_then(Value::as_u64).unwrap_or(0);
248                if let Some(d) = delta {
249                    match d.get("type").and_then(Value::as_str) {
250                        Some("text_delta") => {
251                            let text = d.get("text").and_then(Value::as_str).unwrap_or("").to_string();
252                            Ok(StreamChunk::Text(text))
253                        }
254                        Some("input_json_delta") => {
255                            let partial = d.get("partial_json").and_then(Value::as_str).unwrap_or("").to_string();
256                            Ok(StreamChunk::ToolCall(json!({
257                                "delta": {
258                                    "tool_calls": [{
259                                        "index": idx,
260                                        "function": {
261                                            "arguments": partial,
262                                        }
263                                    }]
264                                }
265                            })))
266                        }
267                        Some("thinking_delta") => {
268                            let thinking = d.get("thinking").and_then(Value::as_str).unwrap_or("").to_string();
269                            Ok(StreamChunk::Thought(thinking))
270                        }
271                        _ => Ok(StreamChunk::Text(String::new())),
272                    }
273                } else {
274                    Ok(StreamChunk::Text(String::new()))
275                }
276            }
277            "content_block_stop" => Ok(StreamChunk::Text(String::new())),
278            "message_delta" => {
279                let output_tokens = data
280                    .get("usage")
281                    .and_then(|u| u.get("output_tokens"))
282                    .and_then(Value::as_u64)
283                    .map(|v| v as u32);
284                Ok(StreamChunk::Usage(UsageInfo {
285                    prompt_tokens: None,
286                    completion_tokens: output_tokens,
287                    total_tokens: None,
288                }))
289            }
290            "message_stop" => Ok(StreamChunk::Stop),
291            "ping" => Ok(StreamChunk::Text(String::new())),
292            _ => Ok(StreamChunk::Text(String::new())),
293        }
294    }
295}
296
297#[async_trait]
298impl LlmClient for AnthropicClient {
299    async fn chat(
300        &self,
301        messages: &[ChatMessage],
302        tools: &[Value],
303        reasoning: Option<&ReasoningConfig>,
304        _response_format: Option<&ResponseFormat>,
305    ) -> AgentResult<Value> {
306        let url = format!("{}/v1/messages", self.base_url);
307        let body = Self::build_body(messages, tools, &self.model, reasoning);
308
309        let response = self
310            .client
311            .post(&url)
312            .header("x-api-key", &self.api_key)
313            .header("anthropic-version", "2023-06-01")
314            .header("Content-Type", "application/json")
315            .json(&body)
316            .send()
317            .await
318            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
319
320        let status = response.status();
321        let res_json: Value = response.json().await
322            .map_err(|e| AgentError::json(format!("Response JSON parse failed: {e}")))?;
323
324        if !status.is_success() {
325            let err_msg = res_json
326                .get("error")
327                .and_then(|e| e.get("message"))
328                .and_then(Value::as_str)
329                .unwrap_or("unknown error");
330            return Err(AgentError::LlmApi {
331                message: err_msg.to_string(),
332            });
333        }
334
335        Ok(res_json)
336    }
337
338    async fn chat_stream(
339        &self,
340        messages: &[ChatMessage],
341        tools: &[Value],
342        reasoning: Option<&ReasoningConfig>,
343        _response_format: Option<&ResponseFormat>,
344    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
345        let url = format!("{}/v1/messages", self.base_url);
346        let mut body = Self::build_body(messages, tools, &self.model, reasoning);
347
348        if let Some(obj) = body.as_object_mut() {
349            obj.insert("stream".to_string(), json!(true));
350        }
351
352        let response = self
353            .client
354            .post(&url)
355            .header("x-api-key", &self.api_key)
356            .header("anthropic-version", "2023-06-01")
357            .header("Content-Type", "application/json")
358            .json(&body)
359            .send()
360            .await
361            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
362
363        if !response.status().is_success() {
364            let err_text = response.text().await
365                .map_err(|e| AgentError::llm(format!("Failed to read error response: {e}")))?;
366            return Err(AgentError::LlmApi { message: err_text });
367        }
368
369        let stream = response
370            .bytes_stream()
371            .eventsource()
372            .filter_map(|event| async move {
373                match event {
374                    Ok(ref ev) if ev.event == "error" => {
375                        let err_msg = ev.data.clone();
376                        Some(Err(AgentError::LlmApi { message: err_msg }))
377                    }
378                    Ok(ev) => {
379                        let event_type = if ev.event.is_empty() { "message_stop" } else { ev.event.as_str() };
380                        match Self::parse_sse(&ev.data, event_type) {
381                            Ok(chunk) => Some(Ok(chunk)),
382                            Err(e) => Some(Err(e)),
383                        }
384                    }
385                    Err(e) => Some(Err(AgentError::LlmStream(format!("SSE Stream error: {e}")))),
386                }
387            });
388
389        Ok(Box::pin(stream))
390    }
391
392    fn capabilities(&self) -> LlmCapabilities {
393        LlmCapabilities {
394            supports_streaming: true,
395            supports_tools: true,
396            supports_vision: true,
397            supports_thinking: true,
398            max_context_tokens: Some(200_000),
399            max_output_tokens: Some(8_192),
400        }
401    }
402}