Skip to main content

rust_genai/
thinking.rs

1//! Thinking support and thought signature validation.
2
3use rust_genai_types::content::{Content, PartKind, Role};
4use rust_genai_types::models::GenerateContentConfig;
5
6use crate::error::{Error, Result};
7
8/// Thought Signature 验证器。
9pub struct ThoughtSignatureValidator {
10    model: String,
11}
12
13impl ThoughtSignatureValidator {
14    /// 创建验证器。
15    pub fn new(model: impl Into<String>) -> Self {
16        Self {
17            model: model.into(),
18        }
19    }
20
21    /// 验证对话历史中的 thought signatures。
22    ///
23    /// # Errors
24    ///
25    /// 当 `thought_signature` 缺失或不符合规则时返回错误。
26    pub fn validate(&self, contents: &[Content]) -> Result<()> {
27        if !is_gemini_3(&self.model) {
28            return Ok(());
29        }
30
31        let current_turn_start = find_current_turn_start(contents);
32
33        for content in &contents[current_turn_start..] {
34            if content.role != Some(Role::Model) {
35                continue;
36            }
37
38            let function_parts: Vec<_> = content
39                .parts
40                .iter()
41                .filter(|part| matches!(part.kind, PartKind::FunctionCall { .. }))
42                .collect();
43
44            if function_parts.is_empty() {
45                continue;
46            }
47
48            if function_parts[0].thought_signature.is_none() {
49                return Err(Error::MissingThoughtSignature {
50                    message: "First function call missing thought_signature".into(),
51                });
52            }
53
54            for part in function_parts.iter().skip(1) {
55                if part.thought_signature.is_some() {
56                    return Err(Error::MissingThoughtSignature {
57                        message: "Only the first function call may include thought_signature"
58                            .into(),
59                    });
60                }
61            }
62        }
63
64        Ok(())
65    }
66}
67
68/// Gemini 3 温度检查。
69///
70/// # Errors
71///
72/// 当前不会返回错误。
73pub fn validate_temperature(model: &str, config: &GenerateContentConfig) -> Result<()> {
74    if !is_gemini_3(model) {
75        return Ok(());
76    }
77
78    if let Some(temperature) = config
79        .generation_config
80        .as_ref()
81        .and_then(|cfg| cfg.temperature)
82    {
83        if temperature < 1.0 {
84            eprintln!(
85                "Warning: Gemini 3 temperature {temperature} < 1.0 may cause looping; use 1.0"
86            );
87        }
88    }
89
90    Ok(())
91}
92
93fn is_gemini_3(model: &str) -> bool {
94    model
95        .rsplit('/')
96        .next()
97        .is_some_and(|name| name.starts_with("gemini-3"))
98}
99
100fn find_current_turn_start(contents: &[Content]) -> usize {
101    for (idx, content) in contents.iter().enumerate().rev() {
102        if content.role != Some(Role::User) {
103            continue;
104        }
105        let has_text = content
106            .parts
107            .iter()
108            .any(|part| matches!(part.kind, PartKind::Text { .. }));
109        if has_text {
110            return idx;
111        }
112    }
113    0
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use rust_genai_types::content::{FunctionCall, Part};
120
121    #[test]
122    fn test_thought_signature_validation_gemini3() {
123        let validator = ThoughtSignatureValidator::new("gemini-3-pro-preview");
124        let contents = vec![
125            Content::user("Check flight AA100"),
126            Content::from_parts(
127                vec![Part::function_call(FunctionCall {
128                    id: None,
129                    name: Some("check_flight".into()),
130                    args: None,
131                    partial_args: None,
132                    will_continue: None,
133                })],
134                Role::Model,
135            ),
136        ];
137
138        assert!(validator.validate(&contents).is_err());
139    }
140
141    #[test]
142    fn test_temperature_warning_gemini3() {
143        let config = GenerateContentConfig {
144            generation_config: Some(rust_genai_types::config::GenerationConfig {
145                temperature: Some(0.5),
146                ..Default::default()
147            }),
148            ..Default::default()
149        };
150        validate_temperature("gemini-3-flash-preview", &config).unwrap();
151    }
152
153    #[test]
154    fn test_thought_signature_validation_non_gemini3_noop() {
155        let validator = ThoughtSignatureValidator::new("gemini-2.0-flash");
156        let contents = vec![Content::from_parts(
157            vec![Part::function_call(FunctionCall {
158                id: None,
159                name: Some("noop".into()),
160                args: None,
161                partial_args: None,
162                will_continue: None,
163            })],
164            Role::Model,
165        )];
166        assert!(validator.validate(&contents).is_ok());
167    }
168
169    #[test]
170    fn test_thought_signature_validation_allows_single_signature() {
171        let validator = ThoughtSignatureValidator::new("gemini-3-pro-preview");
172        let contents = vec![
173            Content::user("Plan"),
174            Content::from_parts(
175                vec![Part::function_call(FunctionCall {
176                    id: None,
177                    name: Some("plan".into()),
178                    args: None,
179                    partial_args: None,
180                    will_continue: None,
181                })
182                .with_thought_signature(vec![1, 2, 3])],
183                Role::Model,
184            ),
185        ];
186        assert!(validator.validate(&contents).is_ok());
187    }
188
189    #[test]
190    fn test_thought_signature_validation_rejects_multiple_signatures() {
191        let validator = ThoughtSignatureValidator::new("gemini-3-pro-preview");
192        let contents = vec![
193            Content::user("Plan"),
194            Content::from_parts(
195                vec![
196                    Part::function_call(FunctionCall {
197                        id: None,
198                        name: Some("step1".into()),
199                        args: None,
200                        partial_args: None,
201                        will_continue: None,
202                    })
203                    .with_thought_signature(vec![1]),
204                    Part::function_call(FunctionCall {
205                        id: None,
206                        name: Some("step2".into()),
207                        args: None,
208                        partial_args: None,
209                        will_continue: None,
210                    })
211                    .with_thought_signature(vec![2]),
212                ],
213                Role::Model,
214            ),
215        ];
216        assert!(validator.validate(&contents).is_err());
217    }
218}