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: Option<String>,
14    pub base64: Option<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/// `guess_mime_type` 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<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 capture = file.get(0).unwrap();
31        let url = file[1].to_string();
32        tasks.push((async |capture: regex::Match<'_>, url: String| {
33            let (mime_type, base64) = if url.starts_with("https://") || url.starts_with("http://") {
34                let response = client.get(&url).send().await;
35                match response {
36                    Ok(mut response) => {
37                        let base64 = response
38                            .body()
39                            .await
40                            .ok()
41                            .map(|bytes| STANDARD.encode(bytes));
42                        let mime_type = match base64 {
43                            Some(_) => response
44                                .mime_type()
45                                .ok()
46                                .flatten()
47                                .map(|mime| mime.to_string())
48                                .or_else(|| Some(guess_mime_type(&url))),
49                            None => None,
50                        };
51                        (mime_type, base64)
52                    }
53                    Err(_) => (None, None),
54                }
55            } else {
56                let base64 = 
57                    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}