Skip to content

[VertexAI] Add minimum, maximum, title, and propertyOrdering to schema #14711

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 8 commits into from
Apr 14, 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
90 changes: 74 additions & 16 deletions FirebaseVertexAI/Sources/Types/Public/Schema.swift
Original file line number Diff line number Diff line change
Expand Up @@ -68,44 +68,77 @@ public final class Schema: Sendable {
/// The format of the data.
public let format: String?

/// A brief description of the parameter.
/// A human-readable explanation of the purpose of the schema or property. While not strictly
/// enforced on the value itself, good descriptions significantly help the model understand the
/// context and generate more relevant and accurate output.
public let description: String?

/// A human-readable name/summary for the schema or a specific property. This helps document the
/// schema's purpose but doesn't typically constrain the generated value. It can subtly guide the
/// model by clarifying the intent of a field.
public let title: String?

/// Indicates if the value may be null.
public let nullable: Bool?

/// Possible values of the element of type "STRING" with "enum" format.
public let enumValues: [String]?

/// Schema of the elements of type `"ARRAY"`.
/// Defines the schema for the elements within the `"ARRAY"`. All items in the generated array
/// must conform to this schema definition. This can be a simple type (like .string) or a complex
/// nested object schema.
public let items: Schema?

/// The minimum number of items (elements) in a schema of type `"ARRAY"`.
/// An integer specifying the minimum number of items the generated `"ARRAY"` must contain.
public let minItems: Int?

/// The maximum number of items (elements) in a schema of type `"ARRAY"`.
/// An integer specifying the maximum number of items the generated `"ARRAY"` must contain.
public let maxItems: Int?

/// Properties of type `"OBJECT"`.
/// The minimum value of a numeric type.
public let minimum: Double?

/// The maximum value of a numeric type.
public let maximum: Double?

/// Defines the members (key-value pairs) expected within an object. It's a dictionary where keys
/// are the property names (strings) and values are nested `Schema` definitions describing each
/// property's type and constraints.
public let properties: [String: Schema]?

/// Required properties of type `"OBJECT"`.
/// 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.
public let requiredProperties: [String]?

/// A specific hint provided to the Gemini model, suggesting the order in which the keys should
/// appear in the generated JSON string. Important: Standard JSON objects are inherently unordered
/// collections of key-value pairs. While the model will try to respect propertyOrdering in its
/// textual JSON output, subsequent parsing into native Swift objects (like Dictionaries or
/// Structs) might not preserve this order. This parameter primarily affects the raw JSON string
/// 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,
properties: [String: Schema]? = nil, requiredProperties: [String]? = nil) {
minItems: Int? = nil, maxItems: Int? = nil, minimum: Double? = nil,
maximum: Double? = nil, properties: [String: Schema]? = nil,
requiredProperties: [String]? = nil, propertyOrdering: [String]? = nil) {
dataType = type
self.format = format
self.description = description
self.title = title
self.nullable = nullable
self.enumValues = enumValues
self.items = items
self.minItems = minItems
self.maxItems = maxItems
self.minimum = minimum
self.maximum = maximum
self.properties = properties
self.requiredProperties = requiredProperties
self.propertyOrdering = propertyOrdering
}

/// Returns a `Schema` representing a string value.
Expand Down Expand Up @@ -184,12 +217,19 @@ public final class Schema: Sendable {
/// use Markdown format.
/// - nullable: If `true`, instructs the model that it may generate `null` instead of a number;
/// defaults to `false`, enforcing that a number is generated.
public static func float(description: String? = nil, nullable: Bool = false) -> Schema {
/// - 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 float(description: String? = nil, nullable: Bool = false,
minimum: Float? = nil, maximum: Float? = nil) -> Schema {
return self.init(
type: .number,
format: "float",
description: description,
nullable: nullable
nullable: nullable,
minimum: minimum.map { Double($0) },
maximum: maximum.map { Double($0) }
)
}

Expand All @@ -203,11 +243,18 @@ public final class Schema: Sendable {
/// use Markdown format.
/// - nullable: If `true`, instructs the model that it may return `null` instead of a number;
/// defaults to `false`, enforcing that a number is returned.
public static func double(description: String? = nil, nullable: Bool = false) -> Schema {
/// - 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 double(description: String? = nil, nullable: Bool = false,
minimum: Double? = nil, maximum: Double? = nil) -> Schema {
return self.init(
type: .number,
description: description,
nullable: nullable
nullable: nullable,
minimum: minimum,
maximum: maximum
)
}

Expand All @@ -232,12 +279,15 @@ public final class Schema: Sendable {
/// formats ``IntegerFormat/int32`` and ``IntegerFormat/int64`` are supported; custom values
/// may be specified using ``IntegerFormat/custom(_:)`` but may be ignored by the model.
public static func integer(description: String? = nil, nullable: Bool = false,
format: IntegerFormat? = nil) -> Schema {
format: IntegerFormat? = nil,
minimum: Int? = nil, maximum: Int? = nil) -> Schema {
return self.init(
type: .integer,
format: format?.rawValue,
description: description,
nullable: nullable
nullable: nullable.self,
minimum: minimum.map { Double($0) },
maximum: maximum.map { Double($0) }
)
}

Expand Down Expand Up @@ -317,7 +367,9 @@ public final class Schema: Sendable {
/// - 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] = [],
description: String? = nil, nullable: Bool = false) -> Schema {
propertyOrdering: [String]? = nil,
description: String? = nil, title: String? = nil,
nullable: Bool = false) -> Schema {
var requiredProperties = Set(properties.keys)
for optionalProperty in optionalProperties {
guard properties.keys.contains(optionalProperty) else {
Expand All @@ -329,9 +381,11 @@ public final class Schema: Sendable {
return self.init(
type: .object,
description: description,
title: title,
nullable: nullable,
properties: properties,
requiredProperties: requiredProperties.sorted()
requiredProperties: requiredProperties.sorted(),
propertyOrdering: propertyOrdering
)
}
}
Expand All @@ -344,12 +398,16 @@ extension Schema: Encodable {
case dataType = "type"
case format
case description
case title
case nullable
case enumValues = "enum"
case items
case minItems
case maxItems
case minimum
case maximum
case properties
case requiredProperties = "required"
case propertyOrdering
}
}
88 changes: 82 additions & 6 deletions FirebaseVertexAI/Tests/TestApp/Tests/Integration/SchemaTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -55,12 +55,12 @@ struct SchemaTests {
generationConfig: GenerationConfig(
responseMIMEType: "application/json",
responseSchema:
.array(
items: .string(description: "The name of the city"),
description: "A list of city names",
minItems: 3,
maxItems: 5
)
.array(
items: .string(description: "The name of the city"),
description: "A list of city names",
minItems: 3,
maxItems: 5
)
),
safetySettings: safetySettings
)
Expand All @@ -72,4 +72,80 @@ struct SchemaTests {
#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)")
}

@Test(arguments: InstanceConfig.allConfigsExceptDeveloperV1)
func generateContentSchemaNumberRange(_ config: InstanceConfig) async throws {
let model = VertexAI.componentInstance(config).generativeModel(
modelName: ModelNames.gemini2FlashLite,
generationConfig: GenerationConfig(
responseMIMEType: "application/json",
responseSchema: .integer(
description: "A number",
minimum: 110,
maximum: 120
)
),
safetySettings: safetySettings
)
let prompt = "Give me a number"
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 decodedNumber = try JSONDecoder().decode(Double.self, from: jsonData)
#expect(decodedNumber >= 110.0, "Expected a number >= 110, but got \(decodedNumber)")
#expect(decodedNumber <= 120.0, "Expected a number <= 120, but got \(decodedNumber)")
}

@Test(arguments: InstanceConfig.allConfigsExceptDeveloperV1)
func generateContentSchemaNumberRangeMultiType(_ config: InstanceConfig) async throws {
struct ProductInfo: Codable {
let productName: String
let rating: Int // Will correspond to .integer in schema
let price: Double // Will correspond to .double in schema
let salePrice: Float // Will correspond to .float in schema
}
let model = VertexAI.componentInstance(config).generativeModel(
modelName: ModelNames.gemini2FlashLite,
generationConfig: GenerationConfig(
responseMIMEType: "application/json",
responseSchema: .object(
properties: [
"productName": .string(description: "The name of the product"),
"price": .double(
description: "A price",
minimum: 10.00,
maximum: 120.00
),
"salePrice": .float(
description: "A sale price",
minimum: 5.00,
maximum: 90.00
),
"rating": .integer(
description: "A rating",
minimum: 1,
maximum: 5
),
],
propertyOrdering: ["salePrice", "rating", "price", "productName"],
title: "ProductInfo"
),
),
safetySettings: safetySettings
)
let prompt = "Describe a premium wireless headphone, including a user rating and price."
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 decodedProduct = try JSONDecoder().decode(ProductInfo.self, from: jsonData)
let price = decodedProduct.price
let salePrice = decodedProduct.salePrice
let rating = decodedProduct.rating
#expect(price >= 10.0, "Expected a price >= 10.00, but got \(price)")
#expect(price <= 120.0, "Expected a price <= 120.00, but got \(price)")
#expect(salePrice >= 5.0, "Expected a salePrice >= 5.00, but got \(salePrice)")
#expect(salePrice <= 90.0, "Expected a salePrice <= 90.00, but got \(salePrice)")
#expect(rating >= 1, "Expected a rating >= 1, but got \(rating)")
#expect(rating <= 5, "Expected a rating <= 5, but got \(rating)")
}
}
Loading