mirror of
https://github.com/firebase/firebase-ios-sdk
synced 2026-11-11 22:55:59 +05:45
433 lines
13 KiB
Swift
433 lines
13 KiB
Swift
// Copyright 2024 Google LLC
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
@testable import FirebaseAILogic
|
|
import Foundation
|
|
import XCTest
|
|
|
|
final class GenerationConfigTests: XCTestCase {
|
|
let encoder = JSONEncoder()
|
|
|
|
override func setUp() {
|
|
encoder.outputFormatting = .init(
|
|
arrayLiteral: .prettyPrinted, .sortedKeys, .withoutEscapingSlashes
|
|
)
|
|
}
|
|
|
|
// MARK: GenerationConfig Encoding
|
|
|
|
func testEncodeGenerationConfig_default() throws {
|
|
let generationConfig = GenerationConfig()
|
|
|
|
let jsonData = try encoder.encode(generationConfig)
|
|
|
|
let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
|
|
XCTAssertEqual(json, """
|
|
{
|
|
|
|
}
|
|
""")
|
|
}
|
|
|
|
func testEncodeGenerationConfig_allOptions() throws {
|
|
let maxOutputTokens = 256
|
|
let stopSequences = ["END", "DONE"]
|
|
let responseMIMEType = "application/json"
|
|
let generationConfig = GenerationConfig(
|
|
maxOutputTokens: maxOutputTokens,
|
|
stopSequences: stopSequences,
|
|
responseMIMEType: responseMIMEType,
|
|
responseSchema: .array(items: .string()),
|
|
responseModalities: [.text, .image]
|
|
)
|
|
|
|
let jsonData = try encoder.encode(generationConfig)
|
|
|
|
let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
|
|
XCTAssertEqual(json, """
|
|
{
|
|
"maxOutputTokens" : \(maxOutputTokens),
|
|
"responseMimeType" : "\(responseMIMEType)",
|
|
"responseModalities" : [
|
|
"TEXT",
|
|
"IMAGE"
|
|
],
|
|
"responseSchema" : {
|
|
"items" : {
|
|
"nullable" : false,
|
|
"type" : "STRING"
|
|
},
|
|
"nullable" : false,
|
|
"type" : "ARRAY"
|
|
},
|
|
"stopSequences" : [
|
|
"END",
|
|
"DONE"
|
|
]
|
|
}
|
|
""")
|
|
}
|
|
|
|
func testEncodeGenerationConfig_jsonResponse() throws {
|
|
let mimeType = "application/json"
|
|
let generationConfig = GenerationConfig(
|
|
responseMIMEType: mimeType,
|
|
responseSchema: .object(properties: [
|
|
"firstName": .string(),
|
|
"middleNames": .array(items: .string(), minItems: 0, maxItems: 3),
|
|
"lastName": .string(),
|
|
"age": .integer(),
|
|
])
|
|
)
|
|
|
|
let jsonData = try encoder.encode(generationConfig)
|
|
|
|
let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
|
|
XCTAssertEqual(json, """
|
|
{
|
|
"responseMimeType" : "\(mimeType)",
|
|
"responseSchema" : {
|
|
"nullable" : false,
|
|
"properties" : {
|
|
"age" : {
|
|
"nullable" : false,
|
|
"type" : "INTEGER"
|
|
},
|
|
"firstName" : {
|
|
"nullable" : false,
|
|
"type" : "STRING"
|
|
},
|
|
"lastName" : {
|
|
"nullable" : false,
|
|
"type" : "STRING"
|
|
},
|
|
"middleNames" : {
|
|
"items" : {
|
|
"nullable" : false,
|
|
"type" : "STRING"
|
|
},
|
|
"maxItems" : 3,
|
|
"minItems" : 0,
|
|
"nullable" : false,
|
|
"type" : "ARRAY"
|
|
}
|
|
},
|
|
"required" : [
|
|
"age",
|
|
"firstName",
|
|
"lastName",
|
|
"middleNames"
|
|
],
|
|
"type" : "OBJECT"
|
|
}
|
|
}
|
|
""")
|
|
}
|
|
|
|
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)"
|
|
}
|
|
""")
|
|
}
|
|
|
|
func testEncodeGenerationConfig_thinkingConfig() throws {
|
|
let testCases: [(ThinkingConfig, String)] = [
|
|
(ThinkingConfig(thinkingBudget: 0), "\"thinkingBudget\" : 0"),
|
|
(ThinkingConfig(thinkingBudget: 1024), "\"thinkingBudget\" : 1024"),
|
|
(ThinkingConfig(thinkingBudget: 1024, includeThoughts: true), """
|
|
"includeThoughts" : true,
|
|
"thinkingBudget" : 1024
|
|
"""),
|
|
(ThinkingConfig(thinkingLevel: .minimal), "\"thinkingLevel\" : \"MINIMAL\""),
|
|
(ThinkingConfig(thinkingLevel: .low), "\"thinkingLevel\" : \"LOW\""),
|
|
(ThinkingConfig(thinkingLevel: .medium), "\"thinkingLevel\" : \"MEDIUM\""),
|
|
(ThinkingConfig(thinkingLevel: .high), "\"thinkingLevel\" : \"HIGH\""),
|
|
(ThinkingConfig(thinkingLevel: .medium, includeThoughts: true), """
|
|
"includeThoughts" : true,
|
|
"thinkingLevel" : \"MEDIUM\"
|
|
"""),
|
|
(ThinkingConfig(thinkingLevel: .medium, includeThoughts: false), """
|
|
"includeThoughts" : false,
|
|
"thinkingLevel" : \"MEDIUM\"
|
|
"""),
|
|
]
|
|
|
|
for (thinkingConfig, expectedJSONSnippet) in testCases {
|
|
let generationConfig = GenerationConfig(thinkingConfig: thinkingConfig)
|
|
let jsonData = try encoder.encode(generationConfig)
|
|
let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
|
|
XCTAssertEqual(json, """
|
|
{
|
|
"thinkingConfig" : {
|
|
\(expectedJSONSnippet)
|
|
}
|
|
}
|
|
""")
|
|
}
|
|
}
|
|
|
|
func testEncodeGenerationConfig_speechConfig() throws {
|
|
let testCases: [(SpeechConfig, String)] = [
|
|
(SpeechConfig(voiceName: "Aoede"), """
|
|
"speechConfig" : {
|
|
"voiceConfig" : {
|
|
"prebuiltVoiceConfig" : {
|
|
"voiceName" : "Aoede"
|
|
}
|
|
}
|
|
}
|
|
"""),
|
|
(SpeechConfig(voiceName: "Aoede", languageCode: "en-US"), """
|
|
"speechConfig" : {
|
|
"languageCode" : "en-US",
|
|
"voiceConfig" : {
|
|
"prebuiltVoiceConfig" : {
|
|
"voiceName" : "Aoede"
|
|
}
|
|
}
|
|
}
|
|
"""),
|
|
(SpeechConfig(multiSpeakerVoiceConfig: MultiSpeakerVoiceConfig(speakerVoiceConfigs: [
|
|
SpeakerVoiceConfig(
|
|
speaker: "Speaker1",
|
|
voiceName: "Puck"
|
|
),
|
|
SpeakerVoiceConfig(
|
|
speaker: "Speaker2",
|
|
voiceName: "Charon"
|
|
),
|
|
])), """
|
|
"speechConfig" : {
|
|
"multiSpeakerVoiceConfig" : {
|
|
"speakerVoiceConfigs" : [
|
|
{
|
|
"speaker" : "Speaker1",
|
|
"voiceConfig" : {
|
|
"prebuiltVoiceConfig" : {
|
|
"voiceName" : "Puck"
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"speaker" : "Speaker2",
|
|
"voiceConfig" : {
|
|
"prebuiltVoiceConfig" : {
|
|
"voiceName" : "Charon"
|
|
}
|
|
}
|
|
}
|
|
]
|
|
}
|
|
}
|
|
"""),
|
|
]
|
|
|
|
for (speechConfig, expectedJSONSnippet) in testCases {
|
|
let generationConfig = GenerationConfig(speechConfig: speechConfig)
|
|
let jsonData = try encoder.encode(generationConfig)
|
|
let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
|
|
XCTAssertEqual(json, """
|
|
{
|
|
\(expectedJSONSnippet)
|
|
}
|
|
""")
|
|
}
|
|
}
|
|
|
|
// MARK: - GenerationConfig Merging
|
|
|
|
func testMerge_bothNil() {
|
|
let result = GenerationConfig.merge(nil, with: nil)
|
|
|
|
XCTAssertNil(result)
|
|
}
|
|
|
|
func testMerge_baseNil() throws {
|
|
let thinkingConfig = ThinkingConfig(thinkingLevel: .high, includeThoughts: true)
|
|
let overrides = GenerationConfig(
|
|
maxOutputTokens: 2048,
|
|
stopSequences: ["stop"],
|
|
responseMIMEType: "application/json",
|
|
responseModalities: [.text],
|
|
thinkingConfig: thinkingConfig
|
|
)
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(nil, with: overrides))
|
|
|
|
XCTAssertEqual(result.maxOutputTokens, 2048)
|
|
XCTAssertEqual(result.stopSequences, ["stop"])
|
|
XCTAssertEqual(result.responseMIMEType, "application/json")
|
|
XCTAssertEqual(result.responseModalities, [.text])
|
|
XCTAssertEqual(result.thinkingConfig, thinkingConfig)
|
|
}
|
|
|
|
func testMerge_overridesNil() throws {
|
|
let thinkingConfig = ThinkingConfig(thinkingLevel: .minimal, includeThoughts: false)
|
|
let overrides = GenerationConfig(
|
|
maxOutputTokens: 1024,
|
|
stopSequences: ["test"],
|
|
responseModalities: [.image],
|
|
thinkingConfig: thinkingConfig
|
|
)
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(nil, with: overrides))
|
|
|
|
XCTAssertEqual(result.maxOutputTokens, 1024)
|
|
XCTAssertEqual(result.stopSequences, ["test"])
|
|
XCTAssertNil(result.responseMIMEType)
|
|
XCTAssertEqual(result.responseModalities, [.image])
|
|
XCTAssertEqual(result.thinkingConfig, thinkingConfig)
|
|
}
|
|
|
|
func testMerge_mergesProperties() throws {
|
|
let base = GenerationConfig(
|
|
responseMIMEType: "text/plain"
|
|
)
|
|
let overrides = GenerationConfig(
|
|
responseMIMEType: "application/json",
|
|
thinkingConfig: ThinkingConfig(thinkingBudget: 1024)
|
|
)
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(base, with: overrides))
|
|
XCTAssertEqual(result.responseMIMEType, "application/json")
|
|
XCTAssertEqual(result.thinkingConfig?.thinkingBudget, 1024)
|
|
}
|
|
|
|
func testMerge_schemaPrecedence_overridesJSONSchema() throws {
|
|
let base = GenerationConfig(responseSchema: .string())
|
|
let overrides = GenerationConfig(
|
|
responseMIMEType: "application/json",
|
|
responseJSONSchema: ["type": .string("string")]
|
|
)
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(base, with: overrides))
|
|
|
|
XCTAssertNil(result.responseSchema)
|
|
XCTAssertEqual(result.responseJSONSchema, ["type": .string("string")])
|
|
XCTAssertEqual(result.responseMIMEType, "application/json")
|
|
}
|
|
|
|
func testMerge_schemaPrecedence_overridesSchema() throws {
|
|
let base = GenerationConfig(
|
|
responseMIMEType: "application/json",
|
|
responseJSONSchema: ["type": .string("string")]
|
|
)
|
|
let overrides = GenerationConfig(responseSchema: .string())
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(base, with: overrides))
|
|
|
|
let schema = try XCTUnwrap(result.responseSchema)
|
|
XCTAssertEqual(schema.type, "STRING")
|
|
XCTAssertEqual(schema.nullable, false)
|
|
XCTAssertNil(result.responseJSONSchema)
|
|
}
|
|
|
|
func testMerge_speechConfig() throws {
|
|
let base = GenerationConfig(speechConfig: SpeechConfig(voiceName: "Kore"))
|
|
let overrides = GenerationConfig(speechConfig: SpeechConfig(
|
|
voiceName: "Puck",
|
|
languageCode: "en-US"
|
|
))
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(base, with: overrides))
|
|
XCTAssertEqual(
|
|
result.speechConfig,
|
|
SpeechConfig(voiceName: "Puck", languageCode: "en-US").speechConfig
|
|
)
|
|
}
|
|
|
|
func testMerge_speechConfig_fallbackToBase() throws {
|
|
let base = GenerationConfig(speechConfig: SpeechConfig(voiceName: "Kore"))
|
|
let overrides = GenerationConfig(maxOutputTokens: 55)
|
|
|
|
let result = try XCTUnwrap(GenerationConfig.merge(base, with: overrides))
|
|
XCTAssertEqual(result.speechConfig, SpeechConfig(voiceName: "Kore").speechConfig)
|
|
}
|
|
}
|