batch/
batch.rs

1use openai_api_rs::v1::api::OpenAIClient;
2use openai_api_rs::v1::batch::CreateBatchRequest;
3use openai_api_rs::v1::file::FileUploadRequest;
4use serde_json::{from_str, to_string_pretty, Value};
5use std::env;
6use std::fs::File;
7use std::io::Write;
8use std::str;
9
10#[tokio::main]
11async fn main() -> Result<(), Box<dyn std::error::Error>> {
12    let api_key = env::var("OPENAI_API_KEY").unwrap().to_string();
13    let mut client = OpenAIClient::builder().with_api_key(api_key).build()?;
14
15    let req = FileUploadRequest::new(
16        "examples/data/batch_request.json".to_string(),
17        "batch".to_string(),
18    );
19
20    let result = client.upload_file(req).await?;
21    println!("File id: {:?}", result.id);
22
23    let input_file_id = result.id;
24    let req = CreateBatchRequest::new(
25        input_file_id.clone(),
26        "/v1/chat/completions".to_string(),
27        "24h".to_string(),
28    );
29
30    let result = client.create_batch(req).await?;
31    println!("Batch id: {:?}", result.id);
32
33    let batch_id = result.id;
34    let result = client.retrieve_batch(batch_id.to_string()).await?;
35    println!("Batch status: {:?}", result.status);
36
37    // sleep 30 seconds
38    println!("Sleeping for 30 seconds...");
39    tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;
40
41    let result = client.retrieve_batch(batch_id.to_string()).await?;
42
43    let file_id = result.output_file_id.unwrap();
44    let result = client.retrieve_file_content(file_id).await?;
45    let s = match str::from_utf8(&result) {
46        Ok(v) => v.to_string(),
47        Err(e) => panic!("Invalid UTF-8 sequence: {}", e),
48    };
49    let json_value: Value = from_str(&s)?;
50    let result_json = to_string_pretty(&json_value)?;
51
52    let output_file_path = "examples/data/batch_result.json";
53    let mut file = File::create(output_file_path)?;
54    file.write_all(result_json.as_bytes())?;
55
56    println!("File writed to {:?}", output_file_path);
57
58    Ok(())
59}
60
61// OPENAI_API_KEY=xxxx cargo run --package openai-api-rs --example batch
batch/batch.rs

batch/
batch.rs
