[Firebase AI] Add internal JSON Schema support in GenerationConfig
#15404
Conversation
# Conflicts: # FirebaseAI/Tests/Unit/GenerationConfigTests.swift
andrewheard
force-pushed
the
ah/ai-json-schema
branch
from
4343ab6 to
397ae19
Compare
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds internal support for response_json_schema in GenerationConfig. The changes include adding the new property, a corresponding internal initializer, and updating the Encodable conformance. The implementation is well-tested with new unit and integration tests that cover various schema types. The changes look good and the new feature is well-implemented. I've left one comment regarding code duplication in the integration tests which could be addressed to improve maintainability.
| func generateContentJSONSchemaItems(_ config: InstanceConfig) async throws { | ||
| let model = FirebaseAI.componentInstance(config).generativeModel( | ||
| modelName: ModelNames.gemini2_5_FlashLite, | ||
| generationConfig: GenerationConfig( | ||
| responseMIMEType: "application/json", | ||
| responseJSONSchema: [ | ||
| "type": .string("array"), | ||
| "description": .string("A list of city names"), | ||
| "items": .object([ | ||
| "type": .string("string"), | ||
| "description": .string("The name of the city"), | ||
| ]), | ||
| "minItems": .number(3), | ||
| "maxItems": .number(5), | ||
| ] | ||
| ), | ||
| safetySettings: safetySettings | ||
| ) | ||
| let prompt = "What are the biggest cities in Canada?" | ||
| let response = try await model.generateContent(prompt) | ||
| let text = try #require(response.text).trimmingCharacters(in: .whitespacesAndNewlines) | ||
| let jsonData = try #require(text.data(using: .utf8)) | ||
| let decodedJSON = try JSONDecoder().decode([String].self, from: jsonData) | ||
| #expect(decodedJSON.count >= 3, "Expected at least 3 cities, but got \(decodedJSON.count)") | ||
| #expect(decodedJSON.count <= 5, "Expected at most 5 cities, but got \(decodedJSON.count)") | ||
| } |
There was a problem hiding this comment.
There's a lot of duplicated code between this test and generateContentSchemaItems. The logic for making the request and asserting the response is identical. Consider extracting this common logic into a private helper function to improve maintainability and reduce code duplication. This new helper could take a GenerativeModel instance as a parameter.
This same pattern of duplication exists for other new tests in this file (generateContentJSONSchemaNumberRange, generateContentJSONSchemaNumberRangeMultiType, generateContentAnyOfJSONSchema), and they would also benefit from a similar refactoring.
1 participant
Added internal support for JSON Schema (
response_json_schema) inGenerationConfig.#no-changelog