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(BatchError::Api(
504                    err.message.unwrap_or_else(|| "unknown error".to_string()),
505                ))
506            } else if let Some(resp_body) = parsed.response {
507                if let Some(inner) = resp_body.body {
508                    let content = inner
509                        .choices
510                        .first()
511                        .and_then(|c| c.message.as_ref())
512                        .and_then(|m| m.content.clone())
513                        .unwrap_or_default();
514
515                    let token_usage = inner.usage.map(|u| crate::language_models::TokenUsage {
516                        prompt_tokens: u.prompt_tokens,
517                        completion_tokens: u.completion_tokens,
518                        total_tokens: u.total_tokens,
519                    });
520
521                    Ok(LLMResult {
522                        content,
523                        model: inner.model.unwrap_or_default(),
524                        token_usage,
525                        tool_calls: None,
526                        thinking_content: None,
527                    })
528                } else {
529                    Err(BatchError::Api("empty response body".to_string()))
530                }
531            } else {
532                Err(BatchError::Api("no response and no error".to_string()))
533            };
534            results.push(BatchResult {
535                custom_id: parsed.custom_id,
536                result,
537            });
538        }
539        Ok(results)
540    }
541
542    async fn download_openai_error_file(
543        &self,
544        error_file_id: &str,
545        headers: &reqwest::header::HeaderMap,
546    ) -> Result<Vec<BatchResult>, BatchError> {
547        let resp = self
548            .http
549            .get(format!("{}/files/{}/content", self.base_url, error_file_id))
550            .headers(headers.clone())
551            .send()
552            .await?;
553
554        let status = resp.status();
555        let text = resp.text().await?;
556
557        if !status.is_success() {
558            return Err(BatchError::Api(format!(
559                "download error file failed ({}): {}",
560                status, text
561            )));
562        }
563
564        // Parse error JSONL the same way as output.
565        self.parse_openai_results_jsonl(&text)
566    }
567
568    async fn results_anthropic(&self, id: &BatchId) -> Result<Vec<BatchResult>, BatchError> {
569        let headers = self.auth_headers()?;
570        let resp = self
571            .http
572            .get(format!(
573                "{}/messages/batches/{}/results",
574                self.base_url, id.0
575            ))
576            .headers(headers)
577            .send()
578            .await?;
579
580        if resp.status() == reqwest::StatusCode::NOT_FOUND {
581            return Err(BatchError::NotFound(id.0.clone()));
582        }
583
584        let status = resp.status();
585        let text = resp.text().await?;
586
587        if !status.is_success() {
588            return Err(BatchError::Api(format!(
589                "fetch results failed ({}): {}",
590                status, text
591            )));
592        }
593
594        self.parse_anthropic_results_jsonl(&text)
595    }
596
597    pub(crate) fn parse_anthropic_results_jsonl(
598        &self,
599        text: &str,
600    ) -> Result<Vec<BatchResult>, BatchError> {
601        let mut results = Vec::new();
602        for line in text.lines() {
603            if line.trim().is_empty() {
604                continue;
605            }
606            let parsed: AnthropicResultLine = serde_json::from_str(line)?;
607            let result = match parsed.result.result_type.as_str() {
608                "succeeded" => {
609                    if let Some(msg) = parsed.result.message {
610                        let content = msg
611                            .content
612                            .iter()
613                            .filter_map(|b| {
614                                if b.block_type == "text" {
615                                    b.text.clone()
616                                } else {
617                                    None
618                                }
619                            })
620                            .collect::<Vec<_>>()
621                            .join("");
622
623                        let token_usage = msg.usage.map(|u| crate::language_models::TokenUsage {
624                            prompt_tokens: u.input_tokens,
625                            completion_tokens: u.output_tokens,
626                            total_tokens: u.input_tokens + u.output_tokens,
627                        });
628
629                        Ok(LLMResult {
630                            content,
631                            model: msg.model,
632                            token_usage,
633                            tool_calls: None,
634                            thinking_content: None,
635                        })
636                    } else {
637                        Err(BatchError::Api(
638                            "succeeded result missing message body".to_string(),
639                        ))
640                    }
641                }
642                "errored" => {
643                    let err_msg = parsed
644                        .result
645                        .error
646                        .and_then(|e| e.message)
647                        .unwrap_or_else(|| "unknown error".to_string());
648                    Err(BatchError::Api(err_msg))
649                }
650                "expired" => Err(BatchError::Api("request expired".to_string())),
651                "canceled" => Err(BatchError::Api("request canceled".to_string())),
652                other => Err(BatchError::Api(format!("unknown result type: {}", other))),
653            };
654            results.push(BatchResult {
655                custom_id: parsed.custom_id,
656                result,
657            });
658        }
659        Ok(results)
660    }
661
662    // -- Cancel -------------------------------------------------------------
663
664    /// Cancel a batch job.
665    pub async fn cancel(&self, id: &BatchId) -> Result<(), BatchError> {
666        match self.provider {
667            BatchProvider::OpenAI => self.cancel_openai(id).await,
668            BatchProvider::Anthropic => Err(BatchError::Api(
669                "Anthropic batch API does not support cancellation".to_string(),
670            )),
671        }
672    }
673
674    async fn cancel_openai(&self, id: &BatchId) -> Result<(), BatchError> {
675        let headers = self.auth_headers()?;
676        let resp = self
677            .http
678            .post(format!("{}/batches/{}/cancel", self.base_url, id.0))
679            .headers(headers)
680            .send()
681            .await?;
682
683        if resp.status() == reqwest::StatusCode::NOT_FOUND {
684            return Err(BatchError::NotFound(id.0.clone()));
685        }
686
687        let status = resp.status();
688        if !status.is_success() {
689            let text = resp.text().await.unwrap_or_default();
690            return Err(BatchError::Api(format!(
691                "cancel failed ({}): {}",
692                status, text
693            )));
694        }
695
696        Ok(())
697    }
698
699    // -- Convenience --------------------------------------------------------
700
701    /// Submit and wait until complete, then return results.
702    ///
703    /// Polls every `poll_interval_ms` milliseconds. Returns
704    /// [`BatchError::Timeout`] if `max_wait_ms` is exceeded.
705    pub async fn submit_and_wait(
706        &self,
707        requests: Vec<BatchRequest>,
708        poll_interval_ms: u64,
709        max_wait_ms: u64,
710    ) -> Result<Vec<BatchResult>, BatchError> {
711        let batch_id = self.submit(requests).await?;
712        let start = std::time::Instant::now();
713        let poll_duration = std::time::Duration::from_millis(poll_interval_ms);
714        let max_duration = std::time::Duration::from_millis(max_wait_ms);
715
716        loop {
717            let status = self.poll(&batch_id).await?;
718
719            match status {
720                BatchStatus::Completed => {
721                    return self.results(&batch_id).await;
722                }
723                BatchStatus::Failed => {
724                    return Err(BatchError::Failed(format!("batch {} failed", batch_id.0)));
725                }
726                BatchStatus::Expired => {
727                    return Err(BatchError::Expired);
728                }
729                BatchStatus::Cancelled => {
730                    return Err(BatchError::Api(format!(
731                        "batch {} was cancelled",
732                        batch_id.0
733                    )));
734                }
735                BatchStatus::InProgress => {
736                    // Continue polling.
737                }
738            }
739
740            if start.elapsed() >= max_duration {
741                return Err(BatchError::Timeout(max_wait_ms));
742            }
743
744            tokio::time::sleep(poll_duration).await;
745        }
746    }
747}
lc_core/batch/client.rs

lc_core/batch/
client.rs
