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: 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 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                match response {
36                    Ok(mut response) => (
37                        response
38                            .mime_type()
39                            .ok()
40                            .flatten()
41                            .map(|mime| mime.to_string())
42                            .unwrap_or_else(|| guess_mime_type(&url)),
43                        response
44                            .body()
45                            .await
46                            .ok()
47                            .map(|bytes| STANDARD.encode(bytes)),
48                    ),
49                    Err(_) => (guess_mime_type(&url), None),
50                }
51            } else {
52                (
53                    guess_mime_type(&url),
54                    tokio::fs::read(url)
55                        .await
56                        .ok()
57                        .map(|bytes| STANDARD.encode(&bytes)),
58                )
59            };
60            MatchedFiles {
61                index: capture.start(),
62                length: capture.len(),
63                mime_type,
64                base64,
65            }
66        })(capture));
67    }
68    join_all(tasks).await
69}