orchestral_runtime/planner/
gemini.rs

1//! Gemini LLM client implementation.
2//!
3//! This module provides a client for Google's Gemini API.
4
5use async_trait::async_trait;
6use reqwest::header::{HeaderMap, HeaderName, HeaderValue, CONTENT_TYPE};
7use serde::{Deserialize, Serialize};
8
9use super::llm::{LlmClient, LlmError, LlmRequest};
10
11/// Gemini client configuration.
12#[derive(Debug, Clone)]
13pub struct GeminiClientConfig {
14    /// API key for authentication.
15    pub api_key: String,
16    /// Model name (e.g., "gemini-1.5-pro", "gemini-1.5-flash").
17    pub model: String,
18    /// Base endpoint URL.
19    pub endpoint: String,
20    /// Temperature for generation (0.0 - 2.0).
21    pub temperature: f32,
22    /// Request timeout in seconds.
23    pub timeout_secs: u64,
24}
25
26impl Default for GeminiClientConfig {
27    fn default() -> Self {
28        Self {
29            api_key: String::new(),
30            model: "gemini-3-flash-preview".to_string(),
31            endpoint: "https://generativelanguage.googleapis.com/v1beta".to_string(),
32            temperature: 0.2,
33            timeout_secs: 30,
34        }
35    }
36}
37
38/// Gemini LLM client.
39pub struct GeminiClient {
40    client: reqwest::Client,
41    config: GeminiClientConfig,
42}
43
44impl GeminiClient {
45    /// Create a new Gemini client.
46    pub fn new(config: GeminiClientConfig) -> Result<Self, LlmError> {
47        let client = reqwest::Client::builder()
48            .timeout(std::time::Duration::from_secs(config.timeout_secs))
49            .build()
50            .map_err(|e| LlmError::Http(e.to_string()))?;
51        Ok(Self { client, config })
52    }
53
54    fn build_url(&self, model: &str) -> String {
55        format!("{}/models/{}:generateContent", self.config.endpoint, model)
56    }
57}
58
59// Gemini API request/response structures
60
61#[derive(Debug, Serialize)]
62struct GeminiRequest {
63    contents: Vec<GeminiContent>,
64    #[serde(rename = "systemInstruction", skip_serializing_if = "Option::is_none")]
65    system_instruction: Option<GeminiSystemInstruction>,
66    #[serde(rename = "generationConfig")]
67    generation_config: GeminiGenerationConfig,
68}
69
70#[derive(Debug, Serialize)]
71struct GeminiContent {
72    role: String,
73    parts: Vec<GeminiPart>,
74}
75
76#[derive(Debug, Serialize)]
77struct GeminiSystemInstruction {
78    parts: Vec<GeminiPart>,
79}
80
81#[derive(Debug, Serialize)]
82struct GeminiPart {
83    text: String,
84}
85
86#[derive(Debug, Serialize)]
87struct GeminiGenerationConfig {
88    temperature: f32,
89}
90
91#[derive(Debug, Deserialize)]
92struct GeminiResponse {
93    candidates: Option<Vec<GeminiCandidate>>,
94    error: Option<GeminiErrorDetail>,
95}
96
97#[derive(Debug, Deserialize)]
98struct GeminiCandidate {
99    content: GeminiContentResponse,
100}
101
102#[derive(Debug, Deserialize)]
103struct GeminiContentResponse {
104    parts: Vec<GeminiPartResponse>,
105}
106
107#[derive(Debug, Deserialize)]
108struct GeminiPartResponse {
109    text: String,
110}
111
112#[derive(Debug, Deserialize)]
113struct GeminiErrorDetail {
114    message: String,
115    #[allow(dead_code)]
116    code: Option<i32>,
117}
118
119#[async_trait]
120impl LlmClient for GeminiClient {
121    async fn complete(&self, request: LlmRequest) -> Result<String, LlmError> {
122        let url = self.build_url(&request.model);
123
124        let mut headers = HeaderMap::new();
125        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
126        let api_key_header = HeaderName::from_static("x-goog-api-key");
127        let api_key_value = HeaderValue::from_str(&self.config.api_key)
128            .map_err(|e| LlmError::Http(format!("invalid Gemini API key header value: {}", e)))?;
129        headers.insert(api_key_header, api_key_value);
130
131        // Build Gemini request
132        let body = GeminiRequest {
133            contents: vec![GeminiContent {
134                role: "user".to_string(),
135                parts: vec![GeminiPart { text: request.user }],
136            }],
137            system_instruction: if request.system.is_empty() {
138                None
139            } else {
140                Some(GeminiSystemInstruction {
141                    parts: vec![GeminiPart {
142                        text: request.system,
143                    }],
144                })
145            },
146            generation_config: GeminiGenerationConfig {
147                temperature: request.temperature,
148            },
149        };
150
151        let response = self
152            .client
153            .post(&url)
154            .headers(headers)
155            .json(&body)
156            .send()
157            .await
158            .map_err(|e| LlmError::Http(e.to_string()))?;
159
160        if !response.status().is_success() {
161            let status = response.status();
162            let text = response.text().await.unwrap_or_default();
163            return Err(LlmError::Response(format!("HTTP {}: {}", status, text)));
164        }
165
166        let text = response
167            .text()
168            .await
169            .map_err(|e| LlmError::Http(e.to_string()))?;
170
171        let parsed: GeminiResponse =
172            serde_json::from_str(&text).map_err(|e| LlmError::Serialization(e.to_string()))?;
173
174        // Check for API error
175        if let Some(error) = parsed.error {
176            return Err(LlmError::Response(format!(
177                "Gemini API error: {}",
178                error.message
179            )));
180        }
181
182        // Extract content from response
183        let content = parsed
184            .candidates
185            .and_then(|c| c.into_iter().next())
186            .and_then(|c| c.content.parts.into_iter().next())
187            .map(|p| p.text)
188            .ok_or_else(|| LlmError::Response("No content in response".to_string()))?;
189
190        Ok(content)
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use crate::planner::LlmRequest;
198
199    #[test]
200    fn test_default_config() {
201        let config = GeminiClientConfig::default();
202        assert_eq!(config.model, "gemini-3-flash-preview");
203        assert!(config
204            .endpoint
205            .contains("generativelanguage.googleapis.com"));
206    }
207
208    #[test]
209    fn test_build_url() {
210        let config = GeminiClientConfig {
211            api_key: "test-key".to_string(),
212            model: "gemini-1.5-pro".to_string(),
213            ..Default::default()
214        };
215        let client = GeminiClient::new(config).unwrap();
216        let url = client.build_url("gemini-1.5-pro");
217        assert!(url.contains("gemini-1.5-pro:generateContent"));
218        assert!(!url.contains("test-key"));
219    }
220
221    #[tokio::test]
222    #[ignore = "requires live GEMINI_API_KEY and network"]
223    async fn test_live_gemini_completion_when_env_set() {
224        let api_key = match std::env::var("GEMINI_API_KEY") {
225            Ok(v) if !v.trim().is_empty() => v,
226            _ => {
227                eprintln!("skipped: GEMINI_API_KEY is not set");
228                return;
229            }
230        };
231
232        let config = GeminiClientConfig {
233            api_key,
234            model: "gemini-3-flash-preview".to_string(),
235            ..Default::default()
236        };
237        let client = GeminiClient::new(config).expect("client should initialize");
238        let request = LlmRequest {
239            system: "You are a concise assistant.".to_string(),
240            user: "Reply with exactly: OK".to_string(),
241            model: "gemini-3-flash-preview".to_string(),
242            temperature: 0.0,
243        };
244
245        let response = client
246            .complete(request)
247            .await
248            .expect("live Gemini completion should succeed");
249        assert!(!response.trim().is_empty());
250    }
251}
orchestral_runtime/planner/gemini.rs

orchestral_runtime/planner/
gemini.rs
