batch_mode_batch_client/
check_and_download.rs

1// ---------------- [ File: src/check_and_download.rs ]
2crate::ix!();
3
4pub async fn check_for_and_download_output_and_error_online(
5    triple: &mut BatchFileTriple,
6    client: &OpenAIClientHandle,
7
8) -> Result<(), BatchDownloadError> {
9
10    loop {
11        match check_batch_status_online(triple,client).await {
12            Ok(status) => {
13                info!("batch online status: {:#?}", status);
14
15                // Download files if available
16                if status.output_file_available() {
17                    download_output_file(triple,client).await?;
18                }
19                if status.error_file_available() {
20                    download_error_file(triple,client).await?;
21                }
22                return Ok(());
23            }
24            Err(BatchDownloadError::BatchStillProcessing { batch_id }) => {
25                // Batch is still processing; decide whether to wait or exit
26                info!("Batch {} is still processing.", batch_id);
27                client.wait_for_batch_completion(&batch_id).await?;
28            }
29            Err(e) => {
30                // Other errors
31                return Err(e);
32            }
33        }
34    }
35}
36
37pub async fn check_batch_status_online(
38    triple: &BatchFileTriple,
39    client: &OpenAIClientHandle,
40
41) -> Result<BatchOnlineStatus, BatchDownloadError> {
42
43    info!("checking batch status online");
44
45    // Load batch metadata to get the batch ID
46    let metadata_filename = triple.metadata_filename_which_maybe_does_not_yet_exist();
47    let mut metadata      = BatchMetadata::load_from_file(&metadata_filename).await?;
48    let batch_id          = metadata.batch_id().to_string();
49
50    // Retrieve batch status from the API
51    let batch = client.retrieve_batch(&batch_id).await?;
52
53    match batch.status {
54        BatchStatus::Completed => {
55
56            // Update metadata with file IDs
57            metadata.set_output_file_id(batch.output_file_id.clone());
58            metadata.set_error_file_id(batch.error_file_id.clone());
59            metadata.save_to_file(&metadata_filename).await?;
60
61            Ok(BatchOnlineStatus::from(&batch))
62        }
63        BatchStatus::Failed => {
64            Err(BatchDownloadError::BatchFailed {
65                batch_id,
66            })
67        }
68        BatchStatus::Validating | BatchStatus::InProgress | BatchStatus::Finalizing => {
69            // Batch is still processing
70            Err(BatchDownloadError::BatchStillProcessing {
71                batch_id,
72            })
73        }
74        _ => {
75            // Handle other statuses if necessary
76            Err(BatchDownloadError::UnknownBatchStatus {
77                batch_id,
78                status: batch.status.clone(),
79            })
80        }
81    }
82}
83
84pub async fn download_output_file(
85    triple: &mut BatchFileTriple,
86    client: &OpenAIClientHandle,
87
88) -> Result<(), BatchDownloadError> {
89
90    info!("downloading batch output file");
91
92    if triple.output().is_some() {
93        return Err(BatchDownloadError::OutputFileAlreadyExists { triple: triple.clone() });
94    }
95
96    let metadata_filename = triple.metadata_filename_which_maybe_does_not_yet_exist();
97    let metadata          = BatchMetadata::load_from_file(&metadata_filename).await?;
98    let output_file_id    = metadata.output_file_id()?;
99
100    // Download the output file content
101    let file_content = client.file_content(output_file_id).await?;
102
103    // Write the content to the output file
104    let output_path = triple.output_filename_which_maybe_does_not_yet_exist().clone();
105
106    assert!(!output_path.exists());
107
108    std::fs::write(&output_path, file_content)?;
109
110    // Update the triple with the output file path
111    triple.set_output_path(Some(output_path));
112
113    Ok(())
114}
115
116pub async fn download_error_file(
117    triple: &mut BatchFileTriple,
118    client: &OpenAIClientHandle,
119
120) -> Result<(), BatchDownloadError> {
121
122    info!("downloading batch error file");
123
124    if triple.error().is_some() {
125        return Err(BatchDownloadError::ErrorFileAlreadyExists { triple: triple.clone() });
126    }
127
128    let metadata_filename = triple.metadata_filename_which_maybe_does_not_yet_exist();
129    let metadata          = BatchMetadata::load_from_file(&metadata_filename).await?;
130    let error_file_id     = metadata.error_file_id()?;
131
132    // Download the error file content
133    let file_content = client.file_content(error_file_id).await?;
134
135    // Write the content to the error file
136    let error_path = triple.error_filename_which_maybe_does_not_yet_exist();
137
138    assert!(!error_path.exists());
139
140    std::fs::write(&error_path, file_content)?;
141
142    // Update the triple with the error file path
143    triple.set_error_path(Some(error_path));
144
145    Ok(())
146}