adding a streamOverwrite option to the response

OvidijusParsiunas · Nov 21, 2023 · 17d8738 · 17d8738
1 parent 12147ca
commit 17d8738
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 21 deletions.
diff --git a/component/src/types/response.ts b/component/src/types/response.ts
@@ -6,5 +6,6 @@ export interface Response {
   html?: string;
   error?: string;
   role?: string;
+  streamOverwrite?: boolean;
   _sessionId?: string;
 }
diff --git a/component/src/utils/HTTP/customHandler.ts b/component/src/utils/HTTP/customHandler.ts
@@ -53,7 +53,7 @@ export class CustomHandler {
       io.streamHandlers.onClose();
       isHandlerActive = false;
     };
-    const onResponse = (result: {text?: string; error?: string, html?: string}) => {
+    const onResponse = (result: Response) => {
       if (!isHandlerActive) return;
       if (!result || typeof result !== 'object'
         || (typeof result.error !== 'string' && typeof result.html !== 'string' && typeof result.text !== 'string')) {

diff --git a/component/src/utils/HTTP/stream.ts b/component/src/utils/HTTP/stream.ts
@@ -39,12 +39,8 @@ export class Stream {
           const response = JSON.parse(message.data) as unknown as OpenAIConverseResult;
           io.extractResultData?.(response)
             .then((textBody?: DResponse) => {
-              if (textBody?.text === undefined) {
-                // strategy - do not to stop the stream on one message failure to give other messages a change to display
-                console.error(`Response data: ${message.data}`);
-              } else {
-                messages.updatedStreamedMessage(textBody);
-              }
+              // do not to stop the stream on one message failure to give other messages a change to display
+              messages.updatedStreamedMessage(textBody);
             })
             .catch((e) => RequestUtils.displayError(messages, e));
         }

diff --git a/component/src/views/chat/messages/stream/messageStream.ts b/component/src/views/chat/messages/stream/messageStream.ts
@@ -28,47 +28,46 @@ export class MessageStream extends MessageBase {
     ElementUtils.scrollToBottom(this.elementRef); // need to scroll down completely
   }
 
-  public updatedStreamedMessage(response: Response, isIncrement = true) {
-    const content = response.text || response.html;
+  public updatedStreamedMessage(response?: Response) {
+    const content = response?.text || response?.html;
     if (!content) return console.error(ErrorMessages.INVALID_STREAM_RESPONSE);
     const isScrollbarAtBottomOfElement = ElementUtils.isScrollbarAtBottomOfElement(this.elementRef);
     const {bubbleElement} = this._messageElementRefs[this._messageElementRefs.length - 1];
-    // WORK - optimize
     if (content.trim().length > 0) {
       const defaultColor = this.messageStyles?.default;
       bubbleElement.style.color = defaultColor?.ai?.bubble?.color || defaultColor?.shared?.bubble?.color || '';
     }
-    this.updateBasedOnType(content, !!response.text, bubbleElement, isIncrement);
+    this.updateBasedOnType(content, !!response.text, bubbleElement, response.streamOverwrite);
     if (isScrollbarAtBottomOfElement) ElementUtils.scrollToBottom(this.elementRef);
   }
 
-  private updateBasedOnType(content: string, isText: boolean, bubbleElement: HTMLElement, isIncrement: boolean) {
+  private updateBasedOnType(content: string, isText: boolean, bubbleElement: HTMLElement, isOverwrite = false) {
     const expectedType = isText ? 'text' : 'html';
     const func = isText ? this.updateText : this.updateHTML;
     if (this._streamType === '') {
       this._streamType = expectedType;
     } else if (this._streamType !== expectedType) {
       return console.error(ErrorMessages.INVALID_STREAM_MIX_RESPONSE);
     }
-    func.bind(this)(content, bubbleElement, isIncrement);
+    func.bind(this)(content, bubbleElement, isOverwrite);
   }
 
-  private updateText(text: string, bubbleElement: HTMLElement, isIncrement: boolean) {
-    this._streamedContent = isIncrement ? this._streamedContent + text : text;
+  private updateText(text: string, bubbleElement: HTMLElement, isOverwrite: boolean) {
+    this._streamedContent = isOverwrite ? text : this._streamedContent + text;
     this._textElementsToText[this._textElementsToText.length - 1][1] = this._streamedContent;
     bubbleElement.innerHTML = this._remarkable.render(this._streamedContent);
   }
 
-  private updateHTML(html: string, bubbleElement: HTMLElement, isIncrement: boolean) {
-    if (isIncrement) {
+  private updateHTML(html: string, bubbleElement: HTMLElement, isOverwrite: boolean) {
+    if (isOverwrite) {
+      this._streamedContent = html;
+      bubbleElement.innerHTML = html;
+    } else {
       if (this._streamedContent === '') bubbleElement.replaceChildren(); // remove '.'
       const wrapper = document.createElement('span');
       wrapper.innerHTML = html;
       bubbleElement.appendChild(wrapper);
       this._streamedContent = MessageStream.HTML_CONTENT_PLACEHOLDER;
-    } else {
-      this._streamedContent = html;
-      bubbleElement.innerHTML = html;
     }
   }
 

diff --git a/component/src/webModel/webModel.ts b/component/src/webModel/webModel.ts
@@ -170,7 +170,7 @@ export class WebModel extends BaseServiceIO {
 
   private static callbackUpdateResponse(messages: Messages, _: number, msg: string) {
     if (!messages.isStreaming()) messages.addNewStreamedMessage();
-    messages.updatedStreamedMessage({text: msg}, false);
+    messages.updatedStreamedMessage({text: msg, streamOverwrite: true});
   }
 
   private async streamResp(messages: Messages, text: string, chat: WebLLM.ChatInterface) {