1use aither_core::{
2 LanguageModel,
3 llm::{
4 Attachment, Event, GenerateError, LLMRequest, Message, ReasoningState, Role, Usage,
5 model::{Ability, Parameters, Profile, ToolChoice},
6 tool::ToolDefinition,
7 },
8};
9use futures_core::Stream;
10use futures_lite::StreamExt;
11use schemars::JsonSchema;
12use serde::de::DeserializeOwned;
13use std::collections::HashMap;
14use tracing::debug;
15
16use crate::PROVIDER_NAME;
17use crate::{
18 client::stream_generate,
19 config::Gemini,
20 error::GeminiError,
21 types::{
22 FunctionCallingConfig, FunctionCallingMode, FunctionDeclaration, GeminiContent, GeminiTool,
23 GenerateContentRequest, GenerationConfig, GoogleSearch, Part, PromptFeedback, SafetyRating,
24 ThinkingConfig, ThinkingLevel, ToolConfig, UsageMetadata,
25 },
26};
27use schemars::schema_for;
28
29impl LanguageModel for Gemini {
30 type Error = GeminiError;
31
32 fn respond(
33 &self,
34 request: LLMRequest,
35 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
36 let cfg = self.config();
37 respond_stream(cfg.clone(), request)
38 }
39
40 fn generate<T: JsonSchema + DeserializeOwned + 'static>(
41 &self,
42 mut request: LLMRequest,
43 ) -> impl core::future::Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
44 let schema = schema_for!(T);
45 let mut params = request.parameters().clone();
46 params.structured_outputs = true;
47 params.response_format = Some(schema);
48 request = request.with_parameters(params);
49
50 let stream = self.respond(request);
51 async move {
52 let text = aither_core::llm::collect_text(stream)
53 .await
54 .map_err(GenerateError::Provider)?;
55 serde_json::from_str::<T>(&text).map_err(|source| GenerateError::Parse {
56 source,
57 response: text.chars().take(500).collect(),
58 })
59 }
60 }
61
62 fn profile(&self) -> impl core::future::Future<Output = Result<Profile, Self::Error>> + Send {
63 let cfg = self.config().clone();
64 async move {
65 let model_name = cfg.text_model.trim_start_matches("models/").to_string();
66
67 let context_length = match crate::client::get_model_info(&cfg, &cfg.text_model).await {
69 Ok(info) => info.input_token_limit,
70 Err(e) => {
71 tracing::debug!("API did not return context_length: {e}");
72 aither_models::lookup(&model_name)
74 .and_then(aither_models::ModelEntry::max_input_tokens)
75 .ok_or_else(|| {
76 GeminiError::Api(format!(
77 "cannot determine context length for model '{model_name}': \
78 model info request failed ({e}) and aither-models has no entry"
79 ))
80 })?
81 }
82 };
83
84 let mut profile = Profile::new(
85 model_name.clone(),
86 "google",
87 model_name,
88 "Gemini Developer API model",
89 context_length,
90 )
91 .with_abilities([Ability::ToolUse, Ability::Vision, Ability::Audio]);
92 for ability in &cfg.native_abilities {
93 if !profile.abilities.contains(ability) {
94 profile.abilities.push(*ability);
95 }
96 }
97 Ok(profile)
98 }
99 }
100}
101
102fn respond_stream(
103 cfg: crate::config::GeminiConfig,
104 request: LLMRequest,
105) -> impl Stream<Item = Result<Event, GeminiError>> + Send {
106 Box::pin(respond_stream_inner(cfg, request))
107}
108
109async fn build_generate_request(
119 cfg: &crate::config::GeminiConfig,
120 request: LLMRequest,
121) -> Result<GenerateContentRequest, GeminiError> {
122 let (messages, parameters, tool_defs) = request.into_parts();
123
124 #[cfg(not(target_arch = "wasm32"))]
125 let messages = crate::attachments::resolve_messages(cfg, messages).await?;
126 #[cfg(target_arch = "wasm32")]
129 let _ = cfg;
130
131 let (system_instruction, contents) = messages_to_gemini(&messages).await?;
132
133 let tool_defs: Vec<ToolDefinition> = match ¶meters.tool_choice {
134 ToolChoice::None => Vec::new(),
135 ToolChoice::Exact(name) => tool_defs
136 .into_iter()
137 .filter(|tool| tool.name() == name)
138 .collect(),
139 ToolChoice::Auto | ToolChoice::Required => tool_defs,
140 };
141 let has_function_tools = !tool_defs.is_empty();
142
143 if let ToolChoice::Exact(name) = ¶meters.tool_choice
144 && !has_function_tools
145 {
146 return Err(GeminiError::Api(format!(
147 "Exact tool choice '{name}' is not present in tool definitions"
148 )));
149 }
150
151 let mut tools: Vec<GeminiTool> = Vec::new();
152 if has_function_tools {
153 tools.push(GeminiTool::FunctionTool {
154 function_declarations: convert_tool_definitions(tool_defs),
155 });
156 }
157
158 let builtins_allowed = !matches!(
161 parameters.tool_choice,
162 ToolChoice::None | ToolChoice::Exact(_)
163 );
164 let native = ¶meters.native_tools.gemini;
165 if builtins_allowed && (parameters.websearch || native.google_search) {
166 tools.push(GeminiTool::GoogleSearchTool {
167 google_search: GoogleSearch {},
168 });
169 }
170 if builtins_allowed && (parameters.code_execution || native.code_execution) {
171 tools.push(GeminiTool::CodeExecutionTool {
172 code_execution: crate::types::CodeExecution {},
173 });
174 }
175 if builtins_allowed && native.url_context {
176 tools.push(GeminiTool::UrlContextTool {
177 url_context: crate::types::UrlContext {},
178 });
179 }
180
181 Ok(GenerateContentRequest {
182 system_instruction,
183 contents,
184 generation_config: build_generation_config(¶meters, None)?,
185 tools,
186 tool_config: build_tool_config(¶meters, has_function_tools),
187 safety_settings: Vec::new(),
188 cached_content: parameters
189 .cache
190 .gemini
191 .as_ref()
192 .map(|cache| cache.cached_content.clone()),
193 })
194}
195
196fn events_from_content(content: &GeminiContent) -> Vec<Event> {
199 let mut events = Vec::new();
200
201 events.extend(content.reasoning_chunks().into_iter().map(Event::Reasoning));
202 events.extend(
203 content
204 .text_chunks()
205 .into_iter()
206 .filter(|text| !text.is_empty())
207 .map(Event::Text),
208 );
209
210 for part in &content.parts {
211 if let Some(code) = &part.executable_code {
212 events.push(Event::BuiltInToolCall {
213 id: String::new(),
214 tool: "code_execution".to_string(),
215 detail: format!("```{}\n{}\n```", code.language.to_lowercase(), code.code),
216 });
217 }
218 if let Some(result) = &part.code_execution_result {
219 events.push(Event::BuiltInToolResult {
220 id: String::new(),
221 tool: "code_execution".to_string(),
222 result: format!("```output\n{}\n```", result.output),
223 });
224 }
225 }
226
227 events.extend(
229 content
230 .function_call_parts()
231 .into_iter()
232 .map(|(call, signature)| {
233 Event::ToolCall(aither_core::llm::ToolCall {
236 id: tool_call_id(),
237 name: call.name,
238 arguments: call.args,
239 reasoning_state: signature
240 .map(|value| ReasoningState::new(PROVIDER_NAME, value)),
241 })
242 }),
243 );
244
245 events
246}
247
248fn respond_stream_inner(
249 cfg: crate::config::GeminiConfig,
250 request: LLMRequest,
251) -> impl Stream<Item = Result<Event, GeminiError>> + Send {
252 async_stream::stream! {
253 let gemini_request = match build_generate_request(&cfg, request).await {
254 Ok(request) => request,
255 Err(e) => {
256 yield Err(e);
257 return;
258 }
259 };
260
261 debug!("Gemini request: {:?}", gemini_request);
262
263 let stream = match stream_generate(&cfg, &cfg.text_model, gemini_request).await {
265 Ok(s) => s,
266 Err(e) => {
267 yield Err(e);
268 return;
269 }
270 };
271 futures_lite::pin!(stream);
272 let mut usage: Option<Usage> = None;
273 let mut finish_reason: Option<String> = None;
274
275 while let Some(result) = stream.next().await {
276 let response = match result {
277 Ok(r) => r,
278 Err(e) => {
279 yield Err(e);
280 continue;
281 }
282 };
283
284 debug!("Gemini stream chunk: {:?}", response);
285
286 if let Some(meta) = &response.usage_metadata {
287 usage = Some(usage_from_metadata(meta));
288 }
289
290 let Some(candidate) = response.primary_candidate() else {
291 if let Some(feedback) = &response.prompt_feedback {
294 yield Err(GeminiError::Api(format_prompt_feedback(feedback)));
295 }
296 continue;
297 };
298
299 let Some(content) = &candidate.content else {
300 if let Some(reason) = candidate.finish_reason.clone() {
301 finish_reason = Some(reason);
302 break;
303 }
304 continue;
305 };
306
307 for event in events_from_content(content) {
308 yield Ok(event);
309 }
310
311 if let Some(reason) = candidate.finish_reason.clone() {
312 finish_reason = Some(reason);
313 break;
314 }
315 }
316
317 if let Some(mut final_usage) = usage {
318 final_usage.stop_reason = finish_reason;
319 yield Ok(Event::Usage(final_usage));
320 } else if let Some(reason) = finish_reason {
321 yield Ok(Event::Usage(Usage {
322 stop_reason: Some(reason),
323 ..Usage::default()
324 }));
325 }
326 }
327}
328
329fn uuid_v4() -> String {
331 use std::time::{SystemTime, UNIX_EPOCH};
332 let timestamp = SystemTime::now()
333 .duration_since(UNIX_EPOCH)
334 .map_or(0, |d| d.as_nanos());
335 format!("{timestamp:032x}")
336}
337
338fn tool_call_id() -> String {
339 format!("gemini_{}", uuid_v4())
340}
341
342fn parse_tool_response(content: &str) -> Option<(String, serde_json::Value, Option<String>)> {
343 let content = content.trim_start();
344 if !content.starts_with('[') {
345 return None;
346 }
347 let end = content.find(']')?;
348 let header = &content[1..end];
349 let mut header_split = header.splitn(2, ':');
350 let id = header_split.next()?.trim();
351 let name = header_split.next()?.trim().to_string();
352 let rest = &content[end + 1..];
353 let output = rest.strip_prefix(' ').unwrap_or(rest);
354 let _ = id;
355 let signature = None;
356 let response_value = match serde_json::from_str::<serde_json::Value>(output) {
357 Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
358 Ok(other) => {
359 let mut map = serde_json::Map::new();
360 map.insert("result".to_string(), other);
361 serde_json::Value::Object(map)
362 }
363 Err(_) => {
364 let mut map = serde_json::Map::new();
365 map.insert(
366 "result".to_string(),
367 serde_json::Value::String(output.to_string()),
368 );
369 serde_json::Value::Object(map)
370 }
371 };
372 Some((name, response_value, signature))
373}
374
375const fn usage_from_metadata(meta: &UsageMetadata) -> Usage {
376 Usage {
377 prompt_tokens: meta.prompt_token_count,
378 completion_tokens: meta.candidates_token_count,
379 total_tokens: meta.total_token_count,
380 reasoning_tokens: meta.thoughts_token_count,
381 cache_read_tokens: meta.cached_content_token_count,
382 cache_write_tokens: None,
383 cost_usd: None,
384 stop_reason: None,
385 }
386}
387
388fn index_tool_calls(messages: &[Message]) -> ToolCallIndex<'_> {
394 let mut names = HashMap::new();
395 let mut signatures = HashMap::new();
396 let mut positions = HashMap::new();
397 for message in messages {
398 for tc in message.tool_calls() {
399 positions.insert(tc.id.as_str(), positions.len());
400 names.insert(tc.id.as_str(), tc.name.as_str());
401 if let Some(signature) = tc
402 .reasoning_state
403 .as_ref()
404 .and_then(|state| state.payload_for(PROVIDER_NAME))
405 {
406 signatures.insert(tc.id.as_str(), signature);
407 }
408 }
409 }
410 ToolCallIndex {
411 names,
412 signatures,
413 positions,
414 }
415}
416
417fn tool_response_value(content: &str) -> serde_json::Value {
419 match serde_json::from_str::<serde_json::Value>(content) {
420 Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
421 Ok(other) => serde_json::json!({ "result": other }),
422 Err(_) => serde_json::json!({ "result": content }),
423 }
424}
425
426struct ToolCallIndex<'a> {
428 names: HashMap<&'a str, &'a str>,
429 signatures: HashMap<&'a str, &'a str>,
430 positions: HashMap<&'a str, usize>,
433}
434
435fn flush_tool_results(contents: &mut Vec<GeminiContent>, pending: &mut Vec<(usize, Part)>) {
441 if pending.is_empty() {
442 return;
443 }
444 pending.sort_by_key(|(position, _)| *position);
445 let parts = pending.drain(..).map(|(_, part)| part).collect();
446 contents.push(GeminiContent::with_parts("user", parts));
447}
448
449async fn messages_to_gemini(
450 messages: &[Message],
451) -> Result<(Option<GeminiContent>, Vec<GeminiContent>), GeminiError> {
452 let mut system_parts = Vec::new();
453 let mut contents = Vec::new();
454
455 let index = index_tool_calls(messages);
456 let mut pending_results: Vec<(usize, Part)> = Vec::new();
457
458 for message in messages {
459 if message.role() != Role::Tool {
460 flush_tool_results(&mut contents, &mut pending_results);
461 }
462 match message.role() {
463 Role::System => system_parts.push(Part::text(message.content())),
464 Role::User => {
465 let attachments = message.attachments();
466 if attachments.is_empty() {
467 contents.push(GeminiContent::text("user", message.content()));
468 } else {
469 let mut parts = Vec::new();
471
472 for attachment in attachments {
474 parts.push(url_to_part(attachment).await?);
475 }
476
477 if !message.content().is_empty() {
479 parts.push(Part::text(message.content()));
480 }
481
482 contents.push(GeminiContent::with_parts("user", parts));
483 }
484 }
485 Role::Tool => {
486 let tool_call_id = message.tool_call_id().unwrap_or("");
488 let base_id = tool_call_id.to_string();
489 let signature = index
490 .signatures
491 .get(tool_call_id)
492 .map(|value| (*value).to_string());
493
494 if let Some(&function_name) = index.names.get(tool_call_id) {
495 pending_results.push((
496 index.positions[tool_call_id],
497 Part::function_response_with_signature(
498 function_name.to_string(),
499 tool_response_value(message.content()),
500 signature,
501 ),
502 ));
503 } else {
504 if let Some((name, response, sig)) = parse_tool_response(message.content()) {
506 pending_results.push((
507 usize::MAX,
508 Part::function_response_with_signature(name, response, sig),
509 ));
510 } else {
511 debug!("Tool message without matching tool_call_id: {}", base_id);
513 flush_tool_results(&mut contents, &mut pending_results);
514 contents.push(GeminiContent::text("user", message.content()));
515 }
516 }
517 }
518 Role::Assistant => {
519 let tool_calls = message.tool_calls();
520 if tool_calls.is_empty() {
521 contents.push(GeminiContent::text("model", message.content()));
523 } else {
524 let mut parts = Vec::new();
526
527 if !message.content().is_empty() {
529 parts.push(Part::text(message.content()));
530 }
531
532 for tc in tool_calls {
534 let signature = tc
535 .reasoning_state
536 .as_ref()
537 .and_then(|state| state.payload_for(PROVIDER_NAME))
538 .map(String::from);
539 parts.push(Part::function_call_with_signature(
540 tc.name.clone(),
541 tc.arguments.clone(),
542 signature,
543 ));
544 }
545
546 contents.push(GeminiContent::with_parts("model", parts));
547 }
548 }
549 }
550 }
551
552 flush_tool_results(&mut contents, &mut pending_results);
553
554 let system_instruction = if system_parts.is_empty() {
555 None
556 } else {
557 Some(GeminiContent::system(system_parts))
558 };
559
560 Ok((system_instruction, contents))
561}
562
563fn format_prompt_feedback(feedback: &PromptFeedback) -> String {
564 let reason = feedback
565 .block_reason
566 .as_deref()
567 .unwrap_or("unknown reason")
568 .to_string();
569 if feedback.safety_ratings.is_empty() {
570 return format!("Gemini response blocked: {reason}");
571 }
572
573 let ratings = feedback
574 .safety_ratings
575 .iter()
576 .map(format_safety_rating)
577 .collect::<Vec<_>>()
578 .join(", ");
579 format!("Gemini response blocked: {reason}; safety: {ratings}")
580}
581
582fn format_safety_rating(rating: &SafetyRating) -> String {
583 let status = rating
584 .blocked
585 .map_or("unspecified", |b| if b { "blocked" } else { "allowed" });
586 let probability = rating.probability.as_deref().unwrap_or("unknown");
587 format!("{} ({status}, probability: {probability})", rating.category)
588}
589
590fn build_generation_config(
595 parameters: &Parameters,
596 modalities: Option<Vec<String>>,
597) -> Result<Option<GenerationConfig>, GeminiError> {
598 let thinking_config = build_thinking_config(parameters)?;
599 let mut config = GenerationConfig {
600 temperature: parameters.temperature,
601 top_p: parameters.top_p,
602 top_k: parameters.top_k,
603 max_output_tokens: parameters.max_tokens.map(
604 #[allow(clippy::cast_possible_wrap)]
605 |value| value as i32,
606 ),
607 stop_sequences: parameters.stop.clone(),
608 response_modalities: modalities,
609 thinking_config,
610 seed: parameters.seed.map(
611 #[allow(clippy::cast_possible_wrap)]
612 |value| value as i32,
613 ),
614 presence_penalty: parameters.presence_penalty,
615 frequency_penalty: parameters.frequency_penalty,
616 response_logprobs: parameters.logprobs,
617 logprobs: parameters.top_logprobs.map(i32::from),
618 ..Default::default()
619 };
620 if let Some(schema) = ¶meters.response_format {
621 config.response_mime_type = Some("application/json".into());
622 config.response_json_schema = Some(schema.clone().to_value());
623 } else if parameters.structured_outputs {
624 config.response_mime_type = Some("application/json".into());
625 }
626 if config.is_meaningful() {
627 Ok(Some(config))
628 } else {
629 Ok(None)
630 }
631}
632
633fn build_tool_config(parameters: &Parameters, has_tools: bool) -> Option<ToolConfig> {
634 if !has_tools {
635 return None;
636 }
637 let (mode, allowed) = match ¶meters.tool_choice {
638 ToolChoice::None => return None,
639 ToolChoice::Auto => (FunctionCallingMode::Auto, None),
640 ToolChoice::Required => (FunctionCallingMode::Any, None),
641 ToolChoice::Exact(name) => (FunctionCallingMode::Any, Some(vec![name.clone()])),
642 };
643 Some(ToolConfig {
644 function_calling_config: Some(FunctionCallingConfig {
645 mode,
646 allowed_function_names: allowed,
647 }),
648 })
649}
650
651fn build_thinking_config(parameters: &Parameters) -> Result<Option<ThinkingConfig>, GeminiError> {
660 let thinking_level = parameters
661 .reasoning_effort
662 .map(ThinkingLevel::try_from)
663 .transpose()?;
664 if !parameters.include_reasoning && thinking_level.is_none() {
665 return Ok(None);
666 }
667 Ok(Some(ThinkingConfig {
668 include_thoughts: parameters.include_reasoning.then_some(true),
669 thinking_level,
670 }))
671}
672
673fn convert_tool_definitions(defs: Vec<ToolDefinition>) -> Vec<FunctionDeclaration> {
674 defs.into_iter()
675 .map(|tool| FunctionDeclaration {
676 name: tool.name().to_string(),
677 description: tool.description().to_string(),
678 parameters: Some(tool.arguments_openai_schema()),
679 })
680 .collect()
681}
682
683async fn url_to_part(attachment: &Attachment) -> Result<Part, GeminiError> {
685 let url = attachment.url();
686 let media_type = attachment.media_type().as_ref();
687 match url.scheme() {
688 "data" => parse_data_url(url.as_str(), media_type),
689 "file" => read_file_to_part(url, media_type).await,
690 "http" | "https" => Ok(Part::from_file(media_type, url.as_str())),
691 scheme => Err(GeminiError::Api(format!(
692 "Gemini does not support attachment URL scheme '{scheme}'"
693 ))),
694 }
695}
696
697fn parse_data_url(url: &str, media_type: &str) -> Result<Part, GeminiError> {
698 use base64::Engine;
699
700 let after_data = url
701 .strip_prefix("data:")
702 .ok_or_else(|| GeminiError::Api("Malformed attachment data URL".to_string()))?;
703 let (header, data) = after_data
704 .split_once(',')
705 .ok_or_else(|| GeminiError::Api("Attachment data URL is missing payload".to_string()))?;
706 let encoded_media_type = header.strip_suffix(";base64").ok_or_else(|| {
707 GeminiError::Api("Attachment data URL must use base64 encoding".to_string())
708 })?;
709 if encoded_media_type != media_type {
710 return Err(GeminiError::Api(format!(
711 "Attachment MIME type '{media_type}' does not match data URL MIME type '{encoded_media_type}'"
712 )));
713 }
714 let bytes = base64::engine::general_purpose::STANDARD.decode(data)?;
715 Ok(Part::inline_media(media_type, bytes))
716}
717
718#[cfg(not(target_arch = "wasm32"))]
719async fn read_file_to_part(url: &url::Url, media_type: &str) -> Result<Part, GeminiError> {
720 let path = url.to_file_path().map_err(|()| {
721 GeminiError::Api("Attachment file URL could not be converted to a path".to_string())
722 })?;
723 let data = async_fs::read(&path).await.map_err(|error| {
724 GeminiError::Api(format!(
725 "Failed to read attachment '{}': {error}",
726 path.display()
727 ))
728 })?;
729 Ok(Part::inline_media(media_type, data))
730}
731
732#[cfg(target_arch = "wasm32")]
733async fn read_file_to_part(_url: &url::Url, _media_type: &str) -> Result<Part, GeminiError> {
734 Err(GeminiError::Api(
735 "file:// attachments are not supported on wasm32".to_string(),
736 ))
737}
738
739#[cfg(test)]
740mod tests {
741 use super::*;
742 use aither_core::llm::model::ReasoningEffort;
743
744 #[tokio::test]
745 async fn data_attachments_preserve_declared_media_types() {
746 for media_type in ["image/png", "audio/wav", "video/mp4", "application/pdf"] {
747 let attachment = Attachment::new(
748 format!("data:{media_type};base64,AA==")
749 .parse()
750 .expect("data URL"),
751 media_type.parse().expect("MIME type"),
752 );
753 let part = url_to_part(&attachment).await.expect("encode attachment");
754 let value = serde_json::to_value(part).expect("serialize part");
755 assert_eq!(value["inlineData"]["mimeType"], media_type);
756 }
757 }
758
759 #[tokio::test]
760 async fn remote_attachment_uses_file_data_with_media_type() {
761 let attachment = Attachment::new(
762 "https://ai.google.dev/static/gemini-api/docs/files/sample.pdf"
763 .parse()
764 .expect("remote URL"),
765 "application/pdf".parse().expect("PDF MIME"),
766 );
767 let part = url_to_part(&attachment).await.expect("encode attachment");
768 let value = serde_json::to_value(part).expect("serialize part");
769 assert_eq!(value["fileData"]["mimeType"], "application/pdf");
770 assert_eq!(
771 value["fileData"]["fileUri"],
772 "https://ai.google.dev/static/gemini-api/docs/files/sample.pdf"
773 );
774 }
775
776 #[test]
777 fn data_url_mime_mismatch_fails() {
778 let error = parse_data_url("data:image/png;base64,AA==", "audio/wav")
779 .expect_err("MIME mismatch must fail");
780 assert!(error.to_string().contains("does not match"));
781 }
782
783 fn thinking_json(params: &Parameters) -> serde_json::Value {
784 let config = build_thinking_config(params)
785 .expect("effort supported by Gemini")
786 .expect("thinking config");
787 serde_json::to_value(config).expect("serialize thinking config")
788 }
789
790 #[test]
793 fn thinking_config_uses_gemini_field_names() {
794 let value = thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::Medium));
795 assert_eq!(value["thinkingLevel"], "medium");
796 assert_eq!(value.get("tokenBudget"), None);
797 assert_eq!(value.get("thinkingBudget"), None);
798 }
799
800 #[test]
803 fn thinking_config_never_pairs_a_budget_with_a_level() {
804 for effort in [
805 ReasoningEffort::Minimal,
806 ReasoningEffort::Low,
807 ReasoningEffort::Medium,
808 ReasoningEffort::High,
809 ] {
810 let value = thinking_json(&Parameters::default().reasoning_effort(effort));
811 let object = value.as_object().expect("thinking config object");
812 assert!(
813 object.keys().all(|key| key != "thinkingBudget"),
814 "effort {effort:?} paired a budget with a level"
815 );
816 assert!(object.contains_key("thinkingLevel"));
817 }
818 }
819
820 #[test]
821 fn minimal_effort_maps_to_the_minimal_level() {
822 let value =
823 thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::Minimal));
824 assert_eq!(value["thinkingLevel"], "minimal");
825 }
826
827 #[test]
830 fn efforts_above_gemini_ceiling_are_rejected() {
831 for effort in [ReasoningEffort::XHigh, ReasoningEffort::Max] {
832 let error = build_thinking_config(&Parameters::default().reasoning_effort(effort))
833 .expect_err("effort above the Gemini ceiling must fail");
834 assert!(
835 error.to_string().contains("tops out at High"),
836 "unexpected error for {effort:?}: {error}"
837 );
838 }
839 }
840
841 #[tokio::test]
845 async fn thought_signature_round_trips_without_mangling_the_id() {
846 let call = aither_core::llm::ToolCall::new("gemini_abc", "lookup", serde_json::json!({}))
847 .with_reasoning_state(ReasoningState::new(PROVIDER_NAME, "sig-xyz"));
848 let message = Message::assistant_with_tool_calls("", vec![call]);
849
850 let (_, contents) = messages_to_gemini(&[message]).await.expect("convert");
851 let value = serde_json::to_value(&contents).expect("serialize contents");
852 let part = &value[0]["parts"][0];
853
854 assert_eq!(part["thoughtSignature"], "sig-xyz");
855 assert_eq!(part["functionCall"]["name"], "lookup");
856 }
857
858 #[tokio::test]
859 async fn generated_tool_call_ids_carry_no_smuggled_signature() {
860 assert!(!tool_call_id().contains("|ts|"));
861 }
862
863 #[tokio::test]
865 async fn foreign_reasoning_state_is_not_sent_as_a_signature() {
866 let call = aither_core::llm::ToolCall::new("gemini_abc", "lookup", serde_json::json!({}))
867 .with_reasoning_state(ReasoningState::new("anthropic", "sig-from-claude"));
868 let message = Message::assistant_with_tool_calls("", vec![call]);
869
870 let (_, contents) = messages_to_gemini(&[message]).await.expect("convert");
871 let value = serde_json::to_value(&contents).expect("serialize contents");
872 assert_eq!(value[0]["parts"][0].get("thoughtSignature"), None);
873 }
874
875 fn call(id: &str, name: &str) -> aither_core::llm::ToolCall {
876 aither_core::llm::ToolCall::new(id, name, serde_json::json!({}))
877 }
878
879 fn responses(content: &serde_json::Value) -> Vec<(String, String)> {
880 content["parts"]
881 .as_array()
882 .expect("parts")
883 .iter()
884 .map(|part| {
885 (
886 part["functionResponse"]["name"]
887 .as_str()
888 .expect("function name")
889 .to_string(),
890 part["functionResponse"]["response"]["result"]
891 .as_str()
892 .expect("result")
893 .to_string(),
894 )
895 })
896 .collect()
897 }
898
899 fn pairs(items: [(&str, &str); 2]) -> Vec<(String, String)> {
900 items
901 .into_iter()
902 .map(|(a, b)| (a.to_string(), b.to_string()))
903 .collect()
904 }
905
906 #[tokio::test]
907 async fn parallel_tool_results_share_one_content() {
908 let messages = [
909 Message::assistant_with_tool_calls("", vec![call("a", "weather"), call("b", "time")]),
910 Message::tool("a", "sunny"),
911 Message::tool("b", "noon"),
912 ];
913 let (_, contents) = messages_to_gemini(&messages).await.expect("convert");
914 let value = serde_json::to_value(&contents).expect("serialize");
915 assert_eq!(value.as_array().expect("contents").len(), 2);
916 assert_eq!(value[1]["role"], "user");
917 assert_eq!(
918 responses(&value[1]),
919 pairs([("weather", "sunny"), ("time", "noon")])
920 );
921 }
922
923 #[tokio::test]
924 async fn same_function_results_stay_distinct_and_ordered() {
925 let messages = [
926 Message::assistant_with_tool_calls("", vec![call("a", "lookup"), call("b", "lookup")]),
927 Message::tool("a", "first"),
928 Message::tool("b", "second"),
929 ];
930 let (_, contents) = messages_to_gemini(&messages).await.expect("convert");
931 let value = serde_json::to_value(&contents).expect("serialize");
932 assert_eq!(value.as_array().expect("contents").len(), 2);
933 assert_eq!(
934 responses(&value[1]),
935 pairs([("lookup", "first"), ("lookup", "second")])
936 );
937 }
938
939 #[tokio::test]
940 async fn out_of_order_results_follow_call_order_and_keep_signatures() {
941 let signed =
942 call("b", "time").with_reasoning_state(ReasoningState::new(PROVIDER_NAME, "sig-b"));
943 let messages = [
944 Message::assistant_with_tool_calls("", vec![call("a", "weather"), signed]),
945 Message::tool("b", "noon"),
946 Message::tool("a", "sunny"),
947 Message::user("thanks"),
948 ];
949 let (_, contents) = messages_to_gemini(&messages).await.expect("convert");
950 let value = serde_json::to_value(&contents).expect("serialize");
951 assert_eq!(value.as_array().expect("contents").len(), 3);
952 assert_eq!(
953 responses(&value[1]),
954 pairs([("weather", "sunny"), ("time", "noon")])
955 );
956 assert_eq!(value[1]["parts"][0].get("thoughtSignature"), None);
957 assert_eq!(value[1]["parts"][1]["thoughtSignature"], "sig-b");
958 assert_eq!(value[2]["parts"][0]["text"], "thanks");
959 }
960
961 #[test]
963 fn none_effort_is_rejected() {
964 let error =
965 build_thinking_config(&Parameters::default().reasoning_effort(ReasoningEffort::None))
966 .expect_err("None effort must fail");
967 assert!(error.to_string().contains("no 'none' thinking level"));
968 }
969
970 #[test]
973 fn effort_and_visibility_are_independent() {
974 let effort_only =
975 thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::High));
976 assert_eq!(effort_only["thinkingLevel"], "high");
977 assert_eq!(effort_only.get("includeThoughts"), None);
978
979 let thoughts_only = thinking_json(&Parameters::default().include_reasoning(true));
980 assert_eq!(thoughts_only["includeThoughts"], true);
981 assert_eq!(thoughts_only.get("thinkingLevel"), None);
982
983 assert!(
984 build_thinking_config(&Parameters::default())
985 .expect("no effort requested")
986 .is_none()
987 );
988 }
989
990 #[test]
993 fn generation_config_forwards_every_sampling_parameter() {
994 let params = Parameters::default()
995 .seed(42)
996 .presence_penalty(0.5)
997 .frequency_penalty(0.25)
998 .logprobs(true)
999 .top_logprobs(3);
1000 let config = build_generation_config(¶ms, None)
1001 .expect("effort supported by Gemini")
1002 .expect("generation config");
1003 let value = serde_json::to_value(config).expect("serialize generation config");
1004 assert_eq!(value["seed"], 42);
1005 assert_eq!(value["presencePenalty"], 0.5);
1006 assert_eq!(value["frequencyPenalty"], 0.25);
1007 assert_eq!(value["responseLogprobs"], true);
1008 assert_eq!(value["logprobs"], 3);
1009 }
1010}