Add system instruction support in Vertex AI (#12749)

firebase · May 1, 2024 · f4fa5f5 · f4fa5f5
1 parent f6dec4e
commit f4fa5f5
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 1 deletion.
diff --git a/FirebaseVertexAI/Sources/GenerateContentRequest.swift b/FirebaseVertexAI/Sources/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/FirebaseVertexAI/Sources/GenerativeModel.swift b/FirebaseVertexAI/Sources/GenerativeModel.swift
@@ -40,6 +40,9 @@ public final class GenerativeModel {
   /// Tool configuration for any `Tool` specified in the request.
   let toolConfig: ToolConfig?
 
+  /// Instructions that direct the model to behave a certain way.
+  let systemInstruction: ModelContent?
+
   /// Configuration parameters for sending requests to the backend.
   let requestOptions: RequestOptions
 
@@ -53,6 +56,8 @@ public final class GenerativeModel {
   ///   - safetySettings: A value describing what types of harmful content your model should allow.
   ///   - tools: A list of ``Tool`` objects that the model may use to generate the next response.
   ///   - toolConfig: Tool configuration for any `Tool` specified in the request.
+  ///   - systemInstruction: Instructions that direct the model to behave a certain way; currently
+  ///     only text content is supported.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
   init(name: String,
@@ -61,6 +66,7 @@ public final class GenerativeModel {
        safetySettings: [SafetySetting]? = nil,
        tools: [Tool]?,
        toolConfig: ToolConfig? = nil,
+       systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
        appCheck: AppCheckInterop?,
        urlSession: URLSession = .shared) {
@@ -74,6 +80,7 @@ public final class GenerativeModel {
     self.safetySettings = safetySettings
     self.tools = tools
     self.toolConfig = toolConfig
+    self.systemInstruction = systemInstruction
     self.requestOptions = requestOptions
 
     Logging.default.info("""
@@ -121,6 +128,7 @@ public final class GenerativeModel {
                                                               safetySettings: safetySettings,
                                                               tools: tools,
                                                               toolConfig: toolConfig,
+                                                              systemInstruction: systemInstruction,
                                                               isStreaming: false,
                                                               options: requestOptions)
       response = try await generativeAIService.loadRequest(request: generateContentRequest)
@@ -194,6 +202,7 @@ public final class GenerativeModel {
                                                         safetySettings: safetySettings,
                                                         tools: tools,
                                                         toolConfig: toolConfig,
+                                                        systemInstruction: systemInstruction,
                                                         isStreaming: true,
                                                         options: requestOptions)
 

diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift
@@ -59,12 +59,15 @@ public class VertexAI: NSObject {
   ///   - safetySettings: A value describing what types of harmful content your model should allow.
   ///   - tools: A list of ``Tool`` objects that the model may use to generate the next response.
   ///   - toolConfig: Tool configuration for any `Tool` specified in the request.
+  ///   - systemInstruction: Instructions that direct the model to behave a certain way; currently
+  ///     only text content is supported.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   public func generativeModel(modelName: String,
                               generationConfig: GenerationConfig? = nil,
                               safetySettings: [SafetySetting]? = nil,
                               tools: [Tool]? = nil,
                               toolConfig: ToolConfig? = nil,
+                              systemInstruction: ModelContent? = nil,
                               requestOptions: RequestOptions = RequestOptions())
     -> GenerativeModel {
     let modelResourceName = modelResourceName(modelName: modelName, location: location)
@@ -80,6 +83,7 @@ public class VertexAI: NSObject {
       safetySettings: safetySettings,
       tools: tools,
       toolConfig: toolConfig,
+      systemInstruction: systemInstruction,
       requestOptions: requestOptions,
       appCheck: appCheck
     )

diff --git a/FirebaseVertexAI/Tests/Unit/VertexAIAPITests.swift b/FirebaseVertexAI/Tests/Unit/VertexAIAPITests.swift
@@ -32,6 +32,7 @@ final class VertexAIAPITests: XCTestCase {
                                   maxOutputTokens: 256,
                                   stopSequences: ["..."])
     let filters = [SafetySetting(harmCategory: .dangerousContent, threshold: .blockOnlyHigh)]
+    let systemInstruction = ModelContent(role: "system", parts: [.text("Talk like a pirate.")])
 
     // Instantiate Vertex AI SDK - Default App
     let vertexAI = VertexAI.vertexAI()
@@ -53,11 +54,17 @@ final class VertexAIAPITests: XCTestCase {
       generationConfig: config
     )
 
+    let _ = vertexAI.generativeModel(
+      modelName: "gemini-1.0-pro",
+      systemInstruction: systemInstruction
+    )
+
     // All arguments passed.
     let genAI = vertexAI.generativeModel(
       modelName: "gemini-1.0-pro",
       generationConfig: config, // Optional
-      safetySettings: filters // Optional
+      safetySettings: filters, // Optional
+      systemInstruction: systemInstruction // Optional
     )
 
     // Full Typed Usage