gemini_client_api/
utils.rs

1use base64::{Engine, engine::general_purpose::STANDARD};
2use futures::future::join_all;
3pub use mime;
4use regex::Regex;
5use reqwest::Client;
6pub use reqwest::header::{HeaderMap, HeaderValue};
7use std::time::Duration;
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) -> mime::Mime,
24    decide_download: fn(headers: &HeaderMap) -> bool,
25    timeout: Duration
26) -> Vec<MatchedFiles> {
27    let client = Client::builder().timeout(timeout).build().unwrap();
28    let mut tasks = Vec::new();
29
30    for file in regex.captures_iter(markdown.as_ref()) {
31        let capture = file.get(0).unwrap();
32        let url = file[1].to_string();
33        tasks.push((async |capture: regex::Match<'_>, url: String| {
34            let (mime_type, base64) = if url.starts_with("https://") || url.starts_with("http://") {
35                let response = client.get(&url).send().await;
36                match response {
37                    Ok(response) if (decide_download)(response.headers()) => {
38                        let mime_type = response
39                            .headers()
40                            .get("Content-Type")
41                            .map(|mime| mime.to_str().ok())
42                            .flatten()
43                            .map(|str| str.to_string());
44
45                        let base64 = response
46                            .bytes()
47                            .await
48                            .ok()
49                            .map(|bytes| STANDARD.encode(bytes));
50                        let mime_type = match base64 {
51                            Some(_) => {
52                                mime_type.or_else(|| Some(guess_mime_type(&url).to_string()))
53                            }
54                            None => None,
55                        };
56                        (mime_type, base64)
57                    }
58                    _ => (None, None),
59                }
60            } else {
61                let base64 = tokio::fs::read(url.clone())
62                    .await
63                    .ok()
64                    .map(|bytes| STANDARD.encode(&bytes));
65                match base64 {
66                    Some(base64) => (Some(guess_mime_type(&url).to_string()), Some(base64)),
67                    None => (None, None),
68                }
69            };
70            MatchedFiles {
71                index: capture.start(),
72                length: capture.len(),
73                mime_type,
74                base64,
75            }
76        })(capture, url));
77    }
78    join_all(tasks).await
79}
gemini_client_api/utils.rs

gemini_client_api/
utils.rs
