Add Grounding with Google Search to Firebase AI Sample App

ai/ai-react-app/src/components/Layout/RightSidebar.tsx

@@ -5,6 +5,7 @@ import {
   AVAILABLE_GENERATIVE_MODELS,
   AVAILABLE_IMAGEN_MODELS,
   defaultFunctionCallingTool,
+  defaultGoogleSearchTool,
 } from "../../services/firebaseAIService";
 import {
   ModelParams,
@@ -158,6 +159,19 @@ const RightSidebar: React.FC<RightSidebarProps> = ({
           nextState.tools = undefined;
           nextState.toolConfig = undefined; // Clear config when turning off
         }
+      } else if (name === "google-search-toggle") {
+        if (checked) {
+          // Turn ON Google Search Grounding
+          nextState.tools = [defaultGoogleSearchTool];
+
+          // Turn OFF JSON mode and Function Calling
+          nextState.generationConfig.responseMimeType = undefined;
+          nextState.generationConfig.responseSchema = undefined;
+          nextState.toolConfig = undefined;
+        } else {
+          // Turn OFF Google Search Grounding
+          nextState.tools = undefined;
+        }
       }
       console.log("[RightSidebar] Updated generative params state:", nextState);
       return nextState;
@@ -219,6 +233,9 @@ const RightSidebar: React.FC<RightSidebarProps> = ({
       generativeParams.toolConfig?.functionCallingConfig?.mode ===
         FunctionCallingMode.ANY) &&
     !!generativeParams.tools?.length;
+  const isGroundingWithGoogleSearchActive = !!generativeParams.tools?.some(
+    (tool) => "googleSearch" in tool,
+  );
 
   return (
     <div className={styles.rightSidebarContainer}>
@@ -360,15 +377,17 @@ const RightSidebar: React.FC<RightSidebarProps> = ({
                   name="structured-output-toggle"
                   checked={isStructuredOutputActive}
                   onChange={handleToggleChange}
-                  disabled={isFunctionCallingActive}
+                  disabled={
+                    isFunctionCallingActive || isGroundingWithGoogleSearchActive
+                  }
                 />
                 <span
-                  className={`${styles.slider} ${isFunctionCallingActive ? styles.disabled : ""}`}
+                  className={`${styles.slider} ${isFunctionCallingActive || isGroundingWithGoogleSearchActive ? styles.disabled : ""}`}
                 ></span>
               </label>
             </div>
             <div
-              className={`${styles.toggleGroup} ${isStructuredOutputActive ? styles.disabledText : ""}`}
+              className={`${styles.toggleGroup} ${isStructuredOutputActive || isGroundingWithGoogleSearchActive ? styles.disabledText : ""}`}
             >
               <label htmlFor="function-call-toggle">Function calling</label>
               <label className={styles.switch}>
@@ -378,13 +397,44 @@ const RightSidebar: React.FC<RightSidebarProps> = ({
                   name="function-call-toggle"
                   checked={isFunctionCallingActive}
                   onChange={handleToggleChange}
-                  disabled={isStructuredOutputActive}
+                  disabled={
+                    isStructuredOutputActive ||
+                    isGroundingWithGoogleSearchActive
+                  }
                 />
                 <span
                   className={`${styles.slider} ${isStructuredOutputActive ? styles.disabled : ""}`}
                 ></span>
               </label>
             </div>
+            <div
+              className={`${styles.toggleGroup} ${
+                isStructuredOutputActive || isFunctionCallingActive
+                  ? styles.disabledText
+                  : ""
+              }`}
+            >
+              <label htmlFor="google-search-toggle">
+                Grounding with Google Search
+              </label>
+              <label className={styles.switch}>
+                <input
+                  type="checkbox"
+                  id="google-search-toggle"
+                  name="google-search-toggle"
+                  checked={isGroundingWithGoogleSearchActive}
+                  onChange={handleToggleChange}
+                  disabled={isStructuredOutputActive || isFunctionCallingActive}
+                />
+                <span
+                  className={`${styles.slider} ${
+                    isStructuredOutputActive || isFunctionCallingActive
+                      ? styles.disabled
+                      : ""
+                  }`}
+                ></span>
+              </label>
+            </div>
           </div>
         </>
       )}

ai/ai-react-app/src/components/Specific/ChatMessage.module.css

@@ -47,3 +47,70 @@
   white-space: pre-wrap;
   word-wrap: break-word;
 }
+
+.sourcesSection {
+  margin-top: 10px;
+  padding-top: 8px;
+  border-top: 1px solid var(--fb-gray-40);
+  font-size: 0.8rem;
+}
+
+.searchEntryPoint {
+  margin-bottom: 8px;
+}
+.searchEntryPoint p {
+  margin: 0 0 4px 0;
+}
+.searchEntryPoint a {
+  color: var(--google-blue);
+  text-decoration: none;
+}
+.searchEntryPoint a:hover {
+  text-decoration: underline;
+}
+
+.sourcesTitle {
+  font-weight: 500;
+  color: var(--fb-gray-70);
+  margin: 0 0 4px 0;
+  font-size: 0.8rem;
+}
+
+.sourcesList {
+  list-style: none;
+  padding-left: 0;
+  margin: 0;
+}
+
+.sourceItem {
+  margin-bottom: 4px;
+}
+
+.sourceItem a {
+  color: var(--google-blue);
+  text-decoration: none;
+  word-break: break-all;
+}
+
+.sourceItem a:hover {
+  text-decoration: underline;
+}
+
+.highlightedSegment {
+  background-color: rgba(
+    var(--google-blue-rgb),
+    0.15
+  );
+  padding: 1px 0;
+  border-radius: 2px;
+  cursor: default;
+}
+
+.sourceSuperscript {
+  font-size: 0.7em;
+  vertical-align: super;
+  color: var(--google-blue);
+  margin-left: 2px;
+  font-weight: bold;
+  user-select: none;
+}

ai/ai-react-app/src/components/Specific/ChatMessage.tsx

@@ -1,9 +1,23 @@
 import React from "react";
-import { Content } from "firebase/ai";
+import {
+  Content,
+  GroundingChunk,
+  GroundingMetadata,
+  GroundingSupport,
+} from "firebase/ai";
 import styles from "./ChatMessage.module.css";
 
 interface ChatMessageProps {
+  /** The message content object containing role and parts. */
   message: Content;
+  groundingMetadata?: GroundingMetadata | null;
+}
+
+interface ProcessedSegment {
+  startIndex: number;
+  endIndex: number;
+  chunkIndices: number[]; // 1-based for display
+  originalSupportIndex: number; // To link back if needed
 }
 
 /**
@@ -23,38 +37,176 @@ const getMessageText = (message: Content): string => {
     .join("");
 };
 
+const renderTextWithInlineHighlighting = (
+  text: string,
+  supports: GroundingSupport[],
+  chunks: GroundingChunk[],
+): React.ReactNode[] => {
+  if (!supports || supports.length === 0 || !text) {
+    return [text];
+  }
+
+  const segmentsToHighlight: ProcessedSegment[] = [];
+
+  supports.forEach((support, supportIndex) => {
+    if (support.segment && support.groundingChunkIndices) {
+      const segment = support.segment;
+      if (segment.partIndex === undefined || segment.partIndex === 0) {
+        segmentsToHighlight.push({
+          startIndex: segment.startIndex,
+          endIndex: segment.endIndex, // API's endIndex is typically exclusive
+          chunkIndices: support.groundingChunkIndices.map((ci) => ci + 1), // 1-based
+          originalSupportIndex: supportIndex,
+        });
+      }
+    }
+  });
+
+  if (segmentsToHighlight.length === 0) {
+    return [text];
+  }
+
+  // Sort segments by start index, then by end index
+  segmentsToHighlight.sort((a, b) => {
+    if (a.startIndex !== b.startIndex) {
+      return a.startIndex - b.startIndex;
+    }
+    return b.endIndex - a.endIndex; // Longer segments first
+  });
+
+  const outputNodes: React.ReactNode[] = [];
+  let lastIndexProcessed = 0;
+
+  segmentsToHighlight.forEach((seg, i) => {
+    // Add un-highlighted text before this segment
+    if (seg.startIndex > lastIndexProcessed) {
+      outputNodes.push(text.substring(lastIndexProcessed, seg.startIndex));
+    }
+
+    // Add the highlighted segment
+    // Ensure we don't re-highlight an already covered portion if a shorter segment comes later
+    const currentSegmentText = text.substring(seg.startIndex, seg.endIndex);
+    const tooltipText = seg.chunkIndices
+      .map((ci) => {
+        const chunk = chunks[ci - 1]; // ci is 1-based
+        return chunk.web?.title || chunk.web?.uri || `Source ${ci}`;
+      })
+      .join("; ");
+
+    outputNodes.push(
+      <span
+        key={`seg-${i}`}
+        className={styles.highlightedSegment}
+        title={`Sources: ${tooltipText}`}
+        data-source-indices={seg.chunkIndices.join(",")}
+      >
+        {currentSegmentText}
+        <sup className={styles.sourceSuperscript}>
+          [{seg.chunkIndices.join(",")}]
+        </sup>
+      </span>,
+    );
+    lastIndexProcessed = Math.max(lastIndexProcessed, seg.endIndex);
+  });
+
+  // Add any remaining un-highlighted text
+  if (lastIndexProcessed < text.length) {
+    outputNodes.push(text.substring(lastIndexProcessed));
+  }
+
+  return outputNodes;
+};
+
 /**
  * Renders a single chat message bubble, styled based on the message role ('user' or 'model').
  * It only renders messages that should be visible in the log (user messages, or model messages
  * containing text). Function role messages or model messages consisting only of function calls
  * are typically not rendered directly as chat bubbles.
  */
-const ChatMessage: React.FC<ChatMessageProps> = ({ message }) => {
+const ChatMessage: React.FC<ChatMessageProps> = ({
+  message,
+  groundingMetadata,
+}) => {
   const text = getMessageText(message);
   const isUser = message.role === "user";
   const isModel = message.role === "model";
 
-  // We render:
-  // 1. User messages (even if they only contain images/files, the 'user' role indicates an entry).
-  // 2. Model messages *only if* they contain actual text content.
-  // We *don't* render:
-  // 1. 'function' role messages (these represent execution results, not direct chat).
-  // 2. 'model' role messages that *only* contain function calls (these are instructions, not display text).
-  // 3. 'system' role messages (handled separately, not usually in chat history display).
-  const shouldRender = isUser || (isModel && text.trim() !== "");
+  const shouldRender =
+    isUser ||
+    (isModel && text.trim() !== "") ||
+    (isModel &&
+      groundingMetadata &&
+      (groundingMetadata.groundingChunks?.length || 0 > 0));
 
   if (!shouldRender) {
     return null;
   }
 
+  let messageContentNodes: React.ReactNode[];
+  if (
+    isModel &&
+    groundingMetadata?.groundingSupports &&
+    groundingMetadata?.groundingChunks
+  ) {
+    messageContentNodes = renderTextWithInlineHighlighting(
+      text,
+      groundingMetadata.groundingSupports,
+      groundingMetadata.groundingChunks,
+    );
+  } else {
+    messageContentNodes = [text];
+  }
+
   return (
     <div
       className={`${styles.messageContainer} ${isUser ? styles.user : styles.model}`}
     >
       <div className={styles.messageBubble}>
-        {/* Use <pre> to preserve whitespace and newlines within the text content.
-                     Handles potential multi-line responses correctly. */}
-        <pre className={styles.messageText}>{text}</pre>
+        <pre className={styles.messageText}>
+          {messageContentNodes.map((node, index) => (
+            <React.Fragment key={index}>{node}</React.Fragment>
+          ))}
+        </pre>
+        {/* Source list rendering */}
+        {isModel &&
+          groundingMetadata &&
+          (groundingMetadata.searchEntryPoint?.renderedContent ||
+          (groundingMetadata.groundingChunks &&
+            groundingMetadata.groundingChunks.length > 0) ? (
+            <div className={styles.sourcesSection}>
+              {groundingMetadata.searchEntryPoint?.renderedContent && (
+                <div
+                  className={styles.searchEntryPoint}
+                  dangerouslySetInnerHTML={{
+                    __html: groundingMetadata.searchEntryPoint.renderedContent,
+                  }}
+                />
+              )}
+              {groundingMetadata.groundingChunks &&
+                groundingMetadata.groundingChunks.length > 0 && (
+                  <>
+                    <h5 className={styles.sourcesTitle}>Sources:</h5>
+                    <ul className={styles.sourcesList}>
+                      {groundingMetadata.groundingChunks.map((chunk, index) => (
+                        <li
+                          key={index}
+                          className={styles.sourceItem}
+                          id={`source-ref-${index + 1}`}
+                        >
+                          <a
+                            href={chunk.web?.uri}
+                            target="_blank"
+                            rel="noopener noreferrer"
+                          >
+                            {`[${index + 1}] ${chunk.web?.title || chunk.web?.uri}`}
+                          </a>
+                        </li>
+                      ))}
+                    </ul>
+                  </>
+                )}
+            </div>
+          ) : null)}
       </div>
     </div>
   );