packages/ai-assistant/src/main.ts

   1 import type { LlmFile } from '@proton/llm/lib/downloader';
   2 import { downloadModel } from '@proton/llm/lib/downloader';
   3 import { isAssistantPostMessage, postMessageIframeToParent } from '@proton/llm/lib/helpers';
   4 import type { AssistantConfig } from '@proton/llm/lib/types';
   5 import { AssistantEvent } from '@proton/llm/lib/types';
   6
   7 window.addEventListener('load', async () => {
   8     let abortController: AbortController | undefined; // defined if status === 'loading'
   9     let parentURL: string;
  10
  11     const handleDownloadModel = async (
  12         assistantConfig: AssistantConfig,
  13         modelVariant: string,
  14         filesToIgnore: LlmFile[]
  15     ) => {
  16         try {
  17             if (!abortController) {
  18                 abortController = new AbortController();
  19                 await downloadModel(modelVariant, assistantConfig, abortController, filesToIgnore, parentURL);
  20                 abortController = undefined;
  21             }
  22         } catch (error: any) {
  23             postMessageIframeToParent(
  24                 {
  25                     type: AssistantEvent.DOWNLOAD_ERROR,
  26                     payload: {
  27                         error,
  28                     },
  29                 },
  30                 parentURL
  31             );
  32         } finally {
  33             abortController = undefined;
  34         }
  35     };
  36
  37     // Request to cancel an ongoing download.
  38     //
  39     // This will return immediately, but the cancellation will be complete when the startDownload() promise is
  40     // resolved, which shouldn't take long.
  41     const cancelDownload = () => {
  42         if (abortController) {
  43             abortController.abort();
  44             abortController = undefined;
  45         }
  46     };
  47
  48     const handleReceivedEvent = async (event: MessageEvent) => {
  49         if (!isAssistantPostMessage(event)) {
  50             return;
  51         }
  52
  53         switch (event.data.type) {
  54             case AssistantEvent.START_DOWNLOAD:
  55                 {
  56                     // If we passed all the validations and receive a start download event,
  57                     // we assume the parent url is the origin of the received event
  58                     parentURL = event.origin;
  59
  60                     postMessageIframeToParent(
  61                         {
  62                             type: AssistantEvent.IFRAME_READY,
  63                         },
  64                         parentURL
  65                     );
  66
  67                     const { config, modelVariant, filesToIgnore } = event.data.payload;
  68                     await handleDownloadModel(config, modelVariant, filesToIgnore);
  69                 }
  70                 break;
  71             case AssistantEvent.PAUSE_DOWNLOAD:
  72                 {
  73                     cancelDownload();
  74                 }
  75                 break;
  76         }
  77     };
  78
  79     window.addEventListener('message', handleReceivedEvent);
  80 });