refactor: move getLocalModels to ModelsManager

packages/backend/src/managers/applicationManager.spec.ts

@@ -1,15 +1,13 @@
 import { type MockInstance, describe, expect, test, vi, beforeEach } from 'vitest';
 import { ApplicationManager } from './applicationManager';
 import type { RecipeStatusRegistry } from '../registries/RecipeStatusRegistry';
-import type { ExtensionContext } from '@podman-desktop/api';
 import type { GitManager } from './gitManager';
 import os from 'os';
-import fs, { Stats, type Dirent } from 'fs';
-import path from 'path';
+import fs from 'fs';
 import type { Recipe } from '@shared/src/models/IRecipe';
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
-import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo';
 import type { RecipeStatusUtils } from '../utils/recipeStatusUtils';
+import type { ModelsManager } from './modelsManager';
 
 const mocks = vi.hoisted(() => {
   return {
@@ -32,70 +30,15 @@ beforeEach(() => {
   vi.resetAllMocks();
 });
 
-test('appUserDirectory should be under home directory', () => {
-  vi.spyOn(os, 'homedir').mockReturnValue('/home/user');
-  const manager = new ApplicationManager({} as GitManager, {} as RecipeStatusRegistry, {} as ExtensionContext);
-  if (process.platform === 'win32') {
-    expect(manager.appUserDirectory).toMatch(/^\\home\\user/);
-  } else {
-    expect(manager.appUserDirectory).toMatch(/^\/home\/user/);
-  }
-});
-
-test('getLocalModels should return models in local directory', () => {
-  vi.spyOn(os, 'homedir').mockReturnValue('/home/user');
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  const readdirSyncMock = vi.spyOn(fs, 'readdirSync') as unknown as MockInstance<
-    [path: string],
-    string[] | fs.Dirent[]
-  >;
-  readdirSyncMock.mockImplementation((dir: string) => {
-    if (dir.endsWith('model-id-1') || dir.endsWith('model-id-2')) {
-      const base = path.basename(dir);
-      return [base + '-model'];
-    } else {
-      return [
-        {
-          isDirectory: () => true,
-          path: '/home/user/appstudio-dir',
-          name: 'model-id-1',
-        },
-        {
-          isDirectory: () => true,
-          path: '/home/user/appstudio-dir',
-          name: 'model-id-2',
-        },
-        {
-          isDirectory: () => false,
-          path: '/home/user/appstudio-dir',
-          name: 'other-file-should-be-ignored.txt',
-        },
-      ] as Dirent[];
-    }
-  });
-  const manager = new ApplicationManager({} as GitManager, {} as RecipeStatusRegistry, {} as ExtensionContext);
-  const models = manager.getLocalModels();
-  expect(models).toEqual([
-    {
-      id: 'model-id-1',
-      file: 'model-id-1-model',
-    },
-    {
-      id: 'model-id-2',
-      file: 'model-id-2-model',
-    },
-  ]);
-});
-
 describe('pullApplication', () => {
   interface mockForPullApplicationOptions {
     recipeFolderExists: boolean;
   }
 
   const setStatusMock = vi.fn();
   const cloneRepositoryMock = vi.fn();
+  const getLocalModelsMock = vi.fn();
   let manager: ApplicationManager;
-  let getLocalModelsSpy: MockInstance<[], LocalModelInfo[]>;
   let downloadModelMainSpy: MockInstance<
     [modelId: string, url: string, taskUtil: RecipeStatusUtils, destFileName?: string],
     Promise<string>
@@ -116,11 +59,11 @@ describe('pullApplication', () => {
     });
     vi.spyOn(fs, 'statSync').mockImplementation((path: string) => {
       if (path.endsWith('recipe1')) {
-        const stat = new Stats();
+        const stat = new fs.Stats();
         stat.isDirectory = () => true;
         return stat;
       } else if (path.endsWith('ai-studio.yaml')) {
-        const stat = new Stats();
+        const stat = new fs.Stats();
         stat.isDirectory = () => false;
         return stat;
       }
@@ -142,16 +85,18 @@ describe('pullApplication', () => {
     mocks.builImageMock.mockResolvedValue(undefined);
 
     manager = new ApplicationManager(
+      '/home/user/aistudio',
       {
         cloneRepository: cloneRepositoryMock,
       } as unknown as GitManager,
       {
         setStatus: setStatusMock,
       } as unknown as RecipeStatusRegistry,
-      {} as ExtensionContext,
+      {
+        getLocalModels: getLocalModelsMock,
+      } as unknown as ModelsManager,
     );
 
-    getLocalModelsSpy = vi.spyOn(manager, 'getLocalModels');
     downloadModelMainSpy = vi.spyOn(manager, 'downloadModelMain');
     downloadModelMainSpy.mockResolvedValue('');
   }
@@ -160,7 +105,7 @@ describe('pullApplication', () => {
     mockForPullApplication({
       recipeFolderExists: false,
     });
-    getLocalModelsSpy.mockReturnValue([]);
+    getLocalModelsMock.mockReturnValue([]);
 
     const recipe: Recipe = {
       id: 'recipe1',
@@ -183,13 +128,9 @@ describe('pullApplication', () => {
 
     await manager.pullApplication(recipe, model);
     if (process.platform === 'win32') {
-      expect(cloneRepositoryMock).toHaveBeenNthCalledWith(
-        1,
-        'repo',
-        '\\home\\user\\podman-desktop\\ai-studio\\recipe1',
-      );
+      expect(cloneRepositoryMock).toHaveBeenNthCalledWith(1, 'repo', '\\home\\user\\aistudio\\recipe1');
     } else {
-      expect(cloneRepositoryMock).toHaveBeenNthCalledWith(1, 'repo', '/home/user/podman-desktop/ai-studio/recipe1');
+      expect(cloneRepositoryMock).toHaveBeenNthCalledWith(1, 'repo', '/home/user/aistudio/recipe1');
     }
     expect(downloadModelMainSpy).toHaveBeenCalledOnce();
     expect(mocks.builImageMock).toHaveBeenCalledOnce();
@@ -199,7 +140,7 @@ describe('pullApplication', () => {
     mockForPullApplication({
       recipeFolderExists: true,
     });
-    getLocalModelsSpy.mockReturnValue([]);
+    getLocalModelsMock.mockReturnValue([]);
 
     const recipe: Recipe = {
       id: 'recipe1',
@@ -228,7 +169,7 @@ describe('pullApplication', () => {
     mockForPullApplication({
       recipeFolderExists: true,
     });
-    getLocalModelsSpy.mockReturnValue([
+    getLocalModelsMock.mockReturnValue([
       {
         id: 'model1',
         file: 'model1.file',

packages/backend/src/managers/applicationManager.ts

@@ -19,22 +19,19 @@
 import type { Recipe } from '@shared/src/models/IRecipe';
 import { arch } from 'node:os';
 import type { GitManager } from './gitManager';
-import os from 'os';
 import fs from 'fs';
 import * as https from 'node:https';
 import * as path from 'node:path';
-import { containerEngine, type ExtensionContext } from '@podman-desktop/api';
+import { containerEngine } from '@podman-desktop/api';
 import type { RecipeStatusRegistry } from '../registries/RecipeStatusRegistry';
 import type { AIConfig } from '../models/AIConfig';
 import { parseYaml } from '../models/AIConfig';
 import type { Task } from '@shared/src/models/ITask';
 import { RecipeStatusUtils } from '../utils/recipeStatusUtils';
 import { getParentDirectory } from '../utils/pathUtils';
-import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo';
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
+import type { ModelsManager } from './modelsManager';
 
-// 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 {
@@ -43,15 +40,12 @@ interface DownloadModelResult {
 }
 
 export class ApplicationManager {
-  readonly appUserDirectory: string; // todo: make configurable
-
   constructor(
+    private appUserDirectory: string,
     private git: GitManager,
     private recipeStatusRegistry: RecipeStatusRegistry,
-    private extensionContext: ExtensionContext,
-  ) {
-    this.appUserDirectory = path.join(os.homedir(), AI_STUDIO_FOLDER);
-  }
+    private modelsManager: ModelsManager,
+  ) {}
 
   async pullApplication(recipe: Recipe, model: ModelInfo) {
     // Create a TaskUtils object to help us
@@ -141,7 +135,7 @@ export class ApplicationManager {
       container => container.arch === undefined || container.arch === arch(),
     );
 
-    const localModels = this.getLocalModels();
+    const localModels = this.modelsManager.getLocalModels();
     if (!localModels.map(m => m.id).includes(model.id)) {
       // Download model
       taskUtil.setTask({
@@ -298,24 +292,4 @@ export class ApplicationManager {
       resp.pipe(file);
     });
   }
-
-  // todo: move somewhere else (dedicated to models)
-  getLocalModels(): LocalModelInfo[] {
-    const result: LocalModelInfo[] = [];
-    const modelsDir = path.join(this.appUserDirectory, '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/managers/modelsManager.spec.ts

@@ -0,0 +1,53 @@
+import { type MockInstance, beforeEach, expect, test, vi } from 'vitest';
+import os from 'os';
+import fs from 'node:fs';
+import path from 'node:path';
+import { ModelsManager } from './modelsManager';
+
+beforeEach(() => {
+  vi.resetAllMocks();
+});
+
+test('getLocalModels should return models in local directory', () => {
+  vi.spyOn(os, 'homedir').mockReturnValue('/home/user');
+  const readdirSyncMock = vi.spyOn(fs, 'readdirSync') as unknown as MockInstance<
+    [path: string],
+    string[] | fs.Dirent[]
+  >;
+  readdirSyncMock.mockImplementation((dir: string) => {
+    if (dir.endsWith('model-id-1') || dir.endsWith('model-id-2')) {
+      const base = path.basename(dir);
+      return [base + '-model'];
+    } else {
+      return [
+        {
+          isDirectory: () => true,
+          path: '/home/user/appstudio-dir',
+          name: 'model-id-1',
+        },
+        {
+          isDirectory: () => true,
+          path: '/home/user/appstudio-dir',
+          name: 'model-id-2',
+        },
+        {
+          isDirectory: () => false,
+          path: '/home/user/appstudio-dir',
+          name: 'other-file-should-be-ignored.txt',
+        },
+      ] as fs.Dirent[];
+    }
+  });
+  const manager = new ModelsManager('/home/user/aistudio');
+  const models = manager.getLocalModels();
+  expect(models).toEqual([
+    {
+      id: 'model-id-1',
+      file: 'model-id-1-model',
+    },
+    {
+      id: 'model-id-2',
+      file: 'model-id-2-model',
+    },
+  ]);
+});

packages/backend/src/managers/modelsManager.ts

@@ -0,0 +1,26 @@
+import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo';
+import fs from 'fs';
+import * as path from 'node:path';
+
+export class ModelsManager {
+  constructor(private appUserDirectory: string) {}
+
+  getLocalModels(): LocalModelInfo[] {
+    const result: LocalModelInfo[] = [];
+    const modelsDir = path.join(this.appUserDirectory, '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.spec.ts

@@ -26,9 +26,10 @@ import type { RecipeStatusRegistry } from './registries/RecipeStatusRegistry';
 import { StudioApiImpl } from './studio-api-impl';
 import type { PlayGroundManager } from './managers/playground';
 import type { Webview } from '@podman-desktop/api';
-
-import * as fs from 'node:fs';
 import { CatalogManager } from './managers/catalogManager';
+import type { ModelsManager } from './managers/modelsManager';
+
+import fs from 'node:fs';
 
 vi.mock('./ai.json', () => {
   return {
@@ -70,12 +71,12 @@ beforeEach(async () => {
 
   // Creating StudioApiImpl
   studioApiImpl = new StudioApiImpl(
-    {
-      appUserDirectory,
-    } as unknown as ApplicationManager,
+    appUserDirectory,
+    {} as unknown as ApplicationManager,
     {} as unknown as RecipeStatusRegistry,
     {} as unknown as PlayGroundManager,
     catalogManager,
+    {} as unknown as ModelsManager,
   );
   vi.resetAllMocks();
   vi.mock('node:fs');