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,38 +37,180 @@ const getMessageText = (message: Content): string => {
2337 . join ( "" ) ;
2438} ;
2539
40+ const renderTextWithInlineHighlighting = (
41+ text : string ,
42+ supports : GroundingSupport [ ] ,
43+ chunks : GroundingChunk [ ] , // Pass chunks for tooltips/links
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+ // Assuming partIndex is 0 for concatenated text.
54+ // If message.parts can have multiple text parts, this needs adjustment.
55+ const segment = support . segment ;
56+ if ( segment . partIndex === undefined || segment . partIndex === 0 ) {
57+ // check if partIndex is specified and is 0 or not specified
58+ segmentsToHighlight . push ( {
59+ startIndex : segment . startIndex ,
60+ endIndex : segment . endIndex , // API's endIndex is typically exclusive
61+ chunkIndices : support . groundingChunkIndices . map ( ( ci ) => ci + 1 ) , // 1-based
62+ originalSupportIndex : supportIndex ,
63+ } ) ;
64+ }
65+ }
66+ } ) ;
67+
68+ if ( segmentsToHighlight . length === 0 ) {
69+ return [ text ] ;
70+ }
71+
72+ // Sort segments by start index, then by end index (longer segments first if they start at same point)
73+ segmentsToHighlight . sort ( ( a , b ) => {
74+ if ( a . startIndex !== b . startIndex ) {
75+ return a . startIndex - b . startIndex ;
76+ }
77+ return b . endIndex - a . endIndex ; // Longer segments first
78+ } ) ;
79+
80+ const outputNodes : React . ReactNode [ ] = [ ] ;
81+ let lastIndexProcessed = 0 ;
82+
83+ segmentsToHighlight . forEach ( ( seg , i ) => {
84+ // Add un-highlighted text before this segment
85+ if ( seg . startIndex > lastIndexProcessed ) {
86+ outputNodes . push ( text . substring ( lastIndexProcessed , seg . startIndex ) ) ;
87+ }
88+
89+ // Add the highlighted segment
90+ // Ensure we don't re-highlight an already covered portion if a shorter segment comes later
91+ const currentSegmentText = text . substring ( seg . startIndex , seg . endIndex ) ;
92+ const tooltipText = seg . chunkIndices
93+ . map ( ( ci ) => {
94+ const chunk = chunks [ ci - 1 ] ; // ci is 1-based
95+ return chunk . web ?. title || chunk . web ?. uri || `Source ${ ci } ;
96+ } )
97+ . join ( "; " ) ;
98+
99+ outputNodes . push (
100+ < span
101+ key = { `seg-${ i } }
102+ className = { styles . highlightedSegment }
103+ title = { `Sources: ${ tooltipText } } // Tooltip for hover
104+ data-source-indices = { seg . chunkIndices . join ( "," ) }
105+ >
106+ { currentSegmentText }
107+ < sup className = { styles . sourceSuperscript } >
108+ [{ seg . chunkIndices . join ( "," ) } ]
109+ </ sup >
110+ </ span > ,
111+ ) ;
112+ lastIndexProcessed = Math . max ( lastIndexProcessed , seg . endIndex ) ;
113+ } ) ;
114+
115+ // Add any remaining un-highlighted text
116+ if ( lastIndexProcessed < text . length ) {
117+ outputNodes . push ( text . substring ( lastIndexProcessed ) ) ;
118+ }
119+
120+ return outputNodes ;
121+ } ;
122+
26123/**
27124 * Renders a single chat message bubble, styled based on the message role ('user' or 'model').
28125 * It only renders messages that should be visible in the log (user messages, or model messages
29126 * containing text). Function role messages or model messages consisting only of function calls
30127 * are typically not rendered directly as chat bubbles.
31128 */
32- const ChatMessage : React . FC < ChatMessageProps > = ( { message } ) => {
129+ const ChatMessage : React . FC < ChatMessageProps > = ( {
130+ message,
131+ groundingMetadata,
132+ } ) => {
33133 const text = getMessageText ( message ) ;
34134 const isUser = message . role === "user" ;
35135 const isModel = message . role === "model" ;
36136
37- // We render:
38- // 1. User messages (even if they only contain images/files, the 'user' role indicates an entry).
39- // 2. Model messages *only if* they contain actual text content.
40- // We *don't* render:
41- // 1. 'function' role messages (these represent execution results, not direct chat).
42- // 2. 'model' role messages that *only* contain function calls (these are instructions, not display text).
43- // 3. 'system' role messages (handled separately, not usually in chat history display).
44- const shouldRender = isUser || ( isModel && text . trim ( ) !== "" ) ;
137+ const shouldRender =
138+ isUser ||
139+ ( isModel && text . trim ( ) !== "" ) ||
140+ ( isModel &&
141+ groundingMetadata &&
142+ ( groundingMetadata . groundingChunks ?. length || 0 > 0 ) ) ;
45143
46144 if ( ! shouldRender ) {
47145 return null ;
48146 }
49147
148+ // Determine content to render
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 */ }
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