gemini_client_api/
utils.rs

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