llm_universal/
gemini.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use serde_json::{json, Value};
5
6use crate::anthropic::Vertex;
7use crate::http::{secret_header, Http};
8use crate::openai::{float_vec, usage, validate_embed_args};
9use crate::stream::{self, parse_data, Accumulator, ChatStream, StreamEvent};
10use crate::{
11    ApiError, ChatRequest, ChatResponse, EmbeddingProvider, Embeddings, ImageSource, LlmError,
12    LlmProvider, Message, ModelInfo, Part, Result, Role, StopReason,
13};
14
15/// Prefix of tool-call ids this crate generates when Gemini returns a call without an id.
16/// Such ids are never sent back to the API.
17const LOCAL_ID_PREFIX: &str = "llm-universal-call-";
18
19/// Client for the Google Gemini API
20/// (`POST {base_url}/models/{model}:generateContent`, key in `x-goog-api-key`),
21/// or for Gemini on Vertex AI via [`vertex`](Self::vertex).
22#[derive(Clone)]
23pub struct GeminiProvider {
24    http: Http,
25    api_key: String,
26    base_url: String,
27    vertex: Option<Vertex>,
28}
29
30impl std::fmt::Debug for GeminiProvider {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("GeminiProvider")
33            .field("api_key", &"<redacted>")
34            .field("base_url", &self.base_url)
35            .field("vertex", &self.vertex)
36            .field("max_retries", &self.http.max_retries)
37            .finish()
38    }
39}
40
41impl GeminiProvider {
42    pub fn new(api_key: impl Into<String>) -> Self {
43        Self {
44            http: Http::default(),
45            api_key: api_key.into(),
46            base_url: "https://generativelanguage.googleapis.com/v1beta".to_string(),
47            vertex: None,
48        }
49    }
50
51    /// Gemini on Google Cloud Vertex AI. `access_token` is an OAuth token
52    /// (e.g. `gcloud auth print-access-token`); `region` is e.g. `global` or `us-central1`.
53    pub fn vertex(
54        project: impl Into<String>,
55        region: impl Into<String>,
56        access_token: impl Into<String>,
57    ) -> Self {
58        let region = region.into();
59        Self {
60            http: Http::default(),
61            api_key: access_token.into(),
62            base_url: Vertex::base_url(&region),
63            vertex: Some(Vertex {
64                project: project.into(),
65                region,
66            }),
67        }
68    }
69
70    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
71        self.base_url = base_url.into().trim_end_matches('/').to_string();
72        self
73    }
74
75    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
76    pub fn with_client(mut self, client: reqwest::Client) -> Self {
77        self.http.client = client;
78        self
79    }
80
81    /// How many times a failed request is retried (default [`DEFAULT_MAX_RETRIES`](crate::DEFAULT_MAX_RETRIES); 0 disables).
82    pub fn with_max_retries(mut self, max_retries: u32) -> Self {
83        self.http.max_retries = max_retries;
84        self
85    }
86
87    fn build_body(req: &ChatRequest) -> Result<Value> {
88        // Gemini needs the function name in `functionResponse`; recover it from the call.
89        let names: HashMap<&str, &str> = req
90            .messages
91            .iter()
92            .flat_map(|m| &m.parts)
93            .filter_map(|p| match p {
94                Part::ToolCall { id, name, .. } => Some((id.as_str(), name.as_str())),
95                _ => None,
96            })
97            .collect();
98
99        let mut system_parts: Vec<Value> =
100            req.system.iter().map(|s| json!({ "text": s })).collect();
101        let mut contents = Vec::new();
102        for message in &req.messages {
103            if message.role == Role::System {
104                system_parts.push(json!({ "text": message.text_content() }));
105                continue;
106            }
107            let role = if message.role == Role::Assistant {
108                "model"
109            } else {
110                "user"
111            };
112            let mut parts = Vec::new();
113            for part in &message.parts {
114                if let Some(p) = to_gemini_part(part, &names)? {
115                    parts.push(p);
116                }
117            }
118            contents.push(json!({ "role": role, "parts": parts }));
119        }
120
121        let mut body = json!({ "contents": contents });
122        if !system_parts.is_empty() {
123            body["systemInstruction"] = json!({ "parts": system_parts });
124        }
125        if !req.tools.is_empty() {
126            let declarations: Vec<Value> = req
127                .tools
128                .iter()
129                .map(|t| {
130                    json!({
131                        "name": t.name,
132                        "description": t.description,
133                        "parameters": t.parameters,
134                    })
135                })
136                .collect();
137            body["tools"] = json!([{ "functionDeclarations": declarations }]);
138        }
139        let mut config = serde_json::Map::new();
140        if let Some(max_tokens) = req.max_tokens {
141            config.insert("maxOutputTokens".into(), json!(max_tokens));
142        }
143        if let Some(temperature) = req.temperature {
144            config.insert("temperature".into(), json!(temperature));
145        }
146        if let Some(format) = &req.response_format {
147            config.insert("responseMimeType".into(), json!("application/json"));
148            config.insert("responseJsonSchema".into(), format.schema.clone());
149        }
150        if let Some(effort) = req.reasoning {
151            // Gemini 2.x sets depth with a token budget; Gemini 3 and later with a level.
152            let mut thinking = json!({ "includeThoughts": true });
153            if Self::model_name(&req.model).starts_with("gemini-2") {
154                thinking["thinkingBudget"] = json!(match effort {
155                    crate::ReasoningEffort::Low => 1024,
156                    crate::ReasoningEffort::Medium => 8192,
157                    crate::ReasoningEffort::High => 24576,
158                });
159            } else {
160                thinking["thinkingLevel"] = json!(effort.as_str().to_uppercase());
161            }
162            config.insert("thinkingConfig".into(), thinking);
163        }
164        if !config.is_empty() {
165            body["generationConfig"] = Value::Object(config);
166        }
167        Ok(body)
168    }
169
170    fn model_url(&self, model: &str, method: &str) -> Result<String> {
171        match &self.vertex {
172            Some(vertex) => vertex.url("GeminiProvider", &self.base_url, "google", model, method),
173            None => {
174                let model = crate::checked_path_segment(
175                    "GeminiProvider",
176                    "model",
177                    Self::model_name(model),
178                    &[],
179                )?;
180                Ok(format!("{}/models/{model}:{method}", self.base_url))
181            }
182        }
183    }
184
185    fn post(&self, url: String) -> Result<reqwest::RequestBuilder> {
186        let builder = self.http.client.post(url);
187        Ok(if self.vertex.is_some() {
188            builder.header(
189                reqwest::header::AUTHORIZATION,
190                secret_header("GeminiProvider", &format!("Bearer {}", self.api_key))?,
191            )
192        } else {
193            builder.header(
194                "x-goog-api-key",
195                secret_header("GeminiProvider", &self.api_key)?,
196            )
197        })
198    }
199
200    fn chat_request(&self, req: &ChatRequest, stream: bool) -> Result<reqwest::RequestBuilder> {
201        crate::validate(
202            "GeminiProvider",
203            &self.api_key,
204            true,
205            &self.base_url,
206            64,
207            req,
208        )?;
209        let method = if stream {
210            "streamGenerateContent?alt=sse"
211        } else {
212            "generateContent"
213        };
214        let url = self.model_url(&req.model, method)?;
215        Ok(self.post(url)?.json(&Self::build_body(req)?))
216    }
217
218    fn model_name(model: &str) -> &str {
219        model.strip_prefix("models/").unwrap_or(model)
220    }
221}
222
223fn to_gemini_part(part: &Part, names: &HashMap<&str, &str>) -> Result<Option<Value>> {
224    let sendable_id =
225        |id: &str| (!id.is_empty() && !id.starts_with(LOCAL_ID_PREFIX)).then(|| id.to_string());
226    Ok(Some(match part {
227        Part::Text { text } if text.is_empty() => return Ok(None),
228        Part::Text { text } => json!({ "text": text }),
229        Part::Image { source } => match source {
230            ImageSource::Base64 { media_type, data } => {
231                json!({ "inlineData": { "mimeType": media_type, "data": data } })
232            }
233            ImageSource::Url { .. } => {
234                return Err(LlmError::InvalidArgument(
235                    "GeminiProvider: image URLs are not supported; use Part::image_base64".into(),
236                ))
237            }
238        },
239        // Thought parts go back exactly as received (they may carry signatures).
240        Part::Reasoning { raw, .. } if raw.is_object() => raw.clone(),
241        Part::Reasoning { .. } => return Ok(None),
242        Part::Other { raw } if raw.is_object() => raw.clone(),
243        Part::Other { .. } => return Ok(None),
244        Part::ToolCall {
245            id,
246            name,
247            arguments,
248            signature,
249        } => {
250            let mut call = json!({ "name": name, "args": arguments });
251            if let Some(id) = sendable_id(id) {
252                call["id"] = json!(id);
253            }
254            let mut part = json!({ "functionCall": call });
255            if let Some(signature) = signature {
256                part["thoughtSignature"] = json!(signature);
257            }
258            part
259        }
260        Part::ToolResult {
261            call_id,
262            content,
263            is_error,
264        } => {
265            // `response` must be an object: use the content if it is a JSON object,
266            // otherwise wrap it the way Google's examples do.
267            let response = match serde_json::from_str::<Value>(content) {
268                Ok(v @ Value::Object(_)) if !is_error => v,
269                _ if *is_error => json!({ "error": content }),
270                _ => json!({ "result": content }),
271            };
272            let mut result = json!({
273                "name": names.get(call_id.as_str()).copied().unwrap_or_default(),
274                "response": response,
275            });
276            if let Some(id) = sendable_id(call_id) {
277                result["id"] = json!(id);
278            }
279            json!({ "functionResponse": result })
280        }
281    }))
282}
283
284fn parse_response(v: &Value, requested_model: &str) -> Result<ChatResponse> {
285    let mut usage = usage(
286        &v["usageMetadata"],
287        "promptTokenCount",
288        "candidatesTokenCount",
289    );
290    // Thinking tokens are billed as output but counted separately.
291    let thoughts = v["usageMetadata"]["thoughtsTokenCount"]
292        .as_u64()
293        .unwrap_or(0);
294    usage.output_tokens = usage
295        .output_tokens
296        .saturating_add(u32::try_from(thoughts).unwrap_or(u32::MAX));
297    let model = v["modelVersion"]
298        .as_str()
299        .unwrap_or(requested_model)
300        .to_string();
301
302    let Some(candidate) = v["candidates"].get(0) else {
303        // The prompt itself was blocked: no candidates, only `promptFeedback.blockReason`.
304        if let Some(reason) = v["promptFeedback"]["blockReason"].as_str() {
305            return Ok(ChatResponse {
306                model,
307                message: Message::new(
308                    Role::Assistant,
309                    vec![Part::text(format!("prompt blocked: {reason}"))],
310                ),
311                stop_reason: StopReason::Refusal,
312                usage,
313            });
314        }
315        return Err(LlmError::InvalidResponse(
316            "missing field `candidates[0]`".into(),
317        ));
318    };
319
320    let mut parts = Vec::new();
321    let mut calls = 0usize;
322    for part in candidate["content"]["parts"]
323        .as_array()
324        .into_iter()
325        .flatten()
326    {
327        if part["thought"].as_bool() == Some(true) {
328            parts.push(Part::Reasoning {
329                text: part["text"].as_str().unwrap_or_default().to_string(),
330                raw: part.clone(),
331            });
332        } else if part["functionCall"].is_object() {
333            let call = &part["functionCall"];
334            let id = match call["id"].as_str() {
335                Some(id) if !id.is_empty() => id.to_string(),
336                _ => format!("{LOCAL_ID_PREFIX}{calls}"),
337            };
338            calls += 1;
339            parts.push(Part::ToolCall {
340                id,
341                name: call["name"].as_str().unwrap_or_default().to_string(),
342                arguments: if call["args"].is_null() {
343                    json!({})
344                } else {
345                    call["args"].clone()
346                },
347                signature: part["thoughtSignature"].as_str().map(str::to_string),
348            });
349        } else if let Some(text) = part["text"].as_str().filter(|t| !t.is_empty()) {
350            parts.push(Part::text(text));
351        } else if is_meaningful_other(part) {
352            // Parts this crate does not model (code execution, an empty text part carrying a
353            // thought signature, …) are kept and sent back unchanged.
354            parts.push(Part::Other { raw: part.clone() });
355        }
356    }
357
358    let stop_reason = match candidate["finishReason"].as_str() {
359        // Gemini reports STOP for function calls as well.
360        Some("STOP") | None if calls > 0 => StopReason::ToolUse,
361        Some("STOP") => StopReason::EndTurn,
362        Some("MAX_TOKENS") => StopReason::MaxTokens,
363        Some(
364            "SAFETY"
365            | "RECITATION"
366            | "BLOCKLIST"
367            | "PROHIBITED_CONTENT"
368            | "SPII"
369            | "IMAGE_SAFETY"
370            | "IMAGE_PROHIBITED_CONTENT"
371            | "IMAGE_RECITATION",
372        ) => StopReason::Refusal,
373        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
374    };
375
376    Ok(ChatResponse {
377        model,
378        message: Message::new(Role::Assistant, parts),
379        stop_reason,
380        usage,
381    })
382}
383
384/// An object part that is not plain empty text.
385fn is_meaningful_other(part: &Value) -> bool {
386    part.as_object()
387        .is_some_and(|o| o.keys().any(|k| k != "text" && k != "thought"))
388}
389
390/// Each streamed chunk is a partial `GenerateContentResponse`; parts are concatenated and the
391/// latest candidate fields, usage and model version are kept separately (the last chunk may
392/// carry only usage).
393struct GeminiStream {
394    model: String,
395    parts: Vec<Value>,
396    candidate: Value,
397    usage: Value,
398    model_version: Value,
399    prompt_feedback: Value,
400    calls: usize,
401    done: bool,
402}
403
404impl GeminiStream {
405    fn new(model: &str) -> Self {
406        Self {
407            model: model.to_string(),
408            parts: Vec::new(),
409            candidate: Value::Null,
410            usage: Value::Null,
411            model_version: Value::Null,
412            prompt_feedback: Value::Null,
413            calls: 0,
414            done: false,
415        }
416    }
417
418    fn push_part(&mut self, part: &Value) {
419        // Merge consecutive plain text (or consecutive thought text) into one part.
420        let mergeable = |p: &Value| {
421            p["text"].is_string()
422                && p.get("thoughtSignature").is_none()
423                && p.get("functionCall").is_none()
424        };
425        if let Some(prev) = self.parts.last_mut() {
426            if mergeable(prev) && mergeable(part) && prev["thought"] == part["thought"] {
427                if let (Value::String(prev), Some(next)) =
428                    (&mut prev["text"], part["text"].as_str())
429                {
430                    prev.push_str(next);
431                }
432                return;
433            }
434        }
435        self.parts.push(part.clone());
436    }
437}
438
439impl Accumulator for GeminiStream {
440    fn on_event(&mut self, _event: &str, data: &str) -> Result<Vec<StreamEvent>> {
441        let chunk = parse_data(data)?;
442        if !chunk.is_object() {
443            return Err(LlmError::InvalidResponse("malformed stream chunk".into()));
444        }
445        if chunk["error"].is_object() {
446            return Err(LlmError::Api(Box::new(ApiError::from_stream_error(
447                &chunk["error"],
448            ))));
449        }
450        if chunk["usageMetadata"].is_object() {
451            self.usage = chunk["usageMetadata"].clone();
452        }
453        if chunk["modelVersion"].is_string() {
454            self.model_version = chunk["modelVersion"].clone();
455        }
456        if chunk["promptFeedback"]["blockReason"].is_string() {
457            self.prompt_feedback = chunk["promptFeedback"].clone();
458            self.done = true;
459        }
460        let candidate = &chunk["candidates"][0];
461        if candidate.is_object() {
462            self.candidate = candidate.clone();
463            if candidate["finishReason"].is_string() {
464                self.done = true;
465            }
466        }
467        let mut out = Vec::new();
468        for part in chunk["candidates"][0]["content"]["parts"]
469            .as_array()
470            .into_iter()
471            .flatten()
472        {
473            if let Some(text) = part["text"].as_str().filter(|t| !t.is_empty()) {
474                out.push(if part["thought"].as_bool() == Some(true) {
475                    StreamEvent::ReasoningDelta(text.to_string())
476                } else {
477                    StreamEvent::TextDelta(text.to_string())
478                });
479            } else if part["functionCall"].is_object() {
480                let call = &part["functionCall"];
481                let name = call["name"].as_str().unwrap_or_default().to_string();
482                // Same id the final response will use.
483                let id = match call["id"].as_str() {
484                    Some(id) if !id.is_empty() => id.to_string(),
485                    _ => format!("{LOCAL_ID_PREFIX}{}", self.calls),
486                };
487                self.calls += 1;
488                out.push(StreamEvent::ToolCallStart {
489                    id: id.clone(),
490                    name,
491                });
492                out.push(StreamEvent::ToolCallDelta {
493                    id,
494                    arguments: call["args"].to_string(),
495                });
496            }
497            if part.is_object() {
498                self.push_part(part);
499            }
500        }
501        Ok(out)
502    }
503
504    fn finish(&mut self) -> Result<ChatResponse> {
505        if !self.done {
506            return Err(crate::stream::ended_early());
507        }
508        let mut response = json!({
509            "usageMetadata": self.usage,
510            "modelVersion": self.model_version,
511            "promptFeedback": self.prompt_feedback,
512        });
513        if self.candidate.is_object() {
514            let mut candidate = std::mem::take(&mut self.candidate);
515            candidate["content"] =
516                json!({ "role": "model", "parts": std::mem::take(&mut self.parts) });
517            response["candidates"] = json!([candidate]);
518        }
519        parse_response(&response, &self.model)
520    }
521}
522
523#[async_trait]
524impl LlmProvider for GeminiProvider {
525    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
526        let builder = self.chat_request(req, false)?;
527        parse_response(
528            &self.http.json(builder).await?,
529            Self::model_name(&req.model),
530        )
531    }
532
533    async fn chat_stream(&self, req: &ChatRequest) -> Result<ChatStream> {
534        let builder = self.chat_request(req, true)?;
535        let response = self.http.send(builder).await?;
536        Ok(stream::drive(
537            response,
538            Box::new(GeminiStream::new(Self::model_name(&req.model))),
539        ))
540    }
541
542    async fn list_models(&self) -> Result<Vec<ModelInfo>> {
543        if self.vertex.is_some() {
544            return Err(LlmError::InvalidArgument(
545                "GeminiProvider: list_models is not available on Vertex AI".into(),
546            ));
547        }
548        crate::validate_endpoint("GeminiProvider", &self.api_key, true, &self.base_url)?;
549        let builder = self
550            .http
551            .client
552            .get(format!("{}/models?pageSize=1000", self.base_url))
553            .header(
554                "x-goog-api-key",
555                secret_header("GeminiProvider", &self.api_key)?,
556            );
557        let v = self.http.json(builder).await?;
558        Ok(v["models"]
559            .as_array()
560            .into_iter()
561            .flatten()
562            .filter_map(|m| {
563                Some(ModelInfo {
564                    id: Self::model_name(m["name"].as_str()?).to_string(),
565                    display_name: m["displayName"].as_str().map(str::to_string),
566                })
567            })
568            .collect())
569    }
570}
571
572#[async_trait]
573impl EmbeddingProvider for GeminiProvider {
574    async fn embed(&self, model: &str, inputs: &[&str]) -> Result<Embeddings> {
575        if self.vertex.is_some() {
576            return Err(LlmError::InvalidArgument(
577                "GeminiProvider: embeddings are not available on Vertex AI in this crate".into(),
578            ));
579        }
580        crate::validate_endpoint("GeminiProvider", &self.api_key, true, &self.base_url)?;
581        validate_embed_args(model, inputs)?;
582        let model = Self::model_name(model);
583        let requests: Vec<Value> = inputs
584            .iter()
585            .map(|text| json!({ "model": format!("models/{model}"), "content": { "parts": [{ "text": text }] } }))
586            .collect();
587        let builder = self
588            .post(self.model_url(model, "batchEmbedContents")?)?
589            .json(&json!({ "requests": requests }));
590        let v = self.http.json(builder).await?;
591        Ok(Embeddings {
592            vectors: v["embeddings"]
593                .as_array()
594                .into_iter()
595                .flatten()
596                .map(|e| float_vec(&e["values"]))
597                .collect(),
598            usage: crate::Usage::default(),
599        })
600    }
601}
602
603#[cfg(test)]
604mod tests {
605    use super::*;
606    use crate::{ReasoningEffort, Tool};
607
608    #[test]
609    fn builds_request_with_tools_and_results() {
610        let req = ChatRequest::new("gemini-2.5-flash")
611            .system("sys")
612            .max_tokens(100)
613            .tool(Tool::new("get_weather", "d", json!({"type": "object"})))
614            .message(Message::user("hi"))
615            .message(Message::new(
616                Role::Assistant,
617                vec![Part::ToolCall {
618                    id: format!("{LOCAL_ID_PREFIX}0"),
619                    name: "get_weather".into(),
620                    arguments: json!({"city": "Berlin"}),
621                    signature: Some("sig".into()),
622                }],
623            ))
624            .message(Message::tool_result(
625                format!("{LOCAL_ID_PREFIX}0"),
626                "18C",
627                false,
628            ));
629        let body = GeminiProvider::build_body(&req).unwrap();
630        assert_eq!(body["systemInstruction"]["parts"][0]["text"], "sys");
631        assert_eq!(body["generationConfig"]["maxOutputTokens"], 100);
632        assert_eq!(
633            body["tools"][0]["functionDeclarations"][0]["name"],
634            "get_weather"
635        );
636        assert_eq!(body["contents"][1]["role"], "model");
637        let call_part = &body["contents"][1]["parts"][0];
638        assert_eq!(call_part["functionCall"]["args"]["city"], "Berlin");
639        assert_eq!(call_part["thoughtSignature"], "sig");
640        assert!(
641            call_part["functionCall"].get("id").is_none(),
642            "local ids are not sent"
643        );
644        let result = &body["contents"][2]["parts"][0]["functionResponse"];
645        assert_eq!(result["name"], "get_weather");
646        assert_eq!(result["response"]["result"], "18C");
647    }
648
649    #[test]
650    fn builds_images_schema_and_thinking() {
651        let req = ChatRequest::new("gemini-2.5-flash")
652            .json_schema("answer", json!({"type": "object"}))
653            .reasoning(ReasoningEffort::Medium)
654            .message(Message::user_with_image(
655                "what?",
656                Part::image_base64("image/jpeg", "AAAA"),
657            ));
658        let body = GeminiProvider::build_body(&req).unwrap();
659        assert_eq!(
660            body["contents"][0]["parts"][1]["inlineData"]["mimeType"],
661            "image/jpeg"
662        );
663        assert_eq!(
664            body["generationConfig"]["responseMimeType"],
665            "application/json"
666        );
667        assert_eq!(
668            body["generationConfig"]["thinkingConfig"]["includeThoughts"],
669            true
670        );
671
672        let url_image = ChatRequest::new("m").message(Message::user_with_image(
673            "x",
674            Part::image_url("https://e.com/a.png"),
675        ));
676        assert!(GeminiProvider::build_body(&url_image).is_err());
677    }
678
679    #[test]
680    fn parses_function_call_with_signature_and_thoughts() {
681        let v = json!({
682            "candidates": [{
683                "content": { "role": "model", "parts": [
684                    { "text": "thinking...", "thought": true },
685                    { "functionCall": { "name": "get_weather", "args": { "city": "Berlin" } }, "thoughtSignature": "abc" }
686                ]},
687                "finishReason": "STOP"
688            }],
689            "usageMetadata": { "promptTokenCount": 4, "candidatesTokenCount": 6 },
690            "modelVersion": "gemini-2.5-flash"
691        });
692        let resp = parse_response(&v, "m").unwrap();
693        assert_eq!(resp.stop_reason, StopReason::ToolUse);
694        assert_eq!(resp.text(), "");
695        assert_eq!(resp.reasoning(), "thinking...");
696        let calls = resp.tool_calls();
697        assert_eq!(calls[0].name, "get_weather");
698        assert!(calls[0].id.starts_with(LOCAL_ID_PREFIX));
699        assert_eq!(resp.usage.output_tokens, 6);
700        assert_eq!(resp.model, "gemini-2.5-flash");
701        // The signature is replayed on the next turn.
702        let replay = to_gemini_part(&resp.message.parts[1], &HashMap::new())
703            .unwrap()
704            .unwrap();
705        assert_eq!(replay["thoughtSignature"], "abc");
706    }
707
708    #[test]
709    fn blocked_prompt_is_a_refusal() {
710        let v = json!({ "promptFeedback": { "blockReason": "SAFETY" } });
711        let resp = parse_response(&v, "m").unwrap();
712        assert_eq!(resp.stop_reason, StopReason::Refusal);
713        assert_eq!(resp.text(), "prompt blocked: SAFETY");
714    }
715
716    #[test]
717    fn stream_merges_text_chunks() {
718        let mut acc = GeminiStream::new("gemini-2.5-flash");
719        acc.on_event(
720            "",
721            r#"{"candidates":[{"content":{"role":"model","parts":[{"text":"Hel"}]}}]}"#,
722        )
723        .unwrap();
724        acc.on_event(
725            "",
726            r#"{"candidates":[{"content":{"role":"model","parts":[{"text":"lo"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3}}"#,
727        )
728        .unwrap();
729        let resp = acc.finish().unwrap();
730        assert_eq!(resp.text(), "Hello");
731        assert_eq!(resp.message.parts.len(), 1);
732        assert_eq!(resp.stop_reason, StopReason::EndTurn);
733        assert_eq!(resp.usage.output_tokens, 3);
734    }
735
736    #[test]
737    fn rejects_model_that_would_change_the_url() {
738        let p = GeminiProvider::new("k");
739        for model in ["../files", "gemini?key=x", "a/b", "gemini pro"] {
740            assert!(p.model_url(model, "generateContent").is_err(), "{model}");
741        }
742        assert!(p
743            .model_url("models/gemini-2.5-flash", "generateContent")
744            .unwrap()
745            .ends_with("/models/gemini-2.5-flash:generateContent"));
746    }
747
748    #[test]
749    fn vertex_url() {
750        let p = GeminiProvider::vertex("proj", "global", "token");
751        assert_eq!(
752            p.model_url("gemini-2.5-flash", "generateContent").unwrap(),
753            "https://aiplatform.googleapis.com/v1/projects/proj/locations/global/publishers/google/models/gemini-2.5-flash:generateContent"
754        );
755    }
756}
757
758#[cfg(test)]
759mod regression_tests {
760    use super::*;
761    use crate::ReasoningEffort;
762
763    #[test]
764    fn effort_maps_to_budget_or_level() {
765        let req = |model: &str| {
766            ChatRequest::new(model)
767                .reasoning(ReasoningEffort::High)
768                .message(Message::user("x"))
769        };
770        let v2 = GeminiProvider::build_body(&req("gemini-2.5-pro")).unwrap();
771        assert_eq!(
772            v2["generationConfig"]["thinkingConfig"]["thinkingBudget"],
773            24576
774        );
775        let v3 = GeminiProvider::build_body(&req("gemini-3-pro")).unwrap();
776        assert_eq!(
777            v3["generationConfig"]["thinkingConfig"]["thinkingLevel"],
778            "HIGH"
779        );
780    }
781
782    #[test]
783    fn thought_tokens_count_as_output() {
784        let v = json!({
785            "candidates": [{ "content": { "parts": [{ "text": "hi" }] }, "finishReason": "STOP" }],
786            "usageMetadata": { "promptTokenCount": 1, "candidatesTokenCount": 2, "thoughtsTokenCount": 30 }
787        });
788        assert_eq!(parse_response(&v, "m").unwrap().usage.output_tokens, 32);
789    }
790
791    #[test]
792    fn unknown_parts_and_bare_signatures_are_replayed() {
793        let v = json!({
794            "candidates": [{ "content": { "parts": [
795                { "text": "answer" },
796                { "text": "", "thoughtSignature": "sig" },
797                { "executableCode": { "code": "print(1)" } },
798                { "text": "" }
799            ] }, "finishReason": "STOP" }]
800        });
801        let resp = parse_response(&v, "m").unwrap();
802        assert_eq!(resp.message.parts.len(), 3, "plain empty text is dropped");
803        let replay: Vec<Value> = resp
804            .message
805            .parts
806            .iter()
807            .filter_map(|p| to_gemini_part(p, &HashMap::new()).unwrap())
808            .collect();
809        assert_eq!(replay[1]["thoughtSignature"], "sig");
810        assert_eq!(replay[2]["executableCode"]["code"], "print(1)");
811    }
812
813    #[test]
814    fn stream_keeps_text_when_last_chunk_has_only_usage() {
815        let mut acc = GeminiStream::new("m");
816        acc.on_event(
817            "",
818            r#"{"candidates":[{"content":{"parts":[{"text":"Hi"}]},"finishReason":"STOP"}]}"#,
819        )
820        .unwrap();
821        acc.on_event(
822            "",
823            r#"{"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1}}"#,
824        )
825        .unwrap();
826        let resp = acc.finish().unwrap();
827        assert_eq!(resp.text(), "Hi");
828        assert_eq!(resp.usage.output_tokens, 1);
829    }
830
831    #[test]
832    fn stream_tool_call_ids_match_the_final_response() {
833        let mut acc = GeminiStream::new("m");
834        let events = acc
835            .on_event(
836                "",
837                r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"f","args":{}}}]},"finishReason":"STOP"}]}"#,
838            )
839            .unwrap();
840        let StreamEvent::ToolCallStart { id, .. } = &events[0] else {
841            panic!()
842        };
843        let resp = acc.finish().unwrap();
844        assert_eq!(&resp.tool_calls()[0].id, id);
845    }
846}
llm_universal/gemini.rs

llm_universal/
gemini.rs
