gemini_client_api/
utils.rs

1use base64::{Engine, engine::general_purpose::STANDARD};
2use futures::future::join_all;
3use regex::Regex;
4use reqwest::Client;
5use std::time::Duration;
6
7const REQ_TIMEOUT: Duration = Duration::from_secs(10);
8
9pub struct MatchedFiles {
10    pub index: usize,
11    pub length: usize,
12    pub mime_type: Option<String>,
13    pub base64: Option<String>,
14}
15/// # Panics
16/// `regex` must have a Regex with atleast 1 capture group with file URL as first capture group, else it PANICS
17/// # Arguments
18/// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
19/// with no "Content-Type" header.
20pub async fn get_file_base64s(
21    markdown: impl AsRef<str>,
22    regex: Regex,
23    guess_mime_type: fn(url: &str) -> String,
24) -> Vec<MatchedFiles> {
25    let client = Client::builder().timeout(REQ_TIMEOUT).build().unwrap();
26    let mut tasks: Vec<_> = Vec::new();
27
28    for file in regex.captures_iter(markdown.as_ref()) {
29        let capture = file.get(0).unwrap();
30        let url = file[1].to_string();
31        tasks.push((async |capture: regex::Match<'_>, url: String| {
32            let (mime_type, base64) = if url.starts_with("https://") || url.starts_with("http://") {
33                let response = client.get(&url).send().await;
34                match response {
35                    Ok(response) => {
36                        let mime_type = response
37                            .headers()
38                            .get("Content-Type")
39                            .map(|mime| mime.to_str().ok())
40                            .flatten()
41                            .map(|str| str.to_string());
42
43                        let base64 = response
44                            .bytes()
45                            .await
46                            .ok()
47                            .map(|bytes| STANDARD.encode(bytes));
48                        let mime_type = match base64 {
49                            Some(_) => mime_type.or_else(|| Some(guess_mime_type(&url))),
50                            None => None,
51                        };
52                        (mime_type, base64)
53                    }
54                    Err(_) => (None, None),
55                }
56            } else {
57                let base64 = tokio::fs::read(url.clone())
58                    .await
59                    .ok()
60                    .map(|bytes| STANDARD.encode(&bytes));
61                match base64 {
62                    Some(base64) => (Some(guess_mime_type(&url)), Some(base64)),
63                    None => (None, None),
64                }
65            };
66            MatchedFiles {
67                index: capture.start(),
68                length: capture.len(),
69                mime_type,
70                base64,
71            }
72        })(capture, url));
73    }
74    join_all(tasks).await
75}