tenant.openagents/omega

1use std::mem;
2
3use anyhow::{Result, anyhow, bail};
4use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
5use http_client::{
6    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
7};
8pub use language_model_core::ModelMode as GoogleModelMode;
9use serde::{Deserialize, Deserializer, Serialize, Serializer};
10pub mod completion;
11
12pub const API_URL: &str = "https://generativelanguage.googleapis.com";
13
14pub async fn stream_generate_content(
15    client: &dyn HttpClient,
16    api_url: &str,
17    api_key: &str,
18    mut request: GenerateContentRequest,
19    extra_headers: &CustomHeaders,
20) -> Result<BoxStream<'static, Result<GenerateContentResponse>>> {
21    let api_key = api_key.trim();
22    validate_generate_content_request(&request)?;
23
24    // The `model` field is emptied as it is provided as a path parameter.
25    let model_id = mem::take(&mut request.model.model_id);
26
27    let uri =
28        format!("{api_url}/v1beta/models/{model_id}:streamGenerateContent?alt=sse&key={api_key}",);
29
30    let request = HttpRequest::builder()
31        .method(Method::POST)
32        .uri(uri)
33        .header("Content-Type", "application/json")
34        .extra_headers(extra_headers)
35        .body(AsyncBody::from(serde_json::to_string(&request)?))?;
36    let mut response = client.send(request).await?;
37    if response.status().is_success() {
38        let reader = BufReader::new(response.into_body());
39        Ok(reader
40            .lines()
41            .filter_map(|line| async move {
42                match line {
43                    Ok(line) => {
44                        if let Some(line) = line.strip_prefix("data: ") {
45                            match serde_json::from_str(line) {
46                                Ok(response) => Some(Ok(response)),
47                                Err(error) => Some(Err(anyhow!(format!(
48                                    "Error parsing JSON: {error:?}\n{line:?}"
49                                )))),
50                            }
51                        } else {
52                            None
53                        }
54                    }
55                    Err(error) => Some(Err(anyhow!(error))),
56                }
57            })
58            .boxed())
59    } else {
60        let mut text = String::new();
61        response.body_mut().read_to_string(&mut text).await?;
62        Err(anyhow!(
63            "error during streamGenerateContent, status code: {:?}, body: {}",
64            response.status(),
65            text
66        ))
67    }
68}
69
70pub fn validate_generate_content_request(request: &GenerateContentRequest) -> Result<()> {
71    if request.model.is_empty() {
72        bail!("Model must be specified");
73    }
74
75    if request.contents.is_empty() {
76        bail!("Request must contain at least one content item");
77    }
78
79    if let Some(user_content) = request
80        .contents
81        .iter()
82        .find(|content| content.role == Role::User)
83        && user_content.parts.is_empty()
84    {
85        bail!("User content must contain at least one part");
86    }
87
88    Ok(())
89}
90
91#[derive(Debug, Serialize, Deserialize)]
92pub enum Task {
93    #[serde(rename = "generateContent")]
94    GenerateContent,
95    #[serde(rename = "streamGenerateContent")]
96    StreamGenerateContent,
97    #[serde(rename = "embedContent")]
98    EmbedContent,
99    #[serde(rename = "batchEmbedContents")]
100    BatchEmbedContents,
101}
102
103#[derive(Debug, Serialize, Deserialize)]
104#[serde(rename_all = "camelCase")]
105pub struct GenerateContentRequest {
106    #[serde(default, skip_serializing_if = "ModelName::is_empty")]
107    pub model: ModelName,
108    pub contents: Vec<Content>,
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub system_instruction: Option<SystemInstruction>,
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub generation_config: Option<GenerationConfig>,
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub safety_settings: Option<Vec<SafetySetting>>,
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub tools: Option<Vec<Tool>>,
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub tool_config: Option<ToolConfig>,
119}
120
121#[derive(Debug, Serialize, Deserialize)]
122#[serde(rename_all = "camelCase")]
123pub struct GenerateContentResponse {
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub candidates: Option<Vec<GenerateContentCandidate>>,
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub prompt_feedback: Option<PromptFeedback>,
128    #[serde(skip_serializing_if = "Option::is_none")]
129    pub usage_metadata: Option<UsageMetadata>,
130}
131
132#[derive(Debug, Serialize, Deserialize)]
133#[serde(rename_all = "camelCase")]
134pub struct GenerateContentCandidate {
135    #[serde(skip_serializing_if = "Option::is_none")]
136    pub index: Option<usize>,
137    pub content: Content,
138    #[serde(skip_serializing_if = "Option::is_none")]
139    pub finish_reason: Option<String>,
140    #[serde(skip_serializing_if = "Option::is_none")]
141    pub finish_message: Option<String>,
142    #[serde(skip_serializing_if = "Option::is_none")]
143    pub safety_ratings: Option<Vec<SafetyRating>>,
144    #[serde(skip_serializing_if = "Option::is_none")]
145    pub citation_metadata: Option<CitationMetadata>,
146}
147
148#[derive(Debug, Serialize, Deserialize)]
149#[serde(rename_all = "camelCase")]
150pub struct Content {
151    #[serde(default)]
152    pub parts: Vec<Part>,
153    pub role: Role,
154}
155
156#[derive(Debug, Serialize, Deserialize)]
157#[serde(rename_all = "camelCase")]
158pub struct SystemInstruction {
159    pub parts: Vec<Part>,
160}
161
162#[derive(Debug, PartialEq, Deserialize, Serialize)]
163#[serde(rename_all = "camelCase")]
164pub enum Role {
165    User,
166    Model,
167}
168
169#[derive(Debug, Serialize, Deserialize)]
170#[serde(untagged)]
171pub enum Part {
172    FunctionCallPart(FunctionCallPart),
173    FunctionResponsePart(FunctionResponsePart),
174    InlineDataPart(InlineDataPart),
175    TextPart(TextPart),
176}
177
178fn is_false(value: &bool) -> bool {
179    !*value
180}
181
182#[derive(Debug, Serialize, Deserialize)]
183#[serde(rename_all = "camelCase")]
184pub struct TextPart {
185    pub text: String,
186    #[serde(default, skip_serializing_if = "is_false")]
187    pub thought: bool,
188    #[serde(skip_serializing_if = "Option::is_none")]
189    pub thought_signature: Option<String>,
190}
191
192#[derive(Debug, Serialize, Deserialize)]
193#[serde(rename_all = "camelCase")]
194pub struct InlineDataPart {
195    pub inline_data: GenerativeContentBlob,
196}
197
198#[derive(Debug, Serialize, Deserialize)]
199#[serde(rename_all = "camelCase")]
200pub struct GenerativeContentBlob {
201    pub mime_type: String,
202    pub data: String,
203}
204
205#[derive(Debug, Serialize, Deserialize)]
206#[serde(rename_all = "camelCase")]
207pub struct FunctionCallPart {
208    pub function_call: FunctionCall,
209    /// Thought signature returned by the model for function calls.
210    /// Only present on the first function call in parallel call scenarios.
211    #[serde(skip_serializing_if = "Option::is_none")]
212    pub thought_signature: Option<String>,
213}
214
215#[derive(Debug, Serialize, Deserialize)]
216#[serde(rename_all = "camelCase")]
217pub struct FunctionResponsePart {
218    pub function_response: FunctionResponse,
219}
220
221#[derive(Debug, Serialize, Deserialize)]
222#[serde(rename_all = "camelCase")]
223pub struct CitationSource {
224    #[serde(skip_serializing_if = "Option::is_none")]
225    pub start_index: Option<usize>,
226    #[serde(skip_serializing_if = "Option::is_none")]
227    pub end_index: Option<usize>,
228    #[serde(skip_serializing_if = "Option::is_none")]
229    pub uri: Option<String>,
230    #[serde(skip_serializing_if = "Option::is_none")]
231    pub license: Option<String>,
232}
233
234#[derive(Debug, Serialize, Deserialize)]
235#[serde(rename_all = "camelCase")]
236pub struct CitationMetadata {
237    pub citation_sources: Vec<CitationSource>,
238}
239
240#[derive(Debug, Serialize, Deserialize)]
241#[serde(rename_all = "camelCase")]
242pub struct PromptFeedback {
243    #[serde(skip_serializing_if = "Option::is_none")]
244    pub block_reason: Option<String>,
245    pub safety_ratings: Option<Vec<SafetyRating>>,
246    #[serde(skip_serializing_if = "Option::is_none")]
247    pub block_reason_message: Option<String>,
248}
249
250#[derive(Debug, Serialize, Deserialize, Default)]
251#[serde(rename_all = "camelCase")]
252pub struct UsageMetadata {
253    #[serde(skip_serializing_if = "Option::is_none")]
254    pub prompt_token_count: Option<u64>,
255    #[serde(skip_serializing_if = "Option::is_none")]
256    pub cached_content_token_count: Option<u64>,
257    #[serde(skip_serializing_if = "Option::is_none")]
258    pub candidates_token_count: Option<u64>,
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub tool_use_prompt_token_count: Option<u64>,
261    #[serde(skip_serializing_if = "Option::is_none")]
262    pub thoughts_token_count: Option<u64>,
263    #[serde(skip_serializing_if = "Option::is_none")]
264    pub total_token_count: Option<u64>,
265}
266
267#[derive(Debug, Serialize, Deserialize, Default)]
268#[serde(rename_all = "camelCase")]
269pub struct ThinkingConfig {
270    #[serde(skip_serializing_if = "Option::is_none")]
271    pub thinking_budget: Option<u32>,
272    #[serde(skip_serializing_if = "Option::is_none")]
273    pub thinking_level: Option<ThinkingLevel>,
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub include_thoughts: Option<bool>,
276}
277
278#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
279#[serde(rename_all = "UPPERCASE")]
280pub enum ThinkingLevel {
281    Minimal,
282    Low,
283    Medium,
284    High,
285}
286
287impl ThinkingLevel {
288    pub fn from_effort(effort: &str) -> Option<Self> {
289        match effort.to_lowercase().as_str() {
290            "minimal" => Some(Self::Minimal),
291            "low" => Some(Self::Low),
292            "medium" => Some(Self::Medium),
293            "high" => Some(Self::High),
294            _ => None,
295        }
296    }
297
298    pub fn name(self) -> &'static str {
299        match self {
300            Self::Minimal => "Minimal",
301            Self::Low => "Low",
302            Self::Medium => "Medium",
303            Self::High => "High",
304        }
305    }
306
307    pub fn value(self) -> &'static str {
308        match self {
309            Self::Minimal => "minimal",
310            Self::Low => "low",
311            Self::Medium => "medium",
312            Self::High => "high",
313        }
314    }
315}
316
317#[derive(Debug, Deserialize, Serialize)]
318#[serde(rename_all = "camelCase")]
319pub struct GenerationConfig {
320    #[serde(skip_serializing_if = "Option::is_none")]
321    pub candidate_count: Option<usize>,
322    #[serde(skip_serializing_if = "Option::is_none")]
323    pub stop_sequences: Option<Vec<String>>,
324    #[serde(skip_serializing_if = "Option::is_none")]
325    pub max_output_tokens: Option<usize>,
326    #[serde(skip_serializing_if = "Option::is_none")]
327    pub temperature: Option<f64>,
328    #[serde(skip_serializing_if = "Option::is_none")]
329    pub top_p: Option<f64>,
330    #[serde(skip_serializing_if = "Option::is_none")]
331    pub top_k: Option<usize>,
332    #[serde(skip_serializing_if = "Option::is_none")]
333    pub thinking_config: Option<ThinkingConfig>,
334}
335
336#[derive(Debug, Serialize, Deserialize)]
337#[serde(rename_all = "camelCase")]
338pub struct SafetySetting {
339    pub category: HarmCategory,
340    pub threshold: HarmBlockThreshold,
341}
342
343#[derive(Debug, Serialize, Deserialize)]
344pub enum HarmCategory {
345    #[serde(rename = "HARM_CATEGORY_UNSPECIFIED")]
346    Unspecified,
347    #[serde(rename = "HARM_CATEGORY_DEROGATORY")]
348    Derogatory,
349    #[serde(rename = "HARM_CATEGORY_TOXICITY")]
350    Toxicity,
351    #[serde(rename = "HARM_CATEGORY_VIOLENCE")]
352    Violence,
353    #[serde(rename = "HARM_CATEGORY_SEXUAL")]
354    Sexual,
355    #[serde(rename = "HARM_CATEGORY_MEDICAL")]
356    Medical,
357    #[serde(rename = "HARM_CATEGORY_DANGEROUS")]
358    Dangerous,
359    #[serde(rename = "HARM_CATEGORY_HARASSMENT")]
360    Harassment,
361    #[serde(rename = "HARM_CATEGORY_HATE_SPEECH")]
362    HateSpeech,
363    #[serde(rename = "HARM_CATEGORY_SEXUALLY_EXPLICIT")]
364    SexuallyExplicit,
365    #[serde(rename = "HARM_CATEGORY_DANGEROUS_CONTENT")]
366    DangerousContent,
367}
368
369#[derive(Debug, Serialize, Deserialize)]
370#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
371pub enum HarmBlockThreshold {
372    #[serde(rename = "HARM_BLOCK_THRESHOLD_UNSPECIFIED")]
373    Unspecified,
374    BlockLowAndAbove,
375    BlockMediumAndAbove,
376    BlockOnlyHigh,
377    BlockNone,
378}
379
380#[derive(Debug, Serialize, Deserialize)]
381#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
382pub enum HarmProbability {
383    #[serde(rename = "HARM_PROBABILITY_UNSPECIFIED")]
384    Unspecified,
385    Negligible,
386    Low,
387    Medium,
388    High,
389}
390
391#[derive(Debug, Serialize, Deserialize)]
392#[serde(rename_all = "camelCase")]
393pub struct SafetyRating {
394    pub category: HarmCategory,
395    pub probability: HarmProbability,
396}
397
398#[derive(Debug, Serialize, Deserialize)]
399pub struct FunctionCall {
400    pub name: String,
401    pub args: serde_json::Value,
402    #[serde(skip_serializing_if = "Option::is_none")]
403    pub id: Option<String>,
404}
405
406#[derive(Debug, Serialize, Deserialize)]
407pub struct FunctionResponse {
408    pub name: String,
409    pub response: serde_json::Value,
410    #[serde(skip_serializing_if = "Option::is_none")]
411    pub id: Option<String>,
412}
413
414#[derive(Debug, Serialize, Deserialize)]
415#[serde(rename_all = "camelCase")]
416pub struct Tool {
417    pub function_declarations: Vec<FunctionDeclaration>,
418}
419
420#[derive(Debug, Serialize, Deserialize)]
421#[serde(rename_all = "camelCase")]
422pub struct ToolConfig {
423    pub function_calling_config: FunctionCallingConfig,
424}
425
426#[derive(Debug, Serialize, Deserialize)]
427#[serde(rename_all = "camelCase")]
428pub struct FunctionCallingConfig {
429    pub mode: FunctionCallingMode,
430    #[serde(skip_serializing_if = "Option::is_none")]
431    pub allowed_function_names: Option<Vec<String>>,
432}
433
434#[derive(Debug, Serialize, Deserialize)]
435#[serde(rename_all = "lowercase")]
436pub enum FunctionCallingMode {
437    Auto,
438    Any,
439    None,
440}
441
442#[derive(Debug, Serialize, Deserialize)]
443pub struct FunctionDeclaration {
444    pub name: String,
445    pub description: String,
446    pub parameters: serde_json::Value,
447}
448
449#[derive(Debug, Default)]
450pub struct ModelName {
451    pub model_id: String,
452}
453
454impl ModelName {
455    pub fn is_empty(&self) -> bool {
456        self.model_id.is_empty()
457    }
458}
459
460const MODEL_NAME_PREFIX: &str = "models/";
461
462impl Serialize for ModelName {
463    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
464    where
465        S: Serializer,
466    {
467        serializer.serialize_str(&format!("{MODEL_NAME_PREFIX}{}", &self.model_id))
468    }
469}
470
471impl<'de> Deserialize<'de> for ModelName {
472    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
473    where
474        D: Deserializer<'de>,
475    {
476        let string = String::deserialize(deserializer)?;
477        if let Some(id) = string.strip_prefix(MODEL_NAME_PREFIX) {
478            Ok(Self {
479                model_id: id.to_string(),
480            })
481        } else {
482            Err(serde::de::Error::custom(format!(
483                "Expected model name to begin with {}, got: {}",
484                MODEL_NAME_PREFIX, string
485            )))
486        }
487    }
488}
489
490#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
491#[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq, Eq, strum::EnumIter)]
492pub enum Model {
493    #[serde(
494        rename = "gemini-2.5-flash-lite",
495        alias = "gemini-2.5-flash-lite-preview-06-17",
496        alias = "gemini-2.0-flash-lite-preview"
497    )]
498    Gemini25FlashLite,
499    #[serde(
500        rename = "gemini-2.5-flash",
501        alias = "gemini-2.0-flash-thinking-exp",
502        alias = "gemini-2.5-flash-preview-04-17",
503        alias = "gemini-2.5-flash-preview-05-20",
504        alias = "gemini-2.5-flash-preview-latest",
505        alias = "gemini-2.0-flash"
506    )]
507    #[default]
508    Gemini25Flash,
509    #[serde(
510        rename = "gemini-2.5-pro",
511        alias = "gemini-2.0-pro-exp",
512        alias = "gemini-2.5-pro-preview-latest",
513        alias = "gemini-2.5-pro-exp-03-25",
514        alias = "gemini-2.5-pro-preview-03-25",
515        alias = "gemini-2.5-pro-preview-05-06",
516        alias = "gemini-2.5-pro-preview-06-05"
517    )]
518    Gemini25Pro,
519    #[serde(rename = "gemini-3.1-flash-lite")]
520    Gemini31FlashLite,
521    #[serde(rename = "gemini-3-flash-preview")]
522    Gemini3Flash,
523    #[serde(rename = "gemini-3.5-flash")]
524    Gemini35Flash,
525    #[serde(rename = "gemini-3.6-flash")]
526    Gemini36Flash,
527    #[serde(rename = "gemini-3.1-pro-preview", alias = "gemini-3-pro-preview")]
528    Gemini31Pro,
529    #[serde(rename = "custom")]
530    Custom {
531        name: String,
532        /// The name displayed in the UI, such as in the agent panel model dropdown menu.
533        display_name: Option<String>,
534        max_tokens: u64,
535        #[serde(default)]
536        mode: GoogleModelMode,
537    },
538}
539
540impl Model {
541    pub fn default_fast() -> Self {
542        Self::Gemini31FlashLite
543    }
544
545    pub fn id(&self) -> &str {
546        match self {
547            Self::Gemini25FlashLite => "gemini-2.5-flash-lite",
548            Self::Gemini25Flash => "gemini-2.5-flash",
549            Self::Gemini25Pro => "gemini-2.5-pro",
550            Self::Gemini31FlashLite => "gemini-3.1-flash-lite",
551            Self::Gemini3Flash => "gemini-3-flash-preview",
552            Self::Gemini35Flash => "gemini-3.5-flash",
553            Self::Gemini36Flash => "gemini-3.6-flash",
554            Self::Gemini31Pro => "gemini-3.1-pro-preview",
555            Self::Custom { name, .. } => name,
556        }
557    }
558    pub fn request_id(&self) -> &str {
559        match self {
560            Self::Gemini25FlashLite => "gemini-2.5-flash-lite",
561            Self::Gemini25Flash => "gemini-2.5-flash",
562            Self::Gemini25Pro => "gemini-2.5-pro",
563            Self::Gemini31FlashLite => "gemini-3.1-flash-lite",
564            Self::Gemini3Flash => "gemini-3-flash-preview",
565            Self::Gemini35Flash => "gemini-3.5-flash",
566            Self::Gemini36Flash => "gemini-3.6-flash",
567            Self::Gemini31Pro => "gemini-3.1-pro-preview",
568            Self::Custom { name, .. } => name,
569        }
570    }
571
572    pub fn display_name(&self) -> &str {
573        match self {
574            Self::Gemini25FlashLite => "Gemini 2.5 Flash-Lite",
575            Self::Gemini25Flash => "Gemini 2.5 Flash",
576            Self::Gemini25Pro => "Gemini 2.5 Pro",
577            Self::Gemini31FlashLite => "Gemini 3.1 Flash-Lite",
578            Self::Gemini3Flash => "Gemini 3 Flash",
579            Self::Gemini35Flash => "Gemini 3.5 Flash",
580            Self::Gemini36Flash => "Gemini 3.6 Flash",
581            Self::Gemini31Pro => "Gemini 3.1 Pro",
582            Self::Custom {
583                name, display_name, ..
584            } => display_name.as_ref().unwrap_or(name),
585        }
586    }
587
588    pub fn max_token_count(&self) -> u64 {
589        match self {
590            Self::Gemini25FlashLite
591            | Self::Gemini25Flash
592            | Self::Gemini25Pro
593            | Self::Gemini31FlashLite
594            | Self::Gemini3Flash
595            | Self::Gemini35Flash
596            | Self::Gemini36Flash
597            | Self::Gemini31Pro => 1_048_576,
598            Self::Custom { max_tokens, .. } => *max_tokens,
599        }
600    }
601
602    pub fn max_output_tokens(&self) -> Option<u64> {
603        match self {
604            Model::Gemini25FlashLite
605            | Model::Gemini25Flash
606            | Model::Gemini25Pro
607            | Model::Gemini31FlashLite
608            | Model::Gemini3Flash
609            | Model::Gemini35Flash
610            | Model::Gemini36Flash
611            | Model::Gemini31Pro => Some(65_536),
612            Model::Custom { .. } => None,
613        }
614    }
615
616    pub fn supports_tools(&self) -> bool {
617        true
618    }
619
620    pub fn supports_images(&self) -> bool {
621        true
622    }
623
624    pub fn supports_thinking(&self) -> bool {
625        matches!(
626            self,
627            Self::Gemini25FlashLite
628                | Self::Gemini25Flash
629                | Self::Gemini25Pro
630                | Self::Gemini31FlashLite
631                | Self::Gemini3Flash
632                | Self::Gemini35Flash
633                | Self::Gemini36Flash
634                | Self::Gemini31Pro
635                | Self::Custom {
636                    mode: GoogleModelMode::Thinking { .. },
637                    ..
638                }
639        )
640    }
641
642    pub fn supported_thinking_levels(&self) -> &'static [ThinkingLevel] {
643        match self {
644            Self::Gemini31FlashLite
645            | Self::Gemini3Flash
646            | Self::Gemini35Flash
647            | Self::Gemini36Flash => &[
648                ThinkingLevel::Minimal,
649                ThinkingLevel::Low,
650                ThinkingLevel::Medium,
651                ThinkingLevel::High,
652            ],
653            Self::Gemini31Pro => &[
654                ThinkingLevel::Low,
655                ThinkingLevel::Medium,
656                ThinkingLevel::High,
657            ],
658            _ => &[],
659        }
660    }
661
662    pub fn default_thinking_level(&self) -> Option<ThinkingLevel> {
663        match self {
664            Self::Gemini31FlashLite => Some(ThinkingLevel::Minimal),
665            Self::Gemini3Flash => Some(ThinkingLevel::High),
666            Self::Gemini35Flash | Self::Gemini36Flash => Some(ThinkingLevel::Medium),
667            Self::Gemini31Pro => Some(ThinkingLevel::High),
668            _ => None,
669        }
670    }
671
672    pub fn mode(&self) -> GoogleModelMode {
673        match self {
674            Self::Gemini25FlashLite | Self::Gemini25Flash | Self::Gemini25Pro => {
675                GoogleModelMode::Thinking {
676                    // By default these models are set to "auto", so we preserve that behavior
677                    // but indicate they are capable of thinking mode
678                    budget_tokens: None,
679                }
680            }
681            Self::Gemini31FlashLite
682            | Self::Gemini3Flash
683            | Self::Gemini35Flash
684            | Self::Gemini36Flash
685            | Self::Gemini31Pro => GoogleModelMode::Thinking {
686                budget_tokens: None,
687            },
688            Self::Custom { mode, .. } => *mode,
689        }
690    }
691}
692
693impl std::fmt::Display for Model {
694    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
695        write!(f, "{}", self.id())
696    }
697}
698
699#[cfg(test)]
700mod tests {
701    use super::*;
702    use serde_json::json;
703
704    #[test]
705    fn test_function_call_part_with_signature_serializes_correctly() {
706        let part = FunctionCallPart {
707            function_call: FunctionCall {
708                name: "test_function".to_string(),
709                args: json!({"arg": "value"}),
710                id: None,
711            },
712            thought_signature: Some("test_signature".to_string()),
713        };
714
715        let serialized = serde_json::to_value(&part).unwrap();
716
717        assert_eq!(serialized["functionCall"]["name"], "test_function");
718        assert_eq!(serialized["functionCall"]["args"]["arg"], "value");
719        assert_eq!(serialized["thoughtSignature"], "test_signature");
720    }
721
722    #[test]
723    fn test_function_call_part_without_signature_omits_field() {
724        let part = FunctionCallPart {
725            function_call: FunctionCall {
726                name: "test_function".to_string(),
727                args: json!({"arg": "value"}),
728                id: None,
729            },
730            thought_signature: None,
731        };
732
733        let serialized = serde_json::to_value(&part).unwrap();
734
735        assert_eq!(serialized["functionCall"]["name"], "test_function");
736        assert_eq!(serialized["functionCall"]["args"]["arg"], "value");
737        // thoughtSignature field should not be present when None
738        assert!(serialized.get("thoughtSignature").is_none());
739    }
740
741    #[test]
742    fn test_function_call_part_deserializes_with_signature() {
743        let json = json!({
744            "functionCall": {
745                "name": "test_function",
746                "args": {"arg": "value"}
747            },
748            "thoughtSignature": "test_signature"
749        });
750
751        let part: FunctionCallPart = serde_json::from_value(json).unwrap();
752
753        assert_eq!(part.function_call.name, "test_function");
754        assert_eq!(part.thought_signature, Some("test_signature".to_string()));
755    }
756
757    #[test]
758    fn test_function_call_part_deserializes_without_signature() {
759        let json = json!({
760            "functionCall": {
761                "name": "test_function",
762                "args": {"arg": "value"}
763            }
764        });
765
766        let part: FunctionCallPart = serde_json::from_value(json).unwrap();
767
768        assert_eq!(part.function_call.name, "test_function");
769        assert_eq!(part.thought_signature, None);
770    }
771
772    #[test]
773    fn test_function_call_part_round_trip() {
774        let original = FunctionCallPart {
775            function_call: FunctionCall {
776                name: "test_function".to_string(),
777                args: json!({"arg": "value", "nested": {"key": "val"}}),
778                id: None,
779            },
780            thought_signature: Some("round_trip_signature".to_string()),
781        };
782
783        let serialized = serde_json::to_value(&original).unwrap();
784        let deserialized: FunctionCallPart = serde_json::from_value(serialized).unwrap();
785
786        assert_eq!(deserialized.function_call.name, original.function_call.name);
787        assert_eq!(deserialized.function_call.args, original.function_call.args);
788        assert_eq!(deserialized.thought_signature, original.thought_signature);
789    }
790
791    #[test]
792    fn test_function_call_part_with_empty_signature_serializes() {
793        let part = FunctionCallPart {
794            function_call: FunctionCall {
795                name: "test_function".to_string(),
796                args: json!({"arg": "value"}),
797                id: None,
798            },
799            thought_signature: Some("".to_string()),
800        };
801
802        let serialized = serde_json::to_value(&part).unwrap();
803
804        // Empty string should still be serialized (normalization happens at a higher level)
805        assert_eq!(serialized["thoughtSignature"], "");
806    }
807}
808