lc_core/batch/
client.rs

1// src/core/batch/client.rs
2//! BatchClient struct and all its method implementations.
3
4use super::types::*;
5use crate::language_models::LLMResult;
6use lc_schema::Message;
7use serde_json::json;
8
9/// Client for submitting, polling, and retrieving batch results.
10pub struct BatchClient {
11    pub(crate) http: reqwest::Client,
12    // M35: API key stored as String — could leak via debug/display.
13    // Consider using a Zeroize wrapper or secret type in a future refactor.
14    pub(crate) api_key: String,
15    pub(crate) provider: BatchProvider,
16    pub(crate) base_url: String,
17}
18
19impl BatchClient {
20    /// Creates a new `BatchClient` for the given provider.
21    pub fn new(provider: BatchProvider, api_key: impl Into<String>) -> Self {
22        let base_url = match provider {
23            BatchProvider::OpenAI => "https://api.openai.com/v1".to_string(),
24            BatchProvider::Anthropic => "https://api.anthropic.com/v1".to_string(),
25        };
26        Self {
27            http: reqwest::Client::new(),
28            api_key: api_key.into(),
29            provider,
30            base_url,
31        }
32    }
33
34    /// Overrides the base URL (useful for proxies or compatible endpoints).
35    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
36        self.base_url = url.into();
37        self
38    }
39
40    // -- Auth header helpers ------------------------------------------------
41
42    fn auth_headers(&self) -> Result<reqwest::header::HeaderMap, BatchError> {
43        let mut headers = reqwest::header::HeaderMap::new();
44        match self.provider {
45            BatchProvider::OpenAI => {
46                let value = format!("Bearer {}", self.api_key);
47                let v = reqwest::header::HeaderValue::from_str(&value).map_err(|e| {
48                    BatchError::Api(format!("invalid API key for Authorization header: {}", e))
49                })?;
50                headers.insert("Authorization", v);
51            }
52            BatchProvider::Anthropic => {
53                let v = reqwest::header::HeaderValue::from_str(&self.api_key).map_err(|e| {
54                    BatchError::Api(format!("invalid API key for x-api-key header: {}", e))
55                })?;
56                headers.insert("x-api-key", v);
57                headers.insert(
58                    "anthropic-version",
59                    reqwest::header::HeaderValue::from_static("2023-06-01"),
60                );
61            }
62        }
63        Ok(headers)
64    }
65
66    // -- Message conversion -------------------------------------------------
67
68    /// Convert a [`Message`] to the OpenAI chat format JSON value.
69    pub(crate) fn message_to_openai(msg: &Message) -> serde_json::Value {
70        match &msg.message_type {
71            lc_schema::MessageType::System => json!({
72                "role": "system",
73                "content": msg.content,
74            }),
75            lc_schema::MessageType::Human => json!({
76                "role": "user",
77                "content": msg.content,
78            }),
79            lc_schema::MessageType::AI => {
80                let mut m = json!({
81                    "role": "assistant",
82                    "content": msg.content,
83                });
84                if let Some(tc) = &msg.tool_calls {
85                    m["tool_calls"] = serde_json::to_value(tc).unwrap_or(serde_json::Value::Null);
86                }
87                m
88            }
89            lc_schema::MessageType::Tool { tool_call_id } => json!({
90                "role": "tool",
91                "tool_call_id": tool_call_id,
92                "content": msg.content,
93            }),
94        }
95    }
96
97    /// Convert a [`Message`] to the Anthropic chat format JSON value.
98    ///
99    /// Anthropic uses a `system` top-level parameter for system messages,
100    /// but in the messages array it only accepts `user` and `assistant` roles.
101    /// System messages are mapped to `user` with a `[System]` prefix to
102    /// preserve the instruction while conforming to the API constraints.
103    /// Tool results are mapped to `user` role per Anthropic's convention.
104    pub(crate) fn message_to_anthropic(msg: &Message) -> serde_json::Value {
105        match &msg.message_type {
106            // H40: System messages should be extracted to top-level `system` parameter
107            // by the caller (submit_anthropic), not mapped as user messages.
108            // Here we map them as system role for identification.
109            lc_schema::MessageType::System => json!({
110                "role": "system",
111                "content": msg.content,
112            }),
113            lc_schema::MessageType::Human => json!({
114                "role": "user",
115                "content": msg.content,
116            }),
117            lc_schema::MessageType::AI => json!({
118                "role": "assistant",
119                "content": msg.content,
120            }),
121            // H40: Tool messages should use Anthropic tool_result content block format
122            lc_schema::MessageType::Tool { tool_call_id } => json!({
123                "role": "user",
124                "content": [{
125                    "type": "tool_result",
126                    "tool_use_id": tool_call_id,
127                    "content": msg.content,
128                }],
129            }),
130        }
131    }
132
133    // -- Submit -------------------------------------------------------------
134
135    /// Submit a batch of requests. Returns the batch ID.
136    ///
137    /// For **OpenAI**: uploads a JSONL file via the Files API, then creates a
138    /// batch via `POST /v1/batches` referencing the uploaded file.
139    ///
140    /// For **Anthropic**: creates a batch via `POST /v1/messages/batches` with
141    /// the requests array inline.
142    pub async fn submit(&self, requests: Vec<BatchRequest>) -> Result<BatchId, BatchError> {
143        match self.provider {
144            BatchProvider::OpenAI => self.submit_openai(requests).await,
145            BatchProvider::Anthropic => self.submit_anthropic(requests).await,
146        }
147    }
148
149    async fn submit_openai(&self, requests: Vec<BatchRequest>) -> Result<BatchId, BatchError> {
150        // Build JSONL content for the input file.
151        let jsonl_lines: Vec<String> = requests
152            .iter()
153            .map(|req| {
154                let openai_msgs: Vec<serde_json::Value> =
155                    req.messages.iter().map(Self::message_to_openai).collect();
156                let mut body = json!({
157                    "model": req.model,
158                    "messages": openai_msgs,
159                });
160                if let Some(t) = req.temperature {
161                    body["temperature"] = json!(t);
162                }
163                if let Some(m) = req.max_tokens {
164                    body["max_tokens"] = json!(m);
165                }
166                let line = json!({
167                    "custom_id": req.custom_id,
168                    "method": "POST",
169                    "url": "/v1/chat/completions",
170                    "body": body,
171                });
172                // H39: Return error instead of empty string on serialization failure
173                serde_json::to_string(&line).map_err(BatchError::Serialization)
174            })
175            .collect::<Result<Vec<String>, BatchError>>()?;
176        let jsonl_content = jsonl_lines.join("\n");
177
178        // Step 1: Upload the JSONL file.
179        let file_id = self.upload_openai_file(&jsonl_content).await?;
180
181        // Step 2: Create the batch.
182        let headers = self.auth_headers()?;
183        let body = json!({
184            "input_file_id": file_id,
185            "endpoint": "/v1/chat/completions",
186            "completion_window": "24h",
187        });
188
189        let resp = self
190            .http
191            .post(format!("{}/batches", self.base_url))
192            .headers(headers)
193            .json(&body)
194            .send()
195            .await?;
196
197        if resp.status() == reqwest::StatusCode::NOT_FOUND {
198            return Err(BatchError::NotFound(
199                "batches endpoint not found".to_string(),
200            ));
201        }
202
203        let status = resp.status();
204        let text = resp.text().await?;
205
206        if !status.is_success() {
207            return Err(BatchError::Api(format!(
208                "batch creation failed ({}): {}",
209                status, text
210            )));
211        }
212
213        let batch: OpenAIBatchResponse = serde_json::from_str(&text)?;
214        Ok(BatchId(batch.id))
215    }
216
217    /// Upload a JSONL file to the OpenAI Files API and return the file ID.
218    async fn upload_openai_file(&self, jsonl_content: &str) -> Result<String, BatchError> {
219        let headers = self.auth_headers()?;
220        let file_part = reqwest::multipart::Part::text(jsonl_content.to_string())
221            .file_name("batch_input.jsonl")
222            .mime_str("application/jsonl")
223            .map_err(|e| BatchError::Api(format!("mime error: {e}")))?;
224
225        let form = reqwest::multipart::Form::new()
226            .text("purpose", "batch")
227            .part("file", file_part);
228
229        let resp = self
230            .http
231            .post(format!("{}/files", self.base_url))
232            .headers(headers)
233            .multipart(form)
234            .send()
235            .await?;
236
237        let status = resp.status();
238        let text = resp.text().await?;
239
240        if !status.is_success() {
241            return Err(BatchError::Api(format!(
242                "file upload failed ({}): {}",
243                status, text
244            )));
245        }
246
247        let file_resp: OpenAIFileResponse = serde_json::from_str(&text)?;
248        Ok(file_resp.id)
249    }
250
251    async fn submit_anthropic(&self, requests: Vec<BatchRequest>) -> Result<BatchId, BatchError> {
252        let req_items: Vec<serde_json::Value> = requests
253            .iter()
254            .map(|req| {
255                // H40: Extract system messages to top-level `system` parameter,
256                // and filter them from the messages array (Anthropic API requirement).
257                let mut system_text = String::new();
258                let anthropic_msgs: Vec<serde_json::Value> = req
259                    .messages
260                    .iter()
261                    .filter_map(|msg| {
262                        if matches!(msg.message_type, lc_schema::MessageType::System) {
263                            if !system_text.is_empty() {
264                                system_text.push('\n');
265                            }
266                            system_text.push_str(&msg.content);
267                            None
268                        } else {
269                            Some(Self::message_to_anthropic(msg))
270                        }
271                    })
272                    .collect();
273                let mut body = json!({
274                    "model": req.model,
275                    "max_tokens": req.max_tokens.unwrap_or(4096),
276                    "messages": anthropic_msgs,
277                });
278                if !system_text.is_empty() {
279                    body["system"] = json!(system_text);
280                }
281                if let Some(t) = req.temperature {
282                    body["temperature"] = json!(t);
283                }
284                json!({
285                    "custom_id": req.custom_id,
286                    "params": body,
287                })
288            })
289            .collect();
290
291        let headers = self.auth_headers()?;
292        let body = json!({
293            "requests": req_items,
294        });
295
296        let resp = self
297            .http
298            .post(format!("{}/messages/batches", self.base_url))
299            .headers(headers)
300            .json(&body)
301            .send()
302            .await?;
303
304        let status = resp.status();
305        let text = resp.text().await?;
306
307        if status == reqwest::StatusCode::NOT_FOUND {
308            return Err(BatchError::NotFound(
309                "messages/batches endpoint not found".to_string(),
310            ));
311        }
312
313        if !status.is_success() {
314            return Err(BatchError::Api(format!(
315                "batch creation failed ({}): {}",
316                status, text
317            )));
318        }
319
320        let batch: AnthropicBatchResponse = serde_json::from_str(&text)?;
321        Ok(BatchId(batch.id))
322    }
323
324    // -- Poll ---------------------------------------------------------------
325
326    /// Poll the status of a batch job.
327    pub async fn poll(&self, id: &BatchId) -> Result<BatchStatus, BatchError> {
328        match self.provider {
329            BatchProvider::OpenAI => self.poll_openai(id).await,
330            BatchProvider::Anthropic => self.poll_anthropic(id).await,
331        }
332    }
333
334    async fn poll_openai(&self, id: &BatchId) -> Result<BatchStatus, BatchError> {
335        let headers = self.auth_headers()?;
336        let resp = self
337            .http
338            .get(format!("{}/batches/{}", self.base_url, id.0))
339            .headers(headers)
340            .send()
341            .await?;
342
343        if resp.status() == reqwest::StatusCode::NOT_FOUND {
344            return Err(BatchError::NotFound(id.0.clone()));
345        }
346
347        let status = resp.status();
348        let text = resp.text().await?;
349
350        if !status.is_success() {
351            return Err(BatchError::Api(format!(
352                "poll failed ({}): {}",
353                status, text
354            )));
355        }
356
357        let batch: OpenAIBatchResponse = serde_json::from_str(&text)?;
358        Ok(match batch.status.as_str() {
359            "in_progress" | "validating" | "finalizing" => BatchStatus::InProgress,
360            "completed" => BatchStatus::Completed,
361            "failed" => BatchStatus::Failed,
362            "expired" => BatchStatus::Expired,
363            "cancelling" | "cancelled" => BatchStatus::Cancelled,
364            other => return Err(BatchError::Api(format!("unknown batch status: {}", other))),
365        })
366    }
367
368    async fn poll_anthropic(&self, id: &BatchId) -> Result<BatchStatus, BatchError> {
369        let headers = self.auth_headers()?;
370        let resp = self
371            .http
372            .get(format!("{}/messages/batches/{}", self.base_url, id.0))
373            .headers(headers)
374            .send()
375            .await?;
376
377        if resp.status() == reqwest::StatusCode::NOT_FOUND {
378            return Err(BatchError::NotFound(id.0.clone()));
379        }
380
381        let status = resp.status();
382        let text = resp.text().await?;
383
384        if !status.is_success() {
385            return Err(BatchError::Api(format!(
386                "poll failed ({}): {}",
387                status, text
388            )));
389        }
390
391        let batch: AnthropicBatchResponse = serde_json::from_str(&text)?;
392        let proc = batch.processing_status.as_deref().unwrap_or("in_progress");
393        Ok(match proc {
394            "in_progress" => BatchStatus::InProgress,
395            "ended" => {
396                // Determine the terminal state from request_counts.
397                if let Some(counts) = batch.request_counts {
398                    if counts.errored > 0 && counts.succeeded == 0 {
399                        BatchStatus::Failed
400                    } else if counts.expired > 0 && counts.succeeded == 0 {
401                        BatchStatus::Expired
402                    } else {
403                        BatchStatus::Completed
404                    }
405                } else {
406                    BatchStatus::Completed
407                }
408            }
409            other => {
410                return Err(BatchError::Api(format!(
411                    "unknown processing_status: {}",
412                    other
413                )))
414            }
415        })
416    }
417
418    // -- Results ------------------------------------------------------------
419
420    /// Retrieve results of a completed batch.
421    pub async fn results(&self, id: &BatchId) -> Result<Vec<BatchResult>, BatchError> {
422        match self.provider {
423            BatchProvider::OpenAI => self.results_openai(id).await,
424            BatchProvider::Anthropic => self.results_anthropic(id).await,
425        }
426    }
427
428    async fn results_openai(&self, id: &BatchId) -> Result<Vec<BatchResult>, BatchError> {
429        // First, get the batch metadata to find the output file ID.
430        let headers = self.auth_headers()?;
431        let resp = self
432            .http
433            .get(format!("{}/batches/{}", self.base_url, id.0))
434            .headers(headers.clone())
435            .send()
436            .await?;
437
438        if resp.status() == reqwest::StatusCode::NOT_FOUND {
439            return Err(BatchError::NotFound(id.0.clone()));
440        }
441
442        let status = resp.status();
443        let text = resp.text().await?;
444
445        if !status.is_success() {
446            return Err(BatchError::Api(format!(
447                "fetch batch metadata failed ({}): {}",
448                status, text
449            )));
450        }
451
452        let batch: OpenAIBatchResponse = serde_json::from_str(&text)?;
453
454        // If the batch has an error file but no output file, report failure.
455        let output_file_id = match batch.output_file_id {
456            Some(fid) => fid,
457            None => {
458                if let Some(err_fid) = batch.error_file_id {
459                    // Download the error file for details.
460                    return self.download_openai_error_file(&err_fid, &headers).await;
461                }
462                return Err(BatchError::Failed(
463                    "batch has no output file and no error file".to_string(),
464                ));
465            }
466        };
467
468        // Download the output JSONL file.
469        let file_resp = self
470            .http
471            .get(format!(
472                "{}/files/{}/content",
473                self.base_url, output_file_id
474            ))
475            .headers(headers)
476            .send()
477            .await?;
478
479        let file_status = file_resp.status();
480        let file_text = file_resp.text().await?;
481
482        if !file_status.is_success() {
483            return Err(BatchError::Api(format!(
484                "download output file failed ({}): {}",
485                file_status, file_text
486            )));
487        }
488
489        self.parse_openai_results_jsonl(&file_text)
490    }
491
492    pub(crate) fn parse_openai_results_jsonl(
493        &self,
494        text: &str,
495    ) -> Result<Vec<BatchResult>, BatchError> {
496        let mut results = Vec::new();
497        for line in text.lines() {
498            if line.trim().is_empty() {
499                continue;
500            }
501            let parsed: OpenAIResultLine = serde_json::from_str(line)?;
502            let result = if let Some(err) = parsed.error {
503                Err(err.message.unwrap_or_else(|| "unknown error".to_string()))
504            } else if let Some(resp_body) = parsed.response {
505                if let Some(inner) = resp_body.body {
506                    let content = inner
507                        .choices
508                        .first()
509                        .and_then(|c| c.message.as_ref())
510                        .and_then(|m| m.content.clone())
511                        .unwrap_or_default();
512
513                    let token_usage = inner.usage.map(|u| crate::language_models::TokenUsage {
514                        prompt_tokens: u.prompt_tokens,
515                        completion_tokens: u.completion_tokens,
516                        total_tokens: u.total_tokens,
517                    });
518
519                    Ok(LLMResult {
520                        content,
521                        model: inner.model.unwrap_or_default(),
522                        token_usage,
523                        tool_calls: None,
524                        thinking_content: None,
525                    })
526                } else {
527                    Err("empty response body".to_string())
528                }
529            } else {
530                Err("no response and no error".to_string())
531            };
532            results.push(BatchResult {
533                custom_id: parsed.custom_id,
534                result,
535            });
536        }
537        Ok(results)
538    }
539
540    async fn download_openai_error_file(
541        &self,
542        error_file_id: &str,
543        headers: &reqwest::header::HeaderMap,
544    ) -> Result<Vec<BatchResult>, BatchError> {
545        let resp = self
546            .http
547            .get(format!("{}/files/{}/content", self.base_url, error_file_id))
548            .headers(headers.clone())
549            .send()
550            .await?;
551
552        let status = resp.status();
553        let text = resp.text().await?;
554
555        if !status.is_success() {
556            return Err(BatchError::Api(format!(
557                "download error file failed ({}): {}",
558                status, text
559            )));
560        }
561
562        // Parse error JSONL the same way as output.
563        self.parse_openai_results_jsonl(&text)
564    }
565
566    async fn results_anthropic(&self, id: &BatchId) -> Result<Vec<BatchResult>, BatchError> {
567        let headers = self.auth_headers()?;
568        let resp = self
569            .http
570            .get(format!(
571                "{}/messages/batches/{}/results",
572                self.base_url, id.0
573            ))
574            .headers(headers)
575            .send()
576            .await?;
577
578        if resp.status() == reqwest::StatusCode::NOT_FOUND {
579            return Err(BatchError::NotFound(id.0.clone()));
580        }
581
582        let status = resp.status();
583        let text = resp.text().await?;
584
585        if !status.is_success() {
586            return Err(BatchError::Api(format!(
587                "fetch results failed ({}): {}",
588                status, text
589            )));
590        }
591
592        self.parse_anthropic_results_jsonl(&text)
593    }
594
595    pub(crate) fn parse_anthropic_results_jsonl(
596        &self,
597        text: &str,
598    ) -> Result<Vec<BatchResult>, BatchError> {
599        let mut results = Vec::new();
600        for line in text.lines() {
601            if line.trim().is_empty() {
602                continue;
603            }
604            let parsed: AnthropicResultLine = serde_json::from_str(line)?;
605            let result = match parsed.result.result_type.as_str() {
606                "succeeded" => {
607                    if let Some(msg) = parsed.result.message {
608                        let content = msg
609                            .content
610                            .iter()
611                            .filter_map(|b| {
612                                if b.block_type == "text" {
613                                    b.text.clone()
614                                } else {
615                                    None
616                                }
617                            })
618                            .collect::<Vec<_>>()
619                            .join("");
620
621                        let token_usage = msg.usage.map(|u| crate::language_models::TokenUsage {
622                            prompt_tokens: u.input_tokens,
623                            completion_tokens: u.output_tokens,
624                            total_tokens: u.input_tokens + u.output_tokens,
625                        });
626
627                        Ok(LLMResult {
628                            content,
629                            model: msg.model,
630                            token_usage,
631                            tool_calls: None,
632                            thinking_content: None,
633                        })
634                    } else {
635                        Err("succeeded result missing message body".to_string())
636                    }
637                }
638                "errored" => {
639                    let err_msg = parsed
640                        .result
641                        .error
642                        .and_then(|e| e.message)
643                        .unwrap_or_else(|| "unknown error".to_string());
644                    Err(err_msg)
645                }
646                "expired" => Err("request expired".to_string()),
647                "canceled" => Err("request canceled".to_string()),
648                other => Err(format!("unknown result type: {}", other)),
649            };
650            results.push(BatchResult {
651                custom_id: parsed.custom_id,
652                result,
653            });
654        }
655        Ok(results)
656    }
657
658    // -- Cancel -------------------------------------------------------------
659
660    /// Cancel a batch job.
661    pub async fn cancel(&self, id: &BatchId) -> Result<(), BatchError> {
662        match self.provider {
663            BatchProvider::OpenAI => self.cancel_openai(id).await,
664            BatchProvider::Anthropic => Err(BatchError::Api(
665                "Anthropic batch API does not support cancellation".to_string(),
666            )),
667        }
668    }
669
670    async fn cancel_openai(&self, id: &BatchId) -> Result<(), BatchError> {
671        let headers = self.auth_headers()?;
672        let resp = self
673            .http
674            .post(format!("{}/batches/{}/cancel", self.base_url, id.0))
675            .headers(headers)
676            .send()
677            .await?;
678
679        if resp.status() == reqwest::StatusCode::NOT_FOUND {
680            return Err(BatchError::NotFound(id.0.clone()));
681        }
682
683        let status = resp.status();
684        if !status.is_success() {
685            let text = resp.text().await.unwrap_or_default();
686            return Err(BatchError::Api(format!(
687                "cancel failed ({}): {}",
688                status, text
689            )));
690        }
691
692        Ok(())
693    }
694
695    // -- Convenience --------------------------------------------------------
696
697    /// Submit and wait until complete, then return results.
698    ///
699    /// Polls every `poll_interval_ms` milliseconds. Returns
700    /// [`BatchError::Timeout`] if `max_wait_ms` is exceeded.
701    pub async fn submit_and_wait(
702        &self,
703        requests: Vec<BatchRequest>,
704        poll_interval_ms: u64,
705        max_wait_ms: u64,
706    ) -> Result<Vec<BatchResult>, BatchError> {
707        let batch_id = self.submit(requests).await?;
708        let start = std::time::Instant::now();
709        let poll_duration = std::time::Duration::from_millis(poll_interval_ms);
710        let max_duration = std::time::Duration::from_millis(max_wait_ms);
711
712        loop {
713            let status = self.poll(&batch_id).await?;
714
715            match status {
716                BatchStatus::Completed => {
717                    return self.results(&batch_id).await;
718                }
719                BatchStatus::Failed => {
720                    return Err(BatchError::Failed(format!("batch {} failed", batch_id.0)));
721                }
722                BatchStatus::Expired => {
723                    return Err(BatchError::Expired);
724                }
725                BatchStatus::Cancelled => {
726                    return Err(BatchError::Api(format!(
727                        "batch {} was cancelled",
728                        batch_id.0
729                    )));
730                }
731                BatchStatus::InProgress => {
732                    // Continue polling.
733                }
734            }
735
736            if start.elapsed() >= max_duration {
737                return Err(BatchError::Timeout(max_wait_ms));
738            }
739
740            tokio::time::sleep(poll_duration).await;
741        }
742    }
743}
lc_core/batch/client.rs

lc_core/batch/
client.rs
