gemini_client_api/
utils.rs

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

gemini_client_api/
utils.rs
