gemini_ai/content_gen/
content.rs

1#[cfg(feature = "sync")]
2use native_tls::TlsConnector;
3#[cfg(feature = "sync")]
4use std::{
5    env,
6    io::{BufReader, Read, Write},
7    net::TcpStream,
8};
9
10#[cfg(feature = "async")]
11use async_std::io::{BufReader, Read, Write};
12#[cfg(feature = "async")]
13use async_std::net::TcpStream;
14#[cfg(feature = "async")]
15use async_tls::TlsConnector;
16
17use dotenv::dotenv;
18
19use crate::{
20    ConfigBuilder, ConfigNotPresent, Default, EnvVariableNotPresent, EnvVariablePresent, Gemini,
21    GeminiContentGen, InstructionNotPresent, Kind, MaxLenNotPresent, MemoryNot, MemoryOK,
22    MemoryType, Memorys, ModelNotPresent, ModelPresent, Models, PropertiesNotPresent,
23    TextNotPresent,
24};
25
26use super::{forgetful::forgetFul, memory::memory};
27
28impl<'gemini>
29    Gemini<
30        'gemini,
31        EnvVariableNotPresent,
32        ModelNotPresent,
33        ConfigNotPresent,
34        TextNotPresent,
35        MaxLenNotPresent,
36        InstructionNotPresent,
37        PropertiesNotPresent,
38        Default,
39    >
40{
41    pub fn new() -> Self {
42        Gemini {
43            model: "",
44            env_variable: "",
45            text: "",
46            instruction: "",
47            max_len: 8192,
48            config: ConfigBuilder {
49                // schema_type: String::new(),
50                r#type: Kind::Text,
51                propertiesstate: std::marker::PhantomData,
52            },
53            memory: MemoryType::NoMemory,
54            envstate: std::marker::PhantomData,
55            maxstate: std::marker::PhantomData,
56            instructionstate: std::marker::PhantomData,
57            modelstate: std::marker::PhantomData,
58            configstate: std::marker::PhantomData,
59            textstate: std::marker::PhantomData,
60            memorystate: std::marker::PhantomData,
61        }
62    }
63}
64
65impl<'gemini>
66    Gemini<
67        'gemini,
68        EnvVariableNotPresent,
69        ModelNotPresent,
70        ConfigNotPresent,
71        TextNotPresent,
72        MaxLenNotPresent,
73        InstructionNotPresent,
74        PropertiesNotPresent,
75        Default,
76    >
77{
78    pub fn env(
79        mut self,
80        env_variable: &'gemini str,
81    ) -> Gemini<
82        'gemini,
83        EnvVariablePresent,
84        ModelNotPresent,
85        ConfigNotPresent,
86        TextNotPresent,
87        MaxLenNotPresent,
88        InstructionNotPresent,
89        PropertiesNotPresent,
90        Default,
91    > {
92        self.env_variable = env_variable;
93        Gemini {
94            env_variable: self.env_variable,
95            model: &self.model,
96            text: self.text,
97            instruction: &self.instruction,
98            max_len: self.max_len,
99            memory: self.memory,
100            config: ConfigBuilder {
101                // schema_type: String::new(),
102                r#type: self.config.r#type,
103                propertiesstate: std::marker::PhantomData,
104            },
105            envstate: std::marker::PhantomData,
106            instructionstate: std::marker::PhantomData,
107            maxstate: std::marker::PhantomData,
108            modelstate: std::marker::PhantomData,
109            configstate: std::marker::PhantomData,
110            textstate: std::marker::PhantomData,
111            memorystate: std::marker::PhantomData,
112        }
113    }
114}
115
116impl<'gemini>
117    Gemini<
118        'gemini,
119        EnvVariablePresent,
120        ModelNotPresent,
121        ConfigNotPresent,
122        TextNotPresent,
123        MaxLenNotPresent,
124        InstructionNotPresent,
125        PropertiesNotPresent,
126        Default,
127    >
128{
129    pub fn model(
130        mut self,
131        model: Models<'gemini>,
132    ) -> Gemini<
133        'gemini,
134        EnvVariablePresent,
135        ModelPresent,
136        ConfigNotPresent,
137        TextNotPresent,
138        MaxLenNotPresent,
139        InstructionNotPresent,
140        PropertiesNotPresent,
141        Default,
142    > {
143        match model {
144            Models::GEMINI_1_0_PRO => self.model = "gemini-1.0-pro",
145            Models::GEMINI_1_5_FLASH => self.model = "gemini-1.5-flash",
146            Models::GEMINI_1_5_FLASH_002 => self.model = "gemini-1.5-flash-002",
147            Models::GEMINI_1_5_FLASH_8B => self.model = "gemini-1.5-flash-8b",
148            Models::GEMINI_1_5_PRO => self.model = "gemini-1.5-pro",
149            Models::GEMINI_1_5_PRO_002 => self.model = "gemini-1.5-pro-002",
150            Models::Custom(model) => self.model = model,
151        }
152        Gemini {
153            env_variable: self.env_variable,
154            model: &self.model,
155            text: self.text,
156            instruction: &self.instruction,
157            max_len: self.max_len,
158            memory: self.memory,
159            config: ConfigBuilder {
160                // schema_type: String::new(),
161                r#type: self.config.r#type,
162                propertiesstate: std::marker::PhantomData,
163            },
164            envstate: std::marker::PhantomData,
165            instructionstate: std::marker::PhantomData,
166            maxstate: std::marker::PhantomData,
167            modelstate: std::marker::PhantomData,
168            configstate: std::marker::PhantomData,
169            textstate: std::marker::PhantomData,
170            memorystate: std::marker::PhantomData,
171        }
172    }
173}
174impl<'gemini>
175    Gemini<
176        'gemini,
177        EnvVariablePresent,
178        ModelPresent,
179        ConfigNotPresent,
180        TextNotPresent,
181        MaxLenNotPresent,
182        InstructionNotPresent,
183        PropertiesNotPresent,
184        Default,
185    >
186{
187    pub fn no_memory(
188        mut self,
189    ) -> Gemini<
190        'gemini,
191        EnvVariablePresent,
192        ModelPresent,
193        ConfigNotPresent,
194        TextNotPresent,
195        MaxLenNotPresent,
196        InstructionNotPresent,
197        PropertiesNotPresent,
198        MemoryNot,
199    > {
200        self.memory = MemoryType::NoMemory;
201        Gemini {
202            env_variable: self.env_variable,
203            model: &self.model,
204            text: self.text,
205            instruction: &self.instruction,
206            max_len: self.max_len,
207            memory: self.memory,
208            config: ConfigBuilder {
209                // schema_type: String::new(),
210                r#type: self.config.r#type,
211                propertiesstate: std::marker::PhantomData,
212            },
213            envstate: std::marker::PhantomData,
214            instructionstate: std::marker::PhantomData,
215            maxstate: std::marker::PhantomData,
216            modelstate: std::marker::PhantomData,
217            configstate: std::marker::PhantomData,
218            textstate: std::marker::PhantomData,
219            memorystate: std::marker::PhantomData,
220        }
221    }
222    pub fn memory(
223        mut self,
224        memory: Memorys,
225    ) -> Gemini<
226        'gemini,
227        EnvVariablePresent,
228        ModelPresent,
229        ConfigNotPresent,
230        TextNotPresent,
231        MaxLenNotPresent,
232        InstructionNotPresent,
233        PropertiesNotPresent,
234        MemoryOK,
235    > {
236        self.memory = MemoryType::Memory(memory);
237        Gemini {
238            env_variable: self.env_variable,
239            model: &self.model,
240            text: self.text,
241            instruction: &self.instruction,
242            max_len: self.max_len,
243            memory: self.memory,
244            config: ConfigBuilder {
245                // schema_type: String::new(),
246                r#type: self.config.r#type,
247                propertiesstate: std::marker::PhantomData,
248            },
249            envstate: std::marker::PhantomData,
250            instructionstate: std::marker::PhantomData,
251            maxstate: std::marker::PhantomData,
252            modelstate: std::marker::PhantomData,
253            configstate: std::marker::PhantomData,
254            textstate: std::marker::PhantomData,
255            memorystate: std::marker::PhantomData,
256        }
257    }
258}
259
260impl<'output> GeminiContentGen<'output> {
261    #[cfg(feature = "sync")]
262    pub fn output(self) -> String {
263        match self.memory {
264            MemoryType::NoMemory => {
265                // println!("{}", "hi rag");
266                forgetFul(&self)
267            }
268            MemoryType::Memory(directory) => memory(directory, &self),
269        }
270        // println!("{:?}", self.config.properties);
271    }
272    #[cfg(feature = "async")]
273    pub async fn output(self) -> String {
274        match self.memory {
275            MemoryType::NoMemory => {
276                // println!("{}", "hi rag");
277                forgetFul(&self).await
278            }
279            MemoryType::Memory(directory) => memory(directory, &self).await,
280        }
281        // println!("{:?}", self.config.properties);
282    }
283}
284#[cfg(feature = "sync")]
285pub(crate) fn gemini(content: String, env: &str, model: &str, mime_type: &str) -> String {
286    dotenv().unwrap();
287
288    let env = env::var(env).expect("Env");
289    let tcp_stream = TcpStream::connect("generativelanguage.googleapis.com:443").unwrap();
290    let connector = TlsConnector::new().unwrap();
291    let mut stream = connector
292        .connect("generativelanguage.googleapis.com", tcp_stream)
293        .unwrap();
294
295    let models = format!(
296        "POST /v1beta/models/{}:generateContent?key={} HTTP/1.1\r\n\
297           Host: generativelanguage.googleapis.com\r\n\
298           Content-Type: {}\r\n\
299           Content-Length: {}\r\n\
300           Connection: close\r\n\r\n{}",
301        model,
302        env,
303        mime_type,
304        content.len(),
305        content
306    );
307    // println!("{}", models);
308
309    stream.write_all(models.as_bytes());
310    stream.flush();
311
312    let mut response = String::new();
313    let mut buffer = BufReader::new(stream);
314    buffer.read_to_string(&mut response);
315    // let mut reader = BufReader::new(stream);
316
317    // loop {
318    //     let mut line = String::new();
319    //     reader.read_line(&mut line);
320
321    //     if line == "\r\n" {
322    //         break;
323    //     }
324    // }
325
326    //     let size = usize::from_str_radix(&line.trim(), 16).unwrap();
327    //     if size == 0 {
328    //         break;
329    //     }
330
331    // }
332    response
333}
334#[cfg(feature = "async")]
335pub(crate) async fn gemini(content: String, env: &str, model: &str, mime_type: &str) -> String {
336    use std::env;
337
338    use async_std::io::{ReadExt, WriteExt};
339
340    dotenv().unwrap();
341    let env = env::var(env).expect("Env");
342    let gemini = TlsConnector::new();
343    let stream = TcpStream::connect("generativelanguage.googleapis.com:443")
344        .await
345        .unwrap();
346    let mut stream = gemini
347        .connect("generativelanguage.googleapis.com", stream)
348        .await
349        .unwrap();
350
351    let models = format!(
352        "POST /v1beta/models/{}:generateContent?key={} HTTP/1.1\r\n\
353           Host: generativelanguage.googleapis.com\r\n\
354           Content-Type: {}\r\n\
355           Content-Length: {}\r\n\
356           Connection: close\r\n\r\n{}",
357        model,
358        env,
359        mime_type,
360        content.len(),
361        content
362    );
363    // println!("{}", models);
364
365    // println!("{}", models);
366    stream.write_all(models.as_bytes()).await;
367    stream.flush().await;
368
369    let mut response = String::new();
370    let mut buffer = BufReader::new(stream);
371    buffer.read_to_string(&mut response).await;
372    // let mut reader = BufReader::new(stream);
373
374    // loop {
375    //     let mut line = String::new();
376    //     reader.read_line(&mut line);
377
378    //     if line == "\r\n" {
379    //         break;
380    //     }
381    // }
382
383    // let mut body = String::new();
384    // loop {
385    //     let mut line = String::new();
386    //     reader.read_line(&mut line);
387    //     // println!("{}", line);
388
389    //     let size = usize::from_str_radix(&line.trim(), 16).unwrap();
390    //     if size == 0 {
391    //         break;
392    //     }
393
394    //     let mut esponse = vec![0; size];
395    //     reader.read_exact(&mut esponse);
396    //     let response = String::from_utf8_lossy(&esponse);
397    //     body.push_str(&response);
398    //     // println!("{}", response);
399
400    //     let mut trail = vec![0; 2];
401    //     reader.read_exact(&mut trail);
402    // }
403    // let mut esponse = vec![0; size];
404    // reader.read_exact(&mut esponse).await?;
405    // let response = String::from_utf8_lossy(&esponse);
406    // body.push_str(&response);
407    // // println!("{}", response);
408
409    // let mut trail = vec![0; 2];
410    // reader.read_exact(&mut trail).await?;
411    response
412    // }
413}