Skip to main content

gemini_client_api/gemini/
utils.rs

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