gemini_client_api/gemini/
types.rs

1use derive_new::new;
2use serde::Serialize;
3use serde_json::Value;
4use std::collections::VecDeque;
5
6#[derive(Serialize)]
7#[allow(non_camel_case_types)]
8pub enum Role {
9    user,
10    developer,
11    model,
12}
13
14#[derive(Serialize, new)]
15pub struct InlineData {
16    mime_type: String,
17    data: String,
18}
19
20#[derive(Serialize)]
21#[allow(non_camel_case_types)]
22pub enum Part {
23    text(String),
24    inline_data(InlineData),
25}
26
27#[derive(Serialize, new)]
28pub struct Chat {
29    role: Role,
30    parts: Vec<Part>,
31}
32
33#[derive(Serialize, new)]
34pub struct SystemInstruction<'a> {
35    parts: &'a [Part],
36}
37
38#[derive(Serialize, new)]
39pub struct GeminiBody<'a> {
40    system_instruction: Option<&'a SystemInstruction<'a>>,
41    contents: &'a [&'a Chat],
42    generation_config: Option<&'a Value>,
43}
44
45#[derive(Serialize)]
46pub struct Session {
47    history: VecDeque<Chat>,
48    history_limit: usize,
49    chat_no: usize,
50}
51impl Session {
52    pub fn new(history_limit: usize) -> Self {
53        Self {
54            history: VecDeque::new(),
55            history_limit,
56            chat_no: 0,
57        }
58    }
59    pub fn get_history_as_vecdeque(&self) -> &VecDeque<Chat> {
60        &self.history
61    }
62    pub(super) fn get_history_as_vecdeque_mut(&mut self) -> &mut VecDeque<Chat> {
63        &mut self.history
64    }
65    pub fn get_history(&self) -> Vec<&Chat> {
66        let (left, right) = self.history.as_slices();
67        left.iter().chain(right.iter()).collect()
68    }
69    pub fn get_history_length(&self) -> usize {
70        self.history.len()
71    }
72    pub fn get_parts_mut(&mut self, chat_previous_no: usize) -> Option<&mut Vec<Part>> {
73        let history_length = self.get_history_length();
74        self.history
75            .get_mut(history_length - chat_previous_no)
76            .map(|chat| &mut chat.parts)
77    }
78    fn add_chat(&mut self, chat: Chat) -> &mut Self {
79        self.history.push_back(chat);
80        self.chat_no += 1;
81        if self.get_history_length() > self.history_limit {
82            self.history.pop_front();
83        }
84        self
85    }
86    pub fn ask(&mut self, parts: Vec<Part>) -> &mut Self {
87        self.add_chat(Chat::new(Role::user, parts))
88    }
89    pub fn ask_string(&mut self, prompt: String) -> &mut Self {
90        self.add_chat(Chat::new(Role::user, vec![Part::text(prompt)]))
91    }
92    pub(super) fn update(&mut self, reply: &str) {
93        let history = &mut self.history;
94        if let Some(chat) = history.back_mut() {
95            if let Role::model = chat.role {
96                if let Some(Part::text(data)) = chat.parts.last_mut() {
97                    data.push_str(reply);
98                } else {
99                    chat.parts.push(Part::text(reply.to_string()));
100                }
101            } else {
102                history.push_back(Chat::new(Role::model, vec![Part::text(reply.to_string())]));
103            }
104        } else {
105            panic!("Cannot update an empty session");
106        }
107    }
108    pub fn last_reply(&self) -> Option<&String> {
109        if let Some(Part::text(reply)) = self
110            .get_history_as_vecdeque()
111            .back()
112            .map(|chat| chat.parts.get(0))
113            .flatten()
114        {
115            Some(reply)
116        } else {
117            None
118        }
119    }
120    pub(super) fn last_reply_mut(&mut self) -> Option<&mut String> {
121        if let Some(Part::text(reply)) = self
122            .get_history_as_vecdeque_mut()
123            .back_mut()
124            .map(|chat| chat.parts.get_mut(0))
125            .flatten()
126        {
127            Some(reply)
128        } else {
129            None
130        }
131    }
132    pub fn forget_last_conversation(&mut self) {
133        self.history.pop_back();
134        if let Some(Chat {
135            role: Role::user,
136            parts: _,
137        }) = self.history.back()
138        {
139            self.history.pop_back();
140        }
141    }
142}
gemini_client_api/gemini/types.rs

gemini_client_api/gemini/
types.rs
