Skip to repository content820 lines · 31.0 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T03:34:02.199Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
completion.rs
1use anyhow::Result;
2use futures::{Stream, StreamExt};
3use language_model_core::{
4 LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelRequest,
5 LanguageModelRequestToolInput, LanguageModelToolChoice, LanguageModelToolUse,
6 LanguageModelToolUseId, LanguageModelToolUseInput, MessageContent, Role, StopReason,
7 TokenUsage,
8};
9use std::pin::Pin;
10use std::sync::Arc;
11use std::sync::atomic::{self, AtomicU64};
12
13use crate::{
14 Content, FunctionCallingConfig, FunctionCallingMode, FunctionDeclaration,
15 GenerateContentResponse, GenerationConfig, GenerativeContentBlob, GoogleModelMode,
16 InlineDataPart, ModelName, Part, SystemInstruction, TextPart, ThinkingConfig, ThinkingLevel,
17 ToolConfig, UsageMetadata,
18};
19
20pub fn into_google(
21 mut request: LanguageModelRequest,
22 model_id: String,
23 mode: GoogleModelMode,
24) -> Result<crate::GenerateContentRequest> {
25 fn map_content(content: Vec<MessageContent>) -> Result<Vec<Part>> {
26 let mut mapped_parts = Vec::new();
27 for content in content {
28 match content {
29 MessageContent::Text(text) => {
30 if !text.is_empty() {
31 mapped_parts.push(Part::TextPart(TextPart {
32 text,
33 thought: false,
34 thought_signature: None,
35 }));
36 }
37 }
38 MessageContent::Thinking {
39 text,
40 signature: Some(signature),
41 } => {
42 if !signature.is_empty() {
43 mapped_parts.push(Part::TextPart(TextPart {
44 text,
45 thought: true,
46 thought_signature: Some(signature),
47 }));
48 }
49 }
50 MessageContent::Thinking { .. } => {}
51 MessageContent::RedactedThinking(_) | MessageContent::Compaction(_) => {}
52 MessageContent::Image(image) => {
53 mapped_parts.push(Part::InlineDataPart(InlineDataPart {
54 inline_data: GenerativeContentBlob {
55 mime_type: "image/png".to_string(),
56 data: image.source.to_string(),
57 },
58 }));
59 }
60 MessageContent::ToolUse(tool_use) => {
61 let thought_signature = tool_use.thought_signature.filter(|s| !s.is_empty());
62 let LanguageModelToolUseInput::Json(input) = tool_use.input else {
63 anyhow::bail!("Google AI does not support custom tool calls");
64 };
65
66 mapped_parts.push(Part::FunctionCallPart(crate::FunctionCallPart {
67 function_call: crate::FunctionCall {
68 name: tool_use.name.to_string(),
69 args: input,
70 id: Some(tool_use.id.to_string()),
71 },
72 thought_signature,
73 }));
74 }
75 MessageContent::ToolResult(tool_result) => {
76 let mut text_output = String::new();
77 let mut images: Vec<InlineDataPart> = Vec::new();
78 for part in tool_result.content {
79 match part {
80 language_model_core::LanguageModelToolResultContent::Text(text) => {
81 text_output.push_str(&text);
82 }
83 language_model_core::LanguageModelToolResultContent::Image(image) => {
84 images.push(InlineDataPart {
85 inline_data: GenerativeContentBlob {
86 mime_type: "image/png".to_string(),
87 data: image.source.to_string(),
88 },
89 });
90 }
91 }
92 }
93 let output = if text_output.is_empty() && !images.is_empty() {
94 "Tool responded with an image".to_string()
95 } else {
96 text_output
97 };
98 mapped_parts.push(Part::FunctionResponsePart(crate::FunctionResponsePart {
99 function_response: crate::FunctionResponse {
100 name: tool_result.tool_name.to_string(),
101 // The API expects a valid JSON object
102 response: serde_json::json!({
103 "output": output
104 }),
105 id: Some(tool_result.tool_use_id.to_string()),
106 },
107 }));
108 mapped_parts.extend(images.into_iter().map(Part::InlineDataPart));
109 }
110 }
111 }
112 Ok(mapped_parts)
113 }
114
115 let thinking_config = thinking_config_for_request(&request, &model_id, mode);
116
117 let system_instructions = if request
118 .messages
119 .first()
120 .is_some_and(|msg| matches!(msg.role, Role::System))
121 {
122 let message = request.messages.remove(0);
123 Some(SystemInstruction {
124 parts: map_content(message.content)?,
125 })
126 } else {
127 None
128 };
129
130 let tools = if request.tools.is_empty() {
131 None
132 } else {
133 Some(vec![crate::Tool {
134 function_declarations: request
135 .tools
136 .into_iter()
137 .map(|tool| match tool.input {
138 LanguageModelRequestToolInput::Function { input_schema, .. } => {
139 Ok(FunctionDeclaration {
140 name: tool.name,
141 description: tool.description,
142 parameters: input_schema,
143 })
144 }
145 LanguageModelRequestToolInput::Custom { .. } => {
146 Err(anyhow::anyhow!("Google AI does not support custom tools"))
147 }
148 })
149 .collect::<Result<_>>()?,
150 }])
151 };
152
153 Ok(crate::GenerateContentRequest {
154 model: ModelName { model_id },
155 system_instruction: system_instructions,
156 contents: request
157 .messages
158 .into_iter()
159 .map(|message| {
160 let parts = map_content(message.content)?;
161 if parts.is_empty() {
162 Ok(None)
163 } else {
164 Ok(Some(Content {
165 parts,
166 role: match message.role {
167 Role::User => crate::Role::User,
168 Role::Assistant => crate::Role::Model,
169 Role::System => crate::Role::User, // Google AI doesn't have a system role
170 },
171 }))
172 }
173 })
174 .collect::<Result<Vec<_>>>()?
175 .into_iter()
176 .flatten()
177 .collect(),
178 generation_config: Some(GenerationConfig {
179 candidate_count: Some(1),
180 stop_sequences: Some(request.stop),
181 max_output_tokens: None,
182 temperature: request.temperature.map(|t| t as f64),
183 thinking_config,
184 top_p: None,
185 top_k: None,
186 }),
187 safety_settings: None,
188 tools,
189 tool_config: request.tool_choice.map(|choice| ToolConfig {
190 function_calling_config: FunctionCallingConfig {
191 mode: match choice {
192 LanguageModelToolChoice::Auto => FunctionCallingMode::Auto,
193 LanguageModelToolChoice::Any => FunctionCallingMode::Any,
194 LanguageModelToolChoice::None => FunctionCallingMode::None,
195 },
196 allowed_function_names: None,
197 },
198 }),
199 })
200}
201
202fn thinking_config_for_request(
203 request: &LanguageModelRequest,
204 model_id: &str,
205 mode: GoogleModelMode,
206) -> Option<ThinkingConfig> {
207 let supports_thinking =
208 matches!(mode, GoogleModelMode::Thinking { .. }) || is_google_thinking_model(model_id);
209 if !supports_thinking {
210 return None;
211 }
212
213 let mut config = ThinkingConfig::default();
214
215 if request.thinking_allowed {
216 config.include_thoughts = Some(true);
217 config.thinking_level = request
218 .thinking_effort
219 .as_deref()
220 .and_then(ThinkingLevel::from_effort);
221
222 if config.thinking_level.is_none()
223 && let GoogleModelMode::Thinking {
224 budget_tokens: Some(budget_tokens),
225 } = mode
226 {
227 config.thinking_budget = Some(budget_tokens);
228 }
229 } else if let Some(thinking_level) = disabled_thinking_level(model_id) {
230 config.thinking_level = Some(thinking_level);
231 } else if supports_thinking_budget_disable(model_id) {
232 config.thinking_budget = Some(0);
233 }
234
235 (!config.is_empty()).then_some(config)
236}
237
238impl ThinkingConfig {
239 fn is_empty(&self) -> bool {
240 self.thinking_budget.is_none()
241 && self.thinking_level.is_none()
242 && self.include_thoughts.is_none()
243 }
244}
245
246fn is_google_thinking_model(model_id: &str) -> bool {
247 model_id.starts_with("gemini-2.5-") || model_id.starts_with("gemini-3")
248}
249
250fn disabled_thinking_level(model_id: &str) -> Option<ThinkingLevel> {
251 match model_id {
252 model_id if model_id.starts_with("gemini-3") && model_id.contains("-pro") => {
253 Some(ThinkingLevel::Low)
254 }
255 model_id if model_id.starts_with("gemini-3") => Some(ThinkingLevel::Minimal),
256 _ => None,
257 }
258}
259
260fn supports_thinking_budget_disable(model_id: &str) -> bool {
261 matches!(
262 model_id,
263 "gemini-2.5-flash"
264 | "gemini-2.5-flash-lite"
265 | "gemini-2.5-flash-preview-latest"
266 | "gemini-2.5-flash-preview-04-17"
267 | "gemini-2.5-flash-preview-05-20"
268 | "gemini-2.5-flash-lite-preview-06-17"
269 )
270}
271
272pub struct GoogleEventMapper {
273 usage: UsageMetadata,
274 stop_reason: StopReason,
275}
276
277impl GoogleEventMapper {
278 pub fn new() -> Self {
279 Self {
280 usage: UsageMetadata::default(),
281 stop_reason: StopReason::EndTurn,
282 }
283 }
284
285 pub fn map_stream(
286 mut self,
287 events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
288 ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
289 {
290 events
291 .map(Some)
292 .chain(futures::stream::once(async { None }))
293 .flat_map(move |event| {
294 futures::stream::iter(match event {
295 Some(Ok(event)) => self.map_event(event),
296 Some(Err(error)) => {
297 vec![Err(LanguageModelCompletionError::from(error))]
298 }
299 None => vec![Ok(LanguageModelCompletionEvent::Stop(self.stop_reason))],
300 })
301 })
302 }
303
304 pub fn map_event(
305 &mut self,
306 event: GenerateContentResponse,
307 ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
308 static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
309
310 let mut events: Vec<_> = Vec::new();
311 let mut wants_to_use_tool = false;
312 if let Some(usage_metadata) = event.usage_metadata {
313 update_usage(&mut self.usage, &usage_metadata);
314 events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(
315 convert_usage(&self.usage),
316 )))
317 }
318
319 if let Some(prompt_feedback) = event.prompt_feedback
320 && let Some(block_reason) = prompt_feedback.block_reason.as_deref()
321 {
322 self.stop_reason = match block_reason {
323 "SAFETY" | "OTHER" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "IMAGE_SAFETY" => {
324 StopReason::Refusal
325 }
326 _ => {
327 log::error!("Unexpected Google block_reason: {block_reason}");
328 StopReason::Refusal
329 }
330 };
331 events.push(Ok(LanguageModelCompletionEvent::Stop(self.stop_reason)));
332
333 return events;
334 }
335
336 if let Some(candidates) = event.candidates {
337 for candidate in candidates {
338 if let Some(finish_reason) = candidate.finish_reason.as_deref() {
339 self.stop_reason = match finish_reason {
340 "STOP" => StopReason::EndTurn,
341 "MAX_TOKENS" => StopReason::MaxTokens,
342 "SAFETY"
343 | "RECITATION"
344 | "LANGUAGE"
345 | "OTHER"
346 | "BLOCKLIST"
347 | "PROHIBITED_CONTENT"
348 | "SPII"
349 | "MALFORMED_FUNCTION_CALL"
350 | "IMAGE_SAFETY"
351 | "IMAGE_PROHIBITED_CONTENT"
352 | "IMAGE_OTHER"
353 | "NO_IMAGE"
354 | "IMAGE_RECITATION"
355 | "UNEXPECTED_TOOL_CALL"
356 | "TOO_MANY_TOOL_CALLS"
357 | "MISSING_THOUGHT_SIGNATURE"
358 | "MALFORMED_RESPONSE" => StopReason::Refusal,
359 _ => {
360 log::error!("Unexpected google finish_reason: {finish_reason}");
361 StopReason::EndTurn
362 }
363 };
364 }
365 candidate
366 .content
367 .parts
368 .into_iter()
369 .for_each(|part| match part {
370 Part::TextPart(text_part) => {
371 let thought_signature =
372 text_part.thought_signature.filter(|s| !s.is_empty());
373 if text_part.thought {
374 if !text_part.text.is_empty() || thought_signature.is_some() {
375 events.push(Ok(LanguageModelCompletionEvent::Thinking {
376 text: text_part.text,
377 signature: thought_signature,
378 }))
379 }
380 } else {
381 if let Some(thought_signature) = thought_signature {
382 events.push(Ok(LanguageModelCompletionEvent::Thinking {
383 text: String::new(),
384 signature: Some(thought_signature),
385 }));
386 }
387 if !text_part.text.is_empty() {
388 events.push(Ok(LanguageModelCompletionEvent::Text(
389 text_part.text,
390 )));
391 }
392 }
393 }
394 Part::InlineDataPart(_) => {}
395 Part::FunctionCallPart(function_call_part) => {
396 wants_to_use_tool = true;
397 let name: Arc<str> = function_call_part.function_call.name.into();
398 let id: LanguageModelToolUseId =
399 if let Some(ref call_id) = function_call_part.function_call.id {
400 call_id.clone().into()
401 } else {
402 let next_tool_id =
403 TOOL_CALL_COUNTER.fetch_add(1, atomic::Ordering::SeqCst);
404 format!("{}-{}", name, next_tool_id).into()
405 };
406
407 // Normalize empty string signatures to None
408 let thought_signature = function_call_part
409 .thought_signature
410 .filter(|s| !s.is_empty());
411
412 events.push(Ok(LanguageModelCompletionEvent::ToolUse(
413 LanguageModelToolUse {
414 id,
415 name,
416 is_input_complete: true,
417 raw_input: function_call_part.function_call.args.to_string(),
418 input: LanguageModelToolUseInput::Json(
419 function_call_part.function_call.args,
420 ),
421 thought_signature,
422 },
423 )));
424 }
425 Part::FunctionResponsePart(_) => {}
426 });
427 }
428 }
429
430 // Even when Gemini wants to use a Tool, the API
431 // responds with `finish_reason: STOP`
432 if wants_to_use_tool {
433 self.stop_reason = StopReason::ToolUse;
434 events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
435 }
436 events
437 }
438}
439
440fn update_usage(usage: &mut UsageMetadata, new: &UsageMetadata) {
441 if let Some(prompt_token_count) = new.prompt_token_count {
442 usage.prompt_token_count = Some(prompt_token_count);
443 }
444 if let Some(cached_content_token_count) = new.cached_content_token_count {
445 usage.cached_content_token_count = Some(cached_content_token_count);
446 }
447 if let Some(candidates_token_count) = new.candidates_token_count {
448 usage.candidates_token_count = Some(candidates_token_count);
449 }
450 if let Some(tool_use_prompt_token_count) = new.tool_use_prompt_token_count {
451 usage.tool_use_prompt_token_count = Some(tool_use_prompt_token_count);
452 }
453 if let Some(thoughts_token_count) = new.thoughts_token_count {
454 usage.thoughts_token_count = Some(thoughts_token_count);
455 }
456 if let Some(total_token_count) = new.total_token_count {
457 usage.total_token_count = Some(total_token_count);
458 }
459}
460
461fn convert_usage(usage: &UsageMetadata) -> TokenUsage {
462 let prompt_tokens = usage.prompt_token_count.unwrap_or(0);
463 let cached_tokens = usage.cached_content_token_count.unwrap_or(0);
464 let input_tokens = prompt_tokens - cached_tokens;
465 let output_tokens = usage.candidates_token_count.unwrap_or(0);
466
467 TokenUsage {
468 input_tokens,
469 output_tokens,
470 cache_read_input_tokens: cached_tokens,
471 cache_creation_input_tokens: 0,
472 }
473}
474
475#[cfg(test)]
476mod tests {
477 use super::*;
478 use crate::{
479 Content, FunctionCall, FunctionCallPart, GenerateContentCandidate, GenerateContentResponse,
480 Part, Role as GoogleRole,
481 };
482 use language_model_core::LanguageModelRequestMessage;
483 use serde_json::json;
484
485 fn text_request() -> LanguageModelRequest {
486 LanguageModelRequest {
487 messages: vec![LanguageModelRequestMessage {
488 role: Role::User,
489 content: vec![MessageContent::Text("Hello".to_string())],
490 cache: false,
491 reasoning_details: None,
492 }],
493 ..Default::default()
494 }
495 }
496
497 #[test]
498 fn into_google_requests_thought_summaries_and_thinking_level() {
499 let mut request = text_request();
500 request.thinking_allowed = true;
501 request.thinking_effort = Some("low".to_string());
502
503 let request = into_google(
504 request,
505 "gemini-3.5-flash".to_string(),
506 GoogleModelMode::Thinking {
507 budget_tokens: None,
508 },
509 )
510 .unwrap();
511
512 let thinking_config = request.generation_config.unwrap().thinking_config.unwrap();
513 assert_eq!(thinking_config.include_thoughts, Some(true));
514 assert_eq!(thinking_config.thinking_level, Some(ThinkingLevel::Low));
515
516 let serialized = serde_json::to_value(thinking_config).unwrap();
517 assert_eq!(serialized["thinkingLevel"], "LOW");
518 assert_eq!(serialized["includeThoughts"], true);
519 }
520
521 #[test]
522 fn into_google_turns_off_budget_thinking_when_supported() {
523 let mut request = text_request();
524 request.thinking_allowed = false;
525
526 let request = into_google(
527 request,
528 "gemini-2.5-flash".to_string(),
529 GoogleModelMode::Thinking {
530 budget_tokens: None,
531 },
532 )
533 .unwrap();
534
535 let thinking_config = request.generation_config.unwrap().thinking_config.unwrap();
536 assert_eq!(thinking_config.thinking_budget, Some(0));
537 assert_eq!(thinking_config.include_thoughts, None);
538 }
539
540 #[test]
541 fn into_google_uses_minimal_level_when_gemini_3_flash_thinking_is_off() {
542 let mut request = text_request();
543 request.thinking_allowed = false;
544
545 let request = into_google(
546 request,
547 "gemini-3.5-flash".to_string(),
548 GoogleModelMode::Thinking {
549 budget_tokens: None,
550 },
551 )
552 .unwrap();
553
554 let thinking_config = request.generation_config.unwrap().thinking_config.unwrap();
555 assert_eq!(thinking_config.thinking_level, Some(ThinkingLevel::Minimal));
556 assert_eq!(thinking_config.include_thoughts, None);
557 }
558
559 #[test]
560 fn into_google_replays_signed_thinking_as_thought_text_part() {
561 let request = LanguageModelRequest {
562 messages: vec![LanguageModelRequestMessage {
563 role: Role::Assistant,
564 content: vec![MessageContent::Thinking {
565 text: "summary".to_string(),
566 signature: Some("signature".to_string()),
567 }],
568 cache: false,
569 reasoning_details: None,
570 }],
571 ..Default::default()
572 };
573
574 let request = into_google(
575 request,
576 "gemini-3.5-flash".to_string(),
577 GoogleModelMode::Thinking {
578 budget_tokens: None,
579 },
580 )
581 .unwrap();
582
583 let Part::TextPart(text_part) = &request.contents[0].parts[0] else {
584 panic!("expected text part");
585 };
586 assert_eq!(text_part.text, "summary");
587 assert!(text_part.thought);
588 assert_eq!(text_part.thought_signature.as_deref(), Some("signature"));
589 }
590
591 #[test]
592 fn thought_text_part_deserializes_and_maps_to_thinking_event() {
593 let part: Part = serde_json::from_value(json!({
594 "text": "checking the constraints",
595 "thought": true,
596 "thoughtSignature": "thought-signature"
597 }))
598 .unwrap();
599
600 let mut mapper = GoogleEventMapper::new();
601 let response = GenerateContentResponse {
602 candidates: Some(vec![GenerateContentCandidate {
603 index: Some(0),
604 content: Content {
605 parts: vec![part],
606 role: GoogleRole::Model,
607 },
608 finish_reason: None,
609 finish_message: None,
610 safety_ratings: None,
611 citation_metadata: None,
612 }]),
613 prompt_feedback: None,
614 usage_metadata: None,
615 };
616
617 let events = mapper.map_event(response);
618 assert_eq!(events.len(), 1);
619 assert!(matches!(
620 &events[0],
621 Ok(LanguageModelCompletionEvent::Thinking { text, signature })
622 if text == "checking the constraints"
623 && signature.as_deref() == Some("thought-signature")
624 ));
625 }
626
627 #[test]
628 fn signed_non_thought_text_part_preserves_signature() {
629 let part: Part = serde_json::from_value(json!({
630 "text": "visible text",
631 "thoughtSignature": "visible-signature"
632 }))
633 .unwrap();
634
635 let Part::TextPart(text_part) = part else {
636 panic!("expected text part");
637 };
638 assert_eq!(text_part.text, "visible text");
639 assert!(!text_part.thought);
640 assert_eq!(
641 text_part.thought_signature.as_deref(),
642 Some("visible-signature")
643 );
644 }
645
646 #[test]
647 fn signed_non_thought_text_part_maps_signature_carrier() {
648 let part: Part = serde_json::from_value(json!({
649 "text": "visible text",
650 "thoughtSignature": "visible-signature"
651 }))
652 .unwrap();
653
654 let mut mapper = GoogleEventMapper::new();
655 let response = GenerateContentResponse {
656 candidates: Some(vec![GenerateContentCandidate {
657 index: Some(0),
658 content: Content {
659 parts: vec![part],
660 role: GoogleRole::Model,
661 },
662 finish_reason: None,
663 finish_message: None,
664 safety_ratings: None,
665 citation_metadata: None,
666 }]),
667 prompt_feedback: None,
668 usage_metadata: None,
669 };
670
671 let events = mapper.map_event(response);
672 assert_eq!(events.len(), 2);
673 assert!(matches!(
674 &events[0],
675 Ok(LanguageModelCompletionEvent::Thinking { text, signature })
676 if text.is_empty() && signature.as_deref() == Some("visible-signature")
677 ));
678 assert!(matches!(
679 &events[1],
680 Ok(LanguageModelCompletionEvent::Text(text)) if text == "visible text"
681 ));
682 }
683
684 #[test]
685 fn safety_finish_reason_is_refusal() {
686 let mut mapper = GoogleEventMapper::new();
687 let response = GenerateContentResponse {
688 candidates: Some(vec![GenerateContentCandidate {
689 index: Some(0),
690 content: Content {
691 parts: Vec::new(),
692 role: GoogleRole::Model,
693 },
694 finish_reason: Some("SAFETY".to_string()),
695 finish_message: None,
696 safety_ratings: None,
697 citation_metadata: None,
698 }]),
699 prompt_feedback: None,
700 usage_metadata: None,
701 };
702
703 mapper.map_event(response);
704 assert_eq!(mapper.stop_reason, StopReason::Refusal);
705 }
706
707 #[test]
708 fn test_function_call_with_signature_creates_tool_use_with_signature() {
709 let mut mapper = GoogleEventMapper::new();
710
711 let response = GenerateContentResponse {
712 candidates: Some(vec![GenerateContentCandidate {
713 index: Some(0),
714 content: Content {
715 parts: vec![Part::FunctionCallPart(FunctionCallPart {
716 function_call: FunctionCall {
717 name: "test_function".to_string(),
718 args: json!({"arg": "value"}),
719 id: None,
720 },
721 thought_signature: Some("test_signature_123".to_string()),
722 })],
723 role: GoogleRole::Model,
724 },
725 finish_reason: None,
726 finish_message: None,
727 safety_ratings: None,
728 citation_metadata: None,
729 }]),
730 prompt_feedback: None,
731 usage_metadata: None,
732 };
733
734 let events = mapper.map_event(response);
735 assert_eq!(events.len(), 2);
736
737 if let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] {
738 assert_eq!(tool_use.name.as_ref(), "test_function");
739 assert_eq!(
740 tool_use.thought_signature.as_deref(),
741 Some("test_signature_123")
742 );
743 } else {
744 panic!("Expected ToolUse event");
745 }
746 }
747
748 #[test]
749 fn test_function_call_without_signature_has_none() {
750 let mut mapper = GoogleEventMapper::new();
751
752 let response = GenerateContentResponse {
753 candidates: Some(vec![GenerateContentCandidate {
754 index: Some(0),
755 content: Content {
756 parts: vec![Part::FunctionCallPart(FunctionCallPart {
757 function_call: FunctionCall {
758 name: "test_function".to_string(),
759 args: json!({"arg": "value"}),
760 id: None,
761 },
762 thought_signature: None,
763 })],
764 role: GoogleRole::Model,
765 },
766 finish_reason: None,
767 finish_message: None,
768 safety_ratings: None,
769 citation_metadata: None,
770 }]),
771 prompt_feedback: None,
772 usage_metadata: None,
773 };
774
775 let events = mapper.map_event(response);
776 assert_eq!(events.len(), 2);
777
778 if let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] {
779 assert!(tool_use.thought_signature.is_none());
780 } else {
781 panic!("Expected ToolUse event");
782 }
783 }
784
785 #[test]
786 fn test_empty_string_signature_normalized_to_none() {
787 let mut mapper = GoogleEventMapper::new();
788
789 let response = GenerateContentResponse {
790 candidates: Some(vec![GenerateContentCandidate {
791 index: Some(0),
792 content: Content {
793 parts: vec![Part::FunctionCallPart(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 role: GoogleRole::Model,
802 },
803 finish_reason: None,
804 finish_message: None,
805 safety_ratings: None,
806 citation_metadata: None,
807 }]),
808 prompt_feedback: None,
809 usage_metadata: None,
810 };
811
812 let events = mapper.map_event(response);
813 if let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] {
814 assert!(tool_use.thought_signature.is_none());
815 } else {
816 panic!("Expected ToolUse event");
817 }
818 }
819}
820