From fd1223da59dfabe186e68cc0727764bdd621be18 Mon Sep 17 00:00:00 2001
From: Andre Bandarra <andreban@gmail.com>
Date: Wed, 27 Nov 2024 16:20:53 +0000
Subject: [PATCH] More builder refactoring

---
 examples/system_instruction.rs       | 14 ++++++++++----
 examples/text-from-text-streaming.rs | 10 ++++++----
 examples/text-from-text.rs           | 10 ++++++----
 src/types/generate_content.rs        | 15 +++------------
 4 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/examples/system_instruction.rs b/examples/system_instruction.rs
index 1b57cec..fc1c63f 100644
--- a/examples/system_instruction.rs
+++ b/examples/system_instruction.rs
@@ -15,12 +15,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         location_id,
     );
 
-    let system_instruction = "Answer as if you were Yoda";
-    let prompt = "What is the airspeed of an unladen swallow?";
+    let system_instruction = Content::builder()
+        .add_text_part("Answer as if you were Yoda")
+        .build();
+
+    let user_prompt = vec![Content::builder()
+        .role(Role::User)
+        .add_text_part("What is the airspeed of an unladen swallow?")
+        .build()];
 
     let request = GenerateContentRequest::builder()
-        .add_text_content(Role::User, prompt)
-        .system_instruction_text(system_instruction)
+        .contents(user_prompt)
+        .system_instruction(system_instruction)
         .build();
 
     let result = gemini
diff --git a/examples/text-from-text-streaming.rs b/examples/text-from-text-streaming.rs
index 5a54c63..06de5f4 100644
--- a/examples/text-from-text-streaming.rs
+++ b/examples/text-from-text-streaming.rs
@@ -14,10 +14,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         location_id,
     );
 
-    let prompt = "Tell me the story of the genesis of the universe as a bedtime story.";
-    let request = GenerateContentRequest::builder()
-        .add_text_content(Role::User, prompt)
-        .build();
+    let prompt = vec![Content::builder()
+        .role(Role::User)
+        .add_text_part("Tell me the story of the genesis of the universe as a bedtime story.")
+        .build()];
+
+    let request = GenerateContentRequest::builder().contents(prompt).build();
 
     let queue = gemini.stream_generate_content(&request, "gemini-pro").await;
 
diff --git a/examples/text-from-text.rs b/examples/text-from-text.rs
index e62329f..fa8e739 100644
--- a/examples/text-from-text.rs
+++ b/examples/text-from-text.rs
@@ -14,10 +14,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         location_id,
     );
 
-    let prompt = "What is the airspeed of an unladen swallow?";
-    let request = GenerateContentRequest::builder()
-        .add_text_content(Role::User, prompt)
-        .build();
+    let prompt = vec![Content::builder()
+        .role(Role::User)
+        .add_text_part("What is the airspeed of an unladen swallow?")
+        .build()];
+
+    let request = GenerateContentRequest::builder().contents(prompt).build();
     let response = gemini.generate_content(&request, "gemini-pro").await?;
     println!("Response: {:?}", response.candidates[0].get_text().unwrap());
 
diff --git a/src/types/generate_content.rs b/src/types/generate_content.rs
index 040952b..e4dc719 100644
--- a/src/types/generate_content.rs
+++ b/src/types/generate_content.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
-use super::{Content, Role, VertexApiError};
+use super::{Content, VertexApiError};
 use crate::error::Result;
 
 #[derive(Clone, Default, Serialize, Deserialize)]
@@ -37,13 +37,8 @@ impl GenerateContentRequestBuilder {
         }
     }
 
-    pub fn add_text_content<T: Into<String>>(self, role: Role, text: T) -> Self {
-        let content = Content::builder().role(role).add_text_part(text).build();
-        self.add_content(content)
-    }
-
-    pub fn add_content(mut self, content: Content) -> Self {
-        self.request.contents.push(content);
+    pub fn contents(mut self, contents: Vec<Content>) -> Self {
+        self.request.contents = contents;
         self
     }
 
@@ -62,10 +57,6 @@ impl GenerateContentRequestBuilder {
         self
     }
 
-    pub fn system_instruction_text<T: Into<String>>(self, text: T) -> Self {
-        self.system_instruction(Content::builder().add_text_part(text).build())
-    }
-
     pub fn system_instruction(mut self, system_instruction: Content) -> Self {
         self.request.system_instruction = Some(system_instruction);
         self