1 files changed, 12 insertions, 7 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 49c9b3292..9e604073d 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -11,7 +11,7 @@ import { ViewBoxAnnotatableComponent } from '../../DocComponent';
 import { FieldView, FieldViewProps } from '../FieldView';
 import './ChatBox.scss';
 import MessageComponentBox from './MessageComponent';
-import { ASSISTANT_ROLE, AssistantMessage, AI_Document, Citation, CHUNK_TYPE, Chunk, getChunkType } from './types';
+import { ASSISTANT_ROLE, AssistantMessage, AI_Document, Citation, CHUNK_TYPE, Chunk, getChunkType, TEXT_TYPE } from './types';
 import { Vectorstore } from './vectorstore/VectorstoreUpload';
 import { Agent } from './Agent';
 import dotenv from 'dotenv';
@@ -58,7 +58,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory);
 
         reaction(
-            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, text_content: msg.text_content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
+            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, content: msg.content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
             serializableHistory => {
                 this.dataDoc.data = JSON.stringify(serializableHistory);
             }
@@ -101,7 +101,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             try {
                 textInput.value = '';
                 runInAction(() => {
-                    this.history.push({ role: ASSISTANT_ROLE.USER, text_content: trimmedText });
+                    this.history.push({ role: ASSISTANT_ROLE.USER, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: trimmedText, citation_ids: null }] });
                     this.isLoading = true;
                 });
 
@@ -113,7 +113,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             } catch (err) {
                 console.error('Error:', err);
                 runInAction(() => {
-                    this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, text_content: 'Sorry, I encountered an error while processing your request.' });
+                    this.history.push({ role: ASSISTANT_ROLE.USER, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }] });
                 });
             } finally {
                 runInAction(() => {
@@ -212,7 +212,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     this.history.push(
                         ...storedHistory.map((msg: AssistantMessage) => ({
                             role: msg.role,
-                            text_content: msg.text_content,
+                            content: msg.content,
                             follow_up_questions: msg.follow_up_questions,
                             citations: msg.citations,
                         }))
@@ -222,7 +222,12 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 console.error('Failed to parse history from dataDoc:', e);
             }
         } else {
-            this.history = [{ role: ASSISTANT_ROLE.ASSISTANT, text_content: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.' }];
+            runInAction(() => {
+                this.history.push({
+                    role: ASSISTANT_ROLE.USER,
+                    content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.', citation_ids: null }],
+                });
+            });
         }
         reaction(
             () => {
@@ -276,7 +281,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     get formattedHistory(): string {
         let history = '<chat_history>\n';
         for (const message of this.history) {
-            history += `<${message.role}>${message.text_content}</${message.role}>\n`;
+            history += `<${message.role}>${message.content.map(content => content.text).join(' ')}</${message.role}>\n`;
         }
         history += '</chat_history>';
         return history;