async_openai/
download.rs

1use std::path::{Path, PathBuf};
2
3use base64::{engine::general_purpose, Engine as _};
4use rand::{distributions::Alphanumeric, Rng};
5use reqwest::Url;
6
7use crate::error::OpenAIError;
8
9fn create_paths<P: AsRef<Path>>(url: &Url, base_dir: P) -> (PathBuf, PathBuf) {
10    let mut dir = PathBuf::from(base_dir.as_ref());
11    let mut path = dir.clone();
12    let segments = url.path_segments().map(|c| c.collect::<Vec<_>>());
13    if let Some(segments) = segments {
14        for (idx, segment) in segments.iter().enumerate() {
15            if idx != segments.len() - 1 {
16                dir.push(segment);
17            }
18            path.push(segment);
19        }
20    }
21
22    (dir, path)
23}
24
25pub(crate) async fn download_url<P: AsRef<Path>>(
26    url: &str,
27    dir: P,
28) -> Result<PathBuf, OpenAIError> {
29    let parsed_url = Url::parse(url).map_err(|e| OpenAIError::FileSaveError(e.to_string()))?;
30    let response = reqwest::get(url)
31        .await
32        .map_err(|e| OpenAIError::FileSaveError(e.to_string()))?;
33
34    if !response.status().is_success() {
35        return Err(OpenAIError::FileSaveError(format!(
36            "couldn't download file, status: {}, url: {url}",
37            response.status()
38        )));
39    }
40
41    let (dir, file_path) = create_paths(&parsed_url, dir);
42
43    tokio::fs::create_dir_all(dir.as_path())
44        .await
45        .map_err(|e| OpenAIError::FileSaveError(format!("{}, dir: {}", e, dir.display())))?;
46
47    tokio::fs::write(
48        file_path.as_path(),
49        response.bytes().await.map_err(|e| {
50            OpenAIError::FileSaveError(format!("{}, file path: {}", e, file_path.display()))
51        })?,
52    )
53    .await
54    .map_err(|e| OpenAIError::FileSaveError(e.to_string()))?;
55
56    Ok(file_path)
57}
58
59pub(crate) async fn save_b64<P: AsRef<Path>>(b64: &str, dir: P) -> Result<PathBuf, OpenAIError> {
60    let filename: String = rand::thread_rng()
61        .sample_iter(&Alphanumeric)
62        .take(10)
63        .map(char::from)
64        .collect();
65
66    let filename = format!("{filename}.png");
67
68    let path = PathBuf::from(dir.as_ref()).join(filename);
69
70    tokio::fs::write(
71        path.as_path(),
72        general_purpose::STANDARD
73            .decode(b64)
74            .map_err(|e| OpenAIError::FileSaveError(e.to_string()))?,
75    )
76    .await
77    .map_err(|e| OpenAIError::FileSaveError(format!("{}, path: {}", e, path.display())))?;
78
79    Ok(path)
80}