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 names an exact tool that was not
116/// supplied, or an attachment cannot be resolved. Cache settings of other
117/// providers are ignored.
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    #[cfg(not(target_arch = "wasm32"))]
125    let messages = crate::attachments::resolve_messages(cfg, messages).await?;
126    // wasm32 has no filesystem, so there is nothing to upload and no use for
127    // the config here.
128    #[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 &parameters.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) = &parameters.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    // Google's own tools cannot be combined with a narrowed tool choice: the
159    // model would be free to answer with a tool the caller did not ask about.
160    let builtins_allowed = !matches!(
161        parameters.tool_choice,
162        ToolChoice::None | ToolChoice::Exact(_)
163    );
164    let native = &parameters.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(&parameters, None)?,
185        tools,
186        tool_config: build_tool_config(&parameters, 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
196/// Turns one candidate's content into the events it represents, in wire order:
197/// reasoning, then text, then built-in tool output, then function calls.
198fn 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    // Tool calls are reported, never executed: the consumer decides.
228    events.extend(
229        content
230            .function_call_parts()
231            .into_iter()
232            .map(|(call, signature)| {
233                // Gemini signs each function call rather than the turn, so the
234                // signature rides on the call it belongs to.
235                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        // Use streaming endpoint for true streaming output
264        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                // Chunks without candidates carry metadata, not content — but a
292                // blocked prompt is reported this way and must not pass silently.
293                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
329/// Simple UUID v4 generator for tool call IDs.
330fn 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
388/// Indexes every tool call by id, yielding its function name and, where the
389/// model signed it, its thought signature.
390///
391/// A tool *result* has to look both up: Gemini scopes the signature to the call
392/// rather than the result, so the result cannot carry it.
393fn index_tool_calls(messages: &[Message]) -> (HashMap<&str, &str>, HashMap<&str, &str>) {
394    let mut names = HashMap::new();
395    let mut signatures = HashMap::new();
396    for message in messages {
397        for tc in message.tool_calls() {
398            names.insert(tc.id.as_str(), tc.name.as_str());
399            if let Some(signature) = tc
400                .reasoning_state
401                .as_ref()
402                .and_then(|state| state.payload_for(PROVIDER_NAME))
403            {
404                signatures.insert(tc.id.as_str(), signature);
405            }
406        }
407    }
408    (names, signatures)
409}
410
411async fn messages_to_gemini(
412    messages: &[Message],
413) -> Result<(Option<GeminiContent>, Vec<GeminiContent>), GeminiError> {
414    let mut system_parts = Vec::new();
415    let mut contents = Vec::new();
416
417    let (tool_call_names, tool_call_signatures) = index_tool_calls(messages);
418
419    for message in messages {
420        match message.role() {
421            Role::System => system_parts.push(Part::text(message.content())),
422            Role::User => {
423                let attachments = message.attachments();
424                if attachments.is_empty() {
425                    contents.push(GeminiContent::text("user", message.content()));
426                } else {
427                    // Build parts with attachments
428                    let mut parts = Vec::new();
429
430                    // Add attachment parts first
431                    for attachment in attachments {
432                        parts.push(url_to_part(attachment).await?);
433                    }
434
435                    // Add text content
436                    if !message.content().is_empty() {
437                        parts.push(Part::text(message.content()));
438                    }
439
440                    contents.push(GeminiContent::with_parts("user", parts));
441                }
442            }
443            Role::Tool => {
444                // Get the function name from the tool_call_id
445                let tool_call_id = message.tool_call_id().unwrap_or("");
446                let base_id = tool_call_id.to_string();
447                let signature = tool_call_signatures
448                    .get(tool_call_id)
449                    .map(|value| (*value).to_string());
450
451                if let Some(&function_name) = tool_call_names.get(tool_call_id) {
452                    // Parse the content as JSON, or wrap it as a string result
453                    let response_value =
454                        match serde_json::from_str::<serde_json::Value>(message.content()) {
455                            Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
456                            Ok(other) => {
457                                let mut map = serde_json::Map::new();
458                                map.insert("result".to_string(), other);
459                                serde_json::Value::Object(map)
460                            }
461                            Err(_) => {
462                                let mut map = serde_json::Map::new();
463                                map.insert(
464                                    "result".to_string(),
465                                    serde_json::Value::String(message.content().to_string()),
466                                );
467                                serde_json::Value::Object(map)
468                            }
469                        };
470                    contents.push(GeminiContent::function_response_with_signature(
471                        function_name.to_string(),
472                        response_value,
473                        signature,
474                    ));
475                } else {
476                    // Fallback: try legacy format or send as user message
477                    if let Some((name, response, sig)) = parse_tool_response(message.content()) {
478                        contents.push(GeminiContent::function_response_with_signature(
479                            name, response, sig,
480                        ));
481                    } else {
482                        // Last resort - wrap as text
483                        debug!("Tool message without matching tool_call_id: {}", base_id);
484                        contents.push(GeminiContent::text("user", message.content()));
485                    }
486                }
487            }
488            Role::Assistant => {
489                let tool_calls = message.tool_calls();
490                if tool_calls.is_empty() {
491                    // Regular text response
492                    contents.push(GeminiContent::text("model", message.content()));
493                } else {
494                    // Assistant message with function calls
495                    let mut parts = Vec::new();
496
497                    // Add text content if present
498                    if !message.content().is_empty() {
499                        parts.push(Part::text(message.content()));
500                    }
501
502                    // Add function call parts with thought signatures extracted from IDs
503                    for tc in tool_calls {
504                        let signature = tc
505                            .reasoning_state
506                            .as_ref()
507                            .and_then(|state| state.payload_for(PROVIDER_NAME))
508                            .map(String::from);
509                        parts.push(Part::function_call_with_signature(
510                            tc.name.clone(),
511                            tc.arguments.clone(),
512                            signature,
513                        ));
514                    }
515
516                    contents.push(GeminiContent::with_parts("model", parts));
517                }
518            }
519        }
520    }
521
522    let system_instruction = if system_parts.is_empty() {
523        None
524    } else {
525        Some(GeminiContent::system(system_parts))
526    };
527
528    Ok((system_instruction, contents))
529}
530
531fn format_prompt_feedback(feedback: &PromptFeedback) -> String {
532    let reason = feedback
533        .block_reason
534        .as_deref()
535        .unwrap_or("unknown reason")
536        .to_string();
537    if feedback.safety_ratings.is_empty() {
538        return format!("Gemini response blocked: {reason}");
539    }
540
541    let ratings = feedback
542        .safety_ratings
543        .iter()
544        .map(format_safety_rating)
545        .collect::<Vec<_>>()
546        .join(", ");
547    format!("Gemini response blocked: {reason}; safety: {ratings}")
548}
549
550fn format_safety_rating(rating: &SafetyRating) -> String {
551    let status = rating
552        .blocked
553        .map_or("unspecified", |b| if b { "blocked" } else { "allowed" });
554    let probability = rating.probability.as_deref().unwrap_or("unknown");
555    format!("{} ({status}, probability: {probability})", rating.category)
556}
557
558/// # Errors
559///
560/// Returns [`GeminiError`] when the requested reasoning effort has no Gemini
561/// equivalent.
562fn build_generation_config(
563    parameters: &Parameters,
564    modalities: Option<Vec<String>>,
565) -> Result<Option<GenerationConfig>, GeminiError> {
566    let thinking_config = build_thinking_config(parameters)?;
567    let mut config = GenerationConfig {
568        temperature: parameters.temperature,
569        top_p: parameters.top_p,
570        top_k: parameters.top_k,
571        max_output_tokens: parameters.max_tokens.map(
572            #[allow(clippy::cast_possible_wrap)]
573            |value| value as i32,
574        ),
575        stop_sequences: parameters.stop.clone(),
576        response_modalities: modalities,
577        thinking_config,
578        seed: parameters.seed.map(
579            #[allow(clippy::cast_possible_wrap)]
580            |value| value as i32,
581        ),
582        presence_penalty: parameters.presence_penalty,
583        frequency_penalty: parameters.frequency_penalty,
584        response_logprobs: parameters.logprobs,
585        logprobs: parameters.top_logprobs.map(i32::from),
586        ..Default::default()
587    };
588    if let Some(schema) = &parameters.response_format {
589        config.response_mime_type = Some("application/json".into());
590        config.response_json_schema = Some(schema.clone().to_value());
591    } else if parameters.structured_outputs {
592        config.response_mime_type = Some("application/json".into());
593    }
594    if config.is_meaningful() {
595        Ok(Some(config))
596    } else {
597        Ok(None)
598    }
599}
600
601fn build_tool_config(parameters: &Parameters, has_tools: bool) -> Option<ToolConfig> {
602    if !has_tools {
603        return None;
604    }
605    let (mode, allowed) = match &parameters.tool_choice {
606        ToolChoice::None => return None,
607        ToolChoice::Auto => (FunctionCallingMode::Auto, None),
608        ToolChoice::Required => (FunctionCallingMode::Any, None),
609        ToolChoice::Exact(name) => (FunctionCallingMode::Any, Some(vec![name.clone()])),
610    };
611    Some(ToolConfig {
612        function_calling_config: Some(FunctionCallingConfig {
613            mode,
614            allowed_function_names: allowed,
615        }),
616    })
617}
618
619/// Builds `thinkingConfig` from the portable reasoning parameters.
620///
621/// Depth (`reasoning_effort`) and visibility (`include_reasoning`) are separate
622/// controls: asking for a deeper think does not require the thoughts to be
623/// returned, and vice versa. Either one alone is enough to emit the config.
624/// # Errors
625///
626/// Returns [`GeminiError`] when the requested effort has no Gemini equivalent.
627fn build_thinking_config(parameters: &Parameters) -> Result<Option<ThinkingConfig>, GeminiError> {
628    let thinking_level = parameters
629        .reasoning_effort
630        .map(ThinkingLevel::try_from)
631        .transpose()?;
632    if !parameters.include_reasoning && thinking_level.is_none() {
633        return Ok(None);
634    }
635    Ok(Some(ThinkingConfig {
636        include_thoughts: parameters.include_reasoning.then_some(true),
637        thinking_level,
638    }))
639}
640
641fn convert_tool_definitions(defs: Vec<ToolDefinition>) -> Vec<FunctionDeclaration> {
642    defs.into_iter()
643        .map(|tool| FunctionDeclaration {
644            name: tool.name().to_string(),
645            description: tool.description().to_string(),
646            parameters: Some(tool.arguments_openai_schema()),
647        })
648        .collect()
649}
650
651/// Convert a typed attachment to a Gemini content part.
652async fn url_to_part(attachment: &Attachment) -> Result<Part, GeminiError> {
653    let url = attachment.url();
654    let media_type = attachment.media_type().as_ref();
655    match url.scheme() {
656        "data" => parse_data_url(url.as_str(), media_type),
657        "file" => read_file_to_part(url, media_type).await,
658        "http" | "https" => Ok(Part::from_file(media_type, url.as_str())),
659        scheme => Err(GeminiError::Api(format!(
660            "Gemini does not support attachment URL scheme '{scheme}'"
661        ))),
662    }
663}
664
665fn parse_data_url(url: &str, media_type: &str) -> Result<Part, GeminiError> {
666    use base64::Engine;
667
668    let after_data = url
669        .strip_prefix("data:")
670        .ok_or_else(|| GeminiError::Api("Malformed attachment data URL".to_string()))?;
671    let (header, data) = after_data
672        .split_once(',')
673        .ok_or_else(|| GeminiError::Api("Attachment data URL is missing payload".to_string()))?;
674    let encoded_media_type = header.strip_suffix(";base64").ok_or_else(|| {
675        GeminiError::Api("Attachment data URL must use base64 encoding".to_string())
676    })?;
677    if encoded_media_type != media_type {
678        return Err(GeminiError::Api(format!(
679            "Attachment MIME type '{media_type}' does not match data URL MIME type '{encoded_media_type}'"
680        )));
681    }
682    let bytes = base64::engine::general_purpose::STANDARD.decode(data)?;
683    Ok(Part::inline_media(media_type, bytes))
684}
685
686#[cfg(not(target_arch = "wasm32"))]
687async fn read_file_to_part(url: &url::Url, media_type: &str) -> Result<Part, GeminiError> {
688    let path = url.to_file_path().map_err(|()| {
689        GeminiError::Api("Attachment file URL could not be converted to a path".to_string())
690    })?;
691    let data = async_fs::read(&path).await.map_err(|error| {
692        GeminiError::Api(format!(
693            "Failed to read attachment '{}': {error}",
694            path.display()
695        ))
696    })?;
697    Ok(Part::inline_media(media_type, data))
698}
699
700#[cfg(target_arch = "wasm32")]
701async fn read_file_to_part(_url: &url::Url, _media_type: &str) -> Result<Part, GeminiError> {
702    Err(GeminiError::Api(
703        "file:// attachments are not supported on wasm32".to_string(),
704    ))
705}
706
707#[cfg(test)]
708mod tests {
709    use super::*;
710    use aither_core::llm::model::ReasoningEffort;
711
712    #[tokio::test]
713    async fn data_attachments_preserve_declared_media_types() {
714        for media_type in ["image/png", "audio/wav", "video/mp4", "application/pdf"] {
715            let attachment = Attachment::new(
716                format!("data:{media_type};base64,AA==")
717                    .parse()
718                    .expect("data URL"),
719                media_type.parse().expect("MIME type"),
720            );
721            let part = url_to_part(&attachment).await.expect("encode attachment");
722            let value = serde_json::to_value(part).expect("serialize part");
723            assert_eq!(value["inlineData"]["mimeType"], media_type);
724        }
725    }
726
727    #[tokio::test]
728    async fn remote_attachment_uses_file_data_with_media_type() {
729        let attachment = Attachment::new(
730            "https://ai.google.dev/static/gemini-api/docs/files/sample.pdf"
731                .parse()
732                .expect("remote URL"),
733            "application/pdf".parse().expect("PDF MIME"),
734        );
735        let part = url_to_part(&attachment).await.expect("encode attachment");
736        let value = serde_json::to_value(part).expect("serialize part");
737        assert_eq!(value["fileData"]["mimeType"], "application/pdf");
738        assert_eq!(
739            value["fileData"]["fileUri"],
740            "https://ai.google.dev/static/gemini-api/docs/files/sample.pdf"
741        );
742    }
743
744    #[test]
745    fn data_url_mime_mismatch_fails() {
746        let error = parse_data_url("data:image/png;base64,AA==", "audio/wav")
747            .expect_err("MIME mismatch must fail");
748        assert!(error.to_string().contains("does not match"));
749    }
750
751    fn thinking_json(params: &Parameters) -> serde_json::Value {
752        let config = build_thinking_config(params)
753            .expect("effort supported by Gemini")
754            .expect("thinking config");
755        serde_json::to_value(config).expect("serialize thinking config")
756    }
757
758    /// Gemini names the field `thinkingLevel`; there is no `tokenBudget`, so a
759    /// misnamed field is accepted by serde and then ignored on the wire.
760    #[test]
761    fn thinking_config_uses_gemini_field_names() {
762        let value = thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::Medium));
763        assert_eq!(value["thinkingLevel"], "medium");
764        assert_eq!(value.get("tokenBudget"), None);
765        assert_eq!(value.get("thinkingBudget"), None);
766    }
767
768    /// `thinkingLevel` and the legacy `thinkingBudget` are mutually exclusive:
769    /// sending both is a 400, so only the level is ever modelled.
770    #[test]
771    fn thinking_config_never_pairs_a_budget_with_a_level() {
772        for effort in [
773            ReasoningEffort::Minimal,
774            ReasoningEffort::Low,
775            ReasoningEffort::Medium,
776            ReasoningEffort::High,
777        ] {
778            let value = thinking_json(&Parameters::default().reasoning_effort(effort));
779            let object = value.as_object().expect("thinking config object");
780            assert!(
781                object.keys().all(|key| key != "thinkingBudget"),
782                "effort {effort:?} paired a budget with a level"
783            );
784            assert!(object.contains_key("thinkingLevel"));
785        }
786    }
787
788    #[test]
789    fn minimal_effort_maps_to_the_minimal_level() {
790        let value =
791            thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::Minimal));
792        assert_eq!(value["thinkingLevel"], "minimal");
793    }
794
795    /// Gemini's ladder stops at High. Clamping `XHigh` or `Max` down to High would
796    /// hand back less reasoning than the caller asked for, silently.
797    #[test]
798    fn efforts_above_gemini_ceiling_are_rejected() {
799        for effort in [ReasoningEffort::XHigh, ReasoningEffort::Max] {
800            let error = build_thinking_config(&Parameters::default().reasoning_effort(effort))
801                .expect_err("effort above the Gemini ceiling must fail");
802            assert!(
803                error.to_string().contains("tops out at High"),
804                "unexpected error for {effort:?}: {error}"
805            );
806        }
807    }
808
809    /// The signature used to be base64'd into the tool-call id behind a `|ts|`
810    /// separator, because core had nowhere to put it. It now rides on the call,
811    /// and the id is a plain identifier again.
812    #[tokio::test]
813    async fn thought_signature_round_trips_without_mangling_the_id() {
814        let call = aither_core::llm::ToolCall::new("gemini_abc", "lookup", serde_json::json!({}))
815            .with_reasoning_state(ReasoningState::new(PROVIDER_NAME, "sig-xyz"));
816        let message = Message::assistant_with_tool_calls("", vec![call]);
817
818        let (_, contents) = messages_to_gemini(&[message]).await.expect("convert");
819        let value = serde_json::to_value(&contents).expect("serialize contents");
820        let part = &value[0]["parts"][0];
821
822        assert_eq!(part["thoughtSignature"], "sig-xyz");
823        assert_eq!(part["functionCall"]["name"], "lookup");
824    }
825
826    #[tokio::test]
827    async fn generated_tool_call_ids_carry_no_smuggled_signature() {
828        assert!(!tool_call_id().contains("|ts|"));
829    }
830
831    /// Signatures are Gemini's; another provider's state must not be replayed.
832    #[tokio::test]
833    async fn foreign_reasoning_state_is_not_sent_as_a_signature() {
834        let call = aither_core::llm::ToolCall::new("gemini_abc", "lookup", serde_json::json!({}))
835            .with_reasoning_state(ReasoningState::new("anthropic", "sig-from-claude"));
836        let message = Message::assistant_with_tool_calls("", vec![call]);
837
838        let (_, contents) = messages_to_gemini(&[message]).await.expect("convert");
839        let value = serde_json::to_value(&contents).expect("serialize contents");
840        assert_eq!(value[0]["parts"][0].get("thoughtSignature"), None);
841    }
842
843    /// Gemini has no "no thinking" level; the caller should omit effort instead.
844    #[test]
845    fn none_effort_is_rejected() {
846        let error =
847            build_thinking_config(&Parameters::default().reasoning_effort(ReasoningEffort::None))
848                .expect_err("None effort must fail");
849        assert!(error.to_string().contains("no 'none' thinking level"));
850    }
851
852    /// Depth and visibility are independent: asking for effort alone must still
853    /// produce a config, and asking for thoughts alone must still produce one.
854    #[test]
855    fn effort_and_visibility_are_independent() {
856        let effort_only =
857            thinking_json(&Parameters::default().reasoning_effort(ReasoningEffort::High));
858        assert_eq!(effort_only["thinkingLevel"], "high");
859        assert_eq!(effort_only.get("includeThoughts"), None);
860
861        let thoughts_only = thinking_json(&Parameters::default().include_reasoning(true));
862        assert_eq!(thoughts_only["includeThoughts"], true);
863        assert_eq!(thoughts_only.get("thinkingLevel"), None);
864
865        assert!(
866            build_thinking_config(&Parameters::default())
867                .expect("no effort requested")
868                .is_none()
869        );
870    }
871
872    /// These reached the wire only if `build_generation_config` populates them;
873    /// they were declared on the struct but never filled.
874    #[test]
875    fn generation_config_forwards_every_sampling_parameter() {
876        let params = Parameters::default()
877            .seed(42)
878            .presence_penalty(0.5)
879            .frequency_penalty(0.25)
880            .logprobs(true)
881            .top_logprobs(3);
882        let config = build_generation_config(&params, None)
883            .expect("effort supported by Gemini")
884            .expect("generation config");
885        let value = serde_json::to_value(config).expect("serialize generation config");
886        assert_eq!(value["seed"], 42);
887        assert_eq!(value["presencePenalty"], 0.5);
888        assert_eq!(value["frequencyPenalty"], 0.25);
889        assert_eq!(value["responseLogprobs"], true);
890        assert_eq!(value["logprobs"], 3);
891    }
892}