From 0ff780b2abf76dc60d59fba422706f78b8c7dcc1 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Wed, 10 Apr 2024 17:57:36 -0400
Subject: [PATCH] Add system instruction support (#129)

---
 Sources/GoogleAI/GenerateContentRequest.swift | 2 ++
 Tests/GoogleAITests/GoogleAITests.swift       | 9 ++++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/Sources/GoogleAI/GenerateContentRequest.swift b/Sources/GoogleAI/GenerateContentRequest.swift
index e0e9c2f..05abadf 100644
--- a/Sources/GoogleAI/GenerateContentRequest.swift
+++ b/Sources/GoogleAI/GenerateContentRequest.swift
@@ -23,6 +23,7 @@ struct GenerateContentRequest {
   let safetySettings: [SafetySetting]?
   let tools: [Tool]?
   let toolConfig: ToolConfig?
+  let systemInstruction: ModelContent?
   let isStreaming: Bool
   let options: RequestOptions
 }
@@ -35,6 +36,7 @@ extension GenerateContentRequest: Encodable {
     case safetySettings
     case tools
     case toolConfig
+    case systemInstruction
   }
 }
 
diff --git a/Tests/GoogleAITests/GoogleAITests.swift b/Tests/GoogleAITests/GoogleAITests.swift
index cbc9252..c0764bb 100644
--- a/Tests/GoogleAITests/GoogleAITests.swift
+++ b/Tests/GoogleAITests/GoogleAITests.swift
@@ -30,17 +30,24 @@ final class GoogleGenerativeAITests: XCTestCase {
                                   maxOutputTokens: 256,
                                   stopSequences: ["..."])
     let filters = [SafetySetting(harmCategory: .dangerousContent, threshold: .blockOnlyHigh)]
+    let systemInstruction = ModelContent(role: "system", parts: [.text("Talk like a pirate.")])
 
     // Permutations without optional arguments.
     let _ = GenerativeModel(name: "gemini-1.0-pro", apiKey: "API_KEY")
     let _ = GenerativeModel(name: "gemini-1.0-pro", apiKey: "API_KEY", safetySettings: filters)
     let _ = GenerativeModel(name: "gemini-1.0-pro", apiKey: "API_KEY", generationConfig: config)
+    let _ = GenerativeModel(
+      name: "gemini-1.0-pro",
+      apiKey: "API_KEY",
+      systemInstruction: systemInstruction
+    )
 
     // All arguments passed.
     let genAI = GenerativeModel(name: "gemini-1.0-pro",
                                 apiKey: "API_KEY",
                                 generationConfig: config, // Optional
-                                safetySettings: filters // Optional
+                                safetySettings: filters, // Optional
+                                systemInstruction: systemInstruction // Optional
     )
     // Full Typed Usage
     let pngData = Data() // ....