Skip to main content

aither_gemini/
llm.rs

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            // Fetch actual context window from API, fallback to models database
68            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                    // Fallback to models database
73                    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
109/// Assembles the Gemini request for a provider-agnostic one.
110///
111/// Async because file attachments are uploaded before the request is sent.
112///
113/// # Errors
114///
115/// Returns [`GeminiError`] if the request carries another provider's cache
116/// settings, names an exact tool that was not supplied, or an attachment cannot
117/// be resolved.
118async 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    if parameters.cache.openai.is_some() || parameters.cache.claude.is_some() {
125        return Err(GeminiError::Api(
126            "Gemini provider only accepts cache.gemini settings".to_string(),
127        ));
128    }
129
130    #[cfg(not(target_arch = "wasm32"))]
131    let messages = crate::attachments::resolve_messages(cfg, messages).await?;
132    // wasm32 has no filesystem, so there is nothing to upload and no use for
133    // the config here.
134    #[cfg(target_arch = "wasm32")]
135    let _ = cfg;
136
137    let (system_instruction, contents) = messages_to_gemini(&messages).await?;
138
139    let tool_defs: Vec<ToolDefinition> = match &parameters.tool_choice {
140        ToolChoice::None => Vec::new(),
141        ToolChoice::Exact(name) => tool_defs
142            .into_iter()
143            .filter(|tool| tool.name() == name)
144            .collect(),
145        ToolChoice::Auto | ToolChoice::Required => tool_defs,
146    };
147    let has_function_tools = !tool_defs.is_empty();
148
149    if let ToolChoice::Exact(name) = &parameters.tool_choice
150        && !has_function_tools
151    {
152        return Err(GeminiError::Api(format!(
153            "Exact tool choice '{name}' is not present in tool definitions"
154        )));
155    }
156
157    let mut tools: Vec<GeminiTool> = Vec::new();
158    if has_function_tools {
159        tools.push(GeminiTool::FunctionTool {
160            function_declarations: convert_tool_definitions(tool_defs),
161        });
162    }
163
164    // Google's own tools cannot be combined with a narrowed tool choice: the
165    // model would be free to answer with a tool the caller did not ask about.
166    let builtins_allowed = !matches!(
167        parameters.tool_choice,
168        ToolChoice::None | ToolChoice::Exact(_)
169    );
170    let native = &parameters.native_tools.gemini;
171    if builtins_allowed && (parameters.websearch || native.google_search) {
172        tools.push(GeminiTool::GoogleSearchTool {
173            google_search: GoogleSearch {},
174        });
175    }
176    if builtins_allowed && (parameters.code_execution || native.code_execution) {
177        tools.push(GeminiTool::CodeExecutionTool {
178            code_execution: crate::types::CodeExecution {},
179        });
180    }
181    if builtins_allowed && native.url_context {
182        tools.push(GeminiTool::UrlContextTool {
183            url_context: crate::types::UrlContext {},
184        });
185    }
186
187    Ok(GenerateContentRequest {
188        system_instruction,
189        contents,
190        generation_config: build_generation_config(&parameters, None)?,
191        tools,
192        tool_config: build_tool_config(&parameters, has_function_tools),
193        safety_settings: Vec::new(),
194        cached_content: parameters
195            .cache
196            .gemini
197            .as_ref()
198            .map(|cache| cache.cached_content.clone()),
199    })
200}
201
202/// Turns one candidate's content into the events it represents, in wire order:
203/// reasoning, then text, then built-in tool output, then function calls.
204fn events_from_content(content: &GeminiContent) -> Vec<Event> {
205    let mut events = Vec::new();
206
207    events.extend(content.reasoning_chunks().into_iter().map(Event::Reasoning));
208    events.extend(
209        content
210            .text_chunks()
211            .into_iter()
212            .filter(|text| !text.is_empty())
213            .map(Event::Text),
214    );
215
216    for part in &content.parts {
217        if let Some(code) = &part.executable_code {
218            events.push(Event::BuiltInToolCall {
219                id: String::new(),
220                tool: "code_execution".to_string(),
221                detail: format!("```{}\n{}\n```", code.language.to_lowercase(), code.code),
222            });
223        }
224        if let Some(result) = &part.code_execution_result {
225            events.push(Event::BuiltInToolResult {
226                id: String::new(),
227                tool: "code_execution".to_string(),
228                result: format!("```output\n{}\n```", result.output),
229            });
230        }
231    }
232
233    // Tool calls are reported, never executed: the consumer decides.
234    events.extend(
235        content
236            .function_call_parts()
237            .into_iter()
238            .map(|(call, signature)| {
239                // Gemini signs each function call rather than the turn, so the
240                // signature rides on the call it belongs to.
241                Event::ToolCall(aither_core::llm::ToolCall {
242                    id: tool_call_id(),
243                    name: call.name,
244                    arguments: call.args,
245                    reasoning_state: signature
246                        .map(|value| ReasoningState::new(PROVIDER_NAME, value)),
247                })
248            }),
249    );
250
251    events
252}
253
254fn respond_stream_inner(
255    cfg: crate::config::GeminiConfig,
256    request: LLMRequest,
257) -> impl Stream<Item = Result<Event, GeminiError>> + Send {
258    async_stream::stream! {
259        let gemini_request = match build_generate_request(&cfg, request).await {
260            Ok(request) => request,
261            Err(e) => {
262                yield Err(e);
263                return;
264            }
265        };
266
267        debug!("Gemini request: {:?}", gemini_request);
268
269        // Use streaming endpoint for true streaming output
270        let stream = match stream_generate(&cfg, &cfg.text_model, gemini_request).await {
271            Ok(s) => s,
272            Err(e) => {
273                yield Err(e);
274                return;
275            }
276        };
277        futures_lite::pin!(stream);
278        let mut usage: Option<Usage> = None;
279        let mut finish_reason: Option<String> = None;
280
281        while let Some(result) = stream.next().await {
282            let response = match result {
283                Ok(r) => r,
284                Err(e) => {
285                    yield Err(e);
286                    continue;
287                }
288            };
289
290            debug!("Gemini stream chunk: {:?}", response);
291
292            if let Some(meta) = &response.usage_metadata {
293                usage = Some(usage_from_metadata(meta));
294            }
295
296            let Some(candidate) = response.primary_candidate() else {
297                // Chunks without candidates carry metadata, not content — but a
298                // blocked prompt is reported this way and must not pass silently.
299                if let Some(feedback) = &response.prompt_feedback {
300                    yield Err(GeminiError::Api(format_prompt_feedback(feedback)));
301                }
302                continue;
303            };
304
305            let Some(content) = &candidate.content else {
306                if let Some(reason) = candidate.finish_reason.clone() {
307                    finish_reason = Some(reason);
308                    break;
309                }
310                continue;
311            };
312
313            for event in events_from_content(content) {
314                yield Ok(event);
315            }
316
317            if let Some(reason) = candidate.finish_reason.clone() {
318                finish_reason = Some(reason);
319                break;
320            }
321        }
322
323        if let Some(mut final_usage) = usage {
324            final_usage.stop_reason = finish_reason;
325            yield Ok(Event::Usage(final_usage));
326        } else if let Some(reason) = finish_reason {
327            yield Ok(Event::Usage(Usage {
328                stop_reason: Some(reason),
329                ..Usage::default()
330            }));
331        }
332    }
333}
334
335/// Simple UUID v4 generator for tool call IDs.
336fn uuid_v4() -> String {
337    use std::time::{SystemTime, UNIX_EPOCH};
338    let timestamp = SystemTime::now()
339        .duration_since(UNIX_EPOCH)
340        .map_or(0, |d| d.as_nanos());
341    format!("{timestamp:032x}")
342}
343
344fn tool_call_id() -> String {
345    format!("gemini_{}", uuid_v4())
346}
347
348fn parse_tool_response(content: &str) -> Option<(String, serde_json::Value, Option<String>)> {
349    let content = content.trim_start();
350    if !content.starts_with('[') {
351        return None;
352    }
353    let end = content.find(']')?;
354    let header = &content[1..end];
355    let mut header_split = header.splitn(2, ':');
356    let id = header_split.next()?.trim();
357    let name = header_split.next()?.trim().to_string();
358    let rest = &content[end + 1..];
359    let output = rest.strip_prefix(' ').unwrap_or(rest);
360    let _ = id;
361    let signature = None;
362    let response_value = match serde_json::from_str::<serde_json::Value>(output) {
363        Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
364        Ok(other) => {
365            let mut map = serde_json::Map::new();
366            map.insert("result".to_string(), other);
367            serde_json::Value::Object(map)
368        }
369        Err(_) => {
370            let mut map = serde_json::Map::new();
371            map.insert(
372                "result".to_string(),
373                serde_json::Value::String(output.to_string()),
374            );
375            serde_json::Value::Object(map)
376        }
377    };
378    Some((name, response_value, signature))
379}
380
381const fn usage_from_metadata(meta: &UsageMetadata) -> Usage {
382    Usage {
383        prompt_tokens: meta.prompt_token_count,
384        completion_tokens: meta.candidates_token_count,
385        total_tokens: meta.total_token_count,
386        reasoning_tokens: meta.thoughts_token_count,
387        cache_read_tokens: meta.cached_content_token_count,
388        cache_write_tokens: None,
389        cost_usd: None,
390        stop_reason: None,
391    }
392}
393
394/// Indexes every tool call by id, yielding its function name and, where the
395/// model signed it, its thought signature.
396///
397/// A tool *result* has to look both up: Gemini scopes the signature to the call
398/// rather than the result, so the result cannot carry it.
399fn index_tool_calls(messages: &[Message]) -> (HashMap<&str, &str>, HashMap<&str, &str>) {
400    let mut names = HashMap::new();
401    let mut signatures = HashMap::new();
402    for message in messages {
403        for tc in message.tool_calls() {
404            names.insert(tc.id.as_str(), tc.name.as_str());
405            if let Some(signature) = tc
406                .reasoning_state
407                .as_ref()
408                .and_then(|state| state.payload_for(PROVIDER_NAME))
409            {
410                signatures.insert(tc.id.as_str(), signature);
411            }
412        }
413    }
414    (names, signatures)
415}
416
417async fn messages_to_gemini(
418    messages: &[Message],
419) -> Result<(Option<GeminiContent>, Vec<GeminiContent>), GeminiError> {
420    let mut system_parts = Vec::new();
421    let mut contents = Vec::new();
422
423    let (tool_call_names, tool_call_signatures) = index_tool_calls(messages);
424
425    for message in messages {
426        match message.role() {
427            Role::System => system_parts.push(Part::text(message.content())),
428            Role::User => {
429                let attachments = message.attachments();
430                if attachments.is_empty() {
431                    contents.push(GeminiContent::text("user", message.content()));
432                } else {
433                    // Build parts with attachments
434                    let mut parts = Vec::new();
435
436                    // Add attachment parts first
437                    for attachment in attachments {
438                        parts.push(url_to_part(attachment).await?);
439                    }
440
441                    // Add text content
442                    if !message.content().is_empty() {
443                        parts.push(Part::text(message.content()));
444                    }
445
446                    contents.push(GeminiContent::with_parts("user", parts));
447                }
448            }
449            Role::Tool => {
450                // Get the function name from the tool_call_id
451                let tool_call_id = message.tool_call_id().unwrap_or("");
452                let base_id = tool_call_id.to_string();
453                let signature = tool_call_signatures
454                    .get(tool_call_id)
455                    .map(|value| (*value).to_string());
456
457                if let Some(&function_name) = tool_call_names.get(tool_call_id) {
458                    // Parse the content as JSON, or wrap it as a string result
459                    let response_value =
460                        match serde_json::from_str::<serde_json::Value>(message.content()) {
461                            Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
462                            Ok(other) => {
463                                let mut map = serde_json::Map::new();
464                                map.insert("result".to_string(), other);
465                                serde_json::Value::Object(map)
466                            }
467                            Err(_) => {
468                                let mut map = serde_json::Map::new();
469                                map.insert(
470                                    "result".to_string(),
471                                    serde_json::Value::String(message.content().to_string()),
472                                );
473                                serde_json::Value::Object(map)
474                            }
475                        };
476                    contents.push(GeminiContent::function_response_with_signature(
477                        function_name.to_string(),
478                        response_value,
479                        signature,
480                    ));
481                } else {
482                    // Fallback: try legacy format or send as user message
483                    if let Some((name, response, sig)) = parse_tool_response(message.content()) {
484                        contents.push(GeminiContent::function_response_with_signature(
485                            name, response, sig,
486                        ));
487                    } else {
488                        // Last resort - wrap as text
489                        debug!("Tool message without matching tool_call_id: {}", base_id);
490                        contents.push(GeminiContent::text("user", message.content()));
491                    }
492                }
493            }
494            Role::Assistant => {
495                let tool_calls = message.tool_calls();
496                if tool_calls.is_empty() {
497                    // Regular text response
498                    contents.push(GeminiContent::text("model", message.content()));
499                } else {
500                    // Assistant message with function calls
501                    let mut parts = Vec::new();
502
503                    // Add text content if present
504                    if !message.content().is_empty() {
505                        parts.push(Part::text(message.content()));
506                    }
507
508                    // Add function call parts with thought signatures extracted from IDs
509                    for tc in tool_calls {
510                        let signature = tc
511                            .reasoning_state
512                            .as_ref()
513                            .and_then(|state| state.payload_for(PROVIDER_NAME))
514                            .map(String::from);
515                        parts.push(Part::function_call_with_signature(
516                            tc.name.clone(),
517                            tc.arguments.clone(),
518                            signature,
519                        ));
520                    }
521
522                    contents.push(GeminiContent::with_parts("model", parts));
523                }
524            }
525        }
526    }
527
528    let system_instruction = if system_parts.is_empty() {
529        None
530    } else {
531        Some(GeminiContent::system(system_parts))
532    };
533
534    Ok((system_instruction, contents))
535}
536
537fn format_prompt_feedback(feedback: &PromptFeedback) -> String {
538    let reason = feedback
539        .block_reason
540        .as_deref()
541        .unwrap_or("unknown reason")
542        .to_string();
543    if feedback.safety_ratings.is_empty() {
544        return format!("Gemini response blocked: {reason}");
545    }
546
547    let ratings = feedback
548        .safety_ratings
549        .iter()
550        .map(format_safety_rating)
551        .collect::<Vec<_>>()
552        .join(", ");
553    format!("Gemini response blocked: {reason}; safety: {ratings}")
554}
555
556fn format_safety_rating(rating: &SafetyRating) -> String {
557    let status = rating
558        .blocked
559        .map_or("unspecified", |b| if b { "blocked" } else { "allowed" });
560    let probability = rating.probability.as_deref().unwrap_or("unknown");
561    format!("{} ({status}, probability: {probability})", rating.category)
562}
563
564/// # Errors
565///
566/// Returns [`GeminiError`] when the requested reasoning effort has no Gemini
567/// equivalent.
568fn build_generation_config(
569    parameters: &Parameters,
570    modalities: Option<Vec<String>>,
571) -> Result<Option<GenerationConfig>, GeminiError> {
572    let thinking_config = build_thinking_config(parameters)?;
573    let mut config = GenerationConfig {
574        temperature: parameters.temperature,
575        top_p: parameters.top_p,
576        top_k: parameters.top_k,
577        max_output_tokens: parameters.max_tokens.map(
578            #[allow(clippy::cast_possible_wrap)]
579            |value| value as i32,
580        ),
581        stop_sequences: parameters.stop.clone(),
582        response_modalities: modalities,
583        thinking_config,
584        seed: parameters.seed.map(
585            #[allow(clippy::cast_possible_wrap)]
586            |value| value as i32,
587        ),
588        presence_penalty: parameters.presence_penalty,
589        frequency_penalty: parameters.frequency_penalty,
590        response_logprobs: parameters.logprobs,
591        logprobs: parameters.top_logprobs.map(i32::from),
592        ..Default::default()
593    };
594    if let Some(schema) = &parameters.response_format {
595        config.response_mime_type = Some("application/json".into());
596        config.response_json_schema = Some(schema.clone().to_value());
597    } else if parameters.structured_outputs {
598        config.response_mime_type = Some("application/json".into());
599    }
600    if config.is_meaningful() {
601        Ok(Some(config))
602    } else {
603        Ok(None)
604    }
605}
606
607fn build_tool_config(parameters: &Parameters, has_tools: bool) -> Option<ToolConfig> {
608    if !has_tools {
609        return None;
610    }
611    let (mode, allowed) = match &parameters.tool_choice {
612        ToolChoice::None => return None,
613        ToolChoice::Auto => (FunctionCallingMode::Auto, None),
614        ToolChoice::Required => (FunctionCallingMode::Any, None),
615        ToolChoice::Exact(name) => (FunctionCallingMode::Any, Some(vec![name.clone()])),
616    };
617    Some(ToolConfig {
618        function_calling_config: Some(FunctionCallingConfig {
619            mode,
620            allowed_function_names: allowed,
621        }),
622    })
623}
624
625/// Builds `thinkingConfig` from the portable reasoning parameters.
626///
627/// Depth (`reasoning_effort`) and visibility (`include_reasoning`) are separate
628/// controls: asking for a deeper think does not require the thoughts to be
629/// returned, and vice versa. Either one alone is enough to emit the config.
630/// # Errors
631///
632/// Returns [`GeminiError`] when the requested effort has no Gemini equivalent.
633fn build_thinking_config(parameters: &Parameters) -> Result<Option<ThinkingConfig>, GeminiError> {
634    let thinking_level = parameters
635        .reasoning_effort
636        .map(ThinkingLevel::try_from)
637        .transpose()?;
638    if !parameters.include_reasoning && thinking_level.is_none() {
639        return Ok(None);
640    }
641    Ok(Some(ThinkingConfig {
642        include_thoughts: parameters.include_reasoning.then_some(true),
643        thinking_level,
644    }))
645}
646
647fn convert_tool_definitions(defs: Vec<ToolDefinition>) -> Vec<FunctionDeclaration> {
648    defs.into_iter()
649        .map(|tool| FunctionDeclaration {
650            name: tool.name().to_string(),
651            description: tool.description().to_string(),
652            parameters: Some(tool.arguments_openai_schema()),
653        })
654        .collect()
655}
656
657/// Convert a typed attachment to a Gemini content part.
658async fn url_to_part(attachment: &Attachment) -> Result<Part, GeminiError> {
659    let url = attachment.url();
660    let media_type = attachment.media_type().as_ref();
661    match url.scheme() {
662        "data" => parse_data_url(url.as_str(), media_type),
663        "file" => read_file_to_part(url, media_type).await,
664        "http" | "https" => Ok(Part::from_file(media_type, url.as_str())),
665        scheme => Err(GeminiError::Api(format!(
666            "Gemini does not support attachment URL scheme '{scheme}'"
667        ))),
668    }
669}
670
671fn parse_data_url(url: &str, media_type: &str) -> Result<Part, GeminiError> {
672    use base64::Engine;
673
674    let after_data = url
675        .strip_prefix("data:")
676        .ok_or_else(|| GeminiError::Api("Malformed attachment data URL".to_string()))?;
677    let (header, data) = after_data
678        .split_once(',')
679        .ok_or_else(|| GeminiError::Api("Attachment data URL is missing payload".to_string()))?;
680    let encoded_media_type = header.strip_suffix(";base64").ok_or_else(|| {
681        GeminiError::Api("Attachment data URL must use base64 encoding".to_string())
682    })?;
683    if encoded_media_type != media_type {
684        return Err(GeminiError::Api(format!(
685            "Attachment MIME type '{media_type}' does not match data URL MIME type '{encoded_media_type}'"
686        )));
687    }
688    let bytes = base64::engine::general_purpose::STANDARD.decode(data)?;
689    Ok(Part::inline_media(media_type, bytes))
690}
691
692#[cfg(not(target_arch = "wasm32"))]
693async fn read_file_to_part(url: &url::Url, media_type: &str) -> Result<Part, GeminiError> {
694    let path = url.to_file_path().map_err(|()| {
695        GeminiError::Api("Attachment file URL could not be converted to a path".to_string())
696    })?;
697    let data = async_fs::read(&path).await.map_err(|error| {
698        GeminiError::Api(format!(
699            "Failed to read attachment '{}': {error}",
700            path.display()
701        ))
702    })?;
703    Ok(Part::inline_media(media_type, data))
704}
705
706#[cfg(target_arch = "wasm32")]
707async fn read_file_to_part(_url: &url::Url, _media_type: &str) -> Result<Part, GeminiError> {
708    Err(GeminiError::Api(
709        "file:// attachments are not supported on wasm32".to_string(),
710    ))
711}
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716    use aither_core::llm::model::ReasoningEffort;
717
718    #[tokio::test]
719    async fn data_attachments_preserve_declared_media_types() {
720        for media_type in ["image/png", "audio/wav", "video/mp4", "application/pdf"] {
721            let attachment = Attachment::new(
722                format!("data:{media_type};base64,AA==")
723                    .parse()
724                    .expect("data URL"),
725                media_type.parse().expect("MIME type"),
726            );
727            let part = url_to_part(&attachment).await.expect("encode attachment");
728            let value = serde_json::to_value(part).expect("serialize part");
729            assert_eq!(value["inlineData"]["mimeType"], media_type);
730        }
731    }
732
733    #[tokio::test]
734    async fn remote_attachment_uses_file_data_with_media_type() {
735        let attachment = Attachment::new(
736            "https://ai.google.dev/static/gemini-api/docs/files/sample.pdf"
737                .parse()
738                .expect("remote URL"),
739            "application/pdf".parse().expect("PDF MIME"),
740        );
741        let part = url_to_part(&attachment).await.expect("encode attachment");
742        let value = serde_json::to_value(part).expect("serialize part");
743        assert_eq!(value["fileData"]["mimeType"], "application/pdf");
744        assert_eq!(
745            value["fileData"]["fileUri"],
746            "https://ai.google.dev/static/gemini-api/docs/files/sample.pdf"
747        );
748    }
749
750    #[test]
751    fn data_url_mime_mismatch_fails() {
752        let error = parse_data_url("data:image/png;base64,AA==", "audio/wav")
753            .expect_err("MIME mismatch must fail");
754        assert!(error.to_string().contains("does not match"));
755    }
756
757    fn thinking_json(params: &Parameters) -> serde_json::Value {
758        let config = build_thinking_config(params)
759            .expect("effort supported by Gemini")
760            .expect("thinking config");
761        serde_json::to_value(config).expect("serialize thinking config")
762    }
763
764    /// Gemini names the field `thinkingLevel`; there is no `tokenBudget`, so a
765    /// misnamed field is accepted by serde and then ignored on the wire.
766    #[test]
767    fn thinking_config_uses_gemini_field_names() {
768        let value = thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::Medium));
769        assert_eq!(value["thinkingLevel"], "medium");
770        assert_eq!(value.get("tokenBudget"), None);
771        assert_eq!(value.get("thinkingBudget"), None);
772    }
773
774    /// `thinkingLevel` and the legacy `thinkingBudget` are mutually exclusive:
775    /// sending both is a 400, so only the level is ever modelled.
776    #[test]
777    fn thinking_config_never_pairs_a_budget_with_a_level() {
778        for effort in [
779            ReasoningEffort::Minimal,
780            ReasoningEffort::Low,
781            ReasoningEffort::Medium,
782            ReasoningEffort::High,
783        ] {
784            let value = thinking_json(&Parameters::default().reasoning_effort(effort));
785            let object = value.as_object().expect("thinking config object");
786            assert!(
787                object.keys().all(|key| key != "thinkingBudget"),
788                "effort {effort:?} paired a budget with a level"
789            );
790            assert!(object.contains_key("thinkingLevel"));
791        }
792    }
793
794    #[test]
795    fn minimal_effort_maps_to_the_minimal_level() {
796        let value =
797            thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::Minimal));
798        assert_eq!(value["thinkingLevel"], "minimal");
799    }
800
801    /// Gemini's ladder stops at High. Clamping `XHigh` or `Max` down to High would
802    /// hand back less reasoning than the caller asked for, silently.
803    #[test]
804    fn efforts_above_gemini_ceiling_are_rejected() {
805        for effort in [ReasoningEffort::XHigh, ReasoningEffort::Max] {
806            let error = build_thinking_config(&Parameters::default().reasoning_effort(effort))
807                .expect_err("effort above the Gemini ceiling must fail");
808            assert!(
809                error.to_string().contains("tops out at High"),
810                "unexpected error for {effort:?}: {error}"
811            );
812        }
813    }
814
815    /// The signature used to be base64'd into the tool-call id behind a `|ts|`
816    /// separator, because core had nowhere to put it. It now rides on the call,
817    /// and the id is a plain identifier again.
818    #[tokio::test]
819    async fn thought_signature_round_trips_without_mangling_the_id() {
820        let call = aither_core::llm::ToolCall::new("gemini_abc", "lookup", serde_json::json!({}))
821            .with_reasoning_state(ReasoningState::new(PROVIDER_NAME, "sig-xyz"));
822        let message = Message::assistant_with_tool_calls("", vec![call]);
823
824        let (_, contents) = messages_to_gemini(&[message]).await.expect("convert");
825        let value = serde_json::to_value(&contents).expect("serialize contents");
826        let part = &value[0]["parts"][0];
827
828        assert_eq!(part["thoughtSignature"], "sig-xyz");
829        assert_eq!(part["functionCall"]["name"], "lookup");
830    }
831
832    #[tokio::test]
833    async fn generated_tool_call_ids_carry_no_smuggled_signature() {
834        assert!(!tool_call_id().contains("|ts|"));
835    }
836
837    /// Signatures are Gemini's; another provider's state must not be replayed.
838    #[tokio::test]
839    async fn foreign_reasoning_state_is_not_sent_as_a_signature() {
840        let call = aither_core::llm::ToolCall::new("gemini_abc", "lookup", serde_json::json!({}))
841            .with_reasoning_state(ReasoningState::new("anthropic", "sig-from-claude"));
842        let message = Message::assistant_with_tool_calls("", vec![call]);
843
844        let (_, contents) = messages_to_gemini(&[message]).await.expect("convert");
845        let value = serde_json::to_value(&contents).expect("serialize contents");
846        assert_eq!(value[0]["parts"][0].get("thoughtSignature"), None);
847    }
848
849    /// Gemini has no "no thinking" level; the caller should omit effort instead.
850    #[test]
851    fn none_effort_is_rejected() {
852        let error =
853            build_thinking_config(&Parameters::default().reasoning_effort(ReasoningEffort::None))
854                .expect_err("None effort must fail");
855        assert!(error.to_string().contains("no 'none' thinking level"));
856    }
857
858    /// Depth and visibility are independent: asking for effort alone must still
859    /// produce a config, and asking for thoughts alone must still produce one.
860    #[test]
861    fn effort_and_visibility_are_independent() {
862        let effort_only =
863            thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::High));
864        assert_eq!(effort_only["thinkingLevel"], "high");
865        assert_eq!(effort_only.get("includeThoughts"), None);
866
867        let thoughts_only = thinking_json(&Parameters::default().include_reasoning(true));
868        assert_eq!(thoughts_only["includeThoughts"], true);
869        assert_eq!(thoughts_only.get("thinkingLevel"), None);
870
871        assert!(
872            build_thinking_config(&Parameters::default())
873                .expect("no effort requested")
874                .is_none()
875        );
876    }
877
878    /// These reached the wire only if `build_generation_config` populates them;
879    /// they were declared on the struct but never filled.
880    #[test]
881    fn generation_config_forwards_every_sampling_parameter() {
882        let params = Parameters::default()
883            .seed(42)
884            .presence_penalty(0.5)
885            .frequency_penalty(0.25)
886            .logprobs(true)
887            .top_logprobs(3);
888        let config = build_generation_config(&params, None)
889            .expect("effort supported by Gemini")
890            .expect("generation config");
891        let value = serde_json::to_value(config).expect("serialize generation config");
892        assert_eq!(value["seed"], 42);
893        assert_eq!(value["presencePenalty"], 0.5);
894        assert_eq!(value["frequencyPenalty"], 0.25);
895        assert_eq!(value["responseLogprobs"], true);
896        assert_eq!(value["logprobs"], 3);
897    }
898}