Vakyartha documentation

1//! Google Gemini adapter (`models/{model}:streamGenerateContent?alt=sse`).
2//!
3//! Wire quirks handled here: roles are `user`/`model`; function responses
4//! ride in a user turn keyed by function NAME (resolved from prior assistant
5//! `functionCall` parts); tool args are JSON objects, not strings; SSE
6//! chunks carry complete `functionCall` objects.
7
8use futures::StreamExt;
9use serde_json::Value;
10use tokio_util::sync::CancellationToken;
11
12use crate::Provider;
13use crate::error::LlmError;
14use crate::gate::ProviderGate;
15use crate::sse::SseDecoder;
16use crate::stream::{EventStream, StreamEvent, channel};
17use crate::turn::{current_turn_boundary, strip_thinking};
18use crate::types::{
19 AssistantMessage, ChatRequest, ContentBlock, Message, Role, StopReason, ToolDefinition,
20};
21
22pub const GOOGLE_DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
23
24#[derive(Debug, Clone)]
25pub struct GoogleConfig {
26 pub api_key: String,
27 pub base_url: String,
28}
29
30#[derive(Clone)]
31pub struct GoogleProvider {
32 http: reqwest::Client,
33 config: GoogleConfig,
34 gate: ProviderGate,
35}
36
37impl GoogleProvider {
38 pub fn new(config: GoogleConfig) -> Result<Self, LlmError> {
39 let http = reqwest::Client::builder()
40 .connect_timeout(std::time::Duration::from_secs(30))
41 .build()
42 .map_err(|e| LlmError::Network(e.to_string()))?;
43 Ok(GoogleProvider {
44 gate: ProviderGate::new(&config.base_url, &config.api_key),
45 http,
46 config,
47 })
48 }
49}
50
51pub fn build_body(request: &ChatRequest) -> Result<Value, LlmError> {
52 let boundary = current_turn_boundary(&request.messages);
53 let mut contents: Vec<Value> = Vec::with_capacity(request.messages.len());
54 // function_call ids → names, resolved while walking history so
55 // functionResponse parts can be keyed correctly.
56 let mut id_to_name: std::collections::HashMap<String, String> =
57 std::collections::HashMap::new();
58
59 for (i, m) in request.messages.iter().enumerate() {
60 let stripped;
61 let m: &Message = if i < boundary {
62 stripped = strip_thinking(m);
63 &stripped
64 } else {
65 m
66 };
67 match m.role {
68 Role::User => {
69 let mut parts: Vec<Value> = Vec::new();
70 let mut text = String::new();
71 let mut images: Vec<&crate::types::ImageSource> = Vec::new();
72 for b in &m.content {
73 match b {
74 ContentBlock::Text { text: t } => {
75 if !text.is_empty() {
76 text.push('\n');
77 }
78 text.push_str(t);
79 }
80 ContentBlock::Image { source } => images.push(source),
81 ContentBlock::ToolResult {
82 tool_use_id,
83 content,
84 ..
85 } => {
86 let name = id_to_name
87 .get(tool_use_id)
88 .cloned()
89 .unwrap_or_else(|| tool_use_id.clone());
90 parts.push(serde_json::json!({
91 "functionResponse": {
92 "name": name,
93 "response": {"result": content},
94 }
95 }));
96 }
97 ContentBlock::ToolUse { .. } => {
98 return Err(LlmError::InvalidRequest(
99 "tool_use blocks must appear in assistant messages".into(),
100 ));
101 }
102 ContentBlock::Thinking { .. } => {}
103 // Only the Anthropic adapter understands
104 // server-side tool search; every other adapter
105 // skips this opaque block (docs/design/68 §5/§12).
106 ContentBlock::Provider { .. } => {}
107 }
108 }
109 if !text.is_empty() {
110 parts.push(serde_json::json!({"text": text}));
111 }
112 for img in images {
113 parts.push(serde_json::json!({
114 "inline_data": {"mime_type": img.media_type, "data": img.data}
115 }));
116 }
117 if !parts.is_empty() {
118 contents.push(serde_json::json!({"role": "user", "parts": parts}));
119 }
120 }
121 Role::Assistant => {
122 let mut parts: Vec<Value> = Vec::new();
123 let mut text = String::new();
124 let mut pending_sig: Option<String> = None;
125 for b in &m.content {
126 match b {
127 ContentBlock::Text { text: t } => {
128 if !text.is_empty() {
129 text.push('\n');
130 }
131 text.push_str(t);
132 }
133 ContentBlock::Thinking { signature, .. } => {
134 if let Some(sig) = signature {
135 pending_sig = Some(sig.clone());
136 }
137 }
138 ContentBlock::ToolUse { id, name, input } => {
139 // Flush pending text first so parts preserve
140 // assistant source order.
141 if !text.is_empty() {
142 let mut text_part = serde_json::json!({"text": text});
143 if let Some(sig) = &pending_sig {
144 text_part["thoughtSignature"] = serde_json::json!(sig);
145 }
146 parts.push(text_part);
147 text = String::new();
148 }
149 id_to_name.insert(id.clone(), name.clone());
150 let mut call_part = serde_json::json!({
151 "functionCall": {
152 "name": name,
153 "args": input,
154 }
155 });
156 if let Some(sig) = &pending_sig {
157 call_part["thoughtSignature"] = serde_json::json!(sig);
158 }
159 parts.push(call_part);
160 }
161 ContentBlock::ToolResult { .. }
162 | ContentBlock::Image { .. }
163 | ContentBlock::Provider { .. } => {}
164 }
165 }
166 if !text.is_empty() {
167 let mut text_part = serde_json::json!({"text": text});
168 if let Some(sig) = &pending_sig {
169 text_part["thoughtSignature"] = serde_json::json!(sig);
170 }
171 parts.push(text_part);
172 }
173 if !parts.is_empty() {
174 contents.push(serde_json::json!({"role": "model", "parts": parts}));
175 }
176 }
177 }
178 }
179
180 let mut body = serde_json::json!({ "contents": contents });
181 if let Some(system) = &request.system {
182 body["systemInstruction"] = serde_json::json!({"parts": [{"text": system}]});
183 }
184 if !request.tools.is_empty() {
185 let decls: Vec<Value> = request
186 .tools
187 .iter()
188 .map(|t: &ToolDefinition| {
189 serde_json::json!({
190 "name": t.name,
191 "description": t.description,
192 "parameters": sanitize_schema(&t.parameters),
193 })
194 })
195 .collect();
196 body["tools"] = serde_json::json!([{ "functionDeclarations": decls }]);
197 }
198 Ok(body)
199}
200
201fn sanitize_schema(val: &Value) -> Value {
202 match val {
203 Value::Object(map) => {
204 let mut cleaned = serde_json::Map::new();
205 for (k, v) in map {
206 if matches!(
207 k.as_str(),
208 "additionalProperties"
209 | "$schema"
210 | "patternProperties"
211 | "definitions"
212 | "$defs"
213 ) {
214 continue;
215 }
216 if k == "type"
217 && let Value::Array(types) = v
218 {
219 let variants: Vec<_> = types
220 .iter()
221 .filter_map(Value::as_str)
222 .map(|kind| serde_json::json!({"type": kind}))
223 .collect();
224 if !variants.is_empty() {
225 cleaned.insert("anyOf".into(), Value::Array(variants));
226 }
227 continue;
228 }
229 cleaned.insert(k.clone(), sanitize_schema(v));
230 }
231 Value::Object(cleaned)
232 }
233 Value::Array(arr) => Value::Array(arr.iter().map(sanitize_schema).collect()),
234 other => other.clone(),
235 }
236}
237
238fn map_status_error(status: u16, body: &str, retry_after: Option<u64>) -> LlmError {
239 let message = serde_json::from_str::<Value>(body)
240 .ok()
241 .and_then(|v| {
242 v.pointer("/error/message")
243 .and_then(|m| m.as_str().map(String::from))
244 })
245 .unwrap_or_else(|| body.chars().take(500).collect());
246 match status {
247 401 | 403 => LlmError::Auth(message),
248 400 => LlmError::classify_400(message),
249 404 | 413 | 422 => LlmError::InvalidRequest(message),
250 429 => LlmError::RateLimit {
251 message,
252 retry_after_secs: retry_after,
253 },
254 503 | 529 => LlmError::Overloaded(message),
255 _ => LlmError::Api { status, message },
256 }
257}
258
259struct Accumulator {
260 message: AssistantMessage,
261 saw_finish: bool,
262}
263
264impl Accumulator {
265 fn new(model: &str) -> Self {
266 Accumulator {
267 message: AssistantMessage::empty(model),
268 saw_finish: false,
269 }
270 }
271
272 fn convert(&mut self, data: &str) -> Result<Option<StreamEvent>, LlmError> {
273 let v: Value =
274 serde_json::from_str(data).map_err(|e| LlmError::Parse(format!("bad chunk: {e}")))?;
275
276 if let Some(err) = v.get("error") {
277 return Err(LlmError::Api {
278 status: 0,
279 message: err
280 .get("message")
281 .and_then(|m| m.as_str())
282 .unwrap_or("unknown gemini error")
283 .to_string(),
284 });
285 }
286
287 if let Some(usage) = v.get("usageMetadata") {
288 // `promptTokenCount` is the whole prompt, cache hits included
289 // (`cachedContentTokenCount` is a subset of it, not an
290 // addition). Normalized `input_tokens` is only the non-cached
291 // remainder, matching every other adapter
292 // (docs/design/68-context-engine.md §1).
293 if let Some(prompt_tokens) = usage.get("promptTokenCount").and_then(|x| x.as_u64()) {
294 let cached_tokens = usage
295 .get("cachedContentTokenCount")
296 .and_then(|x| x.as_u64())
297 .unwrap_or(0);
298 self.message.usage.input_tokens = prompt_tokens.saturating_sub(cached_tokens);
299 self.message.usage.cache_read_input_tokens =
300 (cached_tokens > 0).then_some(cached_tokens);
301 }
302 self.message.usage.output_tokens = usage
303 .get("candidatesTokenCount")
304 .and_then(|x| x.as_u64())
305 .unwrap_or(self.message.usage.output_tokens);
306 }
307
308 let Some(candidate) = v.pointer("/candidates/0") else {
309 return Ok(None);
310 };
311
312 let mut event: Option<StreamEvent> = None;
313
314 if let Some(parts) = candidate
315 .pointer("/content/parts")
316 .and_then(|p| p.as_array())
317 {
318 for part in parts {
319 let sig = part
320 .get("thoughtSignature")
321 .or_else(|| part.get("thought_signature"))
322 .or_else(|| {
323 part.get("functionCall").and_then(|c| {
324 c.get("thoughtSignature")
325 .or_else(|| c.get("thought_signature"))
326 })
327 })
328 .and_then(|s| s.as_str())
329 .map(String::from);
330
331 if let Some(signature) = sig {
332 self.message.content.push(ContentBlock::Thinking {
333 text: String::new(),
334 signature: Some(signature),
335 });
336 }
337
338 if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
339 if text.is_empty() {
340 continue;
341 }
342 append_text_block(&mut self.message.content, text);
343 event = Some(StreamEvent::TextDelta {
344 delta: text.to_string(),
345 partial: self.message.clone(),
346 });
347 }
348 if let Some(call) = part.get("functionCall") {
349 let name = call
350 .get("name")
351 .and_then(|n| n.as_str())
352 .unwrap_or_default()
353 .to_string();
354 let input = call
355 .get("args")
356 .cloned()
357 .unwrap_or(Value::Object(Default::default()));
358 let pos = self.message.content.len();
359 self.message.content.push(ContentBlock::ToolUse {
360 id: format!("gemini-call-{pos}"),
361 name: name.clone(),
362 input,
363 });
364 self.message.stop_reason = StopReason::ToolUse;
365 event = Some(StreamEvent::ToolUseStart {
366 index: pos,
367 id: format!("gemini-call-{pos}"),
368 name,
369 partial: self.message.clone(),
370 });
371 }
372 }
373 }
374
375 if let Some(finish) = candidate.get("finishReason").and_then(|f| f.as_str()) {
376 if self.message.stop_reason != StopReason::ToolUse {
377 self.message.stop_reason = match finish {
378 "MAX_TOKENS" => StopReason::MaxTokens,
379 _ => StopReason::EndTurn,
380 };
381 }
382 self.saw_finish = true;
383 return Ok(Some(StreamEvent::End {
384 message: self.message.clone(),
385 }));
386 }
387
388 Ok(event)
389 }
390}
391
392fn append_text_block(content: &mut Vec<ContentBlock>, text: &str) {
393 if let Some(ContentBlock::Text { text: last }) = content.last_mut() {
394 last.push_str(text);
395 return;
396 }
397 content.push(ContentBlock::text(text));
398}
399
400#[async_trait::async_trait]
401impl Provider for GoogleProvider {
402 fn name(&self) -> &str {
403 "google"
404 }
405
406 fn circuit_key(&self) -> String {
407 crate::gate::route_identity(self.name(), &self.config.base_url, &self.config.api_key)
408 }
409
410 async fn stream(
411 &self,
412 request: ChatRequest,
413 cancel: CancellationToken,
414 ) -> Result<EventStream, LlmError> {
415 let provider_permit = self.gate.acquire(&cancel).await?;
416 let url = format!(
417 "{}/models/{}:streamGenerateContent?alt=sse",
418 self.config.base_url.trim_end_matches('/'),
419 request.model
420 );
421 let body = build_body(&request)?;
422 let send_fut = self
423 .http
424 .post(&url)
425 .header("x-goog-api-key", &self.config.api_key)
426 .json(&body)
427 .send();
428 let response = tokio::select! {
429 _ = cancel.cancelled() => return Err(LlmError::Aborted { partial: None }),
430 r = send_fut => match r {
431 Ok(r) => r,
432 Err(e) => return Err(LlmError::Network(e.to_string())),
433 },
434 };
435
436 let status = response.status();
437 if !status.is_success() {
438 let retry_after = response
439 .headers()
440 .get("retry-after")
441 .and_then(|value| value.to_str().ok())
442 .and_then(|value| value.parse::<u64>().ok());
443 let text = response.text().await.unwrap_or_default();
444 return Err(map_status_error(status.as_u16(), &text, retry_after));
445 }
446
447 let model = request.model.clone();
448 let (mut sink, stream_rx) = channel(256);
449 let mut byte_stream = response.bytes_stream();
450 let mut decoder = SseDecoder::new();
451 let mut acc = Accumulator::new(&model);
452
453 tokio::spawn(async move {
454 loop {
455 tokio::select! {
456 _ = cancel.cancelled() => {
457 let partial = (!acc.message.content.is_empty()).then(|| Box::new(acc.message.clone()));
458 sink.close_error(LlmError::Aborted { partial }).await;
459 return;
460 }
461 chunk = byte_stream.next() => {
462 match chunk {
463 Some(Ok(bytes)) => {
464 decoder.push(&bytes);
465 while let Some(frame) = decoder.next_frame() {
466 match acc.convert(&frame.data) {
467 Ok(Some(event)) => sink.push(event),
468 Ok(None) => {}
469 Err(e) => {
470 sink.close_error(e).await;
471 return;
472 }
473 }
474 }
475 }
476 Some(Err(e)) => {
477 sink.close_error(LlmError::Network(e.to_string())).await;
478 return;
479 }
480 None => {
481 if acc.saw_finish {
482 sink.close_message(acc.message.clone()).await;
483 } else {
484 sink.close_error(LlmError::Parse(
485 "stream closed before finishReason".into(),
486 )).await;
487 }
488 return;
489 }
490 }
491 }
492 }
493 }
494 });
495
496 Ok(stream_rx.with_guard(provider_permit))
497 }
498}
499
500#[cfg(test)]
501mod build_body_tests {
502 #![allow(clippy::unwrap_used, clippy::expect_used)]
503 use super::*;
504 use crate::types::Role;
505
506 #[test]
507 fn provider_tool_schema_represents_json_type_unions_with_any_of() {
508 let input = serde_json::json!({
509 "type": "object",
510 "properties": {"value": {"type": ["string", "number"]}}
511 });
512 let output = sanitize_schema(&input);
513 assert_eq!(
514 output["properties"]["value"]["anyOf"],
515 serde_json::json!([{"type": "string"}, {"type": "number"}])
516 );
517 assert!(output["properties"]["value"].get("type").is_none());
518 }
519
520 #[test]
521 fn thought_signature_round_trips_into_the_request_body() {
522 let mut req = ChatRequest::new("gemini-3-pro");
523 req.messages = vec![
524 Message::user_text("do the thing"),
525 Message::assistant(vec![
526 ContentBlock::Thinking {
527 text: String::new(),
528 signature: Some("thought-sig-abc".into()),
529 },
530 ContentBlock::ToolUse {
531 id: "t1".into(),
532 name: "search".into(),
533 input: serde_json::json!({"q": "x"}),
534 },
535 ]),
536 ];
537 let body = build_body(&req).unwrap();
538 let parts = body["contents"][1]["parts"].as_array().unwrap();
539 let call_part = parts
540 .iter()
541 .find(|p| p.get("functionCall").is_some())
542 .unwrap();
543 assert_eq!(call_part["thoughtSignature"], "thought-sig-abc");
544 }
545
546 #[test]
547 fn thinking_before_the_current_turn_boundary_is_stripped() {
548 let mut req = ChatRequest::new("gemini-3-pro");
549 req.messages = vec![
550 Message::user_text("first"),
551 Message::assistant(vec![
552 ContentBlock::Thinking {
553 text: String::new(),
554 signature: Some("old-sig".into()),
555 },
556 ContentBlock::ToolUse {
557 id: "t1".into(),
558 name: "search".into(),
559 input: serde_json::json!({}),
560 },
561 ]),
562 Message::user_text("second, a fresh directive"),
563 ];
564 let body = build_body(&req).unwrap();
565 let parts = body["contents"][1]["parts"].as_array().unwrap();
566 assert!(
567 parts.iter().all(|p| p.get("thoughtSignature").is_none()),
568 "thinking from a closed turn must not carry a thought signature"
569 );
570 }
571
572 #[test]
573 fn tool_result_only_message_does_not_start_a_new_turn() {
574 let mut req = ChatRequest::new("gemini-3-pro");
575 req.messages = vec![
576 Message::user_text("do the thing"),
577 Message::assistant(vec![
578 ContentBlock::Thinking {
579 text: String::new(),
580 signature: Some("sig".into()),
581 },
582 ContentBlock::ToolUse {
583 id: "t1".into(),
584 name: "search".into(),
585 input: serde_json::json!({}),
586 },
587 ]),
588 Message {
589 role: Role::User,
590 content: vec![ContentBlock::tool_result("t1", "result")],
591 },
592 ];
593 let body = build_body(&req).unwrap();
594 let parts = body["contents"][1]["parts"].as_array().unwrap();
595 let call_part = parts
596 .iter()
597 .find(|p| p.get("functionCall").is_some())
598 .unwrap();
599 assert_eq!(call_part["thoughtSignature"], "sig");
600 }
601}
602

1//! Google Gemini adapter (`models/{model}:streamGenerateContent?alt=sse`).
2//!
3//! Wire quirks handled here: roles are `user`/`model`; function responses
4//! ride in a user turn keyed by function NAME (resolved from prior assistant
5//! `functionCall` parts); tool args are JSON objects, not strings; SSE
6//! chunks carry complete `functionCall` objects.
7
8use futures::StreamExt;
9use serde_json::Value;
10use tokio_util::sync::CancellationToken;
11
12use crate::Provider;
13use crate::error::LlmError;
14use crate::gate::ProviderGate;
15use crate::sse::SseDecoder;
16use crate::stream::{EventStream, StreamEvent, channel};
17use crate::turn::{current_turn_boundary, strip_thinking};
18use crate::types::{
19 AssistantMessage, ChatRequest, ContentBlock, Message, Role, StopReason, ToolDefinition,
20};
21
22pub const GOOGLE_DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
23
24#[derive(Debug, Clone)]
25pub struct GoogleConfig {
26 pub api_key: String,
27 pub base_url: String,
28}
29
30#[derive(Clone)]
31pub struct GoogleProvider {
32 http: reqwest::Client,
33 config: GoogleConfig,
34 gate: ProviderGate,
35}
36
37impl GoogleProvider {
38 pub fn new(config: GoogleConfig) -> Result<Self, LlmError> {
39 let http = reqwest::Client::builder()
40 .connect_timeout(std::time::Duration::from_secs(30))
41 .build()
42 .map_err(|e| LlmError::Network(e.to_string()))?;
43 Ok(GoogleProvider {
44 gate: ProviderGate::new(&config.base_url, &config.api_key),
45 http,
46 config,
47 })
48 }
49}
50
51pub fn build_body(request: &ChatRequest) -> Result<Value, LlmError> {
52 let boundary = current_turn_boundary(&request.messages);
53 let mut contents: Vec<Value> = Vec::with_capacity(request.messages.len());
54 // function_call ids → names, resolved while walking history so
55 // functionResponse parts can be keyed correctly.
56 let mut id_to_name: std::collections::HashMap<String, String> =
57 std::collections::HashMap::new();
58
59 for (i, m) in request.messages.iter().enumerate() {
60 let stripped;
61 let m: &Message = if i < boundary {
62 stripped = strip_thinking(m);
63 &stripped
64 } else {
65 m
66 };
67 match m.role {
68 Role::User => {
69 let mut parts: Vec<Value> = Vec::new();
70 let mut text = String::new();
71 let mut images: Vec<&crate::types::ImageSource> = Vec::new();
72 for b in &m.content {
73 match b {
74 ContentBlock::Text { text: t } => {
75 if !text.is_empty() {
76 text.push('\n');
77 }
78 text.push_str(t);
79 }
80 ContentBlock::Image { source } => images.push(source),
81 ContentBlock::ToolResult {
82 tool_use_id,
83 content,
84 ..
85 } => {
86 let name = id_to_name
87 .get(tool_use_id)
88 .cloned()
89 .unwrap_or_else(|| tool_use_id.clone());
90 parts.push(serde_json::json!({
91 "functionResponse": {
92 "name": name,
93 "response": {"result": content},
94 }
95 }));
96 }
97 ContentBlock::ToolUse { .. } => {
98 return Err(LlmError::InvalidRequest(
99 "tool_use blocks must appear in assistant messages".into(),
100 ));
101 }
102 ContentBlock::Thinking { .. } => {}
103 // Only the Anthropic adapter understands
104 // server-side tool search; every other adapter
105 // skips this opaque block (docs/design/68 §5/§12).
106 ContentBlock::Provider { .. } => {}
107 }
108 }
109 if !text.is_empty() {
110 parts.push(serde_json::json!({"text": text}));
111 }
112 for img in images {
113 parts.push(serde_json::json!({
114 "inline_data": {"mime_type": img.media_type, "data": img.data}
115 }));
116 }
117 if !parts.is_empty() {
118 contents.push(serde_json::json!({"role": "user", "parts": parts}));
119 }
120 }
121 Role::Assistant => {
122 let mut parts: Vec<Value> = Vec::new();
123 let mut text = String::new();
124 let mut pending_sig: Option<String> = None;
125 for b in &m.content {
126 match b {
127 ContentBlock::Text { text: t } => {
128 if !text.is_empty() {
129 text.push('\n');
130 }
131 text.push_str(t);
132 }
133 ContentBlock::Thinking { signature, .. } => {
134 if let Some(sig) = signature {
135 pending_sig = Some(sig.clone());
136 }
137 }
138 ContentBlock::ToolUse { id, name, input } => {
139 // Flush pending text first so parts preserve
140 // assistant source order.
141 if !text.is_empty() {
142 let mut text_part = serde_json::json!({"text": text});
143 if let Some(sig) = &pending_sig {
144 text_part["thoughtSignature"] = serde_json::json!(sig);
145 }
146 parts.push(text_part);
147 text = String::new();
148 }
149 id_to_name.insert(id.clone(), name.clone());
150 let mut call_part = serde_json::json!({
151 "functionCall": {
152 "name": name,
153 "args": input,
154 }
155 });
156 if let Some(sig) = &pending_sig {
157 call_part["thoughtSignature"] = serde_json::json!(sig);
158 }
159 parts.push(call_part);
160 }
161 ContentBlock::ToolResult { .. }
162 | ContentBlock::Image { .. }
163 | ContentBlock::Provider { .. } => {}
164 }
165 }
166 if !text.is_empty() {
167 let mut text_part = serde_json::json!({"text": text});
168 if let Some(sig) = &pending_sig {
169 text_part["thoughtSignature"] = serde_json::json!(sig);
170 }
171 parts.push(text_part);
172 }
173 if !parts.is_empty() {
174 contents.push(serde_json::json!({"role": "model", "parts": parts}));
175 }
176 }
177 }
178 }
179
180 let mut body = serde_json::json!({ "contents": contents });
181 if let Some(system) = &request.system {
182 body["systemInstruction"] = serde_json::json!({"parts": [{"text": system}]});
183 }
184 if !request.tools.is_empty() {
185 let decls: Vec<Value> = request
186 .tools
187 .iter()
188 .map(|t: &ToolDefinition| {
189 serde_json::json!({
190 "name": t.name,
191 "description": t.description,
192 "parameters": sanitize_schema(&t.parameters),
193 })
194 })
195 .collect();
196 body["tools"] = serde_json::json!([{ "functionDeclarations": decls }]);
197 }
198 Ok(body)
199}
200
201fn sanitize_schema(val: &Value) -> Value {
202 match val {
203 Value::Object(map) => {
204 let mut cleaned = serde_json::Map::new();
205 for (k, v) in map {
206 if matches!(
207 k.as_str(),
208 "additionalProperties"
209 | "$schema"
210 | "patternProperties"
211 | "definitions"
212 | "$defs"
213 ) {
214 continue;
215 }
216 if k == "type"
217 && let Value::Array(types) = v
218 {
219 let variants: Vec<_> = types
220 .iter()
221 .filter_map(Value::as_str)
222 .map(|kind| serde_json::json!({"type": kind}))
223 .collect();
224 if !variants.is_empty() {
225 cleaned.insert("anyOf".into(), Value::Array(variants));
226 }
227 continue;
228 }
229 cleaned.insert(k.clone(), sanitize_schema(v));
230 }
231 Value::Object(cleaned)
232 }
233 Value::Array(arr) => Value::Array(arr.iter().map(sanitize_schema).collect()),
234 other => other.clone(),
235 }
236}
237
238fn map_status_error(status: u16, body: &str, retry_after: Option<u64>) -> LlmError {
239 let message = serde_json::from_str::<Value>(body)
240 .ok()
241 .and_then(|v| {
242 v.pointer("/error/message")
243 .and_then(|m| m.as_str().map(String::from))
244 })
245 .unwrap_or_else(|| body.chars().take(500).collect());
246 match status {
247 401 | 403 => LlmError::Auth(message),
248 400 => LlmError::classify_400(message),
249 404 | 413 | 422 => LlmError::InvalidRequest(message),
250 429 => LlmError::RateLimit {
251 message,
252 retry_after_secs: retry_after,
253 },
254 503 | 529 => LlmError::Overloaded(message),
255 _ => LlmError::Api { status, message },
256 }
257}
258
259struct Accumulator {
260 message: AssistantMessage,
261 saw_finish: bool,
262}
263
264impl Accumulator {
265 fn new(model: &str) -> Self {
266 Accumulator {
267 message: AssistantMessage::empty(model),
268 saw_finish: false,
269 }
270 }
271
272 fn convert(&mut self, data: &str) -> Result<Option<StreamEvent>, LlmError> {
273 let v: Value =
274 serde_json::from_str(data).map_err(|e| LlmError::Parse(format!("bad chunk: {e}")))?;
275
276 if let Some(err) = v.get("error") {
277 return Err(LlmError::Api {
278 status: 0,
279 message: err
280 .get("message")
281 .and_then(|m| m.as_str())
282 .unwrap_or("unknown gemini error")
283 .to_string(),
284 });
285 }
286
287 if let Some(usage) = v.get("usageMetadata") {
288 // `promptTokenCount` is the whole prompt, cache hits included
289 // (`cachedContentTokenCount` is a subset of it, not an
290 // addition). Normalized `input_tokens` is only the non-cached
291 // remainder, matching every other adapter
292 // (docs/design/68-context-engine.md §1).
293 if let Some(prompt_tokens) = usage.get("promptTokenCount").and_then(|x| x.as_u64()) {
294 let cached_tokens = usage
295 .get("cachedContentTokenCount")
296 .and_then(|x| x.as_u64())
297 .unwrap_or(0);
298 self.message.usage.input_tokens = prompt_tokens.saturating_sub(cached_tokens);
299 self.message.usage.cache_read_input_tokens =
300 (cached_tokens > 0).then_some(cached_tokens);
301 }
302 self.message.usage.output_tokens = usage
303 .get("candidatesTokenCount")
304 .and_then(|x| x.as_u64())
305 .unwrap_or(self.message.usage.output_tokens);
306 }
307
308 let Some(candidate) = v.pointer("/candidates/0") else {
309 return Ok(None);
310 };
311
312 let mut event: Option<StreamEvent> = None;
313
314 if let Some(parts) = candidate
315 .pointer("/content/parts")
316 .and_then(|p| p.as_array())
317 {
318 for part in parts {
319 let sig = part
320 .get("thoughtSignature")
321 .or_else(|| part.get("thought_signature"))
322 .or_else(|| {
323 part.get("functionCall").and_then(|c| {
324 c.get("thoughtSignature")
325 .or_else(|| c.get("thought_signature"))
326 })
327 })
328 .and_then(|s| s.as_str())
329 .map(String::from);
330
331 if let Some(signature) = sig {
332 self.message.content.push(ContentBlock::Thinking {
333 text: String::new(),
334 signature: Some(signature),
335 });
336 }
337
338 if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
339 if text.is_empty() {
340 continue;
341 }
342 append_text_block(&mut self.message.content, text);
343 event = Some(StreamEvent::TextDelta {
344 delta: text.to_string(),
345 partial: self.message.clone(),
346 });
347 }
348 if let Some(call) = part.get("functionCall") {
349 let name = call
350 .get("name")
351 .and_then(|n| n.as_str())
352 .unwrap_or_default()
353 .to_string();
354 let input = call
355 .get("args")
356 .cloned()
357 .unwrap_or(Value::Object(Default::default()));
358 let pos = self.message.content.len();
359 self.message.content.push(ContentBlock::ToolUse {
360 id: format!("gemini-call-{pos}"),
361 name: name.clone(),
362 input,
363 });
364 self.message.stop_reason = StopReason::ToolUse;
365 event = Some(StreamEvent::ToolUseStart {
366 index: pos,
367 id: format!("gemini-call-{pos}"),
368 name,
369 partial: self.message.clone(),
370 });
371 }
372 }
373 }
374
375 if let Some(finish) = candidate.get("finishReason").and_then(|f| f.as_str()) {
376 if self.message.stop_reason != StopReason::ToolUse {
377 self.message.stop_reason = match finish {
378 "MAX_TOKENS" => StopReason::MaxTokens,
379 _ => StopReason::EndTurn,
380 };
381 }
382 self.saw_finish = true;
383 return Ok(Some(StreamEvent::End {
384 message: self.message.clone(),
385 }));
386 }
387
388 Ok(event)
389 }
390}
391
392fn append_text_block(content: &mut Vec<ContentBlock>, text: &str) {
393 if let Some(ContentBlock::Text { text: last }) = content.last_mut() {
394 last.push_str(text);
395 return;
396 }
397 content.push(ContentBlock::text(text));
398}
399
400#[async_trait::async_trait]
401impl Provider for GoogleProvider {
402 fn name(&self) -> &str {
403 "google"
404 }
405
406 fn circuit_key(&self) -> String {
407 crate::gate::route_identity(self.name(), &self.config.base_url, &self.config.api_key)
408 }
409
410 async fn stream(
411 &self,
412 request: ChatRequest,
413 cancel: CancellationToken,
414 ) -> Result<EventStream, LlmError> {
415 let provider_permit = self.gate.acquire(&cancel).await?;
416 let url = format!(
417 "{}/models/{}:streamGenerateContent?alt=sse",
418 self.config.base_url.trim_end_matches('/'),
419 request.model
420 );
421 let body = build_body(&request)?;
422 let send_fut = self
423 .http
424 .post(&url)
425 .header("x-goog-api-key", &self.config.api_key)
426 .json(&body)
427 .send();
428 let response = tokio::select! {
429 _ = cancel.cancelled() => return Err(LlmError::Aborted { partial: None }),
430 r = send_fut => match r {
431 Ok(r) => r,
432 Err(e) => return Err(LlmError::Network(e.to_string())),
433 },
434 };
435
436 let status = response.status();
437 if !status.is_success() {
438 let retry_after = response
439 .headers()
440 .get("retry-after")
441 .and_then(|value| value.to_str().ok())
442 .and_then(|value| value.parse::<u64>().ok());
443 let text = response.text().await.unwrap_or_default();
444 return Err(map_status_error(status.as_u16(), &text, retry_after));
445 }
446
447 let model = request.model.clone();
448 let (mut sink, stream_rx) = channel(256);
449 let mut byte_stream = response.bytes_stream();
450 let mut decoder = SseDecoder::new();
451 let mut acc = Accumulator::new(&model);
452
453 tokio::spawn(async move {
454 loop {
455 tokio::select! {
456 _ = cancel.cancelled() => {
457 let partial = (!acc.message.content.is_empty()).then(|| Box::new(acc.message.clone()));
458 sink.close_error(LlmError::Aborted { partial }).await;
459 return;
460 }
461 chunk = byte_stream.next() => {
462 match chunk {
463 Some(Ok(bytes)) => {
464 decoder.push(&bytes);
465 while let Some(frame) = decoder.next_frame() {
466 match acc.convert(&frame.data) {
467 Ok(Some(event)) => sink.push(event),
468 Ok(None) => {}
469 Err(e) => {
470 sink.close_error(e).await;
471 return;
472 }
473 }
474 }
475 }
476 Some(Err(e)) => {
477 sink.close_error(LlmError::Network(e.to_string())).await;
478 return;
479 }
480 None => {
481 if acc.saw_finish {
482 sink.close_message(acc.message.clone()).await;
483 } else {
484 sink.close_error(LlmError::Parse(
485 "stream closed before finishReason".into(),
486 )).await;
487 }
488 return;
489 }
490 }
491 }
492 }
493 }
494 });
495
496 Ok(stream_rx.with_guard(provider_permit))
497 }
498}
499
500#[cfg(test)]
501mod build_body_tests {
502 #![allow(clippy::unwrap_used, clippy::expect_used)]
503 use super::*;
504 use crate::types::Role;
505
506 #[test]
507 fn provider_tool_schema_represents_json_type_unions_with_any_of() {
508 let input = serde_json::json!({
509 "type": "object",
510 "properties": {"value": {"type": ["string", "number"]}}
511 });
512 let output = sanitize_schema(&input);
513 assert_eq!(
514 output["properties"]["value"]["anyOf"],
515 serde_json::json!([{"type": "string"}, {"type": "number"}])
516 );
517 assert!(output["properties"]["value"].get("type").is_none());
518 }
519
520 #[test]
521 fn thought_signature_round_trips_into_the_request_body() {
522 let mut req = ChatRequest::new("gemini-3-pro");
523 req.messages = vec![
524 Message::user_text("do the thing"),
525 Message::assistant(vec![
526 ContentBlock::Thinking {
527 text: String::new(),
528 signature: Some("thought-sig-abc".into()),
529 },
530 ContentBlock::ToolUse {
531 id: "t1".into(),
532 name: "search".into(),
533 input: serde_json::json!({"q": "x"}),
534 },
535 ]),
536 ];
537 let body = build_body(&req).unwrap();
538 let parts = body["contents"][1]["parts"].as_array().unwrap();
539 let call_part = parts
540 .iter()
541 .find(|p| p.get("functionCall").is_some())
542 .unwrap();
543 assert_eq!(call_part["thoughtSignature"], "thought-sig-abc");
544 }
545
546 #[test]
547 fn thinking_before_the_current_turn_boundary_is_stripped() {
548 let mut req = ChatRequest::new("gemini-3-pro");
549 req.messages = vec![
550 Message::user_text("first"),
551 Message::assistant(vec![
552 ContentBlock::Thinking {
553 text: String::new(),
554 signature: Some("old-sig".into()),
555 },
556 ContentBlock::ToolUse {
557 id: "t1".into(),
558 name: "search".into(),
559 input: serde_json::json!({}),
560 },
561 ]),
562 Message::user_text("second, a fresh directive"),
563 ];
564 let body = build_body(&req).unwrap();
565 let parts = body["contents"][1]["parts"].as_array().unwrap();
566 assert!(
567 parts.iter().all(|p| p.get("thoughtSignature").is_none()),
568 "thinking from a closed turn must not carry a thought signature"
569 );
570 }
571
572 #[test]
573 fn tool_result_only_message_does_not_start_a_new_turn() {
574 let mut req = ChatRequest::new("gemini-3-pro");
575 req.messages = vec![
576 Message::user_text("do the thing"),
577 Message::assistant(vec![
578 ContentBlock::Thinking {
579 text: String::new(),
580 signature: Some("sig".into()),
581 },
582 ContentBlock::ToolUse {
583 id: "t1".into(),
584 name: "search".into(),
585 input: serde_json::json!({}),
586 },
587 ]),
588 Message {
589 role: Role::User,
590 content: vec![ContentBlock::tool_result("t1", "result")],
591 },
592 ];
593 let body = build_body(&req).unwrap();
594 let parts = body["contents"][1]["parts"].as_array().unwrap();
595 let call_part = parts
596 .iter()
597 .find(|p| p.get("functionCall").is_some())
598 .unwrap();
599 assert_eq!(call_part["thoughtSignature"], "sig");
600 }
601}
602

Indexing the workspace…