feat: add system prompts to req && add unit tests (#17)

Author: wiseaidev

.circleci/config.yml

@@ -0,0 +1,20 @@
+# Rust CircleCI 2.1 configuration file
+version: 2.1
+aliases:
+  - &rust_container
+    docker:
+      - image: cimg/rust:1.86.0
+jobs:
+  testing:
+    <<: *rust_container
+    steps:
+      - checkout
+      - run:
+          name: Run Tests
+          command: cargo test --all-features
+
+workflows:
+  version: 2
+  test:
+    jobs:
+      - testing

.github/workflows/ci.yml

@@ -2,21 +2,18 @@ name: CI
 
 on:
   push:
-    branches: [ "main" ]
+    branches: ["main"]
   pull_request:
-    branches: [ "main" ]
+    branches: ["main"]
 
 env:
   CARGO_TERM_COLOR: always
 
 jobs:
   build:
-
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v3
-    - name: Build
-      run: cargo build --verbose
-    - name: Run tests
-      run: cargo test --verbose
+      - uses: actions/checkout@v3
+      - name: Build
+        run: cargo build --all-features --release

src/chat.rs

@@ -16,6 +16,8 @@ use reqwest::Method;
 pub struct Chat {
     pub model: Model,
     pub messages: Vec<Message>,
+    #[builder(setter(into, strip_option), default)]
+    pub system: Option<Vec<Message>>,
 }
 
 #[derive(Clone)]
@@ -29,9 +31,14 @@ impl Chats {
             parts: params.messages.iter().map(|msg| msg.to_part()).collect(),
         };
 
+        let system_instruction = params.system.as_ref().map(|messages| Content {
+            parts: messages.iter().map(|msg| msg.to_part()).collect(),
+        });
+
         let request_body = GeminiRequest {
             model: params.model.to_string(),
             contents: vec![content],
+            system_instruction,
             config: None,
         };
 

src/client.rs

@@ -18,7 +18,7 @@ use std::sync::{Arc, RwLock};
 
 const GEMINI_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta/models";
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 #[allow(dead_code)]
 pub struct Client {
     http_client: Arc<HttpClient>,

src/imagen.rs

@@ -20,16 +20,24 @@ pub struct Images {
 pub struct ImageGen {
     pub model: Model,
     pub input: Message,
+    #[builder(setter(into, strip_option), default)]
+    pub system: Option<Vec<Message>>,
 }
 
 impl Images {
     pub async fn generate(&self, params: ImageGen) -> Result<Vec<u8>> {
         let content = Content {
             parts: vec![params.input.to_part()],
         };
+
+        let system_instruction = params.system.as_ref().map(|messages| Content {
+            parts: messages.iter().map(|msg| msg.to_part()).collect(),
+        });
+
         let request_body = GeminiRequest {
             model: params.model.to_string(),
             contents: vec![content],
+            system_instruction,
             config: Some(GenerationConfig {
                 response_modalities: vec!["Text".into(), "Image".into()],
             }),

src/messages.rs

@@ -1,6 +1,6 @@
 use crate::requests::Part;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum Content {
     Text(String),
 }
@@ -11,7 +11,7 @@ impl Default for Content {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum Message {
     User {
         content: Content,

src/models.rs

@@ -7,7 +7,7 @@ use derive_builder::Builder;
 use reqwest::Method;
 use std::str::FromStr;
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone, Default, PartialEq)]
 pub enum Model {
     Pro25Preview,
     #[default]

src/requests.rs

@@ -11,6 +11,10 @@ pub struct GeminiRequest {
 
     #[serde(rename = "generationConfig", skip_serializing_if = "Option::is_none")]
     pub config: Option<GenerationConfig>,
+
+    /// Optional system-level instruction.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub system_instruction: Option<Content>,
 }
 
 /// Request structure for content embedding.

src/responses.rs

@@ -16,7 +16,7 @@ pub struct BatchEmbedContentsResponse {
 }
 
 /// Structure representing embedding information.
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, Default, Clone)]
 pub struct Embedding {
     /// List of values for the embedding.
     pub values: Vec<f64>,

src/stream.rs

@@ -19,15 +19,22 @@ pub struct Streaming {
 pub struct Stream {
     pub model: Model,
     pub input: Message,
+    #[builder(setter(into, strip_option), default)]
+    pub system: Option<Vec<Message>>,
 }
 
 impl Streaming {
     pub async fn generate(&self, params: Stream) -> Result<Response> {
+        let system_instruction = params.system.as_ref().map(|messages| Content {
+            parts: messages.iter().map(|msg| msg.to_part()).collect(),
+        });
+
         let request_body = GeminiRequest {
             model: params.model.to_string(),
             contents: vec![Content {
                 parts: vec![params.input.to_part()],
             }],
+            system_instruction,
             config: None,
         };
 

src/tokens.rs

@@ -19,15 +19,21 @@ pub struct Tokens {
 pub struct Token {
     model: Model,
     input: Message,
+    system: Vec<Message>,
 }
 
 impl Tokens {
     pub async fn count(&self, params: Token) -> Result<usize> {
+        let system_instruction = Content {
+            parts: params.system.iter().map(|msg| msg.to_part()).collect(),
+        };
+
         let request_body = GeminiRequest {
             model: params.model.to_string(),
             contents: vec![Content {
                 parts: vec![params.input.to_part()],
             }],
+            system_instruction: Some(system_instruction),
             config: None,
         };
 

src/utils.rs

@@ -106,8 +106,7 @@ pub fn extract_image_or_text(parts: &[Part]) -> Result<Vec<u8>> {
     }) {
         let image_bytes = STANDARD.decode(&base64_data)?;
         Ok(image_bytes)
-    }
-    else if let Some(text) = parts.iter().find_map(|part| match part {
+    } else if let Some(text) = parts.iter().find_map(|part| match part {
         Part::Text { text } => Some(text.clone()),
         _ => None,
     }) {

src/vision.rs

@@ -1,6 +1,7 @@
 use crate::messages::Content;
 use crate::messages::Message;
 use crate::models::Model;
+use crate::requests::Content as ReqContent;
 use crate::requests::GeminiRequest;
 use crate::requests::ImageContent;
 use crate::requests::Part;
@@ -22,12 +23,18 @@ pub struct Visions {
 pub struct Vision {
     pub input: Message,
     pub image: Message,
+    #[builder(setter(into, strip_option), default)]
+    pub system: Option<Vec<Message>>,
 }
 
 impl Visions {
     pub async fn generate(&self, params: Vision) -> Result<String> {
         let input_part = params.input.to_part();
 
+        let system_instruction = params.system.as_ref().map(|messages| ReqContent {
+            parts: messages.iter().map(|msg| msg.to_part()).collect(),
+        });
+
         let image_data = match &params.image {
             Message::Tool { content } => content.clone(),
             Message::User { content, .. }
@@ -47,6 +54,7 @@ impl Visions {
             contents: vec![crate::requests::Content {
                 parts: vec![input_part, image_part],
             }],
+            system_instruction,
             config: None,
         };