11import React from "react" ;
2- import { Content } from "firebase/ai" ;
2+ import {
3+ Content ,
4+ GroundingChunk ,
5+ GroundingMetadata ,
6+ GroundingSupport ,
7+ } from "firebase/ai" ;
38import styles from "./ChatMessage.module.css" ;
49
510interface ChatMessageProps {
11+ /** The message content object containing role and parts. */
612 message : Content ;
13+ groundingMetadata ?: GroundingMetadata | null ;
14+ }
15+
16+ interface ProcessedSegment {
17+ startIndex : number ;
18+ endIndex : number ;
19+ chunkIndices : number [ ] ; // 1-based for display
20+ originalSupportIndex : number ; // To link back if needed
721}
822
923/**
@@ -23,13 +37,96 @@ const getMessageText = (message: Content): string => {
2337 . join ( "" ) ;
2438} ;
2539
40+ const renderTextWithInlineHighlighting = (
41+ text : string ,
42+ supports : GroundingSupport [ ] ,
43+ chunks : GroundingChunk [ ] ,
44+ ) : React . ReactNode [ ] => {
45+ if ( ! supports || supports . length === 0 || ! text ) {
46+ return [ text ] ;
47+ }
48+
49+ const segmentsToHighlight : ProcessedSegment [ ] = [ ] ;
50+
51+ supports . forEach ( ( support , supportIndex ) => {
52+ if ( support . segment && support . groundingChunkIndices ) {
53+ const segment = support . segment ;
54+ if ( segment . partIndex === undefined || segment . partIndex === 0 ) {
55+ segmentsToHighlight . push ( {
56+ startIndex : segment . startIndex ,
57+ endIndex : segment . endIndex , // API's endIndex is typically exclusive
58+ chunkIndices : support . groundingChunkIndices . map ( ( ci ) => ci + 1 ) , // 1-based
59+ originalSupportIndex : supportIndex ,
60+ } ) ;
61+ }
62+ }
63+ } ) ;
64+
65+ if ( segmentsToHighlight . length === 0 ) {
66+ return [ text ] ;
67+ }
68+
69+ // Sort segments by start index, then by end index
70+ segmentsToHighlight . sort ( ( a , b ) => {
71+ if ( a . startIndex !== b . startIndex ) {
72+ return a . startIndex - b . startIndex ;
73+ }
74+ return b . endIndex - a . endIndex ; // Longer segments first
75+ } ) ;
76+
77+ const outputNodes : React . ReactNode [ ] = [ ] ;
78+ let lastIndexProcessed = 0 ;
79+
80+ segmentsToHighlight . forEach ( ( seg , i ) => {
81+ // Add un-highlighted text before this segment
82+ if ( seg . startIndex > lastIndexProcessed ) {
83+ outputNodes . push ( text . substring ( lastIndexProcessed , seg . startIndex ) ) ;
84+ }
85+
86+ // Add the highlighted segment
87+ // Ensure we don't re-highlight an already covered portion if a shorter segment comes later
88+ const currentSegmentText = text . substring ( seg . startIndex , seg . endIndex ) ;
89+ const tooltipText = seg . chunkIndices
90+ . map ( ( ci ) => {
91+ const chunk = chunks [ ci - 1 ] ; // ci is 1-based
92+ return chunk . web ?. title || chunk . web ?. uri || `Source ${ ci } ;
93+ } )
94+ . join ( "; " ) ;
95+
96+ outputNodes . push (
97+ < span
98+ key = { `seg-${ i } }
99+ className = { styles . highlightedSegment }
100+ title = { `Sources: ${ tooltipText } }
101+ data-source-indices = { seg . chunkIndices . join ( "," ) }
102+ >
103+ { currentSegmentText }
104+ < sup className = { styles . sourceSuperscript } >
105+ [{ seg . chunkIndices . join ( "," ) } ]
106+ </ sup >
107+ </ span > ,
108+ ) ;
109+ lastIndexProcessed = Math . max ( lastIndexProcessed , seg . endIndex ) ;
110+ } ) ;
111+
112+ // Add any remaining un-highlighted text
113+ if ( lastIndexProcessed < text . length ) {
114+ outputNodes . push ( text . substring ( lastIndexProcessed ) ) ;
115+ }
116+
117+ return outputNodes ;
118+ } ;
119+
26120/**
27121 * Renders a single chat message bubble, styled based on the message role ('user' or 'model').
28122 * It only renders messages that should be visible in the log (user messages, or model messages
29123 * containing text). Function role messages or model messages consisting only of function calls
30124 * are typically not rendered directly as chat bubbles.
31125 */
32- const ChatMessage : React . FC < ChatMessageProps > = ( { message } ) => {
126+ const ChatMessage : React . FC < ChatMessageProps > = ( {
127+ message,
128+ groundingMetadata,
129+ } ) => {
33130 const text = getMessageText ( message ) ;
34131 const isUser = message . role === "user" ;
35132 const isModel = message . role === "model" ;
@@ -41,20 +138,79 @@ const ChatMessage: React.FC<ChatMessageProps> = ({ message }) => {
41138 // 1. 'function' role messages (these represent execution results, not direct chat).
42139 // 2. 'model' role messages that *only* contain function calls (these are instructions, not display text).
43140 // 3. 'system' role messages (handled separately, not usually in chat history display).
44- const shouldRender = isUser || ( isModel && text . trim ( ) !== "" ) ;
141+ const shouldRender =
142+ isUser ||
143+ ( isModel && text . trim ( ) !== "" ) ;
45144
46145 if ( ! shouldRender ) {
47146 return null ;
48147 }
49148
149+ let messageContentNodes : React . ReactNode [ ] ;
150+ if (
151+ isModel &&
152+ groundingMetadata ?. groundingSupports &&
153+ groundingMetadata ?. groundingChunks
154+ ) {
155+ messageContentNodes = renderTextWithInlineHighlighting (
156+ text ,
157+ groundingMetadata . groundingSupports ,
158+ groundingMetadata . groundingChunks ,
159+ ) ;
160+ } else {
161+ messageContentNodes = [ text ] ;
162+ }
163+
50164 return (
51165 < div
52166 className = { `${ styles . messageContainer } ${ isUser ? styles . user : styles . model } }
53167 >
54168 < div className = { styles . messageBubble } >
55- { /* Use <pre> to preserve whitespace and newlines within the text content.
56- Handles potential multi-line responses correctly. */ }
57- < pre className = { styles . messageText } > { text } </ pre >
169+ < pre className = { styles . messageText } >
170+ { messageContentNodes . map ( ( node , index ) => (
171+ < React . Fragment key = { index } > { node } </ React . Fragment >
172+ ) ) }
173+ </ pre >
174+ { /* Source list rendering for grounded results. This display must comply with the display requirements in the Service Terms. */ }
175+ { isModel &&
176+ groundingMetadata &&
177+ ( groundingMetadata . searchEntryPoint ?. renderedContent ||
178+ ( groundingMetadata . groundingChunks &&
179+ groundingMetadata . groundingChunks . length > 0 ) ? (
180+ < div className = { styles . sourcesSection } >
181+ { groundingMetadata . searchEntryPoint ?. renderedContent && (
182+ < div
183+ className = { styles . searchEntryPoint }
184+ dangerouslySetInnerHTML = { {
185+ __html : groundingMetadata . searchEntryPoint . renderedContent ,
186+ } }
187+ />
188+ ) }
189+ { groundingMetadata . groundingChunks &&
190+ groundingMetadata . groundingChunks . length > 0 && (
191+ < >
192+ < h5 className = { styles . sourcesTitle } > Sources:</ h5 >
193+ < ul className = { styles . sourcesList } >
194+ { groundingMetadata . groundingChunks . map ( ( chunk , index ) => (
195+ < li
196+ key = { index }
197+ className = { styles . sourceItem }
198+ id = { `source-ref-${ index + 1 } }
199+ >
200+ < a
201+ href = { chunk . web ?. uri }
202+ target = "_blank"
203+ rel = "noopener noreferrer"
204+ >
205+ { `[${ index + 1 } ${ chunk . web ?. title || chunk . web ?. uri } }
206+ </ a >
207+ </ li >
208+ ) ) }
209+ </ ul >
210+ </ >
211+ ) }
212+ </ div >
213+ ) : null ) }
58214 </ div >
59215 </ div >
60216 ) ;
0 commit comments