Skip to main content

gemini_client_api/gemini/
utils.rs

1use super::types::request::*;
2use crate::utils::{self, MatchedFiles};
3use getset::Getters;
4use regex::Regex;
5use reqwest::header::HeaderMap;
6use std::time::Duration;
7mod macros;
8pub use gemini_proc_macros::{execute_function_calls, gemini_function, gemini_schema};
9pub use macros::GeminiSchema;
10
11const REQ_TIMEOUT: Duration = Duration::from_secs(10);
12
13pub struct MarkdownToPartsBuilder {
14    regex: Option<Regex>,
15    guess_mime_type: Option<fn(url: &str) -> mime::Mime>,
16    decide_download: Option<fn(headers: &HeaderMap) -> bool>,
17    timeout: Option<Duration>,
18}
19impl MarkdownToPartsBuilder {
20    ///# Panics
21    ///`regex` must have a Regex with only 1 capture group with file URL as first capture
22    ///group, else it PANICS when `.build()` is called.
23    pub fn regex(mut self, regex: Regex) -> Self {
24        self.regex = Some(regex);
25        self
26    }
27    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
28    /// with no "Content-Type" header.
29    pub fn guess_mime_type(mut self, guess_mime_type: fn(url: &str) -> mime::Mime) -> Self {
30        self.guess_mime_type = Some(guess_mime_type);
31        self
32    }
33    /// `decide_download` is used to decide if to download. If it returns false, resource will not
34    /// be fetched and won't be in `parts`
35    pub fn decide_download(mut self, decide_download: fn(headers: &HeaderMap) -> bool) -> Self {
36        self.decide_download = Some(decide_download);
37        self
38    }
39    pub fn timeout(mut self, timeout: Duration) -> Self {
40        self.timeout = Some(timeout);
41        self
42    }
43    pub async fn build<'a>(self, markdown: &'a str) -> MarkdownToParts<'a> {
44        MarkdownToParts {
45            markdown,
46            base64s: utils::get_file_base64s(
47                markdown,
48                self.regex
49                    .unwrap_or(Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap()),
50                self.guess_mime_type.unwrap_or(|_| mime::IMAGE_PNG),
51                |_| true,
52                self.timeout.unwrap_or(REQ_TIMEOUT),
53            )
54            .await,
55        }
56    }
57}
58#[derive(Getters, Clone)]
59///Converts markdown to parts considering `![image](link)` means Gemini will be see the images too. `link` can be URL or file path.  
60pub struct MarkdownToParts<'a> {
61    markdown: &'a str,
62    #[get = "pub"]
63    base64s: Vec<MatchedFiles>,
64}
65impl<'a> MarkdownToParts<'a> {
66    pub fn builder() -> MarkdownToPartsBuilder {
67        MarkdownToPartsBuilder {
68            regex: None,
69            guess_mime_type: None,
70            decide_download: None,
71            timeout: None,
72        }
73    }
74    ///# Panics
75    ///`regex` must have a Regex with only 1 capture group with file URL as first capture
76    ///group, else it PANICS.
77    /// # Arguments
78    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
79    /// with no "Content-Type" header.
80    /// `decide_download` is used to decide if to download. If it returns false, resource will not
81    /// be fetched and won't be in `parts`
82    /// # Example
83    /// ```ignore
84    /// from_regex("Your markdown string...", Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap(), |_| mime::IMAGE_PNG, |_| true)
85    /// ```
86    pub async fn from_regex_checked(
87        markdown: &'a str,
88        regex: Regex,
89        guess_mime_type: fn(url: &str) -> mime::Mime,
90        decide_download: fn(headers: &HeaderMap) -> bool,
91    ) -> Self {
92        Self {
93            base64s: utils::get_file_base64s(
94                markdown,
95                regex,
96                guess_mime_type,
97                decide_download,
98                REQ_TIMEOUT,
99            )
100            .await,
101            markdown,
102        }
103    }
104    ///# Panics
105    /// `regex` must have a Regex with only 1 capture group with file URL as first capture group, else it PANICS.
106    /// # Arguments
107    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
108    /// with no "Content-Type" header.
109    /// # Example
110    /// ```ignore
111    /// from_regex("Your markdown string...", Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap(), |_|
112    /// mime::IMAGE_PNG)
113    /// ```
114    pub async fn from_regex(
115        markdown: &'a str,
116        regex: Regex,
117        guess_mime_type: fn(url: &str) -> mime::Mime,
118    ) -> Self {
119        Self::from_regex_checked(markdown, regex, guess_mime_type, |_| true).await
120    }
121    /// # Arguments
122    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
123    /// with no "Content-Type" header.
124    /// `decide_download` is used to decide if to download. If it returns false, resource will not
125    /// be fetched and won't be in `parts`
126    /// # Example
127    /// ```ignore
128    /// new("Your markdown string...", |_| mime::IMAGE_PNG, |_| true)
129    /// ```
130    pub async fn new_checked(
131        markdown: &'a str,
132        guess_mime_type: fn(url: &str) -> mime::Mime,
133        decide_download: fn(headers: &HeaderMap) -> bool,
134    ) -> Self {
135        let image_regex = Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap();
136        Self::from_regex_checked(markdown, image_regex, guess_mime_type, decide_download).await
137    }
138    /// # Arguments
139    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
140    /// with no "Content-Type" header.
141    /// # Example
142    /// ```ignore
143    /// new("Your markdown string...", |_| mime::IMAGE_PNG)
144    /// ```
145    pub async fn new(markdown: &'a str, guess_mime_type: fn(url: &str) -> mime::Mime) -> Self {
146        Self::new_checked(markdown, guess_mime_type, |_| true).await
147    }
148    pub fn process(mut self) -> Vec<Part> {
149        let mut parts: Vec<Part> = Vec::new();
150        let mut removed_length = 0;
151        for file in self.base64s {
152            if let MatchedFiles {
153                index,
154                length,
155                mime_type: Some(mime_type),
156                base64: Some(base64),
157            } = file
158            {
159                let end = index + length - removed_length;
160                let text = &self.markdown[..end];
161                parts.push(Part::text(text.into()));
162                parts.push(Part::inline_data(InlineData::new(mime_type, base64)));
163
164                self.markdown = &self.markdown[end..];
165                removed_length += end;
166            }
167        }
168        if self.markdown.len() != 0 {
169            parts.push(Part::text(self.markdown.into()));
170        }
171        parts
172    }
173}