Merge pull request #11 from dieharders/bugfix/HBAI-118-2

HBAI-118-2

main/electron.js

@@ -154,12 +154,41 @@ ipcMain.handle('api', async (event, eventName, options) => {
   const modelCard = options?.modelCard
   const id = modelCard?.id
   let dlService = downloaders[id]
-  if (!dlService) {
+  const canWritePath = () => {
+    // Check we can write to path
+    const fs = require('fs')
+    const exists = fs.existsSync(filePath)
+    if (exists) {
+      try {
+        fs.accessSync(filePath, fs.constants.W_OK | fs.constants.R_OK)
+        return true
+      } catch (e) {
+        console.log('[Electron] Error: Cannot access directory')
+        return false
+      }
+    } else {
+      try {
+        fs.mkdirSync(filePath)
+        return true
+      } catch (e) {
+        if (e.code == 'EACCESS') {
+          console.log('[Electron] Error: Cannot create directory, access denied.')
+        } else {
+          console.log(`[Electron] Error: ${e.code}`)
+        }
+        return false
+      }
+    }
+  }
+  // Create downloader instance
+  const createDownloaderInstance = () => {
+    if (dlService) return
     dlService = downloader({ config, modelCard, event, filePath })
     downloaders[id] = dlService
     console.log('[Electron] New downloader created.')
   }
 
+  // Handle api events
   switch (eventName) {
     case 'showConfirmDialog':
       return dialog.showMessageBoxSync(mainWindow, options)
@@ -170,16 +199,27 @@ ipcMain.handle('api', async (event, eventName, options) => {
     case 'getAppPath':
       return app.getAppPath()
     case 'delete_file': {
-      const success = await dlService.onDelete()
-      if (dlService) delete downloaders[id]
-      return success
+      try {
+        createDownloaderInstance()
+        const success = await dlService.onDelete()
+        if (downloaders[id]) delete downloaders[id]
+        return success
+      } catch (err) {
+        console.log('[Electron] Error deleting:', err)
+        return false
+      }
     }
     case 'pause_download':
       return dlService.onPause()
+    // Initiate the previous download
     case 'resume_download':
+      if (!canWritePath()) return
+      createDownloaderInstance()
       return dlService.onStart(true)
     // Start service and return a config
     case 'start_download':
+      if (!canWritePath()) return
+      createDownloaderInstance()
       return dlService.onStart()
     default:
       return

main/utils/downloader.js

@@ -24,6 +24,7 @@ const axios = require('axios')
  * @returns
  */
 const fetchTotalSize = async url => {
+  console.log('[Downloader] Fetching headers...')
   const response = await axios({
     url,
     method: 'HEAD',
@@ -134,6 +135,7 @@ const downloader = payload => {
   let progress = config?.progress ?? 0
   let endChunk = config?.endChunk
   // Other
+  let fileStream
   let hash // object used to create checksum from chunks
   let state = progress > 0 ? EProgressState.Idle : EProgressState.None
   const ipcEvent = payload?.event
@@ -188,7 +190,8 @@ const downloader = payload => {
       downloadUrl,
     )
     const options = startChunk > 0 ? { flags: 'a' } : null
-    const fileStream = fs.createWriteStream(writePath, options)
+    // Open handler
+    fileStream = fs.createWriteStream(writePath, options)
     // Create crypto hash object and update with each chunk.
     // Dont create a chunked hash if we are resuming from cold boot.
     const shouldCreateChunkedHashing =
@@ -334,6 +337,8 @@ const downloader = payload => {
     } catch (err) {
       console.log('[Downloader] Failed writing file to disk', err)
       updateProgressState(EProgressState.Errored)
+      // Close the file
+      fileStream && fileStream.end()
       return false
     }
   }

renderer/components/text-inference-config/TextInferenceConfigMenu.tsx

@@ -47,6 +47,8 @@ const StartEngine = ({ isStarted, setIsStarted, currentTextModelId, ip }: IProps
       // Get installed/stored model configs list and combine
       const storedConfig = getTextModelConfig(currentTextModelId)
       const config = textModels.find(model => model.id === currentTextModelId)
+      if (!storedConfig || !config) throw Error('Cannot find text model config data')
+
       const modelConfig = {
         id: config?.id,
         name: config?.name,
@@ -65,8 +67,6 @@ const StartEngine = ({ isStarted, setIsStarted, currentTextModelId, ip }: IProps
         promptTemplate: config?.promptTemplate || '{{PROMPT}}',
       }
 
-      if (!storedConfig || !config) throw Error('Cannot find text model config data')
-
       const options = { modelConfig }
 
       const response = await fetch(`${ip}/v1/text/start`, {

renderer/components/text-model-browser/ModelBrowser.tsx

@@ -1,16 +1,17 @@
 'use client'
 
-import { Dispatch, SetStateAction, useCallback, useMemo } from 'react'
+import { Dispatch, SetStateAction, useCallback, useEffect, useState } from 'react'
 import ModelCard from './ModelCard'
 import { IModelCard } from '@/models/models'
 import createConfig, { IModelConfig } from './configs'
 import { getTextModelConfig, setUpdateTextModelConfig } from '@/utils/localStorage'
+import { CURRENT_DOWNLOAD_PATH } from '@/components/text-inference-config/TextInferenceConfigMenu'
 
 interface IProps {
   data: Array<IModelCard>
   currentTextModel: string
-  savePath: string
   setCurrentTextModel: Dispatch<SetStateAction<string>>
+  loadTextModelAction: (payload: { modelId: string; pathToModel: string }) => void
 }
 
 // LocalStorage keys
@@ -20,17 +21,28 @@ export const ITEM_CURRENT_MODEL = 'current-text-model'
 /**
  * List of curated text inference models
  */
-const ModelBrowser = ({ data, currentTextModel, savePath, setCurrentTextModel }: IProps) => {
+const ModelBrowser = ({
+  data,
+  currentTextModel,
+  setCurrentTextModel,
+  loadTextModelAction,
+}: IProps) => {
+  const [runOnce, setRunOnce] = useState(false)
+
   // Handlers
   const onSelectTextModel = useCallback(
     (id: string) => {
       console.log('[UI] Set current text model:', id)
-      if (id) {
+      const savePath = localStorage.getItem(CURRENT_DOWNLOAD_PATH)
+      if (id && savePath) {
         setCurrentTextModel(id)
+        // Tell backend which model to load
+        const payload = { modelId: id, pathToModel: savePath }
+        loadTextModelAction(payload)
         localStorage.setItem(ITEM_CURRENT_MODEL, id)
-      }
+      } else console.log('[UI] Error: No id or savePath')
     },
-    [setCurrentTextModel],
+    [loadTextModelAction, setCurrentTextModel],
   )
   const onDownloadComplete = useCallback(() => {
     console.log('[UI] File saved successfully!')
@@ -46,17 +58,18 @@ const ModelBrowser = ({ data, currentTextModel, savePath, setCurrentTextModel }:
   }
 
   // Components
-  const cards = useMemo(() => {
-    return data?.map(item => {
+  const [cards, setCards] = useState<JSX.Element[] | null>(null)
+
+  const createCard = useCallback(
+    (item: IModelCard) => {
       return (
         <ModelCard
           key={item.id}
           modelCard={item}
-          saveToPath={savePath}
           isLoaded={currentTextModel === item.id}
           loadModelConfig={() => {
             try {
-              // Look up the installed model.
+              // Look up the installed model if exists
               return getTextModelConfig(item.id)
             } catch (err) {
               // Error cant load model. `localStorage` possibly undefined.
@@ -68,8 +81,16 @@ const ModelBrowser = ({ data, currentTextModel, savePath, setCurrentTextModel }:
           onDownloadComplete={onDownloadComplete}
         ></ModelCard>
       )
-    })
-  }, [currentTextModel, data, onDownloadComplete, onSelectTextModel, savePath])
+    },
+    [currentTextModel, onDownloadComplete, onSelectTextModel],
+  )
+
+  useEffect(() => {
+    if (runOnce || !data) return
+    const ref = data.map(createCard)
+    setCards(ref)
+    setRunOnce(true)
+  }, [createCard, data, runOnce])
 
   return (
     <div className="z-5 mt-16 flex h-full w-full flex-col justify-center gap-8 rounded-xl border-b border-gray-300 bg-gray-200 p-6 backdrop-blur-sm dark:border-neutral-800 dark:bg-zinc-800/30 lg:mb-0 lg:mt-52 lg:w-10/12 2xl:w-3/6">

renderer/components/text-model-browser/ModelCard.tsx

@@ -3,6 +3,7 @@
 import useDownloader, { EProgressState } from './useDownloader'
 import { IModelCard } from '@/models/models'
 import { EValidationState, IModelConfig } from './configs'
+import { CURRENT_DOWNLOAD_PATH } from '@/components/text-inference-config/TextInferenceConfigMenu'
 import {
   CancelDownloadButton,
   CheckHardware,
@@ -16,7 +17,6 @@ import {
 interface IProps {
   modelCard: IModelCard
   isLoaded: boolean
-  saveToPath: string
   onSelectModel: (modelId: string) => void
   onDownloadComplete: () => void
   loadModelConfig: () => IModelConfig | undefined
@@ -25,7 +25,6 @@ interface IProps {
 
 const ModelCard = ({
   modelCard,
-  saveToPath,
   isLoaded,
   onSelectModel,
   onDownloadComplete,
@@ -34,6 +33,7 @@ const ModelCard = ({
 }: IProps) => {
   // Vars
   const { id, name, description, fileSize, ramSize, licenses, provider } = modelCard
+  const saveToPath = localStorage.getItem(CURRENT_DOWNLOAD_PATH) || ''
 
   // Downloader Hook
   const {

renderer/components/text-model-browser/useDownloader.ts

@@ -70,20 +70,21 @@ const useDownloader = ({ modelCard, saveToPath, loadModelConfig, saveModelConfig
   // Remove config record
   const deleteConfig = useCallback(() => {
     removeTextModelConfig(modelId)
-    console.log(`[Downloader] File ${modelId} removed successfully!`)
+    console.log(`[Downloader] Config ${modelId} removed successfully!`)
   }, [modelId])
   // Delete file
   const deleteFile = useCallback(async () => {
     const success = await window.electron.api('delete_file', apiPayload)
 
     if (success) {
+      console.log('[Downloader] Model file removed successfully!')
       deleteConfig()
       // Set the state
       setModelConfig(undefined)
       return true
     }
 
-    console.log('[Downloader] File removal failed!', success)
+    console.log('[Downloader] File removal failed!')
     return false
   }, [apiPayload, deleteConfig])
   /**
@@ -160,16 +161,18 @@ const useDownloader = ({ modelCard, saveToPath, loadModelConfig, saveModelConfig
       }
     }
 
-    window.electron.message.on(handler)
+    window.electron?.message?.on(handler)
 
     return () => {
-      window.electron.message.off(handler)
+      window.electron?.message?.off(handler)
     }
   }, [modelId])
 
   // Load and update model config from storage whenever progress state changes
   useEffect(() => {
     const c = loadModelConfig()
+    if (!c) return
+
     const progress = c?.progress ?? 0
     setModelConfig(c)
     // We shouldnt have to do this here but the backend has no access to initial `config` state.

renderer/models/models.ts

@@ -251,16 +251,16 @@ const models: Array<IModelCard> = [
     \n
     ASSISTANT:`,
   },
-  {
-    id: 'example-cat-anim',
-    name: 'Example Cute Cat Animation',
-    provider: 'giphy',
-    licenses: [LicenseType.Academic, LicenseType.Commercial, LicenseType.Other],
-    description: 'This is a test file (gif) for testing download behavior.',
-    fileSize: 0.03, // 3060203 bytes
-    fileName: 'cute-cat-anim.gif',
-    downloadUrl: 'https://media.giphy.com/media/04uUJdw2DliDjsNOZV/giphy.gif',
-  },
+  // {
+  //   id: 'example-cat-anim',
+  //   name: 'Example Cute Cat Animation',
+  //   provider: 'giphy',
+  //   licenses: [LicenseType.Academic, LicenseType.Commercial, LicenseType.Other],
+  //   description: 'This is a test file (gif) for testing download behavior.',
+  //   fileSize: 0.03, // 3060203 bytes
+  //   fileName: 'cute-cat-anim.gif',
+  //   downloadUrl: 'https://media.giphy.com/media/04uUJdw2DliDjsNOZV/giphy.gif',
+  // },
 ]
 
 export default models

renderer/pages/home.tsx

@@ -1,6 +1,6 @@
 'use client'
 
-import { useEffect, useState } from 'react'
+import { useCallback, useEffect, useState } from 'react'
 import constants from '@/shared/constants.json'
 import textModels from '@/models/models'
 import AppsBrowser from '@/components/AppsBrowser'
@@ -25,6 +25,23 @@ export default function Home() {
   const [isStarted, setIsStarted] = useState(false)
   const [savePath, setSavePath] = useState<string>('')
   const [currentTextModel, setCurrentTextModel] = useState<string>('')
+
+  const loadTextModelAction = useCallback(
+    (payload: any) => {
+      fetch(`${HOMEBREW_BASE_PATH}/v1/text/load`, {
+        method: 'POST',
+        headers: {
+          Accept: 'application/json',
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify(payload),
+      }).catch(err => {
+        console.log('[UI] Error loading model:', err)
+      })
+    },
+    [HOMEBREW_BASE_PATH],
+  )
+
   // Company credits (built by)
   const renderCredits = () => {
     return (
@@ -79,9 +96,13 @@ export default function Home() {
     storedPath && setSavePath(storedPath)
     const currModel = localStorage.getItem(ITEM_CURRENT_MODEL)
     currModel && setCurrentTextModel(currModel)
+    // Load stored model from storage if found
+    if (storedPath && currModel)
+      loadTextModelAction({ modelId: currModel, pathToModel: storedPath })
 
+    // Set defaults if none found
     if (!storedPath) saveDefaultPath()
-  }, [])
+  }, [loadTextModelAction])
 
   return (
     <div className="xs:p-0 mb-32 flex min-h-screen flex-col items-center justify-between overflow-x-hidden lg:mb-0 lg:p-24">
@@ -122,8 +143,8 @@ export default function Home() {
         <TextModelBrowserMenu
           data={textModels}
           currentTextModel={currentTextModel}
-          savePath={savePath}
           setCurrentTextModel={setCurrentTextModel}
+          loadTextModelAction={loadTextModelAction}
         />
       )}
     </div>

renderer/utils/localStorage.ts

@@ -12,7 +12,7 @@ const getTextModelConfig = (id: string) => {
   const config = modelConfigs.find((item: IModelConfig) => item.id === id)
 
   if (!config) {
-    console.log('[localStorage] Cannot find text model config data')
+    console.log('[localStorage] Cannot find text model config data for', id)
     return undefined
   }