use std::{fs::File, io::Write};
use serde_json::json;
use zai_rs::{batches::*, file::*};
fn make_jsonl_line(custom_id: &str, user_content: &str) -> String {
let v = json!({
"custom_id": custom_id,
"method": "POST",
"url": "/v4/chat/completions",
"body": {
"model": "glm-4",
"messages": [
{"role": "system", "content": "你是一个意图分类器."},
{"role": "user", "content": user_content}
]
}
});
v.to_string()
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let _ = env_logger::try_init();
let key = std::env::var("ZHIPU_API_KEY").expect("Please set ZHIPU_API_KEY env var");
std::fs::create_dir_all("data")?;
let path = "data/batch_demo.jsonl";
let mut f = File::create(path)?;
let lines = vec![
make_jsonl_line("request-1", "订单处理速度太慢,等了很久才发货。"),
make_jsonl_line(
"request-2",
"商品有点小瑕疵,不过客服处理得很快,总体满意。",
),
make_jsonl_line("request-3", "这款产品性价比很高,非常满意。"),
make_jsonl_line("request-4", "说明书写得不清楚,看了半天也不知道怎么用。"),
];
for line in lines {
writeln!(f, "{}", line)?;
}
let upload = FileUploadRequest::new(key.clone(), FilePurpose::Batch, path)
.with_content_type("application/jsonl");
let file: FileObject = upload.send().await?;
let file_id = file
.id
.ok_or_else(|| Box::<dyn std::error::Error>::from("missing file id"))?;
let create = CreateBatchRequest::new(key.clone(), file_id, BatchEndpoint::ChatCompletions)
.with_auto_delete_input_file(true);
let batch: CreateBatchResponse = create.send().await?;
println!(
"created batch: id={:?} status={:?} input_file_id={:?}",
batch.id, batch.status, batch.input_file_id
);
Ok(())
}