diff --git a/FirebaseAI/Sources/Tool.swift b/FirebaseAI/Sources/Tool.swift index 16c05b3a2e4..907e8214c8f 100644 --- a/FirebaseAI/Sources/Tool.swift +++ b/FirebaseAI/Sources/Tool.swift @@ -63,6 +63,11 @@ public struct GoogleSearch: Sendable { public init() {} } +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public struct URLContext: Sendable { + public init() {} +} + /// A helper tool that the model may use when generating responses. /// /// A `Tool` is a piece of code that enables the system to interact with external systems to perform @@ -74,14 +79,24 @@ public struct Tool: Sendable { /// Specifies the Google Search configuration. let googleSearch: GoogleSearch? + let urlContext: URLContext? + init(functionDeclarations: [FunctionDeclaration]?) { self.functionDeclarations = functionDeclarations googleSearch = nil + urlContext = nil } init(googleSearch: GoogleSearch) { self.googleSearch = googleSearch functionDeclarations = nil + urlContext = nil + } + + init(urlContext: URLContext) { + self.urlContext = urlContext + functionDeclarations = nil + googleSearch = nil } /// Creates a tool that allows the model to perform function calling. @@ -126,6 +141,10 @@ public struct Tool: Sendable { public static func googleSearch(_ googleSearch: GoogleSearch = GoogleSearch()) -> Tool { return self.init(googleSearch: googleSearch) } + + public static func urlContext(_ urlContext: URLContext = URLContext()) -> Tool { + return self.init(urlContext: urlContext) + } } /// Configuration for specifying function calling behavior. @@ -214,5 +233,8 @@ extension FunctionCallingConfig.Mode: Encodable {} @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) extension GoogleSearch: Encodable {} +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension URLContext: Encodable {} + @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) extension ToolConfig: Encodable {} diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift index d83c300623d..bd3d8b04203 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift @@ -291,6 +291,44 @@ 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 groundingMetadata = try #require(candidate.groundingMetadata) + #expect(!groundingMetadata.groundingChunks.isEmpty) + #expect(!groundingMetadata.groundingSupports.isEmpty) + + for chunk in groundingMetadata.groundingChunks { + #expect(chunk.web != nil) + } + + for support in groundingMetadata.groundingSupports { + let segment = support.segment + #expect(segment.endIndex > segment.startIndex) + #expect(!segment.text.isEmpty) + #expect(!support.groundingChunkIndices.isEmpty) + + // Ensure indices point to valid chunks + for index in support.groundingChunkIndices { + #expect(index < groundingMetadata.groundingChunks.count) + } + } + } + // MARK: Streaming Tests @Test(arguments: [