Skip to content

Commit

Permalink
refactor: move getLocalModels to ModelsManager
Browse files Browse the repository at this point in the history
  • Loading branch information
feloy committed Jan 23, 2024
1 parent 51d9a65 commit 0ef4670
Show file tree
Hide file tree
Showing 7 changed files with 131 additions and 119 deletions.
87 changes: 14 additions & 73 deletions packages/backend/src/managers/applicationManager.spec.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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>
Expand All @@ -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;
}
Expand All @@ -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('');
}
Expand All @@ -160,7 +105,7 @@ describe('pullApplication', () => {
mockForPullApplication({
recipeFolderExists: false,
});
getLocalModelsSpy.mockReturnValue([]);
getLocalModelsMock.mockReturnValue([]);

const recipe: Recipe = {
id: 'recipe1',
Expand All @@ -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();
Expand All @@ -199,7 +140,7 @@ describe('pullApplication', () => {
mockForPullApplication({
recipeFolderExists: true,
});
getLocalModelsSpy.mockReturnValue([]);
getLocalModelsMock.mockReturnValue([]);

const recipe: Recipe = {
id: 'recipe1',
Expand Down Expand Up @@ -228,7 +169,7 @@ describe('pullApplication', () => {
mockForPullApplication({
recipeFolderExists: true,
});
getLocalModelsSpy.mockReturnValue([
getLocalModelsMock.mockReturnValue([
{
id: 'model1',
file: 'model1.file',
Expand Down
38 changes: 6 additions & 32 deletions packages/backend/src/managers/applicationManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -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;
}
}
53 changes: 53 additions & 0 deletions packages/backend/src/managers/modelsManager.spec.ts
Original file line number Diff line number Diff line change
@@ -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',
},
]);
});
26 changes: 26 additions & 0 deletions packages/backend/src/managers/modelsManager.ts
Original file line number Diff line number Diff line change
@@ -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;
}
}
9 changes: 5 additions & 4 deletions packages/backend/src/studio-api-impl.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 { CatalogManager } from './managers/catalogManager';
import type { ModelsManager } from './managers/modelsManager';

import * as fs from 'node:fs';
import { CatalogManager } from './managers/catalogManager';

vi.mock('./ai.json', () => {
return {
Expand Down Expand Up @@ -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');
Expand Down
Loading

0 comments on commit 0ef4670

Please sign in to comment.