[Firebase AI] Add internal JSON Schema support in GenerationConfig

FirebaseAI/Sources/GenerationConfig.swift

@@ -48,6 +48,11 @@ public struct GenerationConfig: Sendable {
   /// Output schema of the generated candidate text.
   let responseSchema: Schema?
 
+  /// Output schema of the generated response in [JSON Schema](https://json-schema.org/) format.
+  ///
+  /// If set, `responseSchema` must be omitted and `responseMIMEType` is required.
+  let responseJSONSchema: JSONObject?
+
   /// Supported modalities of the response.
   let responseModalities: [ResponseModality]?
 
@@ -175,6 +180,26 @@ public struct GenerationConfig: Sendable {
     self.stopSequences = stopSequences
     self.responseMIMEType = responseMIMEType
     self.responseSchema = responseSchema
+    responseJSONSchema = nil
+    self.responseModalities = responseModalities
+    self.thinkingConfig = thinkingConfig
+  }
+
+  init(temperature: Float? = nil, topP: Float? = nil, topK: Int? = nil, candidateCount: Int? = nil,
+       maxOutputTokens: Int? = nil, presencePenalty: Float? = nil, frequencyPenalty: Float? = nil,
+       stopSequences: [String]? = nil, responseMIMEType: String, responseJSONSchema: JSONObject,
+       responseModalities: [ResponseModality]? = nil, thinkingConfig: ThinkingConfig? = nil) {
+    self.temperature = temperature
+    self.topP = topP
+    self.topK = topK
+    self.candidateCount = candidateCount
+    self.maxOutputTokens = maxOutputTokens
+    self.presencePenalty = presencePenalty
+    self.frequencyPenalty = frequencyPenalty
+    self.stopSequences = stopSequences
+    self.responseMIMEType = responseMIMEType
+    responseSchema = nil
+    self.responseJSONSchema = responseJSONSchema
     self.responseModalities = responseModalities
     self.thinkingConfig = thinkingConfig
   }
@@ -195,6 +220,7 @@ extension GenerationConfig: Encodable {
     case stopSequences
     case responseMIMEType = "responseMimeType"
     case responseSchema
+    case responseJSONSchema = "responseJsonSchema"
     case responseModalities
     case thinkingConfig
   }

FirebaseAI/Sources/Tool.swift

@@ -29,6 +29,11 @@ public struct FunctionDeclaration: Sendable {
   /// Describes the parameters to this function; must be of type `DataType.object`.
   let parameters: Schema?
 
+  /// Describes the parameters to the function in JSON Schema format.
+  ///
+  /// This field is mutually exclusive with `parameters`.
+  let parametersJSONSchema: JSONObject?
+
   /// Constructs a new `FunctionDeclaration`.
   ///
   /// - Parameters:
@@ -47,6 +52,14 @@ public struct FunctionDeclaration: Sendable {
       optionalProperties: optionalParameters,
       nullable: false
     )
+    parametersJSONSchema = nil
+  }
+
+  init(name: String, description: String, parametersJSONSchema: JSONObject) {
+    self.name = name
+    self.description = description
+    parameters = nil
+    self.parametersJSONSchema = parametersJSONSchema
   }
 }
 
@@ -215,13 +228,15 @@ extension FunctionDeclaration: Encodable {
     case name
     case description
     case parameters
+    case parametersJSONSchema = "parametersJsonSchema"
   }
 
   public func encode(to encoder: Encoder) throws {
     var container = encoder.container(keyedBy: CodingKeys.self)
     try container.encode(name, forKey: .name)
     try container.encode(description, forKey: .description)
-    try container.encode(parameters, forKey: .parameters)
+    try container.encodeIfPresent(parameters, forKey: .parameters)
+    try container.encodeIfPresent(parametersJSONSchema, forKey: .parametersJSONSchema)
   }
 }
 

FirebaseAI/Tests/Unit/GenerationConfigTests.swift

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import FirebaseAI
+@testable import FirebaseAI
 import Foundation
 import XCTest
 
@@ -153,4 +153,85 @@ final class GenerationConfigTests: XCTestCase {
     }
     """)
   }
+
+  func testEncodeGenerationConfig_responseJSONSchema() throws {
+    let mimeType = "application/json"
+    let responseJSONSchema: JSONObject = [
+      "type": .string("object"),
+      "title": .string("Person"),
+      "properties": .object([
+        "firstName": .object(["type": .string("string")]),
+        "middleNames": .object([
+          "type": .string("array"),
+          "items": .object(["type": .string("string")]),
+          "minItems": .number(0),
+          "maxItems": .number(3),
+        ]),
+        "lastName": .object(["type": .string("string")]),
+        "age": .object(["type": .string("integer")]),
+      ]),
+      "required": .array([
+        .string("firstName"),
+        .string("middleNames"),
+        .string("lastName"),
+        .string("age"),
+      ]),
+      "propertyOrdering": .array([
+        .string("firstName"),
+        .string("middleNames"),
+        .string("lastName"),
+        .string("age"),
+      ]),
+      "additionalProperties": .bool(false),
+    ]
+    let generationConfig = GenerationConfig(
+      responseMIMEType: mimeType,
+      responseJSONSchema: responseJSONSchema
+    )
+
+    let jsonData = try encoder.encode(generationConfig)
+
+    let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(json, """
+    {
+      "responseJsonSchema" : {
+        "additionalProperties" : false,
+        "properties" : {
+          "age" : {
+            "type" : "integer"
+          },
+          "firstName" : {
+            "type" : "string"
+          },
+          "lastName" : {
+            "type" : "string"
+          },
+          "middleNames" : {
+            "items" : {
+              "type" : "string"
+            },
+            "maxItems" : 3,
+            "minItems" : 0,
+            "type" : "array"
+          }
+        },
+        "propertyOrdering" : [
+          "firstName",
+          "middleNames",
+          "lastName",
+          "age"
+        ],
+        "required" : [
+          "firstName",
+          "middleNames",
+          "lastName",
+          "age"
+        ],
+        "title" : "Person",
+        "type" : "object"
+      },
+      "responseMimeType" : "\(mimeType)"
+    }
+    """)
+  }
 }

FirebaseAI/Tests/Unit/Types/ToolTests.swift

@@ -89,4 +89,49 @@ final class ToolTests: XCTestCase {
     }
     """)
   }
+
+  func testEncodeTool_functionDeclarations_parametersJSONSchema() throws {
+    let functionDecl = FunctionDeclaration(
+      name: "test_function",
+      description: "A test function.",
+      parametersJSONSchema: [
+        "type": .string("object"),
+        "properties": .object([
+          "param1": .object(["type": .string("string")]),
+        ]),
+        "required": .array([.string("param1")]),
+        "propertyOrdering": .array([.string("param1")]),
+        "additionalProperties": .bool(false),
+      ]
+    )
+    let tool = Tool.functionDeclarations([functionDecl])
+    let jsonData = try encoder.encode(tool)
+
+    let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(jsonString, """
+    {
+      "functionDeclarations" : [
+        {
+          "description" : "A test function.",
+          "name" : "test_function",
+          "parametersJsonSchema" : {
+            "additionalProperties" : false,
+            "properties" : {
+              "param1" : {
+                "type" : "string"
+              }
+            },
+            "propertyOrdering" : [
+              "param1"
+            ],
+            "required" : [
+              "param1"
+            ],
+            "type" : "object"
+          }
+        }
+      ]
+    }
+    """)
+  }
 }