Skip to content

[Vertex AI] Add anyOf support to Schema #14708

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 16 commits into from
Apr 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions FirebaseVertexAI/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
# Unreleased
# 11.12.0
- [added] **Public Preview**: Added support for specifying response modalities
in `GenerationConfig`. This includes **public experimental** support for image
generation using Gemini 2.0 Flash (`gemini-2.0-flash-exp`). (#14658)
<br /><br />
Note: This feature is in Public Preview and relies on experimental models,
which means that it is not subject to any SLA or deprecation policy and could
change in backwards-incompatible ways.
- [added] Added support for specifying the minimum and maximum number of items
(`minItems` / `maxItems`) to generate in an array `Schema`. (#14671)
- [added] Added support for more `Schema` fields: `minItems`/`maxItems` (array
size limits), `title` (schema name), `minimum`/`maximum` (numeric ranges),
`anyOf` (select from sub-schemas), and `propertyOrdering` (JSON key order). (#14647)
- [fixed] Fixed an issue where network requests would fail in the iOS 18.4
simulator due to a `URLSession` bug introduced in Xcode 16.3. (#14677)

Expand Down
70 changes: 62 additions & 8 deletions FirebaseVertexAI/Sources/Types/Public/Schema.swift
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,11 @@ public final class Schema: Sendable {
let rawValue: String
}

let dataType: DataType
// May only be nil for `anyOf` schemas, which do not have an explicit `type` in the OpenAPI spec.
let dataType: DataType?

/// The data type.
public var type: String { dataType.rawValue }
public var type: String { dataType?.rawValue ?? "UNSPECIFIED" }

/// The format of the data.
public let format: String?
Expand Down Expand Up @@ -106,6 +107,16 @@ public final class Schema: Sendable {
/// property's type and constraints.
public let properties: [String: Schema]?

/// An array of `Schema` objects. The generated data must be valid against *any* (one or more)
/// of the schemas listed in this array. This allows specifying multiple possible structures or
/// types for a single field.
///
/// For example, a value could be either a `String` or an `Integer`:
/// ```
/// Schema.anyOf(schemas: [.string(), .integer()])
/// ```
public let anyOf: [Schema]?

/// An array of strings, where each string is the name of a property defined in the `properties`
/// dictionary that must be present in the generated object. If a property is listed here, the
/// model must include it in the output.
Expand All @@ -119,12 +130,14 @@ public final class Schema: Sendable {
/// serialization.
public let propertyOrdering: [String]?

required init(type: DataType, format: String? = nil, description: String? = nil,
title: String? = nil,
nullable: Bool = false, enumValues: [String]? = nil, items: Schema? = nil,
minItems: Int? = nil, maxItems: Int? = nil, minimum: Double? = nil,
maximum: Double? = nil, properties: [String: Schema]? = nil,
requiredProperties: [String]? = nil, propertyOrdering: [String]? = nil) {
required init(type: DataType?, format: String? = nil, description: String? = nil,
title: String? = nil, nullable: Bool? = nil, enumValues: [String]? = nil,
items: Schema? = nil, minItems: Int? = nil, maxItems: Int? = nil,
minimum: Double? = nil, maximum: Double? = nil, anyOf: [Schema]? = nil,
properties: [String: Schema]? = nil, requiredProperties: [String]? = nil,
propertyOrdering: [String]? = nil) {
precondition(type != nil || anyOf != nil,
"A schema must have either a `type` or an `anyOf` array of sub-schemas.")
dataType = type
self.format = format
self.description = description
Expand All @@ -136,6 +149,7 @@ public final class Schema: Sendable {
self.maxItems = maxItems
self.minimum = minimum
self.maximum = maximum
self.anyOf = anyOf
self.properties = properties
self.requiredProperties = requiredProperties
self.propertyOrdering = propertyOrdering
Expand Down Expand Up @@ -278,6 +292,10 @@ public final class Schema: Sendable {
/// - format: An optional modifier describing the expected format of the integer. Currently the
/// formats ``IntegerFormat/int32`` and ``IntegerFormat/int64`` are supported; custom values
/// may be specified using ``IntegerFormat/custom(_:)`` but may be ignored by the model.
/// - minimum: If specified, instructs the model that the value should be greater than or
/// equal to the specified minimum.
/// - maximum: If specified, instructs the model that the value should be less than or equal
/// to the specified maximum.
public static func integer(description: String? = nil, nullable: Bool = false,
format: IntegerFormat? = nil,
minimum: Int? = nil, maximum: Int? = nil) -> Schema {
Expand Down Expand Up @@ -362,8 +380,11 @@ public final class Schema: Sendable {
/// - optionalProperties: A list of property names that may be be omitted in objects generated
/// by the model; these names must correspond to the keys provided in the `properties`
/// dictionary and may be an empty list.
/// - propertyOrdering: An optional hint to the model suggesting the order for keys in the
/// generated JSON string. See ``propertyOrdering`` for details.
/// - description: An optional description of what the object should contain or represent; may
/// use Markdown format.
/// - title: An optional human-readable name/summary for the object schema.
/// - nullable: If `true`, instructs the model that it may return `null` instead of an object;
/// defaults to `false`, enforcing that an object is returned.
public static func object(properties: [String: Schema], optionalProperties: [String] = [],
Expand All @@ -388,6 +409,38 @@ public final class Schema: Sendable {
propertyOrdering: propertyOrdering
)
}

/// Returns a `Schema` representing a value that must conform to *any* (one or more) of the
/// provided sub-schemas.
///
/// This schema instructs the model to produce data that is valid against at least one of the
/// schemas listed in the `schemas` array. This is useful when a field can accept multiple
/// distinct types or structures.
///
/// **Example:** A field that can hold either a simple user ID (integer) or a detailed user
/// object.
/// ```
/// Schema.anyOf(schemas: [
/// .integer(description: "User ID"),
/// .object(properties: [
/// "userId": .integer(),
/// "userName": .string()
/// ], description: "Detailed User Object")
/// ])
/// ```
/// The generated data could be decoded based on which schema it matches.
///
/// - Parameters:
/// - schemas: An array of `Schema` objects. The generated data must be valid against at least
/// one of these schemas. The array must not be empty.
public static func anyOf(schemas: [Schema]) -> Schema {
if schemas.isEmpty {
VertexLog.error(code: .invalidSchemaFormat, "The `anyOf` schemas array cannot be empty.")
}
// Note: The 'type' for an 'anyOf' schema is implicitly defined by the presence of the
// 'anyOf' keyword and doesn't have a specific explicit type like "OBJECT" or "STRING".
return self.init(type: nil, anyOf: schemas)
}
}

// MARK: - Codable Conformance
Expand All @@ -406,6 +459,7 @@ extension Schema: Encodable {
case maxItems
case minimum
case maximum
case anyOf
case properties
case requiredProperties = "required"
case propertyOrdering
Expand Down
1 change: 1 addition & 0 deletions FirebaseVertexAI/Sources/VertexLog.swift
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ enum VertexLog {
// Generative Model Configuration
case generativeModelInitialized = 1000
case unsupportedGeminiModel = 1001
case invalidSchemaFormat = 1002

// Imagen Model Configuration
case unsupportedImagenModel = 1200
Expand Down
84 changes: 84 additions & 0 deletions FirebaseVertexAI/Tests/TestApp/Tests/Integration/SchemaTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -148,4 +148,88 @@ struct SchemaTests {
#expect(rating >= 1, "Expected a rating >= 1, but got \(rating)")
#expect(rating <= 5, "Expected a rating <= 5, but got \(rating)")
}

@Test(arguments: InstanceConfig.allConfigsExceptDeveloperV1)
func generateContentAnyOfSchema(_ config: InstanceConfig) async throws {
struct MailingAddress: Decodable {
let streetAddress: String
let city: String

// Canadian-specific
let province: String?
let postalCode: String?

// U.S.-specific
let state: String?
let zipCode: String?

var isCanadian: Bool {
return province != nil && postalCode != nil && state == nil && zipCode == nil
}

var isAmerican: Bool {
return province == nil && postalCode == nil && state != nil && zipCode != nil
}
}

let streetSchema = Schema.string(description:
"The civic number and street name, for example, '123 Main Street'.")
let citySchema = Schema.string(description: "The name of the city.")
let canadianAddressSchema = Schema.object(
properties: [
"streetAddress": streetSchema,
"city": citySchema,
"province": .string(description:
"The 2-letter province or territory code, for example, 'ON', 'QC', or 'NU'."),
"postalCode": .string(description: "The postal code, for example, 'A1A 1A1'."),
],
description: "A Canadian mailing address"
)
let americanAddressSchema = Schema.object(
properties: [
"streetAddress": streetSchema,
"city": citySchema,
"state": .string(description:
"The 2-letter U.S. state or territory code, for example, 'CA', 'NY', or 'TX'."),
"zipCode": .string(description: "The 5-digit ZIP code, for example, '12345'."),
],
description: "A U.S. mailing address"
)
let model = VertexAI.componentInstance(config).generativeModel(
modelName: ModelNames.gemini2Flash,
generationConfig: GenerationConfig(
temperature: 0.0,
topP: 0.0,
topK: 1,
responseMIMEType: "application/json",
responseSchema: .array(items: .anyOf(
schemas: [canadianAddressSchema, americanAddressSchema]
))
),
safetySettings: safetySettings
)
let prompt = """
What are the mailing addresses for the University of Waterloo, UC Berkeley and Queen's U?
"""
let response = try await model.generateContent(prompt)
let text = try #require(response.text)
let jsonData = try #require(text.data(using: .utf8))
let decodedAddresses = try JSONDecoder().decode([MailingAddress].self, from: jsonData)
try #require(decodedAddresses.count == 3, "Expected 3 JSON addresses, got \(text).")
let waterlooAddress = decodedAddresses[0]
#expect(
waterlooAddress.isCanadian,
"Expected Canadian University of Waterloo address, got \(waterlooAddress)."
)
let berkeleyAddress = decodedAddresses[1]
#expect(
berkeleyAddress.isAmerican,
"Expected American UC Berkeley address, got \(berkeleyAddress)."
)
let queensAddress = decodedAddresses[2]
#expect(
queensAddress.isCanadian,
"Expected Canadian Queen's University address, got \(queensAddress)."
)
}
}
Loading
Loading