diff --git a/FirebaseAI/CHANGELOG.md b/FirebaseAI/CHANGELOG.md index d1a00c824c3..8c0a0068fde 100644 --- a/FirebaseAI/CHANGELOG.md +++ b/FirebaseAI/CHANGELOG.md @@ -1,4 +1,6 @@ # Unreleased +- [feature] Added support for the URL context tool, which allows the model to access content + from provided public web URLs to inform and enhance its responses. (#15221) - [changed] Using Firebase AI Logic with the Gemini Developer API is now Generally Available (GA). - [changed] Using Firebase AI Logic with the Imagen generation APIs is now Generally Available (GA). diff --git a/FirebaseAI/Sources/AILog.swift b/FirebaseAI/Sources/AILog.swift index fe04716384a..03232ff23df 100644 --- a/FirebaseAI/Sources/AILog.swift +++ b/FirebaseAI/Sources/AILog.swift @@ -66,6 +66,7 @@ enum AILog { case codeExecutionResultUnrecognizedOutcome = 3015 case executableCodeUnrecognizedLanguage = 3016 case fallbackValueUsed = 3017 + case urlMetadataUnrecognizedURLRetrievalStatus = 3018 // SDK State Errors case generateContentResponseNoCandidates = 4000 diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index 015d5dae56c..c13d3d53c8a 100644 --- a/FirebaseAI/Sources/GenerateContentResponse.swift +++ b/FirebaseAI/Sources/GenerateContentResponse.swift @@ -26,6 +26,9 @@ public struct GenerateContentResponse: Sendable { /// The total number of tokens across the generated response candidates. public let candidatesTokenCount: Int + /// The number of tokens used by tools. + public let toolUsePromptTokenCount: Int + /// The number of tokens used by the model's internal "thinking" process. /// /// For models that support thinking (like Gemini 2.5 Pro and Flash), this represents the actual @@ -39,11 +42,15 @@ public struct GenerateContentResponse: Sendable { /// The total number of tokens in both the request and response. public let totalTokenCount: Int - /// The breakdown, by modality, of how many tokens are consumed by the prompt + /// The breakdown, by modality, of how many tokens are consumed by the prompt. public let promptTokensDetails: [ModalityTokenCount] /// The breakdown, by modality, of how many tokens are consumed by the candidates public let candidatesTokensDetails: [ModalityTokenCount] + + /// The breakdown, by modality, of how many tokens were consumed by the tools used to process + /// the request. + public let toolUsePromptTokensDetails: [ModalityTokenCount] } /// A list of candidate response content, ordered from best to worst. @@ -154,14 +161,19 @@ public struct Candidate: Sendable { public let groundingMetadata: GroundingMetadata? + /// Metadata related to the ``URLContext`` tool. + public let urlContextMetadata: URLContextMetadata? + /// Initializer for SwiftUI previews or tests. public init(content: ModelContent, safetyRatings: [SafetyRating], finishReason: FinishReason?, - citationMetadata: CitationMetadata?, groundingMetadata: GroundingMetadata? = nil) { + citationMetadata: CitationMetadata?, groundingMetadata: GroundingMetadata? = nil, + urlContextMetadata: URLContextMetadata? = nil) { self.content = content self.safetyRatings = safetyRatings self.finishReason = finishReason self.citationMetadata = citationMetadata self.groundingMetadata = groundingMetadata + self.urlContextMetadata = urlContextMetadata } // Returns `true` if the candidate contains no information that a developer could use. @@ -469,10 +481,12 @@ extension GenerateContentResponse.UsageMetadata: Decodable { enum CodingKeys: CodingKey { case promptTokenCount case candidatesTokenCount + case toolUsePromptTokenCount case thoughtsTokenCount case totalTokenCount case promptTokensDetails case candidatesTokensDetails + case toolUsePromptTokensDetails } public init(from decoder: any Decoder) throws { @@ -480,6 +494,8 @@ extension GenerateContentResponse.UsageMetadata: Decodable { promptTokenCount = try container.decodeIfPresent(Int.self, forKey: .promptTokenCount) ?? 0 candidatesTokenCount = try container.decodeIfPresent(Int.self, forKey: .candidatesTokenCount) ?? 0 + toolUsePromptTokenCount = + try container.decodeIfPresent(Int.self, forKey: .toolUsePromptTokenCount) ?? 0 thoughtsTokenCount = try container.decodeIfPresent(Int.self, forKey: .thoughtsTokenCount) ?? 0 totalTokenCount = try container.decodeIfPresent(Int.self, forKey: .totalTokenCount) ?? 0 promptTokensDetails = @@ -488,6 +504,9 @@ extension GenerateContentResponse.UsageMetadata: Decodable { [ModalityTokenCount].self, forKey: .candidatesTokensDetails ) ?? [] + toolUsePromptTokensDetails = try container.decodeIfPresent( + [ModalityTokenCount].self, forKey: .toolUsePromptTokensDetails + ) ?? [] } } @@ -499,6 +518,7 @@ extension Candidate: Decodable { case finishReason case citationMetadata case groundingMetadata + case urlContextMetadata } /// Initializes a response from a decoder. Used for decoding server responses; not for public @@ -540,6 +560,14 @@ extension Candidate: Decodable { GroundingMetadata.self, forKey: .groundingMetadata ) + + if let urlContextMetadata = + try container.decodeIfPresent(URLContextMetadata.self, forKey: .urlContextMetadata), + !urlContextMetadata.urlMetadata.isEmpty { + self.urlContextMetadata = urlContextMetadata + } else { + urlContextMetadata = nil + } } } diff --git a/FirebaseAI/Sources/Tool.swift b/FirebaseAI/Sources/Tool.swift index 78dc8ef9443..53e0ee8b49e 100644 --- a/FirebaseAI/Sources/Tool.swift +++ b/FirebaseAI/Sources/Tool.swift @@ -76,12 +76,15 @@ public struct Tool: Sendable { let googleSearch: GoogleSearch? let codeExecution: CodeExecution? + let urlContext: URLContext? init(functionDeclarations: [FunctionDeclaration]? = nil, googleSearch: GoogleSearch? = nil, + urlContext: URLContext? = nil, codeExecution: CodeExecution? = nil) { self.functionDeclarations = functionDeclarations self.googleSearch = googleSearch + self.urlContext = urlContext self.codeExecution = codeExecution } @@ -128,6 +131,15 @@ public struct Tool: Sendable { return self.init(googleSearch: googleSearch) } + /// Creates a tool that allows you to provide additional context to the models in the form of + /// public web URLs. + /// + /// By including URLs in your request, the Gemini model will access the content from those pages + /// to inform and enhance its response. + public static func urlContext() -> Tool { + return self.init(urlContext: URLContext()) + } + /// Creates a tool that allows the model to execute code. /// /// For more details, see ``CodeExecution``. diff --git a/FirebaseAI/Sources/Types/Internal/Tools/URLContext.swift b/FirebaseAI/Sources/Types/Internal/Tools/URLContext.swift new file mode 100644 index 00000000000..2033bb940f1 --- /dev/null +++ b/FirebaseAI/Sources/Types/Internal/Tools/URLContext.swift @@ -0,0 +1,18 @@ +// 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. + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +struct URLContext: Sendable, Encodable { + init() {} +} diff --git a/FirebaseAI/Sources/Types/Public/URLContextMetadata.swift b/FirebaseAI/Sources/Types/Public/URLContextMetadata.swift new file mode 100644 index 00000000000..5ff671f68eb --- /dev/null +++ b/FirebaseAI/Sources/Types/Public/URLContextMetadata.swift @@ -0,0 +1,34 @@ +// 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. + +/// Metadata related to the ``Tool/urlContext()`` tool. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public struct URLContextMetadata: Sendable, Hashable { + /// List of URL metadata used to provide context to the Gemini model. + public let urlMetadata: [URLMetadata] +} + +// MARK: - Codable Conformances + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension URLContextMetadata: Decodable { + enum CodingKeys: CodingKey { + case urlMetadata + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + urlMetadata = try container.decodeIfPresent([URLMetadata].self, forKey: .urlMetadata) ?? [] + } +} diff --git a/FirebaseAI/Sources/Types/Public/URLMetadata.swift b/FirebaseAI/Sources/Types/Public/URLMetadata.swift new file mode 100644 index 00000000000..50ee16a7e86 --- /dev/null +++ b/FirebaseAI/Sources/Types/Public/URLMetadata.swift @@ -0,0 +1,85 @@ +// 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 Foundation + +/// Metadata for a single URL retrieved by the ``Tool/urlContext()`` tool. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public struct URLMetadata: Sendable, Hashable { + /// Status of the URL retrieval. + public struct URLRetrievalStatus: DecodableProtoEnum, Hashable { + enum Kind: String { + case unspecified = "URL_RETRIEVAL_STATUS_UNSPECIFIED" + case success = "URL_RETRIEVAL_STATUS_SUCCESS" + case error = "URL_RETRIEVAL_STATUS_ERROR" + case paywall = "URL_RETRIEVAL_STATUS_PAYWALL" + case unsafe = "URL_RETRIEVAL_STATUS_UNSAFE" + } + + /// Internal only - Unspecified retrieval status. + static let unspecified = URLRetrievalStatus(kind: .unspecified) + + /// The URL retrieval was successful. + public static let success = URLRetrievalStatus(kind: .success) + + /// The URL retrieval failed. + public static let error = URLRetrievalStatus(kind: .error) + + /// The URL retrieval failed because the content is behind a paywall. + public static let paywall = URLRetrievalStatus(kind: .paywall) + + /// The URL retrieval failed because the content is unsafe. + public static let unsafe = URLRetrievalStatus(kind: .unsafe) + + /// Returns the raw string representation of the `URLRetrievalStatus` value. + public let rawValue: String + + static let unrecognizedValueMessageCode = + AILog.MessageCode.urlMetadataUnrecognizedURLRetrievalStatus + } + + /// The retrieved URL. + public let retrievedURL: URL? + + /// The status of the URL retrieval. + public let retrievalStatus: URLRetrievalStatus +} + +// MARK: - Codable Conformances + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension URLMetadata: Decodable { + enum CodingKeys: String, CodingKey { + case retrievedURL = "retrievedUrl" + case retrievalStatus = "urlRetrievalStatus" + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + + if let retrievedURLString = try container.decodeIfPresent(String.self, forKey: .retrievedURL), + let retrievedURL = URL(string: retrievedURLString) { + self.retrievedURL = retrievedURL + } else { + retrievedURL = nil + } + let retrievalStatus = try container.decodeIfPresent( + URLMetadata.URLRetrievalStatus.self, forKey: .retrievalStatus + ) + + self.retrievalStatus = AILog.safeUnwrap( + retrievalStatus, fallback: URLMetadata.URLRetrievalStatus(kind: .unspecified) + ) + } +} diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift index d2fb589a432..747b1dc5bea 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift @@ -424,6 +424,33 @@ struct GenerateContentIntegrationTests { } } + @Test( + "generateContent with URL Context", + arguments: InstanceConfig.allConfigs + ) + func generateContent_withURLContext_succeeds(_ config: InstanceConfig) async throws { + let model = FirebaseAI.componentInstance(config).generativeModel( + modelName: ModelNames.gemini2_5_Flash, + tools: [.urlContext()] + ) + let prompt = """ + Write a one paragraph summary of this blog post: \ + https://developers.googleblog.com/en/introducing-gemma-3-270m/ + """ + + let response = try await model.generateContent(prompt) + + let candidate = try #require(response.candidates.first) + let urlContextMetadata = try #require(candidate.urlContextMetadata) + #expect(urlContextMetadata.urlMetadata.count == 1) + let urlMetadata = try #require(urlContextMetadata.urlMetadata.first) + let retrievedURL = try #require(urlMetadata.retrievedURL) + #expect( + retrievedURL == URL(string: "https://developers.googleblog.com/en/introducing-gemma-3-270m/") + ) + #expect(urlMetadata.retrievalStatus == .success) + } + @Test(arguments: InstanceConfig.allConfigs) func generateContent_codeExecution_succeeds(_ config: InstanceConfig) async throws { let model = FirebaseAI.componentInstance(config).generativeModel( diff --git a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift index c6335142959..59e1581a638 100644 --- a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift +++ b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift @@ -333,6 +333,55 @@ final class GenerativeModelGoogleAITests: XCTestCase { 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 { @@ -642,4 +691,27 @@ final class GenerativeModelGoogleAITests: XCTestCase { 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) + } } diff --git a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift index 847f5a8e643..1d2498f07e5 100644 --- a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift +++ b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift @@ -487,6 +487,73 @@ final class GenerativeModelVertexAITests: XCTestCase { XCTAssertEqual( textPart2.text, "The sum of the first 5 prime numbers (2, 3, 5, 7, and 11) is 28." ) + let usageMetadata = try XCTUnwrap(response.usageMetadata) + XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 371) + } + + func testGenerateContent_success_urlContext() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-url-context", + withExtension: "json", + subdirectory: vertexSubdirectory + ) + + 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, 34) + XCTAssertEqual(usageMetadata.thoughtsTokenCount, 36) + } + + func testGenerateContent_success_urlContext_mixedValidity() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-url-context-mixed-validity", + withExtension: "json", + subdirectory: vertexSubdirectory + ) + + 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_success_urlContext_retrievedURLPresentOnErrorStatus() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-url-context-missing-retrievedurl", + withExtension: "json", + subdirectory: vertexSubdirectory + ) + + let response = try await model.generateContent(testPrompt) + + let candidate = try XCTUnwrap(response.candidates.first) + let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata) + let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first) + let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL) + XCTAssertEqual(retrievedURL.absoluteString, "https://example.com/8") + XCTAssertEqual(urlMetadata.retrievalStatus, .error) } func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws { @@ -1718,6 +1785,29 @@ final class GenerativeModelVertexAITests: XCTestCase { XCTAssertEqual(responses, 1) } + func testGenerateContentStream_success_urlContext() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "streaming-success-url-context", + withExtension: "txt", + subdirectory: vertexSubdirectory + ) + + 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) + } + // MARK: - Count Tokens func testCountTokens_succeeds() async throws { diff --git a/FirebaseAI/Tests/Unit/MockURLProtocol.swift b/FirebaseAI/Tests/Unit/MockURLProtocol.swift index 5385b164015..6db227d5cfb 100644 --- a/FirebaseAI/Tests/Unit/MockURLProtocol.swift +++ b/FirebaseAI/Tests/Unit/MockURLProtocol.swift @@ -21,6 +21,7 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable { URLResponse, AsyncLineSequence? ))? + override class func canInit(with request: URLRequest) -> Bool { #if os(watchOS) print("MockURLProtocol cannot be used on watchOS.") @@ -33,13 +34,14 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable { override class func canonicalRequest(for request: URLRequest) -> URLRequest { return request } override func startLoading() { - guard let requestHandler = MockURLProtocol.requestHandler else { - fatalError("`requestHandler` is nil.") - } guard let client = client else { fatalError("`client` is nil.") } + guard let requestHandler = MockURLProtocol.requestHandler else { + fatalError("No request handler set.") + } + Task { let (response, stream) = try requestHandler(self.request) client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) diff --git a/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift b/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift index a53d215359f..276308f63aa 100644 --- a/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift +++ b/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift @@ -17,6 +17,8 @@ import XCTest @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) final class GenerateContentResponseTests: XCTestCase { + let jsonDecoder = JSONDecoder() + // MARK: - GenerateContentResponse Computed Properties func testGenerateContentResponse_inlineDataParts_success() throws { @@ -106,4 +108,53 @@ final class GenerateContentResponseTests: XCTestCase { "functionCalls should be empty when there are no candidates." ) } + + // MARK: - Decoding Tests + + func testDecodeCandidate_emptyURLMetadata_urlContextMetadataIsNil() throws { + let json = """ + { + "content": { "role": "model", "parts": [ { "text": "Some text." } ] }, + "finishReason": "STOP", + "urlContextMetadata": { "urlMetadata": [] } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData) + + XCTAssertNil( + candidate.urlContextMetadata, + "urlContextMetadata should be nil if the `urlMetadata` array is empty in the candidate." + ) + XCTAssertEqual(candidate.content.role, "model") + let part = try XCTUnwrap(candidate.content.parts.first) + let textPart = try XCTUnwrap(part as? TextPart) + XCTAssertEqual(textPart.text, "Some text.") + XCTAssertFalse(textPart.isThought) + XCTAssertEqual(candidate.finishReason, .stop) + } + + func testDecodeCandidate_missingURLMetadata_urlContextMetadataIsNil() throws { + let json = """ + { + "content": { "role": "model", "parts": [ { "text": "Some text." } ] }, + "finishReason": "STOP" + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData) + + XCTAssertNil( + candidate.urlContextMetadata, + "urlContextMetadata should be nil if `urlMetadata` is not provided in the candidate." + ) + XCTAssertEqual(candidate.content.role, "model") + let part = try XCTUnwrap(candidate.content.parts.first) + let textPart = try XCTUnwrap(part as? TextPart) + XCTAssertEqual(textPart.text, "Some text.") + XCTAssertFalse(textPart.isThought) + XCTAssertEqual(candidate.finishReason, .stop) + } }