limit_llm/
client.rs

1use crate::error::LlmError;
2use crate::providers::{LlmProvider, ProviderResponseChunk};
3
4use crate::types::{Message, Tool, Usage};
5use async_stream::stream;
6use futures::{Stream, StreamExt};
7use reqwest::Client;
8use serde_json::Value;
9use std::boxed::Box;
10use std::pin::Pin;
11use std::time::Duration;
12use tracing::{debug, error, info, instrument, trace, warn};
13
14pub struct AnthropicClient {
15    api_key: String,
16    client: Client,
17    base_url: String,
18    model: String,
19    max_tokens: u32,
20}
21
22#[derive(Debug)]
23struct SseEvent {
24    data: String,
25}
26
27impl Clone for AnthropicClient {
28    fn clone(&self) -> Self {
29        Self {
30            api_key: self.api_key.clone(),
31            client: self.client.clone(),
32            base_url: self.base_url.clone(),
33            model: self.model.clone(),
34            max_tokens: self.max_tokens,
35        }
36    }
37}
38
39#[async_trait::async_trait]
40impl LlmProvider for AnthropicClient {
41    async fn send(
42        &self,
43        messages: Vec<Message>,
44        tools: Vec<Tool>,
45    ) -> Result<
46        Pin<Box<dyn Stream<Item = Result<ProviderResponseChunk, LlmError>> + Send + '_>>,
47        LlmError,
48    > {
49        Ok(self.send(messages, tools).await)
50    }
51
52    fn provider_name(&self) -> &str {
53        "anthropic"
54    }
55
56    fn model_name(&self) -> &str {
57        &self.model
58    }
59
60    fn clone_box(&self) -> Box<dyn LlmProvider> {
61        Box::new(self.clone())
62    }
63}
64
65impl AnthropicClient {
66    pub fn new(
67        api_key: String,
68        base_url: Option<&str>,
69        timeout: u64,
70        model: &str,
71        max_tokens: u32,
72    ) -> Self {
73        let client = Client::builder()
74            .timeout(Duration::from_secs(timeout))
75            .connect_timeout(Duration::from_secs(30))
76            .build()
77            .expect("Failed to build HTTP client");
78
79        Self {
80            api_key,
81            client,
82            base_url: base_url
83                .unwrap_or("https://api.anthropic.com/v1/messages")
84                .to_string(),
85            model: model.to_string(),
86            max_tokens,
87        }
88    }
89
90    #[instrument(skip(self, messages, tools))]
91    pub async fn send(
92        &self,
93        messages: Vec<Message>,
94        tools: Vec<Tool>,
95    ) -> Pin<Box<dyn Stream<Item = Result<ProviderResponseChunk, LlmError>> + Send + '_>> {
96        let api_key = self.api_key.clone();
97        let base_url = self.base_url.clone();
98        let model = self.model.clone();
99        let max_tokens = self.max_tokens;
100        let messages_cloned = messages.clone();
101        let tools_cloned = tools.clone();
102        let client_clone = self.client.clone();
103
104        Box::pin(stream! {
105            info!("API request: model={}, max_tokens={}", self.model, self.max_tokens);
106
107            let request_body = match build_request_body(&messages_cloned, &tools_cloned, &model, max_tokens) {
108                Ok(body) => body,
109                Err(e) => {
110                    error!("API error: {}", e);
111                    yield Err(e);
112                    return;
113                }
114            };
115
116            for attempt in 0..3 {
117                let delay = Duration::from_secs(2_u64.pow(attempt));
118
119                match do_request(&client_clone, &api_key, &base_url, &request_body).await {
120                    Ok(mut stream) => {
121                        while let Some(chunk) = stream.next().await {
122                            yield chunk;
123                        }
124                        return;
125                    }
126                    Err(e) => {
127                        if attempt == 2 {
128                            error!("API error: {}", e);
129                            yield Err(e);
130                            return;
131                        }
132                        warn!("API retry: attempt={}, delay_ms={}", attempt, delay.as_millis());
133                        tokio::time::sleep(delay).await;
134                    }
135                }
136            }
137        })
138    }
139}
140
141#[instrument(skip_all)]
142#[allow(clippy::type_complexity)]
143async fn do_request(
144    client: &Client,
145    api_key: &str,
146    base_url: &str,
147    request_body: &Value,
148) -> Result<
149    Pin<Box<dyn Stream<Item = Result<ProviderResponseChunk, LlmError>> + Send + 'static>>,
150    LlmError,
151> {
152    let response = client
153        .post(base_url)
154        .header("x-api-key", api_key)
155        .header("anthropic-version", "2023-06-01")
156        .header("content-type", "application/json")
157        .json(request_body)
158        .send()
159        .await
160        .map_err(|e| LlmError::NetworkError(e.to_string()))?;
161
162    let status = response.status();
163    debug!("API response received: status={}", status.as_u16());
164
165    if status.is_client_error() || status.is_server_error() {
166        let error_text = response
167            .text()
168            .await
169            .unwrap_or_else(|_| "Unknown error".to_string());
170
171        if status.as_u16() == 429 {
172            error!("API error: Rate limited");
173            return Err(LlmError::ApiError(format!("Rate limited: {}", error_text)));
174        }
175        error!("API error: HTTP {}: {}", status, error_text);
176        return Err(LlmError::ApiError(format!(
177            "HTTP {}: {}",
178            status, error_text
179        )));
180    }
181
182    let byte_stream = response.bytes_stream();
183    let stream = parse_sse_stream(byte_stream);
184    Ok(stream)
185}
186
187fn build_request_body(
188    messages: &[Message],
189    tools: &[Tool],
190    model: &str,
191    max_tokens: u32,
192) -> Result<Value, LlmError> {
193    let mut request = serde_json::json!({
194        "model": model,
195        "max_tokens": max_tokens,
196        "messages": messages,
197        "stream": true
198    });
199
200    if !tools.is_empty() {
201        request["tools"] = serde_json::to_value(tools)
202            .map_err(|e| LlmError::ApiError(format!("Failed to serialize tools: {}", e)))?;
203    }
204
205    Ok(request)
206}
207/// Parse potentially incomplete JSON during streaming.
208/// Returns empty object if parsing fails.
209fn parse_partial_json(json: &str) -> serde_json::Value {
210    if json.trim().is_empty() {
211        return serde_json::json!({});
212    }
213
214    // Try standard parsing first
215    if let Ok(value) = serde_json::from_str::<serde_json::Value>(json) {
216        return value;
217    }
218
219    // If parsing fails, return empty object
220    // (In future, could use partial-json crate for better handling)
221    serde_json::json!({})
222}
223
224fn parse_sse_stream(
225    byte_stream: impl Stream<Item = reqwest::Result<bytes::Bytes>> + Send + Unpin + 'static,
226) -> Pin<Box<dyn Stream<Item = Result<ProviderResponseChunk, LlmError>> + Send + 'static>> {
227    Box::pin(stream! {
228            let mut buffer = String::new();
229            let mut tool_calls_by_id: std::collections::HashMap<u64, (String, String)> = std::collections::HashMap::new();
230
231            let mut tool_partial_json: std::collections::HashMap<u64, String> = std::collections::HashMap::new();
232
233            let mut lines = byte_stream
234                .map(|chunk| chunk.map_err(|e| LlmError::NetworkError(e.to_string())));
235
236            while let Some(chunk_result) = lines.next().await {
237                let chunk = match chunk_result {
238                    Ok(c) => c,
239                    Err(e) => {
240                        yield Err(e);
241                        continue;
242                    }
243                };
244
245                let text = String::from_utf8_lossy(&chunk);
246
247                buffer.push_str(&text);
248
249                while let Some(event) = parse_sse_line(&mut buffer) {
250
251                    if event.data == "[DONE]" {
252                        return;
253                    }
254
255                    if let Ok(parsed) = serde_json::from_str::<Value>(&event.data) {
256                        trace!("SSE: {}", &event.data.chars().take(200).collect::<String>());
257                        let chunk_type = parsed.get("type").and_then(|v| v.as_str()).unwrap_or("");
258
259                        match chunk_type {
260                            "content_block_delta" => {
261                                if let Some(delta) = parsed.get("delta") {
262                                    // Handle text deltas
263                                    if let Some(text) = delta.get("text").and_then(|v| v.as_str()) {
264                                        yield Ok(ProviderResponseChunk::ContentDelta(text.to_string()));
265                                    }
266
267                                    // Handle tool argument deltas (input_json_delta)
268                                    let delta_type = delta.get("type").and_then(|v| v.as_str());
269                                    if delta_type == Some("input_json_delta") {
270                                        if let Some(partial_json) = delta.get("partial_json").and_then(|v| v.as_str()) {
271                                            // Get tool call index
272                                            if let Some(index) = parsed.get("index").and_then(|v| v.as_u64()) {
273                                                // Accumulate partial JSON
274    tool_partial_json.entry(index)
275                                                    .or_default()
276                                                    .push_str(partial_json);
277
278                                                // Look up tool call metadata and parse accumulated JSON
279                                                if let Some((id, name)) = tool_calls_by_id.get(&index) {
280                                                    let accumulated = tool_partial_json.get(&index).unwrap();
281                                                    let args = parse_partial_json(accumulated);
282
283                                                    yield Ok(ProviderResponseChunk::ToolCallDelta {
284                                                        id: id.clone(),
285                                                        name: name.clone(),
286                                                        arguments: args,
287                                                    });
288                                                }
289                                            }
290                                        }
291                                    }
292                                }
293                            }
294                            "content_block_start" => {
295                                if let Some(content_block) = parsed.get("content_block") {
296                                    let block_type = content_block.get("type").and_then(|v| v.as_str());
297                                    if block_type == Some("tool_use") {
298                                        let id = content_block.get("id")
299                                            .and_then(|v| v.as_str())
300                                            .unwrap_or("")
301                                            .to_string();
302                                        let name = content_block.get("name")
303                                            .and_then(|v| v.as_str())
304                                            .unwrap_or("")
305                                            .to_string();
306
307                                        // Track tool call by index
308                                        if let Some(index) = parsed.get("index").and_then(|v| v.as_u64()) {
309                                            tool_calls_by_id.insert(index, (id.clone(), name.clone()));
310                                        }
311
312                                        yield Ok(ProviderResponseChunk::ToolCallDelta {
313                                            id,
314                                            name,
315                                            arguments: serde_json::json!({}),
316                                        });
317                                    }
318                                }
319                            }
320                            "content_block_stop" => {
321                                // Tool call completed
322                            }
323                            "message_delta" => {
324                                if let Some(delta) = parsed.get("delta") {
325                                    if let Some(stop_reason) = delta.get("stop_reason").and_then(|v| v.as_str()) {
326                                        debug!("stop_reason: {}", stop_reason);
327                                        if stop_reason == "end_turn" || stop_reason == "tool_use" {
328                                            if let Some(usage) = parsed.get("usage") {
329                                                if let Ok(usage_obj) = serde_json::from_value::<Usage>(usage.clone()) {
330                                                    yield Ok(ProviderResponseChunk::Done(usage_obj));
331                                                    return;
332                                                }
333                                            }
334                                        }
335                                    }
336                                }
337                            }
338                            _ => {
339                                debug!("Unknown chunk_type: {}", chunk_type);
340                            }
341                        }
342                    }
343                }
344            }
345        })
346}
347
348fn parse_sse_line(buffer: &mut String) -> Option<SseEvent> {
349    loop {
350        let newline_pos = buffer.find('\n')?;
351        let line = buffer[..newline_pos].trim().to_string();
352        *buffer = buffer[newline_pos + 1..].to_string();
353
354        // Skip empty lines and comments
355        if line.is_empty() || line.starts_with(':') {
356            continue;
357        }
358
359        // Skip event: lines (we only care about data)
360        if line.starts_with("event:") {
361            continue;
362        }
363
364        // Parse data: lines
365        if let Some(data_pos) = line.find("data: ") {
366            let data = line[data_pos + 6..].trim();
367            return Some(SseEvent {
368                data: data.to_string(),
369            });
370        }
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377    use mockito::Server;
378
379    #[tokio::test]
380    async fn test_streaming() {
381        let mut server = Server::new_async().await;
382        let mock = server
383            .mock("POST", "/v1/messages")
384            .with_status(200)
385            .with_header("content-type", "text/event-stream")
386            .with_chunked_body(|w| {
387                w.write_all(b"data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"Hello\"}}\n\n")?;
388                w.write_all(b"data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\" world\"}}\n\n")?;
389                w.write_all(b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"input_tokens\":10,\"output_tokens\":5}}\n\n")?;
390                Ok::<(), std::io::Error>(())
391            })
392            .create_async()
393            .await;
394
395        let client = AnthropicClient::new(
396            "test-key".to_string(),
397            None,
398            300,
399            "claude-3-5-sonnet-20241022",
400            4096,
401        );
402        let messages = vec![Message {
403            role: crate::types::Role::User,
404            content: Some("Hello".to_string()),
405            tool_calls: None,
406            tool_call_id: None,
407        }];
408
409        let base_url = format!("{}/v1/messages", server.url());
410        let client_with_url = AnthropicClient {
411            api_key: "test-key".to_string(),
412            client: client.client,
413            base_url,
414            model: "claude-3-5-sonnet-20241022".to_string(),
415            max_tokens: 4096,
416        };
417
418        let stream = client_with_url.send(messages, vec![]).await;
419        let chunks: Vec<_> = stream.collect().await;
420        assert!(chunks.len() >= 3);
421
422        mock.assert_async().await;
423    }
424
425    #[tokio::test]
426    async fn test_retry_on_429() {
427        let mut server = Server::new_async().await;
428        let mock = server
429            .mock("POST", "/v1/messages")
430            .with_status(429)
431            .with_header("content-type", "application/json")
432            .with_body(r#"{"error":{"type":"rate_limit_error","message":"Rate limited"}}"#)
433            .expect(2)
434            .create_async()
435            .await;
436
437        let success_mock = server
438            .mock("POST", "/v1/messages")
439            .with_status(200)
440            .with_header("content-type", "text/event-stream")
441            .with_chunked_body(|w| {
442                w.write_all(b"data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"Hello\"}}\n\n")?;
443                w.write_all(b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"input_tokens\":10,\"output_tokens\":5}}\n\n")?;
444                Ok::<(), std::io::Error>(())
445            })
446            .expect(1)
447            .create_async()
448            .await;
449
450        let client = AnthropicClient::new(
451            "test-key".to_string(),
452            None,
453            300,
454            "claude-3-5-sonnet-20241022",
455            4096,
456        );
457        let messages = vec![Message {
458            role: crate::types::Role::User,
459            content: Some("Hello".to_string()),
460            tool_calls: None,
461            tool_call_id: None,
462        }];
463
464        let base_url = format!("{}/v1/messages", server.url());
465        let client_with_url = AnthropicClient {
466            api_key: "test-key".to_string(),
467            client: client.client,
468            base_url,
469            model: "claude-3-5-sonnet-20241022".to_string(),
470            max_tokens: 4096,
471        };
472
473        let stream = client_with_url.send(messages, vec![]).await;
474        let chunks: Vec<_> = stream.collect().await;
475        assert!(!chunks.is_empty());
476
477        mock.assert_async().await;
478        success_mock.assert_async().await;
479    }
480
481    #[tokio::test]
482    async fn test_timeout() {
483        let mut server = Server::new_async().await;
484        let _mock = server
485            .mock("POST", "/v1/messages")
486            .with_status(200)
487            .with_header("content-type", "text/event-stream")
488            .with_chunked_body(|w| {
489                // Sleep to simulate slow response
490                std::thread::sleep(std::time::Duration::from_millis(500));
491                w.write_all(
492                    b"data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"Hello\"}}\n\n",
493                )?;
494                Ok::<(), std::io::Error>(())
495            })
496            .create_async()
497            .await;
498
499        let client = AnthropicClient::new(
500            "test-key".to_string(),
501            None,
502            300,
503            "claude-3-5-sonnet-20241022",
504            4096,
505        );
506        let messages = vec![Message {
507            role: crate::types::Role::User,
508            content: Some("Hello".to_string()),
509            tool_calls: None,
510            tool_call_id: None,
511        }];
512
513        let base_url = format!("{}/v1/messages", server.url());
514        let client_with_url = AnthropicClient {
515            api_key: "test-key".to_string(),
516            client: client.client,
517            base_url,
518            model: "claude-3-5-sonnet-20241022".to_string(),
519            max_tokens: 4096,
520        };
521
522        // The test should pass since timeout is 300s
523        let stream = client_with_url.send(messages, vec![]).await;
524        let chunks: Vec<_> = stream.collect().await;
525        assert!(!chunks.is_empty());
526    }
527
528    #[tokio::test]
529    async fn test_tool_call_streaming() {
530        let mut server = Server::new_async().await;
531        let mock = server
532            .mock("POST", "/v1/messages")
533            .with_status(200)
534            .with_header("content-type", "text/event-stream")
535            .with_chunked_body(|w| {
536                w.write_all(b"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_123\",\"name\":\"test_tool\"}}\n\n")?;
537                w.write_all(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"partial_json\":\"{\\\"arg\\\":\\\"value\\\"}\"}}\n\n")?;
538                w.write_all(b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"input_tokens\":15,\"output_tokens\":20}}\n\n")?;
539                Ok::<(), std::io::Error>(())
540            })
541            .create_async()
542            .await;
543
544        let client = AnthropicClient::new(
545            "test-key".to_string(),
546            None,
547            300,
548            "claude-3-5-sonnet-20241022",
549            4096,
550        );
551        let messages = vec![Message {
552            role: crate::types::Role::User,
553            content: Some("Use test_tool".to_string()),
554            tool_calls: None,
555            tool_call_id: None,
556        }];
557
558        let tools = vec![Tool {
559            tool_type: "function".to_string(),
560            function: crate::types::ToolFunction {
561                name: "test_tool".to_string(),
562                description: "A test tool".to_string(),
563                parameters: serde_json::json!({"type": "object"}),
564            },
565        }];
566
567        let base_url = format!("{}/v1/messages", server.url());
568        let client_with_url = AnthropicClient {
569            api_key: "test-key".to_string(),
570            client: client.client,
571            base_url,
572            model: "claude-3-5-sonnet-20241022".to_string(),
573            max_tokens: 4096,
574        };
575
576        let stream = client_with_url.send(messages, tools).await;
577        let chunks: Vec<_> = stream.collect().await;
578        assert!(!chunks.is_empty());
579
580        mock.assert_async().await;
581    }
582
583    #[test]
584    fn test_parse_sse_line() {
585        let mut buffer = String::from("data: {\"type\":\"test\"}\n\nother data");
586        let event = parse_sse_line(&mut buffer);
587        assert!(event.is_some());
588        assert_eq!(event.unwrap().data, "{\"type\":\"test\"}");
589        assert_eq!(buffer, "\nother data");
590    }
591
592    #[test]
593    fn test_parse_sse_line_empty() {
594        let mut buffer = String::from("\n\ndata: test");
595        let event = parse_sse_line(&mut buffer);
596        assert!(event.is_none());
597        assert_eq!(buffer, "data: test");
598    }
599
600    #[test]
601    fn test_parse_sse_line_comment() {
602        let mut buffer = String::from(": comment\n\ndata: test");
603        let event = parse_sse_line(&mut buffer);
604        assert!(event.is_none());
605    }
606
607    #[test]
608    fn test_parse_sse_line_zai_format() {
609        let mut buffer = String::from("event: content_block_start\ndata: {\"type\":\"test\"}\n\n");
610        let event = parse_sse_line(&mut buffer);
611        assert!(event.is_some());
612        assert_eq!(event.unwrap().data, "{\"type\":\"test\"}");
613    }
614}
limit_llm/client.rs

limit_llm/
client.rs
