1 files changed, 24 insertions, 46 deletions

diff --git a/src/client/views/nodes/ChatBox/AnswerParser.ts b/src/client/views/nodes/ChatBox/AnswerParser.ts
index 4b6c817fd..dd7ec3499 100644
--- a/src/client/views/nodes/ChatBox/AnswerParser.ts
+++ b/src/client/views/nodes/ChatBox/AnswerParser.ts
@@ -17,71 +17,49 @@ export class AnswerParser {
         }
 
         const rawTextContent = answerMatch[1].trim();
-        let textContent: AssistantMessage['content'] = [];
+        let content: AssistantMessage['content'] = [];
         let citations: Citation[] = [];
         let contentIndex = 0;
 
         // Parse citations
         let citationMatch;
+        const citationMap = new Map<string, string>();
         while ((citationMatch = citationRegex.exec(rawTextContent)) !== null) {
             const [_, index, chunk_id, type, direct_text] = citationMatch;
+            const citation_id = uuid();
+            citationMap.set(index, citation_id);
             citations.push({
                 direct_text: direct_text.trim(),
                 type: getChunkType(type),
                 chunk_id,
-                citation_id: uuid(),
+                citation_id,
             });
         }
 
-        // Parse text content (normal and grounded)
-        let lastIndex = 0;
-        let matches = [];
+        // Parse grounded text content
+        const parseGroundedText = (text: string): AssistantMessage['content'] => {
+            const result: AssistantMessage['content'] = [];
+            let lastIndex = 0;
+            let match;
 
-        // Find all grounded text matches
-        let groundedTextMatch;
-        while ((groundedTextMatch = groundedTextRegex.exec(rawTextContent)) !== null) {
-            matches.push({
-                type: 'grounded',
-                index: groundedTextMatch.index,
-                length: groundedTextMatch[0].length,
-                citationIndexes: groundedTextMatch[1],
-                text: groundedTextMatch[2],
-            });
-        }
-
-        // Sort matches by their index in the original text
-        matches.sort((a, b) => a.index - b.index);
-
-        // Process normal and grounded text in order
-        for (let i = 0; i <= matches.length; i++) {
-            const currentMatch = matches[i];
-            const nextMatchIndex = currentMatch ? currentMatch.index : rawTextContent.length;
-
-            // Add normal text before the current grounded text (or end of content)
-            if (nextMatchIndex > lastIndex) {
-                const normalText = rawTextContent.slice(lastIndex, nextMatchIndex).trim();
-                if (normalText) {
-                    textContent.push({
-                        index: contentIndex++,
-                        type: TEXT_TYPE.NORMAL,
-                        text: normalText,
-                        citation_ids: null,
-                    });
-                }
-            }
+            while ((match = groundedTextRegex.exec(text)) !== null) {
+                const [fullMatch, citationIndex, groundedText] = match;
+                const citation_ids = citationIndex.split(',').map(index => citationMap.get(index) || '');
 
-            // Add grounded text if there's a match
-            if (currentMatch) {
-                const citationIds = currentMatch.citationIndexes.split(',').map(index => citations[parseInt(index) - 1].citation_id);
-                textContent.push({
+                result.push({
                     index: contentIndex++,
                     type: TEXT_TYPE.GROUNDED,
-                    text: currentMatch.text.trim(),
-                    citation_ids: citationIds,
+                    text: groundedText.trim(),
+                    citation_ids,
                 });
-                lastIndex = currentMatch.index + currentMatch.length;
+
+                lastIndex = match.index + fullMatch.length;
             }
-        }
+
+            return result;
+        };
+
+        content = parseGroundedText(rawTextContent);
 
         let followUpQuestions: string[] = [];
         if (followUpQuestionsMatch) {
@@ -94,7 +72,7 @@ export class AnswerParser {
 
         const assistantResponse: AssistantMessage = {
             role: ASSISTANT_ROLE.ASSISTANT,
-            content: textContent,
+            content,
             follow_up_questions: followUpQuestions,
             citations,
         };