Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: move getLocalModels to ModelsManager #112

Merged
merged 2 commits into from
Jan 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 'node: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[] {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just a question.
It seems it's a lazy initialization and without cache, do we need to read the files each time we're called or you're planning to read it eagerly or have a cache

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could watch the folder for changes

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, I created #116

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
Loading