llm_universal/
gemini.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use serde_json::{json, Value};
5
6use crate::{
7    send_json, ChatRequest, ChatResponse, LlmError, LlmProvider, Message, Part, Result, Role,
8    StopReason, Usage,
9};
10
11/// Prefix of tool-call ids this crate generates when Gemini returns a call without an id.
12/// Such ids are never sent back to the API.
13const LOCAL_ID_PREFIX: &str = "llm-universal-call-";
14
15/// Client for the Google Gemini API
16/// (`POST {base_url}/models/{model}:generateContent`, key in `x-goog-api-key`).
17#[derive(Clone)]
18pub struct GeminiProvider {
19    client: reqwest::Client,
20    api_key: String,
21    base_url: String,
22}
23
24impl std::fmt::Debug for GeminiProvider {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        f.debug_struct("GeminiProvider")
27            .field("api_key", &"<redacted>")
28            .field("base_url", &self.base_url)
29            .finish()
30    }
31}
32
33impl GeminiProvider {
34    pub fn new(api_key: impl Into<String>) -> Self {
35        Self {
36            client: crate::http_client(),
37            api_key: api_key.into(),
38            base_url: "https://generativelanguage.googleapis.com/v1beta".to_string(),
39        }
40    }
41
42    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
43        self.base_url = base_url.into().trim_end_matches('/').to_string();
44        self
45    }
46
47    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
48    pub fn with_client(mut self, client: reqwest::Client) -> Self {
49        self.client = client;
50        self
51    }
52
53    fn build_body(req: &ChatRequest) -> Value {
54        // Gemini needs the function name in `functionResponse`; recover it from the call.
55        let names: HashMap<&str, &str> = req
56            .messages
57            .iter()
58            .flat_map(|m| &m.parts)
59            .filter_map(|p| match p {
60                Part::ToolCall { id, name, .. } => Some((id.as_str(), name.as_str())),
61                _ => None,
62            })
63            .collect();
64
65        let mut system_parts: Vec<Value> =
66            req.system.iter().map(|s| json!({ "text": s })).collect();
67        let mut contents = Vec::new();
68        for message in &req.messages {
69            if message.role == Role::System {
70                system_parts.push(json!({ "text": message.text_content() }));
71                continue;
72            }
73            let role = if message.role == Role::Assistant {
74                "model"
75            } else {
76                "user"
77            };
78            let parts: Vec<Value> = message
79                .parts
80                .iter()
81                .filter_map(|p| to_gemini_part(p, &names))
82                .collect();
83            contents.push(json!({ "role": role, "parts": parts }));
84        }
85
86        let mut body = json!({ "contents": contents });
87        if !system_parts.is_empty() {
88            body["systemInstruction"] = json!({ "parts": system_parts });
89        }
90        if !req.tools.is_empty() {
91            let declarations: Vec<Value> = req
92                .tools
93                .iter()
94                .map(|t| {
95                    json!({
96                        "name": t.name,
97                        "description": t.description,
98                        "parameters": t.parameters,
99                    })
100                })
101                .collect();
102            body["tools"] = json!([{ "functionDeclarations": declarations }]);
103        }
104        let mut config = serde_json::Map::new();
105        if let Some(max_tokens) = req.max_tokens {
106            config.insert("maxOutputTokens".into(), json!(max_tokens));
107        }
108        if let Some(temperature) = req.temperature {
109            config.insert("temperature".into(), json!(temperature));
110        }
111        if !config.is_empty() {
112            body["generationConfig"] = Value::Object(config);
113        }
114        body
115    }
116}
117
118fn to_gemini_part(part: &Part, names: &HashMap<&str, &str>) -> Option<Value> {
119    let sendable_id =
120        |id: &str| (!id.is_empty() && !id.starts_with(LOCAL_ID_PREFIX)).then(|| id.to_string());
121    match part {
122        Part::Text { text } if text.is_empty() => None,
123        Part::Text { text } => Some(json!({ "text": text })),
124        Part::ToolCall {
125            id,
126            name,
127            arguments,
128        } => {
129            let mut call = json!({ "name": name, "args": arguments });
130            if let Some(id) = sendable_id(id) {
131                call["id"] = json!(id);
132            }
133            Some(json!({ "functionCall": call }))
134        }
135        Part::ToolResult {
136            call_id,
137            content,
138            is_error,
139        } => {
140            // `response` must be an object: use the content if it is a JSON object,
141            // otherwise wrap it the way Google's examples do.
142            let response = match serde_json::from_str::<Value>(content) {
143                Ok(v @ Value::Object(_)) if !is_error => v,
144                _ if *is_error => json!({ "error": content }),
145                _ => json!({ "result": content }),
146            };
147            let mut result = json!({
148                "name": names.get(call_id.as_str()).copied().unwrap_or_default(),
149                "response": response,
150            });
151            if let Some(id) = sendable_id(call_id) {
152                result["id"] = json!(id);
153            }
154            Some(json!({ "functionResponse": result }))
155        }
156    }
157}
158
159fn parse_response(v: &Value, requested_model: &str) -> Result<ChatResponse> {
160    let tokens = |field: &str| {
161        u32::try_from(v["usageMetadata"][field].as_u64().unwrap_or(0)).unwrap_or(u32::MAX)
162    };
163    let usage = Usage {
164        input_tokens: tokens("promptTokenCount"),
165        output_tokens: tokens("candidatesTokenCount"),
166    };
167    let model = v["modelVersion"]
168        .as_str()
169        .unwrap_or(requested_model)
170        .to_string();
171
172    let Some(candidate) = v["candidates"].get(0) else {
173        // The prompt itself was blocked: no candidates, only `promptFeedback.blockReason`.
174        if let Some(reason) = v["promptFeedback"]["blockReason"].as_str() {
175            let text = format!("prompt blocked: {reason}");
176            return Ok(ChatResponse {
177                model,
178                message: Message::new(Role::Assistant, vec![Part::Text { text }]),
179                stop_reason: StopReason::Refusal,
180                usage,
181            });
182        }
183        return Err(LlmError::InvalidResponse(
184            "missing field `candidates[0]`".into(),
185        ));
186    };
187
188    let mut parts = Vec::new();
189    let mut calls = 0usize;
190    for part in candidate["content"]["parts"]
191        .as_array()
192        .into_iter()
193        .flatten()
194    {
195        // Thought summaries (`thought: true`) are not part of the answer.
196        if part["thought"].as_bool() == Some(true) {
197            continue;
198        }
199        if let Some(text) = part["text"].as_str() {
200            if !text.is_empty() {
201                parts.push(Part::Text {
202                    text: text.to_string(),
203                });
204            }
205        } else if part["functionCall"].is_object() {
206            let call = &part["functionCall"];
207            let id = match call["id"].as_str() {
208                Some(id) if !id.is_empty() => id.to_string(),
209                _ => format!("{LOCAL_ID_PREFIX}{calls}"),
210            };
211            calls += 1;
212            parts.push(Part::ToolCall {
213                id,
214                name: call["name"].as_str().unwrap_or_default().to_string(),
215                arguments: if call["args"].is_null() {
216                    json!({})
217                } else {
218                    call["args"].clone()
219                },
220            });
221        }
222    }
223
224    let stop_reason = match candidate["finishReason"].as_str() {
225        // Gemini reports STOP for function calls as well.
226        Some("STOP") if calls > 0 => StopReason::ToolUse,
227        Some("STOP") => StopReason::EndTurn,
228        Some("MAX_TOKENS") => StopReason::MaxTokens,
229        Some(
230            "SAFETY"
231            | "RECITATION"
232            | "BLOCKLIST"
233            | "PROHIBITED_CONTENT"
234            | "SPII"
235            | "IMAGE_SAFETY"
236            | "IMAGE_PROHIBITED_CONTENT"
237            | "IMAGE_RECITATION",
238        ) => StopReason::Refusal,
239        None if calls > 0 => StopReason::ToolUse,
240        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
241    };
242
243    Ok(ChatResponse {
244        model,
245        message: Message::new(Role::Assistant, parts),
246        stop_reason,
247        usage,
248    })
249}
250
251/// The model goes into the URL path, so only plain model-name characters are allowed.
252fn checked_model(model: &str) -> Result<&str> {
253    let model = model.strip_prefix("models/").unwrap_or(model);
254    if model
255        .chars()
256        .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
257    {
258        Ok(model)
259    } else {
260        Err(LlmError::InvalidArgument(format!(
261            "GeminiProvider: model `{model}` may only contain a-z, A-Z, 0-9, '-', '_' and '.'"
262        )))
263    }
264}
265
266#[async_trait]
267impl LlmProvider for GeminiProvider {
268    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
269        crate::validate(
270            "GeminiProvider",
271            &self.api_key,
272            true,
273            &self.base_url,
274            64,
275            req,
276        )?;
277        let model = checked_model(&req.model)?;
278        let url = format!("{}/models/{model}:generateContent", self.base_url);
279        let builder = self
280            .client
281            .post(url)
282            .header(
283                "x-goog-api-key",
284                crate::secret_header("GeminiProvider", &self.api_key)?,
285            )
286            .json(&Self::build_body(req));
287        let body = send_json(builder).await?;
288        parse_response(&body, model)
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use crate::Tool;
296
297    #[test]
298    fn builds_request_with_tools_and_results() {
299        let req = ChatRequest::new("gemini-2.5-flash")
300            .system("sys")
301            .max_tokens(100)
302            .tool(Tool::new("get_weather", "d", json!({"type": "object"})))
303            .message(Message::user("hi"))
304            .message(Message::new(
305                Role::Assistant,
306                vec![Part::ToolCall {
307                    id: format!("{LOCAL_ID_PREFIX}0"),
308                    name: "get_weather".into(),
309                    arguments: json!({"city": "Berlin"}),
310                }],
311            ))
312            .message(Message::tool_result(
313                format!("{LOCAL_ID_PREFIX}0"),
314                "18C",
315                false,
316            ));
317        let body = GeminiProvider::build_body(&req);
318        assert_eq!(body["systemInstruction"]["parts"][0]["text"], "sys");
319        assert_eq!(body["generationConfig"]["maxOutputTokens"], 100);
320        assert_eq!(
321            body["tools"][0]["functionDeclarations"][0]["name"],
322            "get_weather"
323        );
324        assert_eq!(body["contents"][1]["role"], "model");
325        let call = &body["contents"][1]["parts"][0]["functionCall"];
326        assert_eq!(call["args"]["city"], "Berlin");
327        assert!(call.get("id").is_none(), "local ids are not sent");
328        let result = &body["contents"][2]["parts"][0]["functionResponse"];
329        assert_eq!(result["name"], "get_weather");
330        assert_eq!(result["response"]["result"], "18C");
331    }
332
333    #[test]
334    fn parses_function_call_as_tool_use() {
335        let v = json!({
336            "candidates": [{
337                "content": { "role": "model", "parts": [
338                    { "text": "thinking...", "thought": true },
339                    { "functionCall": { "name": "get_weather", "args": { "city": "Berlin" } } }
340                ]},
341                "finishReason": "STOP"
342            }],
343            "usageMetadata": { "promptTokenCount": 4, "candidatesTokenCount": 6 },
344            "modelVersion": "gemini-2.5-flash"
345        });
346        let resp = parse_response(&v, "m").unwrap();
347        assert_eq!(resp.stop_reason, StopReason::ToolUse);
348        assert_eq!(resp.text(), "");
349        let calls = resp.tool_calls();
350        assert_eq!(calls[0].1, "get_weather");
351        assert!(calls[0].0.starts_with(LOCAL_ID_PREFIX));
352        assert_eq!(resp.usage.output_tokens, 6);
353        assert_eq!(resp.model, "gemini-2.5-flash");
354    }
355
356    #[test]
357    fn blocked_prompt_is_a_refusal() {
358        let v = json!({ "promptFeedback": { "blockReason": "SAFETY" } });
359        let resp = parse_response(&v, "m").unwrap();
360        assert_eq!(resp.stop_reason, StopReason::Refusal);
361        assert_eq!(resp.text(), "prompt blocked: SAFETY");
362    }
363
364    #[test]
365    fn rejects_model_that_would_change_the_url() {
366        for model in ["../files", "gemini?key=x", "a/b", "gemini pro"] {
367            assert!(checked_model(model).is_err(), "{model}");
368        }
369        assert_eq!(
370            checked_model("models/gemini-2.5-flash").unwrap(),
371            "gemini-2.5-flash"
372        );
373    }
374}
llm_universal/gemini.rs

llm_universal/
gemini.rs
