Merge pull request #11 from projectatomic/feat/download-model

feat: downloadModel method

packages/backend/src/ai.json

@@ -13,28 +13,13 @@
       "readme": "# Locallm\n\nThis repo contains artifacts that can be used to build and run LLM (Large Language Model) services locally on your Mac using podman. These containerized LLM services can be used to help developers quickly prototype new LLM based applications, without the need for relying on any other externally hosted services. Since they are already containerized, it also helps developers move from their prototype to production quicker.     \n\n## Current Locallm Services: \n\n* [Chatbot](#chatbot)\n* [Text Summarization](#text-summarization)\n* [Fine-tuning](#fine-tuning)\n\n### Chatbot\n\nA simple chatbot using the gradio UI. Learn how to build and run this model service here: [Chatbot](/chatbot/).\n\n### Text Summarization\n\nAn LLM app that can summarize arbitrarily long text inputs. Learn how to build and run this model service here: [Text Summarization](/summarizer/).\n\n### Fine Tuning \n\nThis application allows a user to select a model and a data set they'd like to fine-tune that model on. Once the application finishes, it outputs a new fine-tuned model for the user to apply to other LLM services. Learn how to build and run this model training job here: [Fine-tuning](/finetune/).\n\n## Architecture\n![](https://raw.githubusercontent.com/MichaelClifford/locallm/main/assets/arch.jpg)\n\nThe diagram above indicates the general architecture for each of the individual model services contained in this repo. The core code available here is the \"LLM Task Service\" and the \"API Server\", bundled together under `model_services`. With an appropriately chosen model downloaded onto your host,`model_services/builds` contains the Containerfiles required to build an ARM or an x86 (with CUDA) image depending on your need. These model services are intended to be light-weight and run with smaller hardware footprints (given the Locallm name), but they can be run on any hardware that supports containers and scaled up if needed.\n\nWe also provide demo \"AI Applications\" under `ai_applications` for each model service to provide an example of how a developers could interact with the model service for their own needs. ",
       "models": [
         {
-          "id": "stable-diffusion-xl-base-1.0",
-          "name": "stable diffusion xl base 1.0",
+          "id": "llama-2-7b-chat.Q5_K_S",
+          "name": "Llama-2-7B-Chat-GGUF",
           "hw": "CPU",
           "registry": "Hugging Face",
           "popularity": 3,
-          "license": "openrail++"
-        },
-        {
-          "id": "albedobase-xl-1.3",
-          "name": "AlbedoBase XL 1.3",
-          "hw": "CPU",
-          "registry": "Civital",
-          "popularity": 3,
-          "license": "openrail++"
-        },
-        {
-          "id": "sdxl-turbo",
-          "name": "SDXL Turbo",
-          "hw": "CPU",
-          "registry": "Hugging Face",
-          "popularity": 3,
-          "license": "sai-c-community"
+          "license": "?",
+          "url": "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q5_K_S.gguf"
         }
       ]
     }

packages/backend/src/managers/applicationManager.ts

@@ -2,22 +2,31 @@ import { Recipe } from '@shared/models/IRecipe';
 import { arch } from 'node:os';
 import { GitManager } from './gitManager';
 import os from 'os';
-import path from 'path';
 import fs from 'fs';
-import { containerEngine, provider } from '@podman-desktop/api';
+import * as https from 'node:https';
+import * as path from 'node:path';
+import { containerEngine, ExtensionContext, provider } from '@podman-desktop/api';
 import { RecipeStatusRegistry } from '../registries/RecipeStatusRegistry';
 import { AIConfig, parseYaml } from '../models/AIConfig';
 import { Task } from '@shared/models/ITask';
 import { TaskUtils } from '../utils/taskUtils';
 import { getParentDirectory } from '../utils/pathUtils';
+import { a } from 'vitest/dist/suite-dF4WyktM';
+import type { LocalModelInfo } from '@shared/models/ILocalModelInfo';
 
 // TODO: Need to be configured
 export const AI_STUDIO_FOLDER = path.join('podman-desktop', 'ai-studio');
 export const CONFIG_FILENAME = "ai-studio.yaml";
+
+interface DownloadModelResult {
+  result: 'ok' | 'failed';
+  error?: string;
+}
+
 export class ApplicationManager {
   private readonly homeDirectory: string; // todo: make configurable
 
-  constructor(private git: GitManager, private recipeStatusRegistry: RecipeStatusRegistry,) {
+  constructor(private git: GitManager, private recipeStatusRegistry: RecipeStatusRegistry, private extensionContext: ExtensionContext) {
     this.homeDirectory = os.homedir();
   }
 
@@ -109,13 +118,25 @@ export class ApplicationManager {
     const filteredContainers = aiConfig.application.containers
       .filter((container) => container.arch === undefined || container.arch === arch())
 
+    // Download first model available (if exist)
+    if(recipe.models && recipe.models.length > 0) {
+      const model = recipe.models[0];
+      taskUtil.setTask({
+        id: model.id,
+        state: 'loading',
+        name: `Downloading model ${model.name}`,
+      });
+
+      await this.downloadModelMain(model.id, model.url, taskUtil)
+    }
+
     filteredContainers.forEach((container) => {
-     taskUtil.setTask({
-       id: container.name,
-       state: 'loading',
-       name: `Building ${container.name}`,
-     })
-    })
+      taskUtil.setTask({
+        id: container.name,
+        state: 'loading',
+        name: `Building ${container.name}`,
+      })
+    });
 
     // Promise all the build images
     return Promise.all(
@@ -151,4 +172,90 @@ export class ApplicationManager {
       )
     )
   }
+
+
+  downloadModelMain(modelId: string, url: string, taskUtil: TaskUtils, destFileName?: string): Promise<string> {
+    return new Promise((resolve, reject) => {
+      const downloadCallback = (result: DownloadModelResult) => {
+        if (result.result) {
+          taskUtil.setTaskState(modelId, 'success');
+          resolve('');
+        } else {
+          taskUtil.setTaskState(modelId, 'error');
+          reject(result.error)
+        }
+      }
+
+      this.downloadModel(modelId, url, taskUtil, downloadCallback, destFileName)
+    })
+  }
+
+  downloadModel(modelId: string, url: string, taskUtil: TaskUtils, callback: (message: DownloadModelResult) => void, destFileName?: string) {
+    const destDir = path.join(this.homeDirectory, AI_STUDIO_FOLDER, 'models', modelId);
+    if (!fs.existsSync(destDir)) {
+      fs.mkdirSync(destDir, { recursive: true });
+    }
+    if (!destFileName) {
+      destFileName =  path.basename(url);
+    }
+    const destFile = path.resolve(destDir, destFileName);
+    const file = fs.createWriteStream(destFile);
+    let totalFileSize = 0;
+    let progress = 0;
+    https.get(url, (resp) => {
+      if (resp.headers.location) {
+        this.downloadModel(modelId, resp.headers.location, taskUtil, callback, destFileName);
+        return;
+      } else {
+        if (totalFileSize === 0 && resp.headers['content-length']) {
+          totalFileSize = parseFloat(resp.headers['content-length']);
+        }
+      }
+
+      resp.on('data', (chunk) => {
+        progress += chunk.length;
+        const progressValue = progress * 100 / totalFileSize;
+
+        taskUtil.setTaskProgress(modelId, progressValue);
+
+        // send progress in percentage (ex. 1.2%, 2.6%, 80.1%) to frontend
+        //this.sendProgress(progressValue);
+        if (progressValue === 100) {
+          callback({
+            result: 'ok'
+          });
+        }
+      });
+      file.on('finish', () => {
+        file.close();
+      });
+      file.on('error', (e) => {
+        callback({
+          result: 'failed',
+          error: e.message,
+        });
+      })
+      resp.pipe(file);
+    });
+  }
+
+  // todo: move somewhere else (dedicated to models)
+  getLocalModels(): LocalModelInfo[] {
+    const result: LocalModelInfo[] = [];
+    const modelsDir = path.join(this.homeDirectory, AI_STUDIO_FOLDER, 'models');
+    const entries = fs.readdirSync(modelsDir, { withFileTypes: true });
+    const dirs = entries.filter(dir => dir.isDirectory());
+    for (const d of dirs) {
+      const modelEntries = fs.readdirSync(path.resolve(d.path, d.name));
+      if (modelEntries.length != 1) {
+        // we support models with one file only for now
+        continue;
+      }
+      result.push({
+        id: d.name,
+        file: modelEntries[0],
+      })
+    }
+    return result;
+  }
 }

packages/backend/src/studio-api-impl.ts

@@ -5,17 +5,14 @@ import content from './ai.json';
 import { ApplicationManager } from './managers/applicationManager';
 import { RecipeStatusRegistry } from './registries/RecipeStatusRegistry';
 import { RecipeStatus } from '@shared/models/IRecipeStatus';
-import { Task } from '@shared/models/ITask';
 import { ModelInfo } from '@shared/models/IModelInfo';
-import { Studio } from './studio';
 
 export const RECENT_CATEGORY_ID = 'recent-category';
 
 export class StudioApiImpl implements StudioAPI {
   constructor(
     private applicationManager: ApplicationManager,
     private recipeStatusRegistry: RecipeStatusRegistry,
-    private studio: Studio,
   ) {}
 
   async openURL(url: string): Promise<void> {
@@ -69,8 +66,9 @@ export class StudioApiImpl implements StudioAPI {
   }
 
   async getLocalModels(): Promise<ModelInfo[]> {
-    const local = this.studio.getLocalModels();
+    const local = this.applicationManager.getLocalModels();
     const localIds = local.map(l => l.id);
     return content.recipes.flatMap(r => r.models.filter(m => localIds.includes(m.id)));
   }
+
 }

packages/backend/src/studio.ts

@@ -23,7 +23,9 @@ import { StudioApiImpl } from './studio-api-impl';
 import { ApplicationManager } from './managers/applicationManager';
 import { GitManager } from './managers/gitManager';
 import { RecipeStatusRegistry } from './registries/RecipeStatusRegistry';
+
 import * as fs from 'node:fs';
+import * as https from 'node:https';
 import * as path from 'node:path';
 import type { LocalModelInfo } from '@shared/models/ILocalModelInfo';
 
@@ -90,12 +92,12 @@ export class Studio {
     const recipeStatusRegistry = new RecipeStatusRegistry();
     const applicationManager = new ApplicationManager(
       gitManager,
-      recipeStatusRegistry
+      recipeStatusRegistry,
+      this.#extensionContext,
     )
     this.studioApi = new StudioApiImpl(
       applicationManager,
       recipeStatusRegistry,
-      this
     );
     // Register the instance
     this.rpcExtension.registerInstance<StudioApiImpl>(StudioApiImpl, this.studioApi);
@@ -114,23 +116,4 @@ export class Studio {
       localResourceRoots: [Uri.joinPath(extensionUri, 'media')],
     };
   }
-
-  getLocalModels(): LocalModelInfo[] {
-    const result: LocalModelInfo[] = [];
-    const modelsDir = path.resolve(this.#extensionContext.storagePath, 'models');
-    const entries = fs.readdirSync(modelsDir, { withFileTypes: true });
-    const dirs = entries.filter(dir => dir.isDirectory());
-    for (const d of dirs) {
-      const modelEntries = fs.readdirSync(path.resolve(d.path, d.name));
-      if (modelEntries.length != 1) {
-        // we support models with one file only for now
-        continue;
-      }
-      result.push({
-        id: d.name,
-        file: modelEntries[0],
-      })
-    }
-    return result;
-  }
 }

packages/backend/src/utils/taskUtils.ts

@@ -37,6 +37,16 @@ export class TaskUtils {
     })
   }
 
+  setTaskProgress(taskId: string, value: number) {
+    if(!this.tasks.has(taskId))
+      throw new Error('task not found.');
+    const task = this.tasks.get(taskId);
+    this.setTask({
+      ...task,
+      progress: value,
+    })
+  }
+
   toRecipeStatus(): RecipeStatus {
     return {
       recipeId: this.recipeId,

packages/frontend/src/lib/progress/TasksProgress.svelte

@@ -8,26 +8,30 @@ export let tasks: Task[] = [];
 
   {#each tasks as task}
     <li class="flex items-center">
-      {#if task.state === 'success'}
-        <svg class="w-4 h-4 me-2 text-green-500 dark:text-green-400 flex-shrink-0" aria-hidden="true" xmlns="http://www.w3.org/2000/svg" fill="currentColor" viewBox="0 0 20 20">
-          <path d="M10 .5a9.5 9.5 0 1 0 9.5 9.5A9.51 9.51 0 0 0 10 .5Zm3.707 8.207-4 4a1 1 0 0 1-1.414 0l-2-2a1 1 0 0 1 1.414-1.414L9 10.586l3.293-3.293a1 1 0 0 1 1.414 1.414Z"/>
-        </svg>
-      {:else if task.state === 'loading'}
-        <svg aria-hidden="true" class="w-4 h-4 me-2 text-gray-200 animate-spin dark:text-gray-600 fill-purple-500"
-             viewBox="0 0 100 101" fill="none" xmlns="http://www.w3.org/2000/svg">
-          <path
-            d="M100 50.5908C100 78.2051 77.6142 100.591 50 100.591C22.3858 100.591 0 78.2051 0 50.5908C0 22.9766 22.3858 0.59082 50 0.59082C77.6142 0.59082 100 22.9766 100 50.5908ZM9.08144 50.5908C9.08144 73.1895 27.4013 91.5094 50 91.5094C72.5987 91.5094 90.9186 73.1895 90.9186 50.5908C90.9186 27.9921 72.5987 9.67226 50 9.67226C27.4013 9.67226 9.08144 27.9921 9.08144 50.5908Z"
-            fill="currentColor" />
-          <path
-            d="M93.9676 39.0409C96.393 38.4038 97.8624 35.9116 97.0079 33.5539C95.2932 28.8227 92.871 24.3692 89.8167 20.348C85.8452 15.1192 80.8826 10.7238 75.2124 7.41289C69.5422 4.10194 63.2754 1.94025 56.7698 1.05124C51.7666 0.367541 46.6976 0.446843 41.7345 1.27873C39.2613 1.69328 37.813 4.19778 38.4501 6.62326C39.0873 9.04874 41.5694 10.4717 44.0505 10.1071C47.8511 9.54855 51.7191 9.52689 55.5402 10.0491C60.8642 10.7766 65.9928 12.5457 70.6331 15.2552C75.2735 17.9648 79.3347 21.5619 82.5849 25.841C84.9175 28.9121 86.7997 32.2913 88.1811 35.8758C89.083 38.2158 91.5421 39.6781 93.9676 39.0409Z"
-            fill="currentFill" />
-        </svg>
-      {:else}
-        <svg class="flex-shrink-0 inline w-4 h-4 me-3 text-red-600 fe" aria-hidden="true" xmlns="http://www.w3.org/2000/svg" fill="currentColor" viewBox="0 0 20 20">
-          <path d="M10 .5a9.5 9.5 0 1 0 9.5 9.5A9.51 9.51 0 0 0 10 .5ZM9.5 4a1.5 1.5 0 1 1 0 3 1.5 1.5 0 0 1 0-3ZM12 15H8a1 1 0 0 1 0-2h1v-3H8a1 1 0 0 1 0-2h2a1 1 0 0 1 1 1v4h1a1 1 0 0 1 0 2Z"/>
-        </svg>
-      {/if}
-      {task.name}
+      <div class="min-w-4 mr-2">
+        {#if task.state === 'success'}
+          <svg class="w-4 h-4 text-green-500 dark:text-green-400 flex-shrink-0" aria-hidden="true" xmlns="http://www.w3.org/2000/svg" fill="currentColor" viewBox="0 0 20 20">
+            <path d="M10 .5a9.5 9.5 0 1 0 9.5 9.5A9.51 9.51 0 0 0 10 .5Zm3.707 8.207-4 4a1 1 0 0 1-1.414 0l-2-2a1 1 0 0 1 1.414-1.414L9 10.586l3.293-3.293a1 1 0 0 1 1.414 1.414Z"/>
+          </svg>
+        {:else if task.state === 'loading'}
+          <svg aria-hidden="true" class="w-4 h-4text-gray-200 animate-spin dark:text-gray-600 fill-purple-500"
+               viewBox="0 0 100 101" fill="none" xmlns="http://www.w3.org/2000/svg">
+            <path
+              d="M100 50.5908C100 78.2051 77.6142 100.591 50 100.591C22.3858 100.591 0 78.2051 0 50.5908C0 22.9766 22.3858 0.59082 50 0.59082C77.6142 0.59082 100 22.9766 100 50.5908ZM9.08144 50.5908C9.08144 73.1895 27.4013 91.5094 50 91.5094C72.5987 91.5094 90.9186 73.1895 90.9186 50.5908C90.9186 27.9921 72.5987 9.67226 50 9.67226C27.4013 9.67226 9.08144 27.9921 9.08144 50.5908Z"
+              fill="currentColor" />
+            <path
+              d="M93.9676 39.0409C96.393 38.4038 97.8624 35.9116 97.0079 33.5539C95.2932 28.8227 92.871 24.3692 89.8167 20.348C85.8452 15.1192 80.8826 10.7238 75.2124 7.41289C69.5422 4.10194 63.2754 1.94025 56.7698 1.05124C51.7666 0.367541 46.6976 0.446843 41.7345 1.27873C39.2613 1.69328 37.813 4.19778 38.4501 6.62326C39.0873 9.04874 41.5694 10.4717 44.0505 10.1071C47.8511 9.54855 51.7191 9.52689 55.5402 10.0491C60.8642 10.7766 65.9928 12.5457 70.6331 15.2552C75.2735 17.9648 79.3347 21.5619 82.5849 25.841C84.9175 28.9121 86.7997 32.2913 88.1811 35.8758C89.083 38.2158 91.5421 39.6781 93.9676 39.0409Z"
+              fill="currentFill" />
+          </svg>
+        {:else}
+          <svg class="flex-shrink-0 inline w-4 h-4 text-red-600 fe" aria-hidden="true" xmlns="http://www.w3.org/2000/svg" fill="currentColor" viewBox="0 0 20 20">
+            <path d="M10 .5a9.5 9.5 0 1 0 9.5 9.5A9.51 9.51 0 0 0 10 .5ZM9.5 4a1.5 1.5 0 1 1 0 3 1.5 1.5 0 0 1 0-3ZM12 15H8a1 1 0 0 1 0-2h1v-3H8a1 1 0 0 1 0-2h2a1 1 0 0 1 1 1v4h1a1 1 0 0 1 0 2Z"/>
+          </svg>
+        {/if}
+      </div>
+      <span>
+        {task.name} {#if task.progress}({Math.floor(task.progress)}%){/if}
+      </span>
     </li>
   {/each}
 </ul>

packages/shared/models/IModelInfo.ts

@@ -5,4 +5,5 @@ export interface ModelInfo {
   registry: string;
   popularity: number;
   license: string;
+  url: string;
 }

packages/shared/models/ITask.ts

@@ -3,5 +3,6 @@ export type TaskState = 'loading' | 'error' | 'success'
 export interface Task {
   id: string,
   state: TaskState;
+  progress?: number
   name: string;
 }