Add system instruction support (google-gemini#129)

SsomsakTH · Jun 21, 2024 · 0ff780b · 0ff780b
1 parent d529c48
commit 0ff780b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git 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
@@ -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() // ....