Skip to main content

vtcode_llm/providers/gemini/
llm_provider.rs

1use super::helpers::InteractionStreamState;
2use super::*;
3use crate::providers::shared::{StreamAssemblyError, extract_data_payload, next_sse_event};
4
5fn normalize_stream_event(event: LLMStreamEvent, interaction_reasoning: bool) -> Vec<NormalizedStreamEvent> {
6    match event {
7        LLMStreamEvent::Reasoning { delta } if interaction_reasoning => {
8            vec![NormalizedStreamEvent::ReasoningDelta { delta, source: ReasoningSource::ProviderSummary }]
9        }
10        LLMStreamEvent::Completed { response } => normalize_completed_event(response),
11        event => event.into_normalized(),
12    }
13}
14
15fn normalize_completed_event(response: Box<LLMResponse>) -> Vec<NormalizedStreamEvent> {
16    let mut events = Vec::new();
17    if let Some(tool_calls) = response.tool_calls.as_ref() {
18        for tool_call in tool_calls {
19            events.push(NormalizedStreamEvent::ToolCallStart {
20                call_id: tool_call.id.clone(),
21                name: tool_call.tool_name().map(ToOwned::to_owned),
22            });
23            if let Some(arguments) = tool_call
24                .raw_input()
25                .filter(|arguments| !arguments.trim().is_empty() && arguments.trim() != "{}")
26            {
27                events.push(NormalizedStreamEvent::ToolCallDelta {
28                    call_id: tool_call.id.clone(),
29                    delta: arguments.to_string(),
30                });
31            }
32        }
33    }
34    events.extend(LLMStreamEvent::Completed { response }.into_normalized());
35    events
36}
37
38/// Shared transport for `generateContent` and `streamGenerateContent`.
39impl GeminiProvider {
40    async fn post_generate_content(
41        &self,
42        url: &str,
43        body: &GenerateContentRequest,
44    ) -> Result<reqwest::Response, LLMError> {
45        self.http_client
46            .post(url)
47            .header("x-goog-api-key", self.api_key.as_ref())
48            .json(body)
49            .send()
50            .await
51            .map_err(|e| format_network_error("Gemini", &e))
52    }
53
54    /// Send a generate request and return only a successful response.
55    ///
56    /// A stale `cachedContent` name (expired or evicted) is retried once
57    /// without the cache and with the full system instruction/tools resent.
58    /// Both `generate` and `stream` share this path so recovery cannot drift
59    /// between them. The dead slot is dropped before the retry so the next
60    /// turn rebuilds it.
61    async fn send_generate_request_with_cache_recovery(
62        &self,
63        url: &str,
64        gemini_request: &GenerateContentRequest,
65        request: &LLMRequest,
66    ) -> Result<reqwest::Response, LLMError> {
67        let response = self.post_generate_content(url, gemini_request).await?;
68        if response.status().is_success() {
69            return Ok(response);
70        }
71
72        let status = response.status();
73        let error_text = crate::providers::common::read_provider_error_body(response).await;
74        if gemini_request.cached_content.is_none()
75            || !explicit_cache::is_stale_cache_error(status.as_u16(), &error_text)
76        {
77            return Err(Self::handle_http_error(status, &error_text));
78        }
79
80        self.explicit_cache.clear();
81        let full_request = self.convert_to_gemini_request(request)?;
82        let retry = self.post_generate_content(url, &full_request).await?;
83        if retry.status().is_success() {
84            return Ok(retry);
85        }
86
87        let retry_status = retry.status();
88        let retry_error_text = crate::providers::common::read_provider_error_body(retry).await;
89        Err(Self::handle_http_error(retry_status, &retry_error_text))
90    }
91}
92
93#[async_trait]
94impl LLMProvider for GeminiProvider {
95    fn name(&self) -> &str {
96        "gemini"
97    }
98
99    fn supports_streaming(&self) -> bool {
100        true
101    }
102
103    fn supports_non_streaming(&self, _model: &str) -> bool {
104        // Pinned so the stream-timeout fallback cannot silently regress.
105        true
106    }
107
108    fn supports_reasoning(&self, model: &str) -> bool {
109        // Codex-inspired robustness: Setting model_supports_reasoning to false
110        // does NOT disable it for known reasoning models.
111        models::google::REASONING_MODELS.contains(&model)
112            || self
113                .model_behavior
114                .as_ref()
115                .and_then(|b| b.model_supports_reasoning)
116                .unwrap_or(false)
117    }
118
119    fn supports_reasoning_effort(&self, model: &str) -> bool {
120        // Same robustness logic for reasoning effort
121        models::google::REASONING_MODELS.contains(&model)
122            || self
123                .model_behavior
124                .as_ref()
125                .and_then(|b| b.model_supports_reasoning_effort)
126                .unwrap_or(false)
127    }
128
129    fn supports_context_caching(&self, model: &str) -> bool {
130        models::google::CACHING_MODELS.contains(&model)
131    }
132
133    fn effective_context_size(&self, model: &str) -> usize {
134        let fallback = if model.contains("gemini-4-argon") || model.contains("gemini-3.1") {
135            // Gemini 4 Argon: unified 1M limit shared by input, thinking, and output
136            1_048_576
137        } else if model.contains("3") || model.contains("1.5-pro") {
138            2_097_152
139        } else {
140            1_048_576
141        };
142        crate::provider::catalog_context_window("gemini", model, fallback)
143    }
144
145    async fn generate(&self, request: LLMRequest) -> Result<LLMResponse, LLMError> {
146        let model = request.model.clone();
147        if self.should_use_interactions(&request) {
148            let interaction_request = self.convert_to_interaction_request(&request)?;
149            let url = format!("{}/interactions", self.base_url);
150            let response = self
151                .http_client
152                .post(&url)
153                .header("x-goog-api-key", self.api_key.as_ref())
154                .json(&interaction_request)
155                .send()
156                .await
157                .map_err(|e| format_network_error("Gemini", &e))?;
158
159            if !response.status().is_success() {
160                let status = response.status();
161                let error_text = crate::providers::common::read_provider_error_body(response).await;
162                return Err(Self::handle_http_error(status, &error_text));
163            }
164
165            let interaction_response: Interaction =
166                response.json().await.map_err(|e| format_parse_error("Gemini", &e))?;
167
168            return Self::convert_from_interaction_response(interaction_response, model);
169        }
170
171        let mut gemini_request = self.convert_to_gemini_request(&request)?;
172        if let Some(cache_name) = self.ensure_explicit_cache(&request, &gemini_request).await? {
173            gemini_request = self.apply_explicit_cache_to_request(gemini_request, &cache_name);
174        }
175
176        let url = format!("{}/models/{}:generateContent", self.base_url, request.model);
177
178        let response = self
179            .send_generate_request_with_cache_recovery(&url, &gemini_request, &request)
180            .await?;
181
182        let gemini_response: GenerateContentResponse =
183            response.json().await.map_err(|e| format_parse_error("Gemini", &e))?;
184
185        Self::convert_from_gemini_response(gemini_response, model)
186    }
187
188    async fn stream(&self, request: LLMRequest) -> Result<LLMStream, LLMError> {
189        if self.should_use_interactions(&request) {
190            let model = request.model.clone();
191            let interaction_request = self.convert_to_interaction_request(&request)?;
192            let url = format!("{}/interactions?alt=sse", self.base_url);
193            let response = self
194                .http_client
195                .post(&url)
196                .header("x-goog-api-key", self.api_key.as_ref())
197                .json(&interaction_request)
198                .send()
199                .await
200                .map_err(|e| format_network_error("Gemini", &e))?;
201
202            if !response.status().is_success() {
203                let status = response.status();
204                let error_text = crate::providers::common::read_provider_error_body(response).await;
205                return Err(Self::handle_http_error(status, &error_text));
206            }
207
208            let stream = {
209                try_stream! {
210                    let mut body_stream = response.bytes_stream();
211                    let mut buf: Vec<u8> = Vec::new();
212                    let mut offset = 0usize;
213                    let mut decoder = crate::providers::shared::Utf8StreamDecoder::new();
214                    let mut state = InteractionStreamState::default();
215
216                    while let Some(chunk_result) = body_stream.next().await {
217                        let chunk = chunk_result
218                            .map_err(|err| format_network_error("Gemini", &err))?;
219
220                        decoder.push_bytes(&chunk, &mut buf);
221
222                        while let Some(event) = next_sse_event(&buf, &mut offset)
223                            .map_err(|e| {
224                                StreamAssemblyError::InvalidPayload(format!("non-utf-8 stream data: {e}"))
225                                    .into_llm_error("Gemini")
226                            })?
227                        {
228
229                            let Some(data_payload) = extract_data_payload(event) else {
230                                continue;
231                            };
232
233                            let trimmed_payload = data_payload.trim();
234                            if trimmed_payload.is_empty() || trimmed_payload == "[DONE]" {
235                                continue;
236                            }
237
238                            let payload: Value = serde_json::from_str(trimmed_payload)
239                                .map_err(|err| {
240                                    StreamAssemblyError::InvalidPayload(err.to_string())
241                                        .into_llm_error("Gemini")
242                                })?;
243
244                            for stream_event in Self::apply_interaction_stream_payload(&mut state, &payload)? {
245                                yield stream_event;
246                            }
247                        }
248
249                        // Drain the consumed prefix so `buf` stays bounded to
250                        // the unprocessed tail rather than growing for the
251                        // entire stream lifetime.
252                        if offset > 0 {
253                            buf.drain(..offset);
254                            offset = 0;
255                        }
256                    }
257
258                    if !state.completed {
259                        let formatted_error = error_display::format_llm_error(
260                            "Gemini",
261                            "Interactions stream ended without an interaction.complete event",
262                        );
263                        Err(LLMError::Provider {
264                            message: formatted_error,
265                            metadata: None,
266                        })?;
267                    }
268
269                    let response =
270                        Self::finalize_interaction_stream_state(state, model)?;
271                    yield LLMStreamEvent::Completed { response: Box::new(response) };
272                }
273            };
274            return Ok(Box::pin(stream));
275        }
276
277        let model = request.model.clone();
278        let mut gemini_request = self.convert_to_gemini_request(&request)?;
279        if let Some(cache_name) = self.ensure_explicit_cache(&request, &gemini_request).await? {
280            gemini_request = self.apply_explicit_cache_to_request(gemini_request, &cache_name);
281        }
282
283        let url = format!("{}/models/{}:streamGenerateContent", self.base_url, request.model);
284
285        let response = self
286            .send_generate_request_with_cache_recovery(&url, &gemini_request, &request)
287            .await?;
288
289        let (event_tx, event_rx) = mpsc::unbounded_channel::<Result<LLMStreamEvent, LLMError>>();
290        let completion_sender = event_tx.clone();
291
292        let streaming_timeout = self.timeouts.streaming_ceiling_seconds;
293
294        let model_clone = model.clone();
295        tokio::spawn(async move {
296            let config = StreamingConfig::with_total_timeout(streaming_timeout);
297            let mut processor = StreamingProcessor::with_config(config);
298            let event_sender = completion_sender.clone();
299            let mut aggregator = crate::providers::shared::StreamAggregator::new(model_clone.clone());
300
301            let mut on_chunk = |chunk: &str| -> Result<(), StreamingError> {
302                if chunk.is_empty() {
303                    return Ok(());
304                }
305
306                if let Some(delta) = Self::apply_stream_delta(&mut aggregator.content, chunk) {
307                    if delta.is_empty() {
308                        return Ok(());
309                    }
310
311                    for event in aggregator.sanitizer.process_chunk(&delta) {
312                        event_sender.send(Ok(event)).map_err(|_e| StreamingError::StreamingError {
313                            message: "Streaming consumer dropped".to_string(),
314                            partial_content: Some(chunk.to_string()),
315                        })?;
316                    }
317                }
318                Ok(())
319            };
320
321            let result = processor.process_stream(response, &mut on_chunk).await;
322            match result {
323                Ok(mut streaming_response) => {
324                    if streaming_response.candidates.is_empty() && !aggregator.content.trim().is_empty() {
325                        streaming_response.candidates.push(StreamingCandidate {
326                            content: Content {
327                                role: "model".to_string(),
328                                parts: vec![Part::Text {
329                                    text: aggregator.content.clone(),
330                                    thought_signature: None,
331                                }],
332                            },
333                            finish_reason: None,
334                            index: Some(0),
335                        });
336                    }
337
338                    match Self::convert_from_streaming_response(streaming_response, model_clone) {
339                        Ok(mut final_response) => {
340                            let aggregator_response = aggregator.finalize();
341                            if final_response.reasoning.is_none() {
342                                final_response.reasoning = aggregator_response.reasoning;
343                            }
344                            if final_response.content.is_none() {
345                                final_response.content = aggregator_response.content;
346                            }
347
348                            let _ = completion_sender
349                                .send(Ok(LLMStreamEvent::Completed { response: Box::new(final_response) }));
350                        }
351                        Err(err) => {
352                            let _ = completion_sender.send(Err(err));
353                        }
354                    }
355                }
356                Err(error) => {
357                    let mapped = Self::map_streaming_error(error);
358                    let _ = completion_sender.send(Err(mapped));
359                }
360            }
361        });
362
363        drop(event_tx);
364
365        let stream = {
366            let mut receiver = event_rx;
367            try_stream! {
368                while let Some(event) = receiver.recv().await {
369                    yield event?;
370                }
371            }
372        };
373
374        Ok(Box::pin(stream))
375    }
376
377    async fn stream_normalized(&self, request: LLMRequest) -> Result<LLMNormalizedStream, LLMError> {
378        let interaction_reasoning = self.should_use_interactions(&request);
379        let mut legacy_stream = self.stream(request).await?;
380        let stream = try_stream! {
381            while let Some(event) = legacy_stream.next().await {
382                for normalized in normalize_stream_event(event?, interaction_reasoning) {
383                    yield normalized;
384                }
385            }
386        };
387
388        Ok(Box::pin(stream))
389    }
390
391    fn supported_models(&self) -> Vec<String> {
392        models::google::SUPPORTED_MODELS.iter().map(|s| s.to_string()).collect()
393    }
394
395    fn validate_request(&self, request: &LLMRequest) -> Result<(), LLMError> {
396        if GeminiProvider::uses_latest_gemini_api(&request.model) {
397            if request.temperature.is_some() || request.top_p.is_some() || request.top_k.is_some() {
398                tracing::warn!(
399                    model = %request.model,
400                    temperature = ?request.temperature,
401                    top_p = ?request.top_p,
402                    top_k = ?request.top_k,
403                    "Sampling parameters (temperature, top_p, top_k) are deprecated for this Gemini model and will be ignored by the API"
404                );
405            }
406        }
407
408        if request.previous_response_id.is_some() && request.response_store == Some(false) {
409            let formatted_error = error_display::format_llm_error(
410                "Gemini",
411                "Interactions with previous_interaction_id cannot set store=false",
412            );
413            return Err(LLMError::InvalidRequest { message: formatted_error, metadata: None });
414        }
415
416        if !models::google::SUPPORTED_MODELS.iter().any(|m| *m == request.model) {
417            let formatted_error =
418                error_display::format_llm_error("Gemini", &format!("Unsupported model: {}", request.model));
419            return Err(LLMError::InvalidRequest { message: formatted_error, metadata: None });
420        }
421
422        if let Some(max_tokens) = request.max_tokens {
423            let model = request.model.as_str();
424            let max_output_tokens = if model.contains("3") || model.contains("4") {
425                65536
426            } else {
427                8192
428            };
429
430            if max_tokens > max_output_tokens {
431                let formatted_error = error_display::format_llm_error(
432                    "Gemini",
433                    &format!(
434                        "Requested max_tokens ({max_tokens}) exceeds model limit ({max_output_tokens}) for {model}"
435                    ),
436                );
437                return Err(LLMError::InvalidRequest { message: formatted_error, metadata: None });
438            }
439        }
440
441        Ok(())
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::{LLMStreamEvent, NormalizedStreamEvent, ReasoningSource, normalize_stream_event};
448    use crate::provider::{LLMResponse, ToolCall};
449
450    #[test]
451    fn interaction_reasoning_is_marked_as_public_summary() {
452        let events = normalize_stream_event(LLMStreamEvent::Reasoning { delta: "summary".to_string() }, true);
453
454        assert!(matches!(
455            events.as_slice(),
456            [NormalizedStreamEvent::ReasoningDelta { delta, source }]
457                if delta == "summary" && *source == ReasoningSource::ProviderSummary
458        ));
459    }
460
461    #[test]
462    fn standard_reasoning_remains_unclassified() {
463        let events = normalize_stream_event(LLMStreamEvent::Reasoning { delta: "trace".to_string() }, false);
464
465        assert!(matches!(
466            events.as_slice(),
467            [NormalizedStreamEvent::ReasoningDelta { delta, source }]
468                if delta == "trace" && *source == ReasoningSource::Unknown
469        ));
470    }
471
472    #[test]
473    fn completed_tool_calls_become_structured_events() {
474        let events = normalize_stream_event(
475            LLMStreamEvent::Completed {
476                response: Box::new(LLMResponse {
477                    tool_calls: Some(vec![ToolCall::function(
478                        "call_1".to_string(),
479                        "search_workspace".to_string(),
480                        "{\"query\":\"vtcode\"}".to_string(),
481                    )]),
482                    ..Default::default()
483                }),
484            },
485            false,
486        );
487
488        assert!(matches!(
489            events.as_slice(),
490            [
491                NormalizedStreamEvent::ToolCallStart { call_id, name },
492                NormalizedStreamEvent::ToolCallDelta { call_id: delta_call_id, delta },
493                NormalizedStreamEvent::Done { .. }
494            ] if call_id == "call_1"
495                && delta_call_id == "call_1"
496                && name.as_deref() == Some("search_workspace")
497                && delta == "{\"query\":\"vtcode\"}"
498        ));
499    }
500}