1.4.3 - adding a new chatStyle property and refactoring streaming

component/package.json

@@ -1,6 +1,6 @@
 {
   "name": "deep-chat",
-  "version": "1.4.2",
+  "version": "1.4.3",
   "description": "Customizable chat component for AI APIs",
   "main": "./dist/deepChat.js",
   "module": "./dist/deepChat.js",

component/src/deepChat.ts

@@ -19,6 +19,7 @@ import {RequestBodyLimits} from './types/chatLimits';
 import {Property} from './utils/decorators/property';
 import {FireEvents} from './utils/events/fireEvents';
 import {ValidateInput} from './types/validateInput';
+import {WebModel} from './types/webModel/webModel';
 import {DropupStyles} from './types/dropupStyles';
 import {HTMLClassUtilities} from './types/html';
 import {ChatView} from './views/chat/chatView';
@@ -159,7 +160,7 @@ export class DeepChat extends InternalHTML {
   @Property('object')
   demo?: Demo;
 
-  _webModel = false;
+  _webModel: WebModel = false;
 
   _hasBeenRendered = false;
 

component/src/services/serviceIO.ts

@@ -3,6 +3,7 @@ import {IWebsocketHandler} from '../utils/HTTP/customHandler';
 import {Messages} from '../views/chat/messages/messages';
 import {InterfacesUnion} from '../types/utilityTypes';
 import {FetchFunc} from '../utils/HTTP/requestUtils';
+import {ResponseI} from '../types/responseInternal';
 import {FILE_TYPES} from '../types/fileTypes';
 import {Response} from '../types/response';
 import {Request} from '../types/request';
@@ -101,8 +102,12 @@ export interface ServiceIO {
 
   deepChat: DeepChat; // this is used for interceptors as the user may pass them much later after component is initiated
 
+  addMessage?(data: ResponseI): void;
+
   isDirectConnection(): boolean;
 
+  isWebModel(): boolean;
+
   isSubmitProgrammaticallyDisabled?: boolean;
 
   sessionId?: string;

component/src/services/utils/baseServiceIO.ts

@@ -3,6 +3,7 @@ import {MessageContentI} from '../../types/messagesInternal';
 import {Messages} from '../../views/chat/messages/messages';
 import {HTTPRequest} from '../../utils/HTTP/HTTPRequest';
 import {ValidateInput} from '../../types/validateInput';
+import {ResponseI} from '../../types/responseInternal';
 import {MessageLimitUtils} from './messageLimitUtils';
 import {Websocket} from '../../utils/HTTP/websocket';
 import {Legacy} from '../../utils/legacy/legacy';
@@ -34,6 +35,7 @@ export class BaseServiceIO implements ServiceIO {
   recordAudio?: MicrophoneFilesServiceConfig;
   totalMessagesMaxCharLength?: number;
   maxMessages?: number;
+  addMessage?: (data: ResponseI) => void;
   demo?: DemoT;
   // these are placeholders that are later populated in submitButton.ts
   completionsHandlers: CompletionsHandlers = {} as CompletionsHandlers;
@@ -142,4 +144,8 @@ export class BaseServiceIO implements ServiceIO {
   public isDirectConnection() {
     return false;
   }
+
+  public isWebModel() {
+    return false;
+  }
 }

component/src/types/responseInternal.ts

@@ -0,0 +1,3 @@
+import {Response} from './response';
+
+export type ResponseI = Response & {sendUpdate?: boolean};

component/src/types/webLLM/webLLMChatModule.ts → component/src/types/webModel/webLLM/webLLMChatModule.ts

@@ -12,24 +12,9 @@ export declare class ChatModule implements ChatInterface {
   runtimeStatsText(): Promise<string>;
   resetChat(): Promise<void>;
   unload(): Promise<void>;
-  /**
-   * @returns Whether the generation stopped.
-   */
   stopped(): boolean;
-  /**
-   * Get the current generated response.
-   *
-   * @returns The current output message.
-   */
   getMessage(): string;
-  /**
-   * Run a prefill step with a given input.
-   * @param input The input prompt.
-   */
   prefill(input: string): Promise<void>;
-  /**
-   * Run a decode step to decode the next token.
-   */
   decode(): Promise<void>;
   private readonly getPipeline;
   private readonly asyncLoadTokenizer;

component/src/types/webModel/webModel.ts

@@ -0,0 +1,13 @@
+export interface WebModelConfig {
+  load?: {
+    onInit?: boolean;
+    onMessage?: boolean;
+  };
+  initialMessage?: {
+    displayed?: boolean;
+    html?: string;
+    downloadClass?: string;
+  };
+}
+
+export type WebModel = boolean | WebModelConfig;

component/src/utils/HTTP/customHandler.ts

@@ -41,10 +41,9 @@ export class CustomHandler {
   public static stream(io: ServiceIO, body: RequestDetails['body'], messages: Messages) {
     let isHandlerActive = true;
     let isOpen = false;
-    let textElement: HTMLElement | null = null;
     const onOpen = () => {
       if (isOpen || !isHandlerActive) return;
-      textElement = messages.addNewStreamedMessage();
+      messages.addNewStreamedMessage();
       io.streamHandlers.onOpen();
       isOpen = true;
     };
@@ -64,8 +63,8 @@ export class CustomHandler {
         io.streamHandlers.onClose();
         messages.addNewErrorMessage('service', result.error);
         isHandlerActive = false;
-      } else if (result.text && textElement) {
-        messages.updateStreamedMessage(result.text, textElement);
+      } else if (result.text) {
+        messages.updateStreamedMessage(result.text);
       }
     };
     io.streamHandlers.abortStream.abort = () => {

component/src/views/chat/input/buttons/submit/submitButton.ts

@@ -162,7 +162,7 @@ export class SubmitButton extends InputButton<Styles> {
     if (await this._validationHandler?.(isProgrammatic) === false) return;
     this.changeToLoadingIcon();
     await this.addNewMessages(userText, uploadedFilesData);
-    this._messages.addLoadingMessage();
+    if (!this._serviceIO.isWebModel()) this._messages.addLoadingMessage();
     if (!isProgrammatic) TextInputEl.clear(this._inputElementRef); // when uploading a file and placeholder text present
 
     const requestContents = {text: userText === '' ? undefined : userText, files: fileData};

component/src/views/chat/messages/messages.ts

@@ -7,6 +7,7 @@ import {HTMLDeepChatElements} from './html/htmlDeepChatElements';
 import {MessageContentI} from '../../../types/messagesInternal';
 import {RemarkableConfig} from './remarkable/remarkableConfig';
 import {FireEvents} from '../../../utils/events/fireEvents';
+import {ResponseI} from '../../../types/responseInternal';
 import {HTMLClassUtilities} from '../../../types/html';
 import {Demo, DemoResponse} from '../../../types/demo';
 import {MessageStyleUtils} from './messageStyleUtils';
@@ -86,6 +87,7 @@ export class Messages {
     deepChat.clearMessages = this.clearMessages.bind(this, serviceIO);
     deepChat.refreshMessages = this.refreshTextMessages.bind(this);
     deepChat.scrollToBottom = this.scrollToBottom.bind(this);
+    serviceIO.addMessage = this.addNewMessage.bind(this);
     if (demo) this.prepareDemo(demo);
     if (deepChat.textToSpeech) {
       TextToSpeech.processConfig(deepChat.textToSpeech, (processedConfig) => {
@@ -237,7 +239,7 @@ export class Messages {
   }
 
   // this should not be activated by streamed messages
-  public addNewMessage(data: Response, isInitial = false) {
+  public addNewMessage(data: ResponseI, isInitial = false) {
     let isNewMessage = true;
     const message = Messages.processMessageContent(data);
     if (message.text !== undefined && data.text !== null) {
@@ -254,12 +256,12 @@ export class Messages {
       if (HTMLDeepChatElements.isElementTemporary(elements)) delete message.html;
       isNewMessage = !!elements;
     }
-    this.updateStateOnMessage(message, isNewMessage, isInitial);
+    this.updateStateOnMessage(message, isNewMessage, data.sendUpdate, isInitial);
   }
 
-  private updateStateOnMessage(messageContent: MessageContentI, isNewMessage: boolean, isInitial = false) {
+  private updateStateOnMessage(messageContent: MessageContentI, isNewMessage: boolean, update = true, initial = false) {
     if (isNewMessage) this.messages.push(messageContent);
-    this.sendClientUpdate(messageContent, isInitial);
+    if (update) this.sendClientUpdate(messageContent, initial);
   }
 
   private sendClientUpdate(message: MessageContentI, isInitial = false) {

component/src/webModel/webModel.ts

@@ -1,7 +1,9 @@
+import {WebModel as WebModelT} from '../types/webModel/webModel';
+import {MessageUtils} from '../views/chat/messages/messageUtils';
 import {BaseServiceIO} from '../services/utils/baseServiceIO';
 import {MessageContentI} from '../types/messagesInternal';
+import * as WebLLM from '../types/webModel/webLLM/webLLM';
 import {Messages} from '../views/chat/messages/messages';
-import * as WebLLM from '../types/webLLM/webLLM';
 // import * as WebLLM2 from '@mlc-ai/web-llm';
 import config from './webModelConfig';
 import {DeepChat} from '../deepChat';
@@ -23,8 +25,11 @@ declare global {
 export class WebModel extends BaseServiceIO {
   private readonly chat?: WebLLM.ChatInterface;
   // private readonly localChat: WebLLM.ChatInterface;
-  private readonly selectedModel: string = 'Llama-2-7b-chat-hf-q4f32_1';
-  private chatLoaded = false;
+  private static readonly DOWNLOAD_BUTTON_CLASS = 'deep-chat-download-button';
+  private readonly _selectedModel: string = 'Llama-2-7b-chat-hf-q4f32_1';
+  private _isModelLoaded = false;
+  private _isModelLoading = false;
+  private _loadOnFirstMessage = false;
 
   constructor(deepChat: DeepChat) {
     super(deepChat);
@@ -37,66 +42,131 @@ export class WebModel extends BaseServiceIO {
       ? new webLLM.ChatWorkerClient(new Worker(new URL('./worker.ts', import.meta.url), {type: 'module'}))
       : new webLLM.ChatModule();
     // this.localChat = new window.webLLM.ChatRestModule();
-    setTimeout(() => {
-      this.asyncInitChat(chat);
-    }, 2000);
     this.chat = chat;
+    this.canSendMessage = this.canSubmit.bind(this);
+    setTimeout(() => {
+      // in timeout for this.addMessage to be set
+      const isMessageSet = this.addInitialMessage(chat, deepChat);
+      this.configureLoad(chat, isMessageSet, deepChat._webModel);
+    });
+  }
+
+  private addInitialMessage(chat: WebLLM.ChatInterface, deepChat: DeepChat) {
+    const {_webModel} = deepChat;
+    const initialMessage = typeof _webModel === 'object' ? _webModel.initialMessage : undefined;
+    if (!_webModel || initialMessage?.displayed === false) return false;
+    const downloadClass = initialMessage?.downloadClass || WebModel.DOWNLOAD_BUTTON_CLASS;
+    const text = `
+      Download a web model that will run entirely on your browser.
+      <br/> <button style="margin-top: 10px; margin-bottom: 5px; margin-left: 1px"
+        class="${downloadClass} deep-chat-button">Download</button>`;
+    const html = initialMessage?.html || `<div class="deep-chat-update-message">${text}</div>`;
+    this.addMessage?.({role: MessageUtils.AI_ROLE, html, sendUpdate: false});
+    const button = deepChat.shadowRoot?.children[0]?.getElementsByClassName(downloadClass)[0] as HTMLButtonElement;
+    if (button) button.onclick = this.loadModel.bind(this, chat);
+    return true;
   }
 
-  private async asyncInitChat(chat: WebLLM.ChatInterface) {
+  private canSubmit(text?: string) {
+    if (!text?.trim() || this._isModelLoading) return false;
+    if (this._loadOnFirstMessage) return true;
+    return !!this._isModelLoaded;
+  }
+
+  private async configureLoad(chat: WebLLM.ChatInterface, isMessageSet: boolean, webModel?: WebModelT) {
+    if (webModel && typeof webModel !== 'boolean' && webModel.load) {
+      if (webModel.load.onInit) {
+        this.loadModel(chat);
+        return;
+      }
+      if (webModel.load.onMessage) {
+        this._loadOnFirstMessage = true;
+        return;
+      }
+    }
+    if (!isMessageSet) this.loadModel(chat);
+  }
+
+  // prettier-ignore
+  private async loadModel(chat: WebLLM.ChatInterface) {
     // const hasIt = await window.webLLM.hasModelInCache(this.selectedModel, config);
     // console.log(hasIt);
-    if (this.chatLoaded) return;
+    if (this._isModelLoaded) return;
+    this._isModelLoading = true;
     const initProgressCallback = (report: WebLLM.InitProgressReport) => {
-      console.log(report);
-      console.log('initializing');
+      this.addMessage?.({html: `<div class="deep-chat-update-message">${report.text}</div>`, sendUpdate: false});
     };
     chat.setInitProgressCallback(initProgressCallback);
 
     try {
-      await chat.reload(this.selectedModel, {conv_config: {system: 'keep responses to one sentence'}}, config);
+      await chat.reload(this._selectedModel, {conv_config: {system: 'keep responses to one sentence'}}, config);
     } catch (err) {
       console.error('initialzitaion error');
       console.log(err);
       this.unloadChat();
       return;
     }
-    console.log('ready');
-    this.chatLoaded = true;
+    this.addMessage?.({html: `<div class="deep-chat-update-message">Model loaded</div>`, sendUpdate: false});
+    this._isModelLoaded = true;
+    this._isModelLoading = false;
+  }
+
+  private async immediateResp(messages: Messages, text: string, chat: WebLLM.ChatInterface) {
+    const output = await chat.generate(text, undefined, 0); // anything but 1 will not stream
+    messages.addNewMessage({text: output});
+    this.completionsHandlers.onFinish();
   }
 
   private static callbackUpdateResponse(messages: Messages, _: number, msg: string) {
     if (!messages.isStreamingText()) messages.addNewStreamedMessage();
     messages.updateStreamedMessage(msg, false);
   }
 
-  override async callServiceAPI(messages: Messages, pMessages: MessageContentI[], __?: File[]) {
-    if (!this.chat || !this.chatLoaded) return;
-    const message = pMessages[pMessages.length - 1].text as string;
+  private async streamResp(messages: Messages, text: string, chat: WebLLM.ChatInterface) {
+    this.streamHandlers.abortStream.abort = () => {
+      chat.interruptGenerate();
+    };
+    this.streamHandlers.onOpen();
+    await chat.generate(text, WebModel.callbackUpdateResponse.bind(this, messages));
+    if (!messages.isStreamingText()) messages.addNewStreamedMessage(); // needed when early abort clicked
+    messages.finaliseStreamedMessage();
+    this.streamHandlers.onClose();
+  }
+
+  private async generateResp(messages: Messages, pMessages: MessageContentI[], chat: WebLLM.ChatInterface) {
+    const text = pMessages[pMessages.length - 1].text as string;
     try {
       if (this.deepChat.stream) {
-        this.streamHandlers.abortStream.abort = () => {
-          this.chat?.interruptGenerate();
-        };
-        this.streamHandlers.onOpen();
-        await this.chat.generate(message, WebModel.callbackUpdateResponse.bind(this, messages));
-        if (!messages.isStreamingText()) messages.addNewStreamedMessage(); // needed when early abort clicked
-        messages.finaliseStreamedMessage();
-        this.streamHandlers.onClose();
+        this.streamResp(messages, text, chat);
       } else {
-        const output = await this.chat.generate(message, undefined, 0); // anything but 1 will not stream
-        messages.addNewMessage({text: output});
-        this.completionsHandlers.onFinish();
+        this.immediateResp(messages, text, chat);
       }
     } catch (err) {
       console.error('error');
       await this.unloadChat();
     }
   }
 
+  override async callServiceAPI(messages: Messages, pMessages: MessageContentI[], __?: File[]) {
+    if (!this.chat || this._isModelLoading) return;
+    if (!this._isModelLoaded) {
+      if (this._loadOnFirstMessage) {
+        await this.loadModel(this.chat);
+      } else {
+        return;
+      }
+    }
+    messages.addLoadingMessage();
+    this.generateResp(messages, pMessages, this.chat);
+  }
+
   private async unloadChat() {
     if (!this.chat) return;
     await this.chat.unload();
-    this.chatLoaded = false;
+    this._isModelLoaded = false;
+  }
+
+  override isWebModel() {
+    return true;
   }
 }