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::{
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(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    #[get = "pub"]
69    base64s: Vec<MatchedFiles>,
70}
71impl<'a> MarkdownToParts<'a> {
72    pub fn builder() -> MarkdownToPartsBuilder {
73        MarkdownToPartsBuilder {
74            regex: None,
75            guess_mime_type: None,
76            #[cfg(feature = "reqwest")]
77            decide_download: None,
78            timeout: None,
79        }
80    }
81    ///# Panics
82    ///`regex` must have a Regex with only 1 capture group with file URL as first capture
83    ///group, else it PANICS.
84    /// # Arguments
85    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
86    /// with no "Content-Type" header.
87    /// `decide_download` is used to decide if to download. If it returns false, resource will not
88    /// be fetched and won't be in `parts`
89    /// # Example
90    /// ```ignore
91    /// from_regex("Your markdown string...", Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap(), |_| mime::IMAGE_PNG, |_| true)
92    /// ```
93    #[cfg(feature = "reqwest")]
94    pub async fn from_regex_checked(
95        markdown: &'a str,
96        regex: Regex,
97        guess_mime_type: fn(url: &str) -> mime::Mime,
98        decide_download: fn(headers: &HeaderMap) -> bool,
99    ) -> Self {
100        Self {
101            base64s: utils::get_file_base64s(
102                markdown,
103                regex,
104                guess_mime_type,
105                decide_download,
106                REQ_TIMEOUT,
107            )
108            .await,
109            markdown,
110        }
111    }
112    ///# Panics
113    ///`regex` must have a Regex with only 1 capture group with file URL as first capture
114    ///group, else it PANICS.
115    /// # Arguments
116    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
117    /// with no "Content-Type" header.
118    /// # Example
119    /// ```ignore
120    /// from_regex("Your markdown string...", Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap(), |_|
121    /// mime::IMAGE_PNG)
122    /// ```
123    #[cfg(feature = "reqwest")]
124    pub async fn from_regex(
125        markdown: &'a str,
126        regex: Regex,
127        guess_mime_type: fn(url: &str) -> mime::Mime,
128    ) -> Self {
129        Self::from_regex_checked(markdown, regex, guess_mime_type, |_| true).await
130    }
131    /// # Arguments
132    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
133    /// with no "Content-Type" header.
134    /// `decide_download` is used to decide if to download. If it returns false, resource will not
135    /// be fetched and won't be in `parts`
136    /// # Example
137    /// ```ignore
138    /// new("Your markdown string...", |_| mime::IMAGE_PNG, |_| true)
139    /// ```
140    #[cfg(feature = "reqwest")]
141    pub async fn new_checked(
142        markdown: &'a str,
143        guess_mime_type: fn(url: &str) -> mime::Mime,
144        decide_download: fn(headers: &HeaderMap) -> bool,
145    ) -> Self {
146        let image_regex = Regex::new(r"(?s)!\[.*?].?\((.*?)\)").unwrap();
147        Self::from_regex_checked(markdown, image_regex, guess_mime_type, decide_download).await
148    }
149    /// # Arguments
150    /// `guess_mime_type` is used to detect mimi_type of URL pointing to file system or web resource
151    /// with no "Content-Type" header.
152    /// # Example
153    /// ```ignore
154    /// new("Your markdown string...", |_| mime::IMAGE_PNG)
155    /// ```
156    #[cfg(feature = "reqwest")]
157    pub async fn new(markdown: &'a str, guess_mime_type: fn(url: &str) -> mime::Mime) -> Self {
158        Self::new_checked(markdown, guess_mime_type, |_| true).await
159    }
160    pub fn process(mut self) -> Vec<Part> {
161        let mut parts: Vec<Part> = Vec::new();
162        let mut removed_length = 0;
163        for file in self.base64s {
164            if let MatchedFiles {
165                index,
166                length,
167                mime_type: Some(mime_type),
168                base64: Some(base64),
169            } = file
170            {
171                let end = index + length - removed_length;
172                let text = &self.markdown[..end];
173                parts.push(text.into());
174                parts.push(InlineData::new(mime_type, base64).into());
175
176                self.markdown = &self.markdown[end..];
177                removed_length += end;
178            }
179        }
180        if self.markdown.len() != 0 {
181            parts.push(self.markdown.into());
182        }
183        parts
184    }
185}