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        tracing::debug!(model = %self.model, url = %url, body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "Anthropic chat request");
309
310        let response = self
311            .client
312            .post(&url)
313            .header("x-api-key", &self.api_key)
314            .header("anthropic-version", "2023-06-01")
315            .header("Content-Type", "application/json")
316            .json(&body)
317            .send()
318            .await
319            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
320
321        let status = response.status();
322        let res_json: Value = response.json().await
323            .map_err(|e| AgentError::json(format!("Response JSON parse failed: {e}")))?;
324
325        if !status.is_success() {
326            let err_msg = res_json
327                .get("error")
328                .and_then(|e| e.get("message"))
329                .and_then(Value::as_str)
330                .unwrap_or("unknown error");
331            tracing::warn!(status = %status, error = %err_msg, "Anthropic API non-success");
332            return Err(AgentError::LlmApi {
333                message: err_msg.to_string(),
334            });
335        }
336
337        tracing::debug!(status = %status, "Anthropic chat response received");
338        Ok(res_json)
339    }
340
341    async fn chat_stream(
342        &self,
343        messages: &[ChatMessage],
344        tools: &[Value],
345        reasoning: Option<&ReasoningConfig>,
346        _response_format: Option<&ResponseFormat>,
347    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
348        let url = format!("{}/v1/messages", self.base_url);
349        let mut body = Self::build_body(messages, tools, &self.model, reasoning);
350
351        if let Some(obj) = body.as_object_mut() {
352            obj.insert("stream".to_string(), json!(true));
353        }
354        tracing::debug!(model = %self.model, url = %url, body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "Anthropic chat_stream request");
355
356        let response = self
357            .client
358            .post(&url)
359            .header("x-api-key", &self.api_key)
360            .header("anthropic-version", "2023-06-01")
361            .header("Content-Type", "application/json")
362            .json(&body)
363            .send()
364            .await
365            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
366
367        if !response.status().is_success() {
368            let status = response.status();
369            let err_text = response.text().await
370                .map_err(|e| AgentError::llm(format!("Failed to read error response: {e}")))?;
371            tracing::warn!(%status, error = %err_text, "Anthropic API stream non-success");
372            return Err(AgentError::LlmApi { message: err_text });
373        }
374
375        let stream = response
376            .bytes_stream()
377            .eventsource()
378            .filter_map(|event| async move {
379                match event {
380                    Ok(ref ev) if ev.event == "error" => {
381                        let err_msg = ev.data.clone();
382                        Some(Err(AgentError::LlmApi { message: err_msg }))
383                    }
384                    Ok(ev) => {
385                        let event_type = if ev.event.is_empty() { "message_stop" } else { ev.event.as_str() };
386                        match Self::parse_sse(&ev.data, event_type) {
387                            Ok(chunk) => Some(Ok(chunk)),
388                            Err(e) => Some(Err(e)),
389                        }
390                    }
391                    Err(e) => Some(Err(AgentError::LlmStream(format!("SSE Stream error: {e}")))),
392                }
393            });
394
395        Ok(Box::pin(stream))
396    }
397
398    fn capabilities(&self) -> LlmCapabilities {
399        LlmCapabilities {
400            supports_streaming: true,
401            supports_tools: true,
402            supports_vision: true,
403            supports_thinking: true,
404            max_context_tokens: Some(200_000),
405            max_output_tokens: Some(8_192),
406        }
407    }
408}