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