rs_agent/models/
gemini.rs

1use async_trait::async_trait;
2use base64::Engine;
3use futures::stream::{BoxStream, StreamExt};
4use reqwest::Client;
5use serde::{Deserialize, Serialize};
6
7use crate::error::{AgentError, Result};
8use crate::models::LLM;
9use crate::types::{File, GenerationChunk, GenerationResponse, Message, Role};
10
11/// Gemini LLM provider
12pub struct GeminiLLM {
13    client: Client,
14    api_key: String,
15    model: String,
16}
17
18#[derive(Debug, Serialize)]
19struct GeminiRequest {
20    contents: Vec<GeminiContent>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    tools: Option<Vec<serde_json::Value>>,
23}
24
25#[derive(Debug, Serialize)]
26struct GeminiContent {
27    role: String,
28    parts: Vec<GeminiPart>,
29}
30
31#[derive(Debug, Serialize)]
32#[serde(untagged)]
33enum GeminiPart {
34    Text { text: String },
35    InlineData { inline_data: GeminiBlob },
36}
37
38#[derive(Debug, Serialize)]
39struct GeminiBlob {
40    mime_type: String,
41    data: String,
42}
43
44#[derive(Debug, Deserialize)]
45struct GeminiResponse {
46    candidates: Option<Vec<GeminiCandidate>>,
47}
48
49#[derive(Debug, Deserialize)]
50struct GeminiCandidate {
51    content: Option<GeminiContentResponse>,
52}
53
54#[derive(Debug, Deserialize)]
55struct GeminiContentResponse {
56    parts: Option<Vec<GeminiPartResponse>>,
57}
58
59#[derive(Debug, Deserialize)]
60struct GeminiPartResponse {
61    text: Option<String>,
62}
63
64impl GeminiLLM {
65    /// Creates a new Gemini LLM with API key from environment
66    pub fn new(model: impl Into<String>) -> Result<Self> {
67        let api_key = std::env::var("GOOGLE_API_KEY")
68            .or_else(|_| std::env::var("GEMINI_API_KEY"))
69            .map_err(|_| {
70                AgentError::ConfigError(
71                    "GOOGLE_API_KEY or GEMINI_API_KEY environment variable not set".to_string(),
72                )
73            })?;
74
75        Ok(Self {
76            client: Client::new(),
77            api_key,
78            model: model.into(),
79        })
80    }
81
82    /// Creates with explicit API key
83    pub fn with_api_key(api_key: impl Into<String>, model: impl Into<String>) -> Self {
84        Self {
85            client: Client::new(),
86            api_key: api_key.into(),
87            model: model.into(),
88        }
89    }
90
91    fn convert_role(role: &Role) -> String {
92        match role {
93            Role::User => "user".to_string(),
94            Role::Assistant => "model".to_string(),
95            Role::System => "user".to_string(), // Gemini maps system to user or uses specific system instruction
96            Role::Tool => "user".to_string(),
97        }
98    }
99
100    fn prepare_request_body(
101        &self,
102        messages: Vec<Message>,
103        files: Option<Vec<File>>,
104    ) -> GeminiRequest {
105        let mut contents: Vec<GeminiContent> = messages
106            .iter()
107            .map(|m| GeminiContent {
108                role: Self::convert_role(&m.role),
109                parts: vec![GeminiPart::Text {
110                    text: m.content.clone(),
111                }],
112            })
113            .collect();
114
115        // Add file attachments if present to the last message
116        if let Some(files) = files {
117            if let Some(last_content) = contents.last_mut() {
118                for file in files {
119                    last_content.parts.push(GeminiPart::InlineData {
120                        inline_data: GeminiBlob {
121                            mime_type: file.mime_type,
122                            data: base64::engine::general_purpose::STANDARD.encode(&file.data),
123                        },
124                    });
125                }
126            }
127        }
128
129        GeminiRequest {
130            contents,
131            tools: None,
132        }
133    }
134}
135
136#[async_trait]
137impl LLM for GeminiLLM {
138    async fn generate(
139        &self,
140        messages: Vec<Message>,
141        files: Option<Vec<File>>,
142    ) -> Result<GenerationResponse> {
143        let request = self.prepare_request_body(messages, files);
144
145        let url = format!(
146            "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}",
147            self.model, self.api_key
148        );
149
150        let response = self
151            .client
152            .post(&url)
153            .json(&request)
154            .send()
155            .await
156            .map_err(|e| AgentError::ModelError(format!("Gemini API error: {}", e)))?;
157
158        if !response.status().is_success() {
159            let status = response.status();
160            let text = response.text().await.unwrap_or_default();
161            return Err(AgentError::ModelError(format!(
162                "Gemini API error {}: {}",
163                status, text
164            )));
165        }
166
167        let gemini_response: GeminiResponse = response
168            .json()
169            .await
170            .map_err(|e| AgentError::ModelError(format!("Failed to parse response: {}", e)))?;
171
172        let content = gemini_response
173            .candidates
174            .as_ref()
175            .and_then(|c| c.first())
176            .and_then(|c| c.content.as_ref())
177            .and_then(|c| c.parts.as_ref())
178            .and_then(|p| p.first())
179            .and_then(|p| p.text.clone())
180            .ok_or_else(|| AgentError::ModelError("No content in response".to_string()))?;
181
182        Ok(GenerationResponse {
183            content,
184            metadata: None,
185        })
186    }
187
188    async fn stream_generate(
189        &self,
190        messages: Vec<Message>,
191        files: Option<Vec<File>>,
192    ) -> Result<BoxStream<'static, Result<GenerationChunk>>> {
193        let request = self.prepare_request_body(messages, files);
194
195        let url = format!(
196            "https://generativelanguage.googleapis.com/v1beta/models/{}:streamGenerateContent?key={}&alt=sse",
197            self.model, self.api_key
198        );
199
200        let response = self
201            .client
202            .post(&url)
203            .json(&request)
204            .send()
205            .await
206            .map_err(|e| AgentError::ModelError(format!("Gemini API error: {}", e)))?;
207
208        if !response.status().is_success() {
209            let status = response.status();
210            let text = response.text().await.unwrap_or_default();
211            return Err(AgentError::ModelError(format!(
212                "Gemini API error {}: {}",
213                status, text
214            )));
215        }
216
217        let stream = response.bytes_stream();
218        let buffer = Vec::new();
219
220        let s = futures::stream::unfold((stream, buffer), |(mut stream, mut buffer)| async move {
221            loop {
222                // Check if buffer has a full line
223                if let Some(pos) = buffer.iter().position(|&b| b == b'\n') {
224                    let line = buffer.drain(0..=pos).collect::<Vec<u8>>();
225                    let s = String::from_utf8_lossy(&line);
226                    let trimmed = s.trim();
227
228                    if trimmed.starts_with("data: ") {
229                        let json_str = trimmed.trim_start_matches("data: ").trim();
230                        if let Ok(resp) = serde_json::from_str::<GeminiResponse>(json_str) {
231                            // Extract content
232                            let content_opt = resp
233                                .candidates
234                                .as_ref()
235                                .and_then(|c| c.first())
236                                .and_then(|c| c.content.as_ref())
237                                .and_then(|c| c.parts.as_ref())
238                                .and_then(|p| p.first())
239                                .and_then(|p| match p {
240                                    GeminiPartResponse { text: Some(t) } => Some(t.clone()),
241                                    _ => None,
242                                });
243
244                            if let Some(content) = content_opt {
245                                return Some((
246                                    Ok(GenerationChunk {
247                                        content,
248                                        metadata: None,
249                                    }),
250                                    (stream, buffer),
251                                ));
252                            }
253                        }
254                    }
255                    continue;
256                }
257
258                // Need more data
259                match stream.next().await {
260                    Some(Ok(chunk)) => {
261                        buffer.extend_from_slice(&chunk);
262                    }
263                    Some(Err(e)) => {
264                        return Some((Err(AgentError::ModelError(e.to_string())), (stream, buffer)))
265                    }
266                    None => {
267                        return None;
268                    }
269                }
270            }
271        });
272
273        Ok(Box::pin(s))
274    }
275
276    fn model_name(&self) -> &str {
277        &self.model
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    #[tokio::test]
286    #[ignore] // Requires API key
287    async fn test_gemini_generate() {
288        let llm = GeminiLLM::new("gemini-2.0-flash").unwrap();
289        let messages = vec![Message {
290            role: Role::User,
291            content: "Say 'Hello, World!' and nothing else.".to_string(),
292            metadata: None,
293        }];
294
295        let response = llm.generate(messages, None).await.unwrap();
296        assert!(response.content.contains("Hello"));
297    }
298}
rs_agent/models/gemini.rs

rs_agent/models/
gemini.rs
