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