mirror of
https://github.com/firebase/firebase-ios-sdk
synced 2026-11-10 22:45:59 +05:45
718 lines
30 KiB
Swift
718 lines
30 KiB
Swift
// Copyright 2025 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.
|
|
|
|
import FirebaseAppCheckInterop
|
|
import FirebaseAuthInterop
|
|
import FirebaseCore
|
|
import XCTest
|
|
|
|
@testable import FirebaseAILogic
|
|
|
|
@available(macOS 12.0, watchOS 8.0, *)
|
|
final class GenerativeModelGoogleAITests: XCTestCase {
|
|
let testPrompt = "What sorts of questions can I ask you?"
|
|
let safetyRatingsNegligible: [SafetyRating] = [
|
|
.init(
|
|
category: .sexuallyExplicit,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
.init(
|
|
category: .hateSpeech,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
.init(
|
|
category: .harassment,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
.init(
|
|
category: .dangerousContent,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
].sorted()
|
|
let testModelName = "test-model"
|
|
let testModelResourceName = "projects/test-project-id/models/test-model"
|
|
let apiConfig = FirebaseAI.defaultVertexAIAPIConfig
|
|
|
|
let googleAISubdirectory = "mock-responses/googleai"
|
|
|
|
var urlSession: URLSession!
|
|
var model: GenerativeModel!
|
|
|
|
override func setUp() async throws {
|
|
let configuration = URLSessionConfiguration.default
|
|
configuration.protocolClasses = [MockURLProtocol.self]
|
|
urlSession = try XCTUnwrap(URLSession(configuration: configuration))
|
|
model = GenerativeModel(
|
|
modelName: testModelName,
|
|
modelResourceName: testModelResourceName,
|
|
firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(),
|
|
apiConfig: apiConfig,
|
|
tools: nil,
|
|
requestOptions: RequestOptions(),
|
|
urlSession: urlSession
|
|
)
|
|
}
|
|
|
|
override func tearDown() {
|
|
MockURLProtocol.requestHandler = nil
|
|
}
|
|
|
|
// MARK: - Generate Content
|
|
|
|
func testGenerateContent_success_basicReplyLong() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-basic-reply-long",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let finishReason = try XCTUnwrap(candidate.finishReason)
|
|
XCTAssertEqual(finishReason, .stop)
|
|
XCTAssertEqual(candidate.safetyRatings.count, 4)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let part = try XCTUnwrap(candidate.content.parts.first)
|
|
let partText = try XCTUnwrap(part as? TextPart).text
|
|
XCTAssertTrue(partText.hasPrefix("Making professional-quality"))
|
|
XCTAssertEqual(response.text, partText)
|
|
XCTAssertEqual(response.functionCalls, [])
|
|
XCTAssertEqual(response.inlineDataParts, [])
|
|
}
|
|
|
|
func testGenerateContent_success_basicReplyShort() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-basic-reply-short",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let finishReason = try XCTUnwrap(candidate.finishReason)
|
|
XCTAssertEqual(finishReason, .stop)
|
|
XCTAssertEqual(candidate.safetyRatings.sorted(), safetyRatingsNegligible)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let part = try XCTUnwrap(candidate.content.parts.first)
|
|
let textPart = try XCTUnwrap(part as? TextPart)
|
|
XCTAssertTrue(textPart.text.hasPrefix("Google's headquarters"))
|
|
XCTAssertEqual(response.text, textPart.text)
|
|
XCTAssertEqual(response.functionCalls, [])
|
|
XCTAssertEqual(response.inlineDataParts, [])
|
|
}
|
|
|
|
func testGenerateContent_success_citations() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-citations",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let text = try XCTUnwrap(response.text)
|
|
XCTAssertTrue(text.hasPrefix("Okay, let's break down quantum mechanics."))
|
|
let citationMetadata = try XCTUnwrap(candidate.citationMetadata)
|
|
XCTAssertEqual(citationMetadata.citations.count, 4)
|
|
let citationSource1 = try XCTUnwrap(citationMetadata.citations[0])
|
|
XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
|
|
XCTAssertEqual(citationSource1.startIndex, 548)
|
|
XCTAssertEqual(citationSource1.endIndex, 690)
|
|
XCTAssertNil(citationSource1.title)
|
|
XCTAssertEqual(citationSource1.license, "mit")
|
|
XCTAssertNil(citationSource1.publicationDate)
|
|
let citationSource2 = try XCTUnwrap(citationMetadata.citations[1])
|
|
XCTAssertEqual(citationSource2.uri, "https://www.example.com/some-citation-1")
|
|
XCTAssertEqual(citationSource2.startIndex, 1240)
|
|
XCTAssertEqual(citationSource2.endIndex, 1407)
|
|
XCTAssertNil(citationSource2.title, "some-citation-2")
|
|
XCTAssertNil(citationSource2.license)
|
|
XCTAssertNil(citationSource2.publicationDate)
|
|
let citationSource3 = try XCTUnwrap(citationMetadata.citations[2])
|
|
XCTAssertEqual(citationSource3.startIndex, 1942)
|
|
XCTAssertEqual(citationSource3.endIndex, 2149)
|
|
XCTAssertNil(citationSource3.uri)
|
|
XCTAssertNil(citationSource3.license)
|
|
XCTAssertNil(citationSource3.title)
|
|
XCTAssertNil(citationSource3.publicationDate)
|
|
let citationSource4 = try XCTUnwrap(citationMetadata.citations[3])
|
|
XCTAssertEqual(citationSource4.startIndex, 2036)
|
|
XCTAssertEqual(citationSource4.endIndex, 2175)
|
|
XCTAssertNil(citationSource4.uri)
|
|
XCTAssertNil(citationSource4.license)
|
|
XCTAssertNil(citationSource4.title)
|
|
XCTAssertNil(citationSource4.publicationDate)
|
|
}
|
|
|
|
func testGenerateContent_usageMetadata() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-basic-reply-short",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
let usageMetadata = try XCTUnwrap(response.usageMetadata)
|
|
XCTAssertEqual(usageMetadata.promptTokenCount, 7)
|
|
XCTAssertEqual(usageMetadata.promptTokensDetails.count, 1)
|
|
XCTAssertEqual(usageMetadata.promptTokensDetails[0].modality, .text)
|
|
XCTAssertEqual(usageMetadata.promptTokensDetails[0].tokenCount, 7)
|
|
XCTAssertEqual(usageMetadata.candidatesTokenCount, 22)
|
|
XCTAssertEqual(usageMetadata.candidatesTokensDetails.count, 1)
|
|
XCTAssertEqual(usageMetadata.candidatesTokensDetails[0].modality, .text)
|
|
XCTAssertEqual(usageMetadata.candidatesTokensDetails[0].tokenCount, 22)
|
|
}
|
|
|
|
func testGenerateContent_groundingMetadata() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-google-search-grounding",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let groundingMetadata = try XCTUnwrap(candidate.groundingMetadata)
|
|
|
|
XCTAssertEqual(groundingMetadata.webSearchQueries, ["current weather in London"])
|
|
let searchEntryPoint = try XCTUnwrap(groundingMetadata.searchEntryPoint)
|
|
XCTAssertFalse(searchEntryPoint.renderedContent.isEmpty)
|
|
|
|
XCTAssertEqual(groundingMetadata.groundingChunks.count, 2)
|
|
let firstChunk = try XCTUnwrap(groundingMetadata.groundingChunks.first?.web)
|
|
XCTAssertEqual(firstChunk.title, "accuweather.com")
|
|
XCTAssertNotNil(firstChunk.uri)
|
|
XCTAssertNil(firstChunk.domain) // Domain is not supported by Google AI backend
|
|
|
|
XCTAssertEqual(groundingMetadata.groundingSupports.count, 3)
|
|
let firstSupport = try XCTUnwrap(groundingMetadata.groundingSupports.first)
|
|
let segment = try XCTUnwrap(firstSupport.segment)
|
|
XCTAssertEqual(segment.text, "The current weather in London, United Kingdom is cloudy.")
|
|
XCTAssertEqual(segment.startIndex, 0)
|
|
XCTAssertEqual(segment.partIndex, 0)
|
|
XCTAssertEqual(segment.endIndex, 56)
|
|
XCTAssertEqual(firstSupport.groundingChunkIndices, [0])
|
|
}
|
|
|
|
// This test case can be deleted once https://b.corp.google.com/issues/422779395 (internal) is
|
|
// fixed.
|
|
func testGenerateContent_groundingMetadata_emptyGroundingChunks() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-google-search-grounding-empty-grounding-chunks",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let groundingMetadata = try XCTUnwrap(candidate.groundingMetadata)
|
|
XCTAssertNotNil(groundingMetadata.searchEntryPoint)
|
|
XCTAssertEqual(groundingMetadata.webSearchQueries, ["current weather London"])
|
|
|
|
// Chunks exist, but contain no web information.
|
|
XCTAssertEqual(groundingMetadata.groundingChunks.count, 2)
|
|
XCTAssertNil(groundingMetadata.groundingChunks[0].web)
|
|
XCTAssertNil(groundingMetadata.groundingChunks[1].web)
|
|
|
|
XCTAssertEqual(groundingMetadata.groundingSupports.count, 1)
|
|
let support = try XCTUnwrap(groundingMetadata.groundingSupports.first)
|
|
XCTAssertEqual(support.groundingChunkIndices, [0])
|
|
XCTAssertEqual(
|
|
support.segment.text,
|
|
"There is a 0% chance of rain and the humidity is around 41%."
|
|
)
|
|
}
|
|
|
|
func testGenerateContent_success_thinking_thoughtSummary() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-thinking-reply-thought-summary",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.content.parts.count, 2)
|
|
let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
|
|
XCTAssertTrue(thoughtPart.isThought)
|
|
XCTAssertTrue(thoughtPart.text.hasPrefix("**Thinking About Google's Headquarters**"))
|
|
XCTAssertEqual(thoughtPart.text, response.thoughtSummary)
|
|
let textPart = try XCTUnwrap(candidate.content.parts.last as? TextPart)
|
|
XCTAssertFalse(textPart.isThought)
|
|
XCTAssertEqual(textPart.text, "Mountain View")
|
|
XCTAssertEqual(textPart.text, response.text)
|
|
}
|
|
|
|
func testGenerateContent_success_thinking_functionCall_thoughtSummaryAndSignature() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-thinking-function-call-thought-summary-signature",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.finishReason, .stop)
|
|
XCTAssertEqual(candidate.content.parts.count, 2)
|
|
let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
|
|
XCTAssertTrue(thoughtPart.isThought)
|
|
XCTAssertTrue(thoughtPart.text.hasPrefix("**Thinking Through the New Year's Eve Calculation**"))
|
|
let functionCallPart = try XCTUnwrap(candidate.content.parts.last as? FunctionCallPart)
|
|
XCTAssertFalse(functionCallPart.isThought)
|
|
XCTAssertEqual(functionCallPart.name, "now")
|
|
XCTAssertTrue(functionCallPart.args.isEmpty)
|
|
let thoughtSignature = try XCTUnwrap(functionCallPart.thoughtSignature)
|
|
XCTAssertTrue(thoughtSignature.hasPrefix("CtQOAVSoXO74PmYr9AFu"))
|
|
}
|
|
|
|
func testGenerateContent_success_codeExecution() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-code-execution",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let parts = candidate.content.parts
|
|
XCTAssertEqual(candidate.finishReason, .stop)
|
|
XCTAssertEqual(parts.count, 3)
|
|
let executableCodePart = try XCTUnwrap(parts[0] as? ExecutableCodePart)
|
|
XCTAssertFalse(executableCodePart.isThought)
|
|
XCTAssertEqual(executableCodePart.language, .python)
|
|
XCTAssertTrue(executableCodePart.code.starts(with: "prime_numbers = [2, 3, 5, 7, 11]"))
|
|
let codeExecutionResultPart = try XCTUnwrap(parts[1] as? CodeExecutionResultPart)
|
|
XCTAssertFalse(codeExecutionResultPart.isThought)
|
|
XCTAssertEqual(codeExecutionResultPart.outcome, .ok)
|
|
XCTAssertEqual(codeExecutionResultPart.output, "sum_of_primes=28\n")
|
|
let textPart = try XCTUnwrap(parts[2] as? TextPart)
|
|
XCTAssertFalse(textPart.isThought)
|
|
XCTAssertTrue(textPart.text.hasPrefix("The first 5 prime numbers are 2, 3, 5, 7, and 11."))
|
|
let usageMetadata = try XCTUnwrap(response.usageMetadata)
|
|
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 160)
|
|
}
|
|
|
|
func testGenerateContent_success_urlContext() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-url-context",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
|
|
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
|
|
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
|
|
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
|
|
XCTAssertEqual(
|
|
retrievedURL,
|
|
URL(string: "https://berkshirehathaway.com")
|
|
)
|
|
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
|
|
let usageMetadata = try XCTUnwrap(response.usageMetadata)
|
|
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 424)
|
|
}
|
|
|
|
func testGenerateContent_success_urlContext_mixedValidity() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-url-context-mixed-validity",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
|
|
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)
|
|
|
|
let paywallURLMetadata = urlContextMetadata.urlMetadata[0]
|
|
XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)
|
|
|
|
let successURLMetadata = urlContextMetadata.urlMetadata[1]
|
|
XCTAssertEqual(successURLMetadata.retrievalStatus, .success)
|
|
|
|
let errorURLMetadata = urlContextMetadata.urlMetadata[2]
|
|
XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
|
|
}
|
|
|
|
func testGenerateContent_failure_invalidAPIKey() async throws {
|
|
let expectedStatusCode = 400
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-api-key",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory,
|
|
statusCode: expectedStatusCode
|
|
)
|
|
|
|
do {
|
|
_ = try await model.generateContent(testPrompt)
|
|
XCTFail("Should throw GenerateContentError.internalError; no error thrown.")
|
|
} catch let GenerateContentError.internalError(error as BackendError) {
|
|
XCTAssertEqual(error.httpResponseCode, 400)
|
|
XCTAssertEqual(error.status, .invalidArgument)
|
|
XCTAssertEqual(error.message, "API key not valid. Please pass a valid API key.")
|
|
XCTAssertTrue(error.localizedDescription.contains(error.message))
|
|
XCTAssertTrue(error.localizedDescription.contains(error.status.rawValue))
|
|
XCTAssertTrue(error.localizedDescription.contains("\(error.httpResponseCode)"))
|
|
let nsError = error as NSError
|
|
XCTAssertEqual(nsError.domain, "\(Constants.baseErrorDomain).\(BackendError.self)")
|
|
XCTAssertEqual(nsError.code, error.httpResponseCode)
|
|
return
|
|
} catch {
|
|
XCTFail("Should throw GenerateContentError.internalError(RPCError); error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
func testGenerateContent_failure_finishReasonSafety() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-finish-reason-safety",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
do {
|
|
_ = try await model.generateContent(testPrompt)
|
|
XCTFail("Should throw")
|
|
} catch let GenerateContentError.responseStoppedEarly(reason, response) {
|
|
XCTAssertEqual(reason, .safety)
|
|
XCTAssertEqual(response.text, "Safety error incoming in 5, 4, 3, 2...")
|
|
} catch {
|
|
XCTFail("Should throw a responseStoppedEarly")
|
|
}
|
|
}
|
|
|
|
func testGenerateContent_failure_unknownModel() async throws {
|
|
let expectedStatusCode = 404
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-unknown-model",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory,
|
|
statusCode: 404
|
|
)
|
|
|
|
do {
|
|
_ = try await model.generateContent(testPrompt)
|
|
XCTFail("Should throw GenerateContentError.internalError; no error thrown.")
|
|
} catch let GenerateContentError.internalError(underlying: rpcError as BackendError) {
|
|
XCTAssertEqual(rpcError.status, .notFound)
|
|
XCTAssertEqual(rpcError.httpResponseCode, expectedStatusCode)
|
|
XCTAssertTrue(rpcError.message.hasPrefix("models/gemini-5.0-flash is not found"))
|
|
} catch {
|
|
XCTFail("Should throw GenerateContentError.internalError; error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
// MARK: - Generate Content (Streaming)
|
|
|
|
func testGenerateContentStream_successBasicReplyLong() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-basic-reply-long",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = 0
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await content in stream {
|
|
XCTAssertNotNil(content.text)
|
|
responses += 1
|
|
}
|
|
|
|
XCTAssertEqual(responses, 36)
|
|
}
|
|
|
|
func testGenerateContentStream_successBasicReplyShort() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-basic-reply-short",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = 0
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await content in stream {
|
|
XCTAssertNotNil(content.text)
|
|
responses += 1
|
|
}
|
|
|
|
XCTAssertEqual(responses, 3)
|
|
}
|
|
|
|
func testGenerateContentStream_successWithCitations() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-citations",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let stream = try model.generateContentStream("Hi")
|
|
var citations = [Citation]()
|
|
var responses = [GenerateContentResponse]()
|
|
for try await content in stream {
|
|
responses.append(content)
|
|
XCTAssertNotNil(content.text)
|
|
let candidate = try XCTUnwrap(content.candidates.first)
|
|
if let sources = candidate.citationMetadata?.citations {
|
|
citations.append(contentsOf: sources)
|
|
}
|
|
}
|
|
|
|
let lastCandidate = try XCTUnwrap(responses.last?.candidates.first)
|
|
XCTAssertEqual(lastCandidate.finishReason, .stop)
|
|
XCTAssertEqual(citations.count, 1)
|
|
let citation = try XCTUnwrap(citations.first)
|
|
XCTAssertEqual(citation.startIndex, 111)
|
|
XCTAssertEqual(citation.endIndex, 236)
|
|
let citationURI = try XCTUnwrap(citation.uri)
|
|
XCTAssertTrue(citationURI.starts(with: "https://www."))
|
|
XCTAssertNil(citation.license)
|
|
XCTAssertNil(citation.title)
|
|
XCTAssertNil(citation.publicationDate)
|
|
}
|
|
|
|
func testGenerateContentStream_successWithThoughtSummary() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-thinking-reply-thought-summary",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var thoughtSummary = ""
|
|
var text = ""
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await response in stream {
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let textPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
|
|
if textPart.isThought {
|
|
let newThought = try XCTUnwrap(response.thoughtSummary)
|
|
XCTAssertEqual(textPart.text, newThought)
|
|
thoughtSummary.append(newThought)
|
|
} else {
|
|
let newText = try XCTUnwrap(response.text)
|
|
XCTAssertEqual(textPart.text, newText)
|
|
text.append(newText)
|
|
}
|
|
}
|
|
|
|
XCTAssertTrue(thoughtSummary.hasPrefix("**Exploring Sky Color**"))
|
|
XCTAssertTrue(text.hasPrefix("The sky is blue because"))
|
|
}
|
|
|
|
func testGenerateContentStream_success_thinking_functionCall_thoughtSummary_signature() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-thinking-function-call-thought-summary-signature",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var thoughtSummary = ""
|
|
var functionCalls: [FunctionCallPart] = []
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await response in stream {
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let part = try XCTUnwrap(candidate.content.parts.first)
|
|
if part.isThought {
|
|
let textPart = try XCTUnwrap(part as? TextPart)
|
|
let newThought = try XCTUnwrap(response.thoughtSummary)
|
|
XCTAssertEqual(textPart.text, newThought)
|
|
thoughtSummary.append(newThought)
|
|
} else {
|
|
let functionCallPart = try XCTUnwrap(part as? FunctionCallPart)
|
|
XCTAssertEqual(response.functionCalls.count, 1)
|
|
let newFunctionCall = try XCTUnwrap(response.functionCalls.first)
|
|
XCTAssertEqual(functionCallPart, newFunctionCall)
|
|
functionCalls.append(newFunctionCall)
|
|
}
|
|
}
|
|
|
|
XCTAssertTrue(thoughtSummary.hasPrefix("**Calculating the Days**"))
|
|
XCTAssertEqual(functionCalls.count, 1)
|
|
let functionCall = try XCTUnwrap(functionCalls.first)
|
|
XCTAssertEqual(functionCall.name, "now")
|
|
XCTAssertTrue(functionCall.args.isEmpty)
|
|
let thoughtSignature = try XCTUnwrap(functionCall.thoughtSignature)
|
|
XCTAssertTrue(thoughtSignature.hasPrefix("CiIBVKhc7vB+vaaq6rA"))
|
|
}
|
|
|
|
func testGenerateContentStream_success_ignoresEmptyParts() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-empty-parts",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await response in stream {
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertGreaterThan(candidate.content.parts.count, 0)
|
|
let text = response.text
|
|
let inlineData = response.inlineDataParts.first
|
|
XCTAssertTrue(text != nil || inlineData != nil, "Response did not contain text or data")
|
|
}
|
|
}
|
|
|
|
func testGenerateContentStream_success_codeExecution() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-code-execution",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var parts = [any Part]()
|
|
let stream = try model.generateContentStream(testPrompt)
|
|
for try await response in stream {
|
|
if let responseParts = response.candidates.first?.content.parts {
|
|
parts.append(contentsOf: responseParts)
|
|
}
|
|
}
|
|
|
|
let thoughtParts = parts.filter { $0.isThought }
|
|
XCTAssertEqual(thoughtParts.count, 0)
|
|
let textParts = parts.filter { $0 is TextPart }
|
|
XCTAssertGreaterThan(textParts.count, 0)
|
|
let executableCodeParts = parts.compactMap { $0 as? ExecutableCodePart }
|
|
XCTAssertEqual(executableCodeParts.count, 1)
|
|
let executableCodePart = try XCTUnwrap(executableCodeParts.first)
|
|
XCTAssertFalse(executableCodePart.isThought)
|
|
XCTAssertEqual(executableCodePart.language, .python)
|
|
XCTAssertTrue(executableCodePart.code.starts(with: "prime_numbers = [2, 3, 5, 7, 11]"))
|
|
let codeExecutionResultParts = parts.compactMap { $0 as? CodeExecutionResultPart }
|
|
XCTAssertEqual(codeExecutionResultParts.count, 1)
|
|
let codeExecutionResultPart = try XCTUnwrap(codeExecutionResultParts.first)
|
|
XCTAssertFalse(codeExecutionResultPart.isThought)
|
|
XCTAssertEqual(codeExecutionResultPart.outcome, .ok)
|
|
XCTAssertEqual(codeExecutionResultPart.output, "The sum of the first 5 prime numbers is: 28\n")
|
|
}
|
|
|
|
func testGenerateContentStream_failureInvalidAPIKey() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-api-key",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
do {
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await _ in stream {
|
|
XCTFail("No content is there, this shouldn't happen.")
|
|
}
|
|
} catch let GenerateContentError.internalError(error as BackendError) {
|
|
XCTAssertEqual(error.httpResponseCode, 400)
|
|
XCTAssertEqual(error.status, .invalidArgument)
|
|
XCTAssertEqual(error.message, "API key not valid. Please pass a valid API key.")
|
|
XCTAssertTrue(error.localizedDescription.contains(error.message))
|
|
XCTAssertTrue(error.localizedDescription.contains(error.status.rawValue))
|
|
XCTAssertTrue(error.localizedDescription.contains("\(error.httpResponseCode)"))
|
|
let nsError = error as NSError
|
|
XCTAssertEqual(nsError.domain, "\(Constants.baseErrorDomain).\(BackendError.self)")
|
|
XCTAssertEqual(nsError.code, error.httpResponseCode)
|
|
return
|
|
}
|
|
|
|
XCTFail("Should have caught an error.")
|
|
}
|
|
|
|
func testGenerateContentStream_failureFinishRecitation() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-failure-recitation-no-content",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = [GenerateContentResponse]()
|
|
do {
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await response in stream {
|
|
responses.append(response)
|
|
}
|
|
XCTFail("Expected a GenerateContentError.responseStoppedEarly error, but got no error.")
|
|
} catch let GenerateContentError.responseStoppedEarly(reason, response) {
|
|
XCTAssertEqual(reason, .recitation)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.finishReason, reason)
|
|
} catch {
|
|
XCTFail("Expected a GenerateContentError.responseStoppedEarly error, but got error: \(error)")
|
|
}
|
|
|
|
XCTAssertEqual(responses.count, 8)
|
|
let firstResponse = try XCTUnwrap(responses.first)
|
|
XCTAssertEqual(firstResponse.text, "text1")
|
|
let lastResponse = try XCTUnwrap(responses.last)
|
|
XCTAssertEqual(lastResponse.text, "text8")
|
|
}
|
|
|
|
func testGenerateContentStream_success_urlContext() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-url-context",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = [GenerateContentResponse]()
|
|
let stream = try model.generateContentStream(testPrompt)
|
|
for try await response in stream {
|
|
responses.append(response)
|
|
}
|
|
|
|
let firstResponse = try XCTUnwrap(responses.first)
|
|
let candidate = try XCTUnwrap(firstResponse.candidates.first)
|
|
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
|
|
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
|
|
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
|
|
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
|
|
XCTAssertEqual(retrievedURL, URL(string: "https://google.com"))
|
|
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
|
|
}
|
|
}
|