feat(firebaseai): Add support for URL context (#17736)

Author: dlarocque

packages/firebase_ai/firebase_ai/lib/firebase_ai.dart

@@ -108,4 +108,5 @@ export 'src/tool.dart'
         Tool,
         ToolConfig,
         GoogleSearch,
-        CodeExecution;
+        CodeExecution,
+        UrlContext;

packages/firebase_ai/firebase_ai/lib/src/api.dart

@@ -182,13 +182,16 @@ final class PromptFeedback {
 /// Metadata on the generation request's token usage.
 final class UsageMetadata {
   // ignore: public_member_api_docs
-  UsageMetadata._(
-      {this.promptTokenCount,
-      this.candidatesTokenCount,
-      this.totalTokenCount,
-      this.thoughtsTokenCount,
-      this.promptTokensDetails,
-      this.candidatesTokensDetails});
+  UsageMetadata._({
+    this.promptTokenCount,
+    this.candidatesTokenCount,
+    this.totalTokenCount,
+    this.thoughtsTokenCount,
+    this.toolUsePromptTokenCount,
+    this.promptTokensDetails,
+    this.candidatesTokensDetails,
+    this.toolUsePromptTokensDetails,
+  });
 
   /// Number of tokens in the prompt.
   final int? promptTokenCount;
@@ -202,19 +205,26 @@ final class UsageMetadata {
   /// Number of tokens present in thoughts output.
   final int? thoughtsTokenCount;
 
+  /// The number of tokens used by tools.
+  final int? toolUsePromptTokenCount;
+
   /// List of modalities that were processed in the request input.
   final List<ModalityTokenCount>? promptTokensDetails;
 
   /// List of modalities that were returned in the response.
   final List<ModalityTokenCount>? candidatesTokensDetails;
+
+  /// A list of tokens used by tools whose usage was triggered from a prompt,
+  /// broken down by modality.
+  final List<ModalityTokenCount>? toolUsePromptTokensDetails;
 }
 
 /// Response candidate generated from a [GenerativeModel].
 final class Candidate {
   // ignore: public_member_api_docs
   Candidate(this.content, this.safetyRatings, this.citationMetadata,
       this.finishReason, this.finishMessage,
-      {this.groundingMetadata});
+      {this.groundingMetadata, this.urlContextMetadata});
 
   /// Generated content returned from the model.
   final Content content;
@@ -242,6 +252,9 @@ final class Candidate {
   /// Metadata returned to the client when grounding is enabled.
   final GroundingMetadata? groundingMetadata;
 
+  /// Metadata returned to the client when the [UrlContext] tool is enabled.
+  final UrlContextMetadata? urlContextMetadata;
+
   /// The concatenation of the text parts of [content], if any.
   ///
   /// If this candidate was finished for a reason of [FinishReason.recitation]
@@ -417,6 +430,76 @@ final class GroundingMetadata {
   final List<String> webSearchQueries;
 }
 
+/// The status of a URL retrieval.
+///
+/// > Warning: For Firebase AI Logic, URL Context
+/// is in Public Preview, which means that the feature is not subject to any SLA
+/// or deprecation policy and could change in backwards-incompatible ways.
+enum UrlRetrievalStatus {
+  /// Unspecified retrieval status.
+  unspecified('URL_RETRIEVAL_STATUS_UNSPECIFIED'),
+
+  /// The URL retrieval was successful.
+  success('URL_RETRIEVAL_STATUS_SUCCESS'),
+
+  /// The URL retrieval failed due.
+  error('URL_RETRIEVAL_STATUS_ERROR'),
+
+  /// The URL retrieval failed because the content is behind a paywall.
+  paywall('URL_RETRIEVAL_STATUS_PAYWALL'),
+
+  /// The URL retrieval failed because the content is unsafe.
+  unsafe('URL_RETRIEVAL_STATUS_UNSAFE');
+
+  const UrlRetrievalStatus(this._jsonString);
+  final String _jsonString;
+
+  // ignore: public_member_api_docs
+  String toJson() => _jsonString;
+
+  // ignore: unused_element
+  static UrlRetrievalStatus _parseValue(Object jsonObject) {
+    return switch (jsonObject) {
+      'URL_RETRIEVAL_STATUS_UNSPECIFIED' => UrlRetrievalStatus.unspecified,
+      'URL_RETRIEVAL_STATUS_SUCCESS' => UrlRetrievalStatus.success,
+      'URL_RETRIEVAL_STATUS_ERROR' => UrlRetrievalStatus.error,
+      'URL_RETRIEVAL_STATUS_PAYWALL' => UrlRetrievalStatus.paywall,
+      'URL_RETRIEVAL_STATUS_UNSAFE' => UrlRetrievalStatus.unsafe,
+      _ => UrlRetrievalStatus
+          .unspecified, // Default to unspecified for unknown values.
+    };
+  }
+}
+
+/// Metadata for a single URL retrieved by the [UrlContext] tool.
+///
+/// > Warning: For Firebase AI Logic, URL Context
+/// is in Public Preview, which means that the feature is not subject to any SLA
+/// or deprecation policy and could change in backwards-incompatible ways.
+final class UrlMetadata {
+  // ignore: public_member_api_docs
+  UrlMetadata({this.retrievedUrl, required this.urlRetrievalStatus});
+
+  /// The retrieved URL.
+  final Uri? retrievedUrl;
+
+  /// The status of the URL retrieval.
+  final UrlRetrievalStatus urlRetrievalStatus;
+}
+
+/// Metadata related to the [UrlContext] tool.
+///
+/// > Warning: For Firebase AI Logic, URL Context
+/// is in Public Preview, which means that the feature is not subject to any SLA
+/// or deprecation policy and could change in backwards-incompatible ways.
+final class UrlContextMetadata {
+  // ignore: public_member_api_docs
+  UrlContextMetadata({required this.urlMetadata});
+
+  /// List of [UrlMetadata] used to provide context to the Gemini model.
+  final List<UrlMetadata> urlMetadata;
+}
+
 /// Safety rating for a piece of content.
 ///
 /// The safety rating contains the category of harm and the harm probability
@@ -1280,6 +1363,11 @@ Candidate _parseCandidate(Object? jsonObject) {
         {'groundingMetadata': final Object groundingMetadata} =>
           parseGroundingMetadata(groundingMetadata),
         _ => null
+      },
+      urlContextMetadata: switch (jsonObject) {
+        {'urlContextMetadata': final Object urlContextMetadata} =>
+          parseUrlContextMetadata(urlContextMetadata),
+        _ => null
       });
 }
 
@@ -1328,6 +1416,11 @@ UsageMetadata parseUsageMetadata(Object jsonObject) {
     {'thoughtsTokenCount': final int thoughtsTokenCount} => thoughtsTokenCount,
     _ => null,
   };
+  final toolUsePromptTokenCount = switch (jsonObject) {
+    {'toolUsePromptTokenCount': final int toolUsePromptTokenCount} =>
+      toolUsePromptTokenCount,
+    _ => null,
+  };
   final promptTokensDetails = switch (jsonObject) {
     {'promptTokensDetails': final List<Object?> promptTokensDetails} =>
       promptTokensDetails.map(_parseModalityTokenCount).toList(),
@@ -1338,13 +1431,23 @@ UsageMetadata parseUsageMetadata(Object jsonObject) {
       candidatesTokensDetails.map(_parseModalityTokenCount).toList(),
     _ => null,
   };
+  final toolUsePromptTokensDetails = switch (jsonObject) {
+    {
+      'toolUsePromptTokensDetails': final List<Object?>
+          toolUsePromptTokensDetails
+    } =>
+      toolUsePromptTokensDetails.map(_parseModalityTokenCount).toList(),
+    _ => null,
+  };
   return UsageMetadata._(
     promptTokenCount: promptTokenCount,
     candidatesTokenCount: candidatesTokenCount,
     totalTokenCount: totalTokenCount,
     thoughtsTokenCount: thoughtsTokenCount,
+    toolUsePromptTokenCount: toolUsePromptTokenCount,
     promptTokensDetails: promptTokensDetails,
     candidatesTokensDetails: candidatesTokensDetails,
+    toolUsePromptTokensDetails: toolUsePromptTokensDetails,
   );
 }
 
@@ -1523,6 +1626,33 @@ SearchEntryPoint _parseSearchEntryPoint(Object? jsonObject) {
   );
 }
 
+UrlMetadata _parseUrlMetadata(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('UrlMetadata', jsonObject);
+  }
+  final uriString = jsonObject['retrievedUrl'] as String?;
+  return UrlMetadata(
+    retrievedUrl: uriString != null ? Uri.parse(uriString) : null,
+    urlRetrievalStatus:
+        UrlRetrievalStatus._parseValue(jsonObject['urlRetrievalStatus']),
+  );
+}
+
+/// Parses a [UrlContextMetadata] from a JSON object.
+///
+/// This function is used internally to convert URL context metadata from the API
+/// response.
+UrlContextMetadata parseUrlContextMetadata(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('UrlContextMetadata', jsonObject);
+  }
+  return UrlContextMetadata(
+    urlMetadata: (jsonObject['urlMetadata'] as List<Object?>? ?? [])
+        .map(_parseUrlMetadata)
+        .toList(),
+  );
+}
+
 /// Supported programming languages for the generated code.
 enum CodeLanguage {
   /// Unspecified status. This value should not be used.

packages/firebase_ai/firebase_ai/lib/src/developer/api.dart

@@ -29,7 +29,8 @@ import '../api.dart'
         SerializationStrategy,
         parseUsageMetadata,
         parseCitationMetadata,
-        parseGroundingMetadata;
+        parseGroundingMetadata,
+        parseUrlContextMetadata;
 import '../content.dart' show Content, parseContent;
 import '../error.dart';
 import '../tool.dart' show Tool, ToolConfig;
@@ -175,6 +176,11 @@ Candidate _parseCandidate(Object? jsonObject) {
         parseGroundingMetadata(groundingMetadata),
       _ => null
     },
+    urlContextMetadata: switch (jsonObject) {
+      {'urlContextMetadata': final Object urlContextMetadata} =>
+        parseUrlContextMetadata(urlContextMetadata),
+      _ => null
+    },
   );
 }
 

packages/firebase_ai/firebase_ai/lib/src/tool.dart

@@ -21,12 +21,13 @@ import 'schema.dart';
 /// knowledge and scope of the model.
 final class Tool {
   // ignore: public_member_api_docs
-  Tool._(this._functionDeclarations, this._googleSearch, this._codeExecution);
+  Tool._(this._functionDeclarations, this._googleSearch, this._codeExecution,
+      this._urlContext);
 
   /// Returns a [Tool] instance with list of [FunctionDeclaration].
   static Tool functionDeclarations(
       List<FunctionDeclaration> functionDeclarations) {
-    return Tool._(functionDeclarations, null, null);
+    return Tool._(functionDeclarations, null, null, null);
   }
 
   /// Creates a tool that allows the model to use Grounding with Google Search.
@@ -47,13 +48,30 @@ final class Tool {
   ///
   /// Returns a `Tool` configured for Google Search.
   static Tool googleSearch({GoogleSearch googleSearch = const GoogleSearch()}) {
-    return Tool._(null, googleSearch, null);
+    return Tool._(null, googleSearch, null, null);
   }
 
   /// Returns a [Tool] instance that enables the model to use Code Execution.
   static Tool codeExecution(
       {CodeExecution codeExecution = const CodeExecution()}) {
-    return Tool._(null, null, codeExecution);
+    return Tool._(null, null, codeExecution, null);
+  }
+
+  /// 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.
+  ///
+  /// - [urlContext]: Specifies the URL context configuration.
+  ///
+  /// Returns a `Tool` configured for URL context.
+  ///
+  /// > Warning: For Firebase AI Logic, URL Context
+  /// is in Public Preview, which means that the feature is not subject to any SLA
+  /// or deprecation policy and could change in backwards-incompatible ways.
+  static Tool urlContext({UrlContext urlContext = const UrlContext()}) {
+    return Tool._(null, null, null, urlContext);
   }
 
   /// A list of `FunctionDeclarations` available to the model that can be used
@@ -74,6 +92,9 @@ final class Tool {
   /// A tool that allows the model to use Code Execution.
   final CodeExecution? _codeExecution;
 
+  /// A tool that allows providing URL context to the model.
+  final UrlContext? _urlContext;
+
   /// Convert to json object.
   Map<String, Object> toJson() => {
         if (_functionDeclarations case final _functionDeclarations?)
@@ -82,7 +103,9 @@ final class Tool {
         if (_googleSearch case final _googleSearch?)
           'googleSearch': _googleSearch.toJson(),
         if (_codeExecution case final _codeExecution?)
-          'codeExecution': _codeExecution.toJson()
+          'codeExecution': _codeExecution.toJson(),
+        if (_urlContext case final _urlContext?)
+          'urlContext': _urlContext.toJson(),
       };
 }
 
@@ -104,6 +127,22 @@ final class GoogleSearch {
   Map<String, Object> toJson() => {};
 }
 
+/// 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.
+///
+/// > Warning: For Firebase AI Logic, URL Context
+/// is in Public Preview, which means that the feature is not subject to any SLA
+/// or deprecation policy and could change in backwards-incompatible ways.
+final class UrlContext {
+  // ignore: public_member_api_docs
+  const UrlContext();
+
+  /// Convert to json object.
+  Map<String, Object> toJson() => {};
+}
+
 /// A tool that allows the model to use Code Execution.
 final class CodeExecution {
   // ignore: public_member_api_docs