From 89613f1c403352265272b754472e796b69ded9d3 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Fri, 12 Jan 2024 11:51:29 -0500 Subject: [PATCH 1/5] feat: adding TaskRegistry - Adding central place to register tasks - Adding labels to tasks etc. --- .../src/managers/applicationManager.ts | 15 ++-- .../src/registries/RecipeStatusRegistry.ts | 7 ++ .../backend/src/registries/TaskRegistry.ts | 18 ++++ packages/backend/src/studio-api-impl.ts | 8 +- packages/backend/src/studio.ts | 6 +- .../{taskUtils.ts => recipeStatusUtils.ts} | 2 +- packages/frontend/src/pages/Models.svelte | 89 ++++++++++++++----- packages/frontend/src/pages/Recipe.svelte | 2 +- packages/shared/StudioAPI.ts | 8 +- packages/shared/models/ITask.ts | 5 +- 10 files changed, 126 insertions(+), 34 deletions(-) create mode 100644 packages/backend/src/registries/TaskRegistry.ts rename packages/backend/src/utils/{taskUtils.ts => recipeStatusUtils.ts} (97%) diff --git a/packages/backend/src/managers/applicationManager.ts b/packages/backend/src/managers/applicationManager.ts index a26b65334..8c1c6a509 100644 --- a/packages/backend/src/managers/applicationManager.ts +++ b/packages/backend/src/managers/applicationManager.ts @@ -9,9 +9,8 @@ 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 { RecipeStatusUtils } from '../utils/recipeStatusUtils'; import { getParentDirectory } from '../utils/pathUtils'; -import { a } from 'vitest/dist/suite-dF4WyktM'; import type { LocalModelInfo } from '@shared/models/ILocalModelInfo'; // TODO: Need to be configured @@ -32,7 +31,7 @@ export class ApplicationManager { async pullApplication(recipe: Recipe) { // Create a TaskUtils object to help us - const taskUtil = new TaskUtils(recipe.id, this.recipeStatusRegistry); + const taskUtil = new RecipeStatusUtils(recipe.id, this.recipeStatusRegistry); const localFolder = path.join(this.homeDirectory, AI_STUDIO_FOLDER, recipe.id); @@ -41,6 +40,9 @@ export class ApplicationManager { id: 'checkout', name: 'Checkout repository', state: 'loading', + labels: { + 'git': 'checkout', + }, } taskUtil.setTask(checkoutTask); @@ -125,6 +127,9 @@ export class ApplicationManager { id: model.id, state: 'loading', name: `Downloading model ${model.name}`, + labels: { + "model-pulling": model.id, + } }); await this.downloadModelMain(model.id, model.url, taskUtil) @@ -174,7 +179,7 @@ export class ApplicationManager { } - downloadModelMain(modelId: string, url: string, taskUtil: TaskUtils, destFileName?: string): Promise { + downloadModelMain(modelId: string, url: string, taskUtil: RecipeStatusUtils, destFileName?: string): Promise { return new Promise((resolve, reject) => { const downloadCallback = (result: DownloadModelResult) => { if (result.result) { @@ -190,7 +195,7 @@ export class ApplicationManager { }) } - downloadModel(modelId: string, url: string, taskUtil: TaskUtils, callback: (message: DownloadModelResult) => void, destFileName?: string) { + private downloadModel(modelId: string, url: string, taskUtil: RecipeStatusUtils, 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 }); diff --git a/packages/backend/src/registries/RecipeStatusRegistry.ts b/packages/backend/src/registries/RecipeStatusRegistry.ts index ef2a46da8..8f636a6f3 100644 --- a/packages/backend/src/registries/RecipeStatusRegistry.ts +++ b/packages/backend/src/registries/RecipeStatusRegistry.ts @@ -1,9 +1,16 @@ import { RecipeStatus } from '@shared/models/IRecipeStatus'; +import { TaskRegistry } from './TaskRegistry'; export class RecipeStatusRegistry { private statuses: Map = new Map(); + constructor(private taskRegistry: TaskRegistry) { } + setStatus(recipeId: string, status: RecipeStatus) { + // Update the TaskRegistry + if(status.tasks && status.tasks.length > 0) { + status.tasks.map((task) => this.taskRegistry.set(task)); + } this.statuses.set(recipeId, status); } diff --git a/packages/backend/src/registries/TaskRegistry.ts b/packages/backend/src/registries/TaskRegistry.ts new file mode 100644 index 000000000..699fa3a4f --- /dev/null +++ b/packages/backend/src/registries/TaskRegistry.ts @@ -0,0 +1,18 @@ +import { Task } from '@shared/models/ITask'; + +export class TaskRegistry { + private tasks: Map = new Map(); + + getTasksByLabel(label: string): Task[] { + return Array.from(this.tasks.values()) + .filter((task) => label in (task.labels || {})) + } + + set(task: Task) { + this.tasks.set(task.id, task); + } + + delete(taskId: string) { + this.tasks.delete(taskId); + } +} diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts index 5940bd6fd..c63c13a5b 100644 --- a/packages/backend/src/studio-api-impl.ts +++ b/packages/backend/src/studio-api-impl.ts @@ -6,6 +6,8 @@ import { AI_STUDIO_FOLDER, ApplicationManager } from './managers/applicationMana import { RecipeStatusRegistry } from './registries/RecipeStatusRegistry'; import { RecipeStatus } from '@shared/models/IRecipeStatus'; import { ModelInfo } from '@shared/models/IModelInfo'; +import { TaskRegistry } from './registries/TaskRegistry'; +import { Task } from '@shared/models/ITask'; import { Studio } from './studio'; import * as path from 'node:path'; import { ModelResponse } from '@shared/models/IModelResponse'; @@ -16,7 +18,7 @@ export class StudioApiImpl implements StudioAPI { constructor( private applicationManager: ApplicationManager, private recipeStatusRegistry: RecipeStatusRegistry, - private studio: Studio, + private taskRegistry: TaskRegistry, ) {} async openURL(url: string): Promise { @@ -82,6 +84,10 @@ export class StudioApiImpl implements StudioAPI { return content.recipes.flatMap(r => r.models.filter(m => localIds.includes(m.id))); } + async getTasksByLabel(label: string): Promise { + return this.taskRegistry.getTasksByLabel(label); + } + async startPlayground(modelId: string): Promise { const localModelInfo = this.applicationManager.getLocalModels().filter(m => m.id === modelId); if (localModelInfo.length !== 1) { diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index efaff9e48..0b160b293 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -28,6 +28,7 @@ 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'; +import { TaskRegistry } from './registries/TaskRegistry'; import { PlayGroundManager } from './playground'; export class Studio { @@ -92,7 +93,8 @@ export class Studio { // Let's create the api that the front will be able to call this.rpcExtension = new RpcExtension(this.#panel.webview); const gitManager = new GitManager(); - const recipeStatusRegistry = new RecipeStatusRegistry(); + const taskRegistry = new TaskRegistry(); + const recipeStatusRegistry = new RecipeStatusRegistry(taskRegistry); const applicationManager = new ApplicationManager( gitManager, recipeStatusRegistry, @@ -101,7 +103,7 @@ export class Studio { this.studioApi = new StudioApiImpl( applicationManager, recipeStatusRegistry, - this, + taskRegistry ); // Register the instance this.rpcExtension.registerInstance(StudioApiImpl, this.studioApi); diff --git a/packages/backend/src/utils/taskUtils.ts b/packages/backend/src/utils/recipeStatusUtils.ts similarity index 97% rename from packages/backend/src/utils/taskUtils.ts rename to packages/backend/src/utils/recipeStatusUtils.ts index 8269a79f7..93c3ddc46 100644 --- a/packages/backend/src/utils/taskUtils.ts +++ b/packages/backend/src/utils/recipeStatusUtils.ts @@ -3,7 +3,7 @@ import type { Task, TaskState } from '@shared/models/ITask'; import { RecipeStatusRegistry } from '../registries/RecipeStatusRegistry'; -export class TaskUtils { +export class RecipeStatusUtils { private tasks: Map = new Map(); private state: RecipeStatusState = 'loading'; diff --git a/packages/frontend/src/pages/Models.svelte b/packages/frontend/src/pages/Models.svelte index 90fa0a951..b1ed951d2 100644 --- a/packages/frontend/src/pages/Models.svelte +++ b/packages/frontend/src/pages/Models.svelte @@ -1,23 +1,60 @@ @@ -25,10 +62,20 @@
- {#if $localModels && $localModels.length} + {#if tasks.length > 0} +
+ +
+
Downloading models
+ +
+ +
+ {/if} + {#if models.length > 0}
diff --git a/packages/frontend/src/pages/Recipe.svelte b/packages/frontend/src/pages/Recipe.svelte index 98865371e..bdca5e59b 100644 --- a/packages/frontend/src/pages/Recipe.svelte +++ b/packages/frontend/src/pages/Recipe.svelte @@ -73,7 +73,7 @@ const onClickRepository = () => {
-
+
Repository
diff --git a/packages/shared/StudioAPI.ts b/packages/shared/StudioAPI.ts index 32dd909af..c2cb51523 100644 --- a/packages/shared/StudioAPI.ts +++ b/packages/shared/StudioAPI.ts @@ -16,11 +16,17 @@ export abstract class StudioAPI { abstract pullApplication(recipeId: string): Promise; abstract openURL(url: string): Promise; /** - * Get the information of models saved locally into the extension's storage directory + * Get the information of models saved locally into the extension's storage directory */ abstract getLocalModels(): Promise; abstract startPlayground(modelId: string): Promise; abstract askPlayground(modelId: string, prompt: string): Promise; + + /** + * Get task by label + * @param label + */ + abstract getTasksByLabel(label: string): Promise; } diff --git a/packages/shared/models/ITask.ts b/packages/shared/models/ITask.ts index a6aec82e7..c26df28d0 100644 --- a/packages/shared/models/ITask.ts +++ b/packages/shared/models/ITask.ts @@ -1,8 +1,9 @@ export type TaskState = 'loading' | 'error' | 'success' export interface Task { - id: string, + id: string; state: TaskState; - progress?: number + progress?: number; name: string; + labels?: {[id: string]: string} } From 2a5da85f63f1ba21fb0b0ea63c8de10c198acf51 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Fri, 12 Jan 2024 12:01:46 -0500 Subject: [PATCH 2/5] fix: rebase --- packages/backend/src/studio-api-impl.ts | 9 ++++++--- packages/backend/src/studio.ts | 3 ++- packages/shared/StudioAPI.ts | 1 + 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts index c63c13a5b..1ae2fec62 100644 --- a/packages/backend/src/studio-api-impl.ts +++ b/packages/backend/src/studio-api-impl.ts @@ -11,6 +11,7 @@ import { Task } from '@shared/models/ITask'; import { Studio } from './studio'; import * as path from 'node:path'; import { ModelResponse } from '@shared/models/IModelResponse'; +import { PlayGroundManager } from './playground'; export const RECENT_CATEGORY_ID = 'recent-category'; @@ -19,6 +20,7 @@ export class StudioApiImpl implements StudioAPI { private applicationManager: ApplicationManager, private recipeStatusRegistry: RecipeStatusRegistry, private taskRegistry: TaskRegistry, + private playgroundManager: PlayGroundManager, ) {} async openURL(url: string): Promise { @@ -93,9 +95,10 @@ export class StudioApiImpl implements StudioAPI { if (localModelInfo.length !== 1) { throw new Error('model not found'); } - const destDir = path.join(); + const modelPath = path.resolve(this.applicationManager.homeDirectory, AI_STUDIO_FOLDER, 'models', modelId, localModelInfo[0].file); - this.studio.playgroundManager.startPlayground(modelId, modelPath); + + await this.playgroundManager.startPlayground(modelId, modelPath); } askPlayground(modelId: string, prompt: string): Promise { @@ -103,6 +106,6 @@ export class StudioApiImpl implements StudioAPI { if (localModelInfo.length !== 1) { throw new Error('model not found'); } - return this.studio.playgroundManager.askPlayground(localModelInfo[0], prompt); + return this.playgroundManager.askPlayground(localModelInfo[0], prompt); } } diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index 0b160b293..eb13c103b 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -103,7 +103,8 @@ export class Studio { this.studioApi = new StudioApiImpl( applicationManager, recipeStatusRegistry, - taskRegistry + taskRegistry, + this.playgroundManager, ); // Register the instance this.rpcExtension.registerInstance(StudioApiImpl, this.studioApi); diff --git a/packages/shared/StudioAPI.ts b/packages/shared/StudioAPI.ts index c2cb51523..b76fa4f1c 100644 --- a/packages/shared/StudioAPI.ts +++ b/packages/shared/StudioAPI.ts @@ -3,6 +3,7 @@ import type { Category } from '@shared/models/ICategory'; import { RecipeStatus } from '@shared/models/IRecipeStatus'; import { ModelInfo } from '@shared/models/IModelInfo'; import { ModelResponse } from '@shared/models/IModelResponse'; +import { Task } from './models/ITask'; export abstract class StudioAPI { abstract ping(): Promise; From 0c7913a77cef9736914c53d330ac8c986ae9e744 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Fri, 12 Jan 2024 13:01:10 -0500 Subject: [PATCH 3/5] fix: caching --- .../src/managers/applicationManager.ts | 6 ++++ packages/frontend/src/pages/Models.svelte | 32 +++++++++++++++---- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/packages/backend/src/managers/applicationManager.ts b/packages/backend/src/managers/applicationManager.ts index 8c1c6a509..2e8c21fba 100644 --- a/packages/backend/src/managers/applicationManager.ts +++ b/packages/backend/src/managers/applicationManager.ts @@ -191,6 +191,12 @@ export class ApplicationManager { } } + if(fs.existsSync(destFileName)) { + taskUtil.setTaskState(modelId, 'success'); + taskUtil.setTaskProgress(modelId, 100); + return; + } + this.downloadModel(modelId, url, taskUtil, downloadCallback, destFileName) }) } diff --git a/packages/frontend/src/pages/Models.svelte b/packages/frontend/src/pages/Models.svelte index b1ed951d2..6ec1104fc 100644 --- a/packages/frontend/src/pages/Models.svelte +++ b/packages/frontend/src/pages/Models.svelte @@ -17,6 +17,7 @@ import TasksProgress from '/@/lib/progress/TasksProgress.svelte'; import { faRefresh } from '@fortawesome/free-solid-svg-icons'; import Card from '/@/lib/Card.svelte'; import Button from '/@/lib/button/Button.svelte'; +import LinearProgress from '/@/lib/progress/LinearProgress.svelte'; const columns: Column[] = [ new Column('Name', { width: '4fr', renderer: ModelColumnName }), @@ -27,25 +28,39 @@ const columns: Column[] = [ ]; const row = new Row({}); +let loading: boolean = true; let intervalId: ReturnType | undefined = undefined; $: tasks = [] as Task[]; $: models = [] as ModelInfo[]; +$: filteredModels = [] as ModelInfo[]; -function filterModels(tasks: Task[], models: ModelInfo[]): void { - const dict: {[id: number]: Task} = Object.fromEntries(tasks.map((task) => [task.id, task])); - models = models.filter((model) => model.id in dict); +function filterModels(): void { + // Let's collect the models we do not want to show (loading, error). + const modelsId: string[] = tasks.reduce((previousValue, currentValue) => { + if(currentValue.state === 'success') + return previousValue; + + if(currentValue.labels) { + previousValue.push(currentValue.labels["model-pulling"]) + } + return previousValue; + }, [] as string[]); + filteredModels = models.filter((model) => !(model.id in modelsId)); } onMount(() => { // Pulling update intervalId = setInterval(async () => { tasks = await studioClient.getTasksByLabel("model-pulling"); - filterModels(tasks, models); + loading = false; + filterModels(); }, 1000); + // Subscribe to the models store return localModels.subscribe((value) => { - filterModels(tasks, value); + models = value; + filterModels(); }) }); @@ -61,6 +76,9 @@ onDestroy(() => {
+ {#if loading} + + {/if}
{#if tasks.length > 0}
@@ -72,10 +90,10 @@ onDestroy(() => {
{/if} - {#if models.length > 0} + {#if filteredModels.length > 0}
From 4d1a5e7279f2cac895d5b896d56776adffe183be Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Mon, 15 Jan 2024 09:59:21 +0100 Subject: [PATCH 4/5] fix: displaying loading screen properly --- packages/frontend/src/pages/Models.svelte | 44 ++++++++++++----------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/packages/frontend/src/pages/Models.svelte b/packages/frontend/src/pages/Models.svelte index 6ec1104fc..f7b9d5262 100644 --- a/packages/frontend/src/pages/Models.svelte +++ b/packages/frontend/src/pages/Models.svelte @@ -41,8 +41,8 @@ function filterModels(): void { if(currentValue.state === 'success') return previousValue; - if(currentValue.labels) { - previousValue.push(currentValue.labels["model-pulling"]) + if(currentValue.labels !== undefined) { + previousValue.push(currentValue.labels["model-pulling"]); } return previousValue; }, [] as string[]); @@ -80,25 +80,27 @@ onDestroy(() => { {/if}
- {#if tasks.length > 0} -
- -
-
Downloading models
- -
- -
- {/if} - {#if filteredModels.length > 0} - -
- {:else} -
There is no model yet
+ {#if !loading} + {#if tasks.length > 0} +
+ +
+
Downloading models
+ +
+ +
+ {/if} + {#if filteredModels.length > 0} + +
+ {:else} +
There is no model yet
+ {/if} {/if}
From 6045ed56142e38ac2e5deea7ead69748f2bd3019 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Mon, 15 Jan 2024 11:57:59 +0100 Subject: [PATCH 5/5] fix: remove $ variables --- packages/frontend/src/pages/Models.svelte | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/frontend/src/pages/Models.svelte b/packages/frontend/src/pages/Models.svelte index f7b9d5262..49a0e5ebc 100644 --- a/packages/frontend/src/pages/Models.svelte +++ b/packages/frontend/src/pages/Models.svelte @@ -31,9 +31,9 @@ const row = new Row({}); let loading: boolean = true; let intervalId: ReturnType | undefined = undefined; -$: tasks = [] as Task[]; -$: models = [] as ModelInfo[]; -$: filteredModels = [] as ModelInfo[]; +let tasks: Task[] = []; +let models: ModelInfo[] = []; +let filteredModels: ModelInfo[] = []; function filterModels(): void { // Let's collect the models we do not want to show (loading, error).