Skip to content

Commit

Permalink
feat: display all models in the models page (#245)
Browse files Browse the repository at this point in the history 
* feat: display all models in the models page

Fixes #43

Signed-off-by: Jeff MAURY <[email protected]>
  • Loading branch information
@jeffmaury
jeffmaury authored Feb 12, 2024
1 parent d5ea0b1 commit f4c251c
Show file tree
Hide file tree
Showing 14 changed files with 151 additions and 140 deletions.
11 changes: 10 additions & 1 deletion packages/backend/src/managers/applicationManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,16 @@ describe('pullApplication', () => {
mocks.createContainerMock.mockResolvedValue({
id: 'id',
});
modelsManager = new ModelsManager('appdir', {} as Webview, {} as CatalogManager, telemetryLogger);
modelsManager = new ModelsManager(
'appdir',
{} as Webview,
{
getModels(): ModelInfo[] {
return [];
},
} as CatalogManager,
telemetryLogger,
);
manager = new ApplicationManager(
'/home/user/aistudio',
{
Expand Down
139 changes: 74 additions & 65 deletions packages/backend/src/managers/modelsManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ function mockFiles(now: Date) {
});
}

test('getLocalModelsFromDisk should get models in local directory', () => {
test('getModelsInfo should get models in local directory', async () => {
const now = new Date();
mockFiles(now);
let appdir: string;
Expand All @@ -129,27 +129,47 @@ test('getLocalModelsFromDisk should get models in local directory', () => {
} else {
appdir = '/home/user/aistudio';
}
const manager = new ModelsManager(appdir, {} as Webview, {} as CatalogManager, telemetryLogger);
manager.getLocalModelsFromDisk();
expect(manager.getLocalModels()).toEqual([
const manager = new ModelsManager(
appdir,
{
postMessage: vi.fn(),
} as unknown as Webview,
{
getModels(): ModelInfo[] {
return [
{ id: 'model-id-1', name: 'model-id-1-model' } as ModelInfo,
{ id: 'model-id-2', name: 'model-id-2-model' } as ModelInfo,
];
},
} as CatalogManager,
telemetryLogger,
);
await manager.loadLocalModels();
expect(manager.getModelsInfo()).toEqual([
{
id: 'model-id-1',
file: 'model-id-1-model',
size: 32000,
creation: now,
path: path.resolve(dirent[0].path, dirent[0].name),
name: 'model-id-1-model',
file: {
size: 32000,
creation: now,
path: path.resolve(dirent[0].path, dirent[0].name),
file: 'model-id-1-model',
},
},
{
id: 'model-id-2',
file: 'model-id-2-model',
size: 32000,
creation: now,
path: path.resolve(dirent[1].path, dirent[1].name),
name: 'model-id-2-model',
file: {
size: 32000,
creation: now,
path: path.resolve(dirent[1].path, dirent[1].name),
file: 'model-id-2-model',
},
},
]);
});

test('getLocalModelsFromDisk should return an empty array if the models folder does not exist', () => {
test('getModelsInfo should return an empty array if the models folder does not exist', () => {
vi.spyOn(os, 'homedir').mockReturnValue('/home/user');
const existsSyncSpy = vi.spyOn(fs, 'existsSync');
existsSyncSpy.mockReturnValue(false);
Expand All @@ -159,9 +179,18 @@ test('getLocalModelsFromDisk should return an empty array if the models folder d
} else {
appdir = '/home/user/aistudio';
}
const manager = new ModelsManager(appdir, {} as Webview, {} as CatalogManager, telemetryLogger);
const manager = new ModelsManager(
appdir,
{} as Webview,
{
getModels(): ModelInfo[] {
return [];
},
} as CatalogManager,
telemetryLogger,
);
manager.getLocalModelsFromDisk();
expect(manager.getLocalModels()).toEqual([]);
expect(manager.getModelsInfo()).toEqual([]);
if (process.platform === 'win32') {
expect(existsSyncSpy).toHaveBeenCalledWith('C:\\home\\user\\aistudio\\models');
} else {
Expand Down Expand Up @@ -198,13 +227,12 @@ test('loadLocalModels should post a message with the message on disk and on cata
);
await manager.loadLocalModels();
expect(postMessageMock).toHaveBeenNthCalledWith(1, {
id: 'new-local-models-state',
id: 'new-models-state',
body: [
{
file: {
creation: now,
file: 'model-id-1-model',
id: 'model-id-1',
size: 32000,
path: path.resolve(dirent[0].path, dirent[0].name),
},
Expand Down Expand Up @@ -242,37 +270,24 @@ test('deleteLocalModel deletes the model folder', async () => {
} as CatalogManager,
telemetryLogger,
);
manager.getLocalModelsFromDisk();
await manager.loadLocalModels();
await manager.deleteLocalModel('model-id-1');
// check that the model's folder is removed from disk
if (process.platform === 'win32') {
expect(rmSpy).toBeCalledWith('C:\\home\\user\\aistudio\\models\\model-id-1', { recursive: true });
} else {
expect(rmSpy).toBeCalledWith('/home/user/aistudio/models/model-id-1', { recursive: true });
}
expect(postMessageMock).toHaveBeenCalledTimes(2);
// check that a state is sent with the model being deleted
expect(postMessageMock).toHaveBeenCalledWith({
id: 'new-local-models-state',
expect(postMessageMock).toHaveBeenCalledTimes(3);
// check that a new state is sent with the model removed
expect(postMessageMock).toHaveBeenNthCalledWith(3, {
id: 'new-models-state',
body: [
{
file: {
creation: now,
file: 'model-id-1-model',
id: 'model-id-1',
size: 32000,
path: path.resolve(dirent[0].path, dirent[0].name),
},
id: 'model-id-1',
state: 'deleting',
},
],
});
// check that a new state is sent with the model removed
expect(postMessageMock).toHaveBeenCalledWith({
id: 'new-local-models-state',
body: [],
});
expect(mocks.logUsageMock).toHaveBeenNthCalledWith(1, 'model.delete', { 'model.id': 'model-id-1' });
});

Expand Down Expand Up @@ -304,44 +319,20 @@ test('deleteLocalModel fails to delete the model folder', async () => {
} as CatalogManager,
telemetryLogger,
);
manager.getLocalModelsFromDisk();
await manager.loadLocalModels();
await manager.deleteLocalModel('model-id-1');
// check that the model's folder is removed from disk
if (process.platform === 'win32') {
expect(rmSpy).toBeCalledWith('C:\\home\\user\\aistudio\\models\\model-id-1', { recursive: true });
} else {
expect(rmSpy).toBeCalledWith('/home/user/aistudio/models/model-id-1', { recursive: true });
}
expect(postMessageMock).toHaveBeenCalledTimes(2);
// check that a state is sent with the model being deleted
expect(postMessageMock).toHaveBeenCalledWith({
id: 'new-local-models-state',
body: [
{
file: {
creation: now,
file: 'model-id-1-model',
id: 'model-id-1',
size: 32000,
path: path.resolve(dirent[0].path, dirent[0].name),
},
id: 'model-id-1',
state: 'deleting',
},
],
});
expect(postMessageMock).toHaveBeenCalledTimes(3);
// check that a new state is sent with the model non removed
expect(postMessageMock).toHaveBeenCalledWith({
id: 'new-local-models-state',
expect(postMessageMock).toHaveBeenNthCalledWith(3, {
id: 'new-models-state',
body: [
{
file: {
creation: now,
file: 'model-id-1-model',
id: 'model-id-1',
size: 32000,
path: path.resolve(dirent[0].path, dirent[0].name),
},
id: 'model-id-1',
},
],
Expand All @@ -351,7 +342,16 @@ test('deleteLocalModel fails to delete the model folder', async () => {
});

describe('downloadModel', () => {
const manager = new ModelsManager('appdir', {} as Webview, {} as CatalogManager, telemetryLogger);
const manager = new ModelsManager(
'appdir',
{} as Webview,
{
getModels(): ModelInfo[] {
return [];
},
} as CatalogManager,
telemetryLogger,
);
test('download model if not already on disk', async () => {
vi.spyOn(manager, 'isModelOnDisk').mockReturnValue(false);
const doDownloadModelWrapperMock = vi
Expand Down Expand Up @@ -403,7 +403,16 @@ describe('downloadModel', () => {
});

describe('doDownloadModelWrapper', () => {
const manager = new ModelsManager('appdir', {} as Webview, {} as CatalogManager, telemetryLogger);
const manager = new ModelsManager(
'appdir',
{} as Webview,
{
getModels(): ModelInfo[] {
return [];
},
} as CatalogManager,
telemetryLogger,
);
test('returning model path if model has been downloaded', async () => {
vi.spyOn(manager, 'doDownloadModel').mockImplementation(
(
Expand Down
Loading

0 comments on commit f4c251c

Please sign in to comment.