1use crate::llm::protocol::{FromProvider, ProtocolError, ProtocolResult, ToProvider};
26use bamboo_domain::{FunctionCall, ToolCall};
27use bamboo_domain::{FunctionSchema, ToolSchema};
28use bamboo_domain::{Message, Role};
29use serde::{Deserialize, Serialize};
30use serde_json::Value;
31
32pub struct GeminiProtocol;
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct GeminiRequest {
42 pub contents: Vec<GeminiContent>,
44 #[serde(
46 skip_serializing_if = "Option::is_none",
47 rename = "systemInstruction",
48 alias = "system_instruction"
49 )]
50 pub system_instruction: Option<GeminiContent>,
51 #[serde(skip_serializing_if = "Option::is_none")]
53 pub tools: Option<Vec<GeminiTool>>,
54 #[serde(
56 skip_serializing_if = "Option::is_none",
57 rename = "generationConfig",
58 alias = "generation_config"
59 )]
60 pub generation_config: Option<Value>,
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct GeminiContent {
66 pub role: String,
68 pub parts: Vec<GeminiPart>,
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct GeminiPart {
75 #[serde(skip_serializing_if = "Option::is_none")]
77 pub text: Option<String>,
78 #[serde(
80 skip_serializing_if = "Option::is_none",
81 rename = "inlineData",
82 alias = "inline_data"
83 )]
84 pub inline_data: Option<GeminiInlineData>,
85 #[serde(
87 skip_serializing_if = "Option::is_none",
88 rename = "fileData",
89 alias = "file_data"
90 )]
91 pub file_data: Option<GeminiFileData>,
92 #[serde(
94 skip_serializing_if = "Option::is_none",
95 rename = "functionCall",
96 alias = "function_call"
97 )]
98 pub function_call: Option<GeminiFunctionCall>,
99 #[serde(
101 skip_serializing_if = "Option::is_none",
102 rename = "functionResponse",
103 alias = "function_response"
104 )]
105 pub function_response: Option<GeminiFunctionResponse>,
106}
107
108#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct GeminiInlineData {
111 #[serde(rename = "mimeType", alias = "mime_type")]
112 pub mime_type: String,
113 pub data: String,
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct GeminiFileData {
119 #[serde(rename = "fileUri", alias = "file_uri")]
120 pub file_uri: String,
121 #[serde(
122 skip_serializing_if = "Option::is_none",
123 rename = "mimeType",
124 alias = "mime_type"
125 )]
126 pub mime_type: Option<String>,
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct GeminiFunctionCall {
132 pub name: String,
133 pub args: Value,
134}
135
136#[derive(Debug, Clone, Serialize, Deserialize)]
138pub struct GeminiFunctionResponse {
139 pub name: String,
140 pub response: Value,
141}
142
143#[derive(Debug, Clone, Serialize, Deserialize)]
145pub struct GeminiTool {
146 #[serde(rename = "functionDeclarations", alias = "function_declarations")]
147 pub function_declarations: Vec<GeminiFunctionDeclaration>,
148}
149
150#[derive(Debug, Clone, Serialize, Deserialize)]
152pub struct GeminiFunctionDeclaration {
153 pub name: String,
154 #[serde(skip_serializing_if = "Option::is_none")]
155 pub description: Option<String>,
156 pub parameters: Value,
157}
158
159#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct GeminiResponse {
162 pub candidates: Vec<GeminiCandidate>,
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct GeminiCandidate {
168 pub content: GeminiContent,
169 #[serde(skip_serializing_if = "Option::is_none")]
170 pub finish_reason: Option<String>,
171}
172
173impl FromProvider<GeminiContent> for Message {
178 fn from_provider(content: GeminiContent) -> ProtocolResult<Self> {
179 let role = match content.role.as_str() {
180 "user" => Role::User,
181 "model" => Role::Assistant,
182 "system" => Role::System,
183 _ => return Err(ProtocolError::InvalidRole(content.role)),
184 };
185
186 let mut text_parts = Vec::new();
188 let mut content_parts = Vec::new();
189 let mut tool_calls = Vec::new();
190 let mut has_image_parts = false;
191
192 for part in content.parts {
193 if let Some(text) = part.text {
194 text_parts.push(text.clone());
195 content_parts.push(bamboo_domain::MessagePart::Text { text });
196 }
197
198 if let Some(inline_data) = part.inline_data {
199 if let Some(url) = inline_data_to_data_url(&inline_data) {
200 has_image_parts = true;
201 content_parts.push(bamboo_domain::MessagePart::ImageUrl {
202 image_url: bamboo_domain::ImageUrlRef { url, detail: None },
203 });
204 }
205 }
206
207 if let Some(file_data) = part.file_data {
208 let file_uri = file_data.file_uri.trim();
209 if !file_uri.is_empty() {
210 has_image_parts = true;
211 content_parts.push(bamboo_domain::MessagePart::ImageUrl {
212 image_url: bamboo_domain::ImageUrlRef {
213 url: file_uri.to_string(),
214 detail: None,
215 },
216 });
217 }
218 }
219
220 if let Some(func_call) = part.function_call {
221 tool_calls.push(ToolCall {
222 id: format!("gemini_{}", uuid::Uuid::new_v4()), tool_type: "function".to_string(),
224 function: FunctionCall {
225 name: func_call.name,
226 arguments: serde_json::to_string(&func_call.args).unwrap_or_default(),
227 },
228 });
229 }
230
231 if let Some(func_response) = part.function_response {
232 return Ok(Message::tool_result(
234 format!("gemini_tool_{}", func_response.name),
235 serde_json::to_string(&func_response.response).unwrap_or_default(),
236 ));
237 }
238 }
239
240 let content_text = text_parts.join("");
241
242 Ok(Message {
243 id: String::new(),
244 role,
245 content: content_text,
246 reasoning: None,
247 content_parts: has_image_parts.then_some(content_parts),
248 image_ocr: None,
249 phase: None,
250 tool_calls: if tool_calls.is_empty() {
251 None
252 } else {
253 Some(tool_calls)
254 },
255 tool_call_id: None,
256 tool_success: None,
257 compressed: false,
258 compressed_by_event_id: None,
259 never_compress: false,
260 compression_level: 0,
261 created_at: chrono::Utc::now(),
262 metadata: None,
263 })
264 }
265}
266
267impl FromProvider<GeminiTool> for ToolSchema {
268 fn from_provider(tool: GeminiTool) -> ProtocolResult<Self> {
269 let func = tool
272 .function_declarations
273 .into_iter()
274 .next()
275 .ok_or_else(|| ProtocolError::InvalidToolCall("Empty tool declarations".to_string()))?;
276
277 Ok(ToolSchema {
278 schema_type: "function".to_string(),
279 function: FunctionSchema {
280 name: func.name,
281 description: func.description.unwrap_or_default(),
282 parameters: func.parameters,
283 },
284 })
285 }
286}
287
288pub struct GeminiRequestBuilder;
296
297impl ToProvider<GeminiRequest> for Vec<Message> {
298 fn to_provider(&self) -> ProtocolResult<GeminiRequest> {
299 let mut system_texts = Vec::new();
300 let mut contents = Vec::new();
301
302 for msg in self {
303 match msg.role {
304 Role::System => {
305 let trimmed = msg.content.trim();
306 if !trimmed.is_empty() {
307 system_texts.push(trimmed.to_string());
308 }
309 }
310 _ => {
311 contents.push(msg.to_provider()?);
312 }
313 }
314 }
315
316 let system_instruction = if system_texts.is_empty() {
317 None
318 } else {
319 Some(GeminiContent {
320 role: "system".to_string(),
321 parts: vec![GeminiPart {
322 text: Some(system_texts.join("\n\n")),
323 inline_data: None,
324 file_data: None,
325 function_call: None,
326 function_response: None,
327 }],
328 })
329 };
330
331 Ok(GeminiRequest {
332 contents,
333 system_instruction,
334 tools: None,
335 generation_config: None,
336 })
337 }
338}
339
340impl ToProvider<GeminiContent> for Message {
341 fn to_provider(&self) -> ProtocolResult<GeminiContent> {
342 if self.role == Role::Tool {
344 let tool_name = self
345 .tool_call_id
346 .clone()
347 .ok_or_else(|| ProtocolError::MissingField("tool_call_id".to_string()))?;
348
349 return Ok(GeminiContent {
350 role: "user".to_string(),
351 parts: vec![GeminiPart {
352 text: None,
353 inline_data: None,
354 file_data: None,
355 function_call: None,
356 function_response: Some(GeminiFunctionResponse {
357 name: tool_name,
358 response: serde_json::from_str(&self.content)
359 .unwrap_or_else(|_| Value::String(self.content.clone())),
360 }),
361 }],
362 });
363 }
364
365 let role = match self.role {
366 Role::User => "user",
367 Role::Assistant => "model",
368 Role::System => "system",
369 Role::Tool => "user", };
371
372 let mut parts = Vec::new();
373
374 if let Some(content_parts) = self.content_parts.as_ref() {
376 for part in content_parts {
377 if let Some(gemini_part) = message_content_part_to_gemini_part(part) {
378 parts.push(gemini_part);
379 }
380 }
381 }
382
383 if parts.is_empty() && !self.content.is_empty() {
385 parts.push(GeminiPart {
386 text: Some(self.content.clone()),
387 inline_data: None,
388 file_data: None,
389 function_call: None,
390 function_response: None,
391 });
392 }
393
394 if let Some(tool_calls) = &self.tool_calls {
396 for tc in tool_calls {
397 let args: Value = serde_json::from_str(&tc.function.arguments)
398 .unwrap_or_else(|_| Value::Object(serde_json::Map::new()));
399
400 parts.push(GeminiPart {
401 text: None,
402 inline_data: None,
403 file_data: None,
404 function_call: Some(GeminiFunctionCall {
405 name: tc.function.name.clone(),
406 args,
407 }),
408 function_response: None,
409 });
410 }
411 }
412
413 if parts.is_empty() {
415 parts.push(GeminiPart {
416 text: Some(String::new()),
417 inline_data: None,
418 file_data: None,
419 function_call: None,
420 function_response: None,
421 });
422 }
423
424 Ok(GeminiContent {
425 role: role.to_string(),
426 parts,
427 })
428 }
429}
430
431impl ToProvider<GeminiTool> for ToolSchema {
432 fn to_provider(&self) -> ProtocolResult<GeminiTool> {
433 Ok(GeminiTool {
434 function_declarations: vec![GeminiFunctionDeclaration {
435 name: self.function.name.clone(),
436 description: Some(self.function.description.clone()),
437 parameters: self.function.parameters.clone(),
438 }],
439 })
440 }
441}
442
443impl ToProvider<Vec<GeminiTool>> for Vec<ToolSchema> {
448 fn to_provider(&self) -> ProtocolResult<Vec<GeminiTool>> {
449 let declarations: Vec<GeminiFunctionDeclaration> = self
451 .iter()
452 .map(|schema| GeminiFunctionDeclaration {
453 name: schema.function.name.clone(),
454 description: Some(schema.function.description.clone()),
455 parameters: schema.function.parameters.clone(),
456 })
457 .collect();
458
459 if declarations.is_empty() {
460 Ok(vec![])
461 } else {
462 Ok(vec![GeminiTool {
463 function_declarations: declarations,
464 }])
465 }
466 }
467}
468
469fn message_content_part_to_gemini_part(part: &bamboo_domain::MessagePart) -> Option<GeminiPart> {
470 match part {
471 bamboo_domain::MessagePart::Text { text } => Some(GeminiPart {
472 text: Some(text.clone()),
473 inline_data: None,
474 file_data: None,
475 function_call: None,
476 function_response: None,
477 }),
478 bamboo_domain::MessagePart::ImageUrl { image_url } => {
479 image_url_to_gemini_part(&image_url.url)
480 }
481 }
482}
483
484fn image_url_to_gemini_part(url: &str) -> Option<GeminiPart> {
485 let trimmed = url.trim();
486 if trimmed.is_empty() {
487 return None;
488 }
489
490 if let Some((mime_type, data)) = parse_data_url_base64(trimmed) {
491 return Some(GeminiPart {
492 text: None,
493 inline_data: Some(GeminiInlineData { mime_type, data }),
494 file_data: None,
495 function_call: None,
496 function_response: None,
497 });
498 }
499
500 Some(GeminiPart {
501 text: None,
502 inline_data: None,
503 file_data: Some(GeminiFileData {
504 file_uri: trimmed.to_string(),
505 mime_type: None,
506 }),
507 function_call: None,
508 function_response: None,
509 })
510}
511
512fn parse_data_url_base64(url: &str) -> Option<(String, String)> {
513 let rest = url.strip_prefix("data:")?;
514 let (meta, data) = rest.split_once(',')?;
515 let data = data.trim();
516 if data.is_empty() {
517 return None;
518 }
519
520 let mut mime_type = "application/octet-stream";
521 let mut is_base64 = false;
522 for (idx, seg) in meta.split(';').enumerate() {
523 let segment = seg.trim();
524 if idx == 0 && !segment.is_empty() && !segment.eq_ignore_ascii_case("base64") {
525 mime_type = segment;
526 }
527 if segment.eq_ignore_ascii_case("base64") {
528 is_base64 = true;
529 }
530 }
531
532 if !is_base64 {
533 return None;
534 }
535
536 Some((mime_type.to_string(), data.to_string()))
537}
538
539fn inline_data_to_data_url(inline: &GeminiInlineData) -> Option<String> {
540 let mime_type = inline.mime_type.trim();
541 let data = inline.data.trim();
542 if mime_type.is_empty() || data.is_empty() {
543 return None;
544 }
545 Some(format!("data:{mime_type};base64,{data}"))
546}
547
548pub trait GeminiExt: Sized {
554 fn into_internal(self) -> ProtocolResult<Message>;
555 fn to_gemini(&self) -> ProtocolResult<GeminiContent>;
556}
557
558impl GeminiExt for GeminiContent {
559 fn into_internal(self) -> ProtocolResult<Message> {
560 Message::from_provider(self)
561 }
562
563 fn to_gemini(&self) -> ProtocolResult<GeminiContent> {
564 Ok(self.clone())
565 }
566}
567
568impl GeminiExt for Message {
569 fn into_internal(self) -> ProtocolResult<Message> {
570 Ok(self)
571 }
572
573 fn to_gemini(&self) -> ProtocolResult<GeminiContent> {
574 self.to_provider()
575 }
576}
577
578#[cfg(test)]
583mod tests {
584 use super::*;
585 use crate::models::{ContentPart, ImageUrl};
586 use bamboo_domain::MessagePart;
587
588 #[test]
589 fn test_gemini_to_internal_user_message() {
590 let gemini = GeminiContent {
591 role: "user".to_string(),
592 parts: vec![GeminiPart {
593 text: Some("Hello".to_string()),
594 inline_data: None,
595 file_data: None,
596 function_call: None,
597 function_response: None,
598 }],
599 };
600
601 let internal: Message = Message::from_provider(gemini).unwrap();
602
603 assert_eq!(internal.role, Role::User);
604 assert_eq!(internal.content, "Hello");
605 assert!(internal.tool_calls.is_none());
606 }
607
608 #[test]
609 fn test_internal_to_gemini_user_message() {
610 let internal = Message::user("Hello");
611
612 let gemini: GeminiContent = internal.to_provider().unwrap();
613
614 assert_eq!(gemini.role, "user");
615 assert_eq!(gemini.parts.len(), 1);
616 assert_eq!(gemini.parts[0].text, Some("Hello".to_string()));
617 }
618
619 #[test]
620 fn test_internal_to_gemini_with_data_url_image_part() {
621 let internal = Message::user_with_parts(
622 "describe",
623 vec![
624 ContentPart::Text {
625 text: "describe".to_string(),
626 },
627 ContentPart::ImageUrl {
628 image_url: ImageUrl {
629 url: "data:image/png;base64,AAAA".to_string(),
630 detail: None,
631 },
632 },
633 ]
634 .into_iter()
635 .map(Into::into)
636 .collect(),
637 );
638
639 let gemini: GeminiContent = internal.to_provider().unwrap();
640
641 assert_eq!(gemini.parts.len(), 2);
642 assert_eq!(gemini.parts[0].text, Some("describe".to_string()));
643 let inline = gemini.parts[1]
644 .inline_data
645 .as_ref()
646 .expect("inlineData should be present");
647 assert_eq!(inline.mime_type, "image/png");
648 assert_eq!(inline.data, "AAAA");
649 assert!(gemini.parts[1].file_data.is_none());
650 }
651
652 #[test]
653 fn test_gemini_to_internal_model_message() {
654 let gemini = GeminiContent {
655 role: "model".to_string(),
656 parts: vec![GeminiPart {
657 text: Some("Hello there!".to_string()),
658 inline_data: None,
659 file_data: None,
660 function_call: None,
661 function_response: None,
662 }],
663 };
664
665 let internal: Message = Message::from_provider(gemini).unwrap();
666
667 assert_eq!(internal.role, Role::Assistant);
668 assert_eq!(internal.content, "Hello there!");
669 }
670
671 #[test]
672 fn test_gemini_to_internal_with_inline_data_image() {
673 let gemini = GeminiContent {
674 role: "user".to_string(),
675 parts: vec![GeminiPart {
676 text: Some("look".to_string()),
677 inline_data: Some(GeminiInlineData {
678 mime_type: "image/png".to_string(),
679 data: "BBBB".to_string(),
680 }),
681 file_data: None,
682 function_call: None,
683 function_response: None,
684 }],
685 };
686
687 let internal: Message = Message::from_provider(gemini).unwrap();
688 assert_eq!(internal.content, "look");
689 let parts = internal
690 .content_parts
691 .as_ref()
692 .expect("content_parts should preserve image");
693 assert!(parts.iter().any(|part| {
694 matches!(
695 part,
696 MessagePart::ImageUrl { image_url }
697 if image_url.url == "data:image/png;base64,BBBB"
698 )
699 }));
700 }
701
702 #[test]
703 fn test_internal_to_gemini_with_tool_call() {
704 let tool_call = ToolCall {
705 id: "call_1".to_string(),
706 tool_type: "function".to_string(),
707 function: FunctionCall {
708 name: "search".to_string(),
709 arguments: r#"{"q":"test"}"#.to_string(),
710 },
711 };
712
713 let internal = Message::assistant("Let me search", Some(vec![tool_call]));
714
715 let gemini: GeminiContent = internal.to_provider().unwrap();
716
717 assert_eq!(gemini.role, "model");
718 assert_eq!(gemini.parts.len(), 2);
719 assert_eq!(gemini.parts[0].text, Some("Let me search".to_string()));
720 assert!(gemini.parts[1].function_call.is_some());
721
722 let func_call = gemini.parts[1].function_call.as_ref().unwrap();
723 assert_eq!(func_call.name, "search");
724 assert_eq!(func_call.args, serde_json::json!({"q": "test"}));
725 }
726
727 #[test]
728 fn test_gemini_to_internal_with_tool_call() {
729 let gemini = GeminiContent {
730 role: "model".to_string(),
731 parts: vec![GeminiPart {
732 text: None,
733 inline_data: None,
734 file_data: None,
735 function_call: Some(GeminiFunctionCall {
736 name: "search".to_string(),
737 args: serde_json::json!({"q": "test"}),
738 }),
739 function_response: None,
740 }],
741 };
742
743 let internal: Message = Message::from_provider(gemini).unwrap();
744
745 assert_eq!(internal.role, Role::Assistant);
746 assert!(internal.tool_calls.is_some());
747
748 let tool_calls = internal.tool_calls.unwrap();
749 assert_eq!(tool_calls.len(), 1);
750 assert_eq!(tool_calls[0].function.name, "search");
751 }
752
753 #[test]
754 fn test_system_message_extraction() {
755 let messages = vec![Message::system("You are helpful"), Message::user("Hello")];
756
757 let request: GeminiRequest = messages.to_provider().unwrap();
758
759 assert!(request.system_instruction.is_some());
760 let sys = request.system_instruction.unwrap();
761 assert_eq!(sys.role, "system");
762 assert_eq!(sys.parts[0].text, Some("You are helpful".to_string()));
763
764 assert_eq!(request.contents.len(), 1);
765 assert_eq!(request.contents[0].role, "user");
766 }
767
768 #[test]
769 fn test_multiple_system_messages_are_joined() {
770 let messages = vec![
771 Message::system("You are helpful"),
772 Message::system("Use tools when needed"),
773 Message::user("Hello"),
774 ];
775
776 let request: GeminiRequest = messages.to_provider().unwrap();
777
778 let sys = request
779 .system_instruction
780 .expect("system instruction should be present");
781 assert_eq!(sys.role, "system");
782 assert_eq!(
783 sys.parts[0].text.as_deref(),
784 Some("You are helpful\n\nUse tools when needed")
785 );
786 assert_eq!(request.contents.len(), 1);
787 assert_eq!(request.contents[0].role, "user");
788 }
789
790 #[test]
791 fn test_tool_response_conversion() {
792 let internal = Message::tool_result("search_tool", r#"{"result": "ok"}"#);
793
794 let gemini: GeminiContent = internal.to_provider().unwrap();
795
796 assert_eq!(gemini.role, "user");
797 assert!(gemini.parts[0].function_response.is_some());
798
799 let func_resp = gemini.parts[0].function_response.as_ref().unwrap();
800 assert_eq!(func_resp.name, "search_tool");
801 }
802
803 #[test]
804 fn test_tool_schema_conversion() {
805 let gemini_tool = GeminiTool {
806 function_declarations: vec![GeminiFunctionDeclaration {
807 name: "search".to_string(),
808 description: Some("Search the web".to_string()),
809 parameters: serde_json::json!({
810 "type": "object",
811 "properties": {
812 "q": { "type": "string" }
813 }
814 }),
815 }],
816 };
817
818 let internal_schema: ToolSchema = ToolSchema::from_provider(gemini_tool.clone()).unwrap();
820 assert_eq!(internal_schema.function.name, "search");
821
822 let roundtrip: GeminiTool = internal_schema.to_provider().unwrap();
824 assert_eq!(roundtrip.function_declarations.len(), 1);
825 assert_eq!(roundtrip.function_declarations[0].name, "search");
826 }
827
828 #[test]
829 fn test_multiple_tools_grouped() {
830 let tools = vec![
831 ToolSchema {
832 schema_type: "function".to_string(),
833 function: FunctionSchema {
834 name: "search".to_string(),
835 description: "Search".to_string(),
836 parameters: serde_json::json!({"type": "object"}),
837 },
838 },
839 ToolSchema {
840 schema_type: "function".to_string(),
841 function: FunctionSchema {
842 name: "read".to_string(),
843 description: "Read file".to_string(),
844 parameters: serde_json::json!({"type": "object"}),
845 },
846 },
847 ];
848
849 let gemini_tools: Vec<GeminiTool> = tools.to_provider().unwrap();
850
851 assert_eq!(gemini_tools.len(), 1);
853 assert_eq!(gemini_tools[0].function_declarations.len(), 2);
854 assert_eq!(gemini_tools[0].function_declarations[0].name, "search");
855 assert_eq!(gemini_tools[0].function_declarations[1].name, "read");
856 }
857
858 #[test]
859 fn test_roundtrip_conversion() {
860 let original = Message::user("Hello, world!");
861
862 let gemini: GeminiContent = original.to_provider().unwrap();
864
865 let roundtrip: Message = Message::from_provider(gemini).unwrap();
867
868 assert_eq!(roundtrip.role, original.role);
869 assert_eq!(roundtrip.content, original.content);
870 }
871
872 #[test]
873 fn test_invalid_role_error() {
874 let gemini = GeminiContent {
875 role: "invalid_role".to_string(),
876 parts: vec![GeminiPart {
877 text: Some("test".to_string()),
878 inline_data: None,
879 file_data: None,
880 function_call: None,
881 function_response: None,
882 }],
883 };
884
885 let result: ProtocolResult<Message> = Message::from_provider(gemini);
886 assert!(matches!(result, Err(ProtocolError::InvalidRole(_))));
887 }
888
889 #[test]
890 fn test_empty_parts_has_default() {
891 let internal = Message::assistant("", None);
892
893 let gemini: GeminiContent = internal.to_provider().unwrap();
894
895 assert_eq!(gemini.parts.len(), 1);
897 assert_eq!(gemini.parts[0].text, Some(String::new()));
898 }
899}