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.

Sign up for GitHub

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

Load user catalog if file exists or use default one #85

Merged
merged 5 commits into from
Jan 18, 2024
Merged

Load user catalog if file exists or use default one #85

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
609258d
fic: declare a appUserDirectory for all user's files
feloy Jan 18, 2024
6d66dcb
add catalog model
feloy Jan 18, 2024
03e01b8
load catalog from user's directory
feloy Jan 18, 2024
8b96600
review changes
feloy Jan 18, 2024
5b9c1c0
do not log error when user catalog is not present
feloy Jan 18, 2024
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
49 changes: 49 additions & 0 deletions packages/backend/src/ai-user-test.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
{
"recipes": [
{
"id": "recipe 1",
"description" : "Recipe 1",
"name" : "Recipe 1",
"repository": "https://recipe1.example.com",
"icon": "natural-language-processing",
"categories": [
"category1"
],
"config": "chatbot/ai-studio.yaml",
"readme": "Readme for recipe 1",
"models": [
"model1",
"model2"
]
}
],
"models": [
{
"id": "model1",
"name": "Model 1",
"description": "Readme for model 1",
"hw": "CPU",
"registry": "Hugging Face",
"popularity": 3,
"license": "?",
"url": "https://model1.example.com"
},
{
"id": "model2",
"name": "Model 2",
"description": "Readme for model 2",
"hw": "CPU",
"registry": "Civital",
"popularity": 3,
"license": "?",
"url": ""
}
],
"categories": [
{
"id": "category1",
"name": "Category 1",
"description" : "Readme for category 1"
}
]
}
10 changes: 5 additions & 5 deletions packages/backend/src/managers/applicationManager.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,21 +25,21 @@ interface DownloadModelResult {
}

export class ApplicationManager {
readonly homeDirectory: string; // todo: make configurable
readonly appUserDirectory: string; // todo: make configurable

constructor(
private git: GitManager,
private recipeStatusRegistry: RecipeStatusRegistry,
private extensionContext: ExtensionContext,
) {
this.homeDirectory = os.homedir();
this.appUserDirectory = path.join(os.homedir(), AI_STUDIO_FOLDER);
}

async pullApplication(recipe: Recipe, model: ModelInfo) {
// Create a TaskUtils object to help us
const taskUtil = new RecipeStatusUtils(recipe.id, this.recipeStatusRegistry);

const localFolder = path.join(this.homeDirectory, AI_STUDIO_FOLDER, recipe.id);
const localFolder = path.join(this.appUserDirectory, recipe.id);

// Adding checkout task
const checkoutTask: Task = {
Expand Down Expand Up @@ -218,7 +218,7 @@ export class ApplicationManager {
callback: (message: DownloadModelResult) => void,
destFileName?: string,
) {
const destDir = path.join(this.homeDirectory, AI_STUDIO_FOLDER, 'models', modelId);
const destDir = path.join(this.appUserDirectory, 'models', modelId);
if (!fs.existsSync(destDir)) {
fs.mkdirSync(destDir, { recursive: true });
}
Expand Down Expand Up @@ -269,7 +269,7 @@ export class ApplicationManager {
// todo: move somewhere else (dedicated to models)
getLocalModels(): LocalModelInfo[] {
const result: LocalModelInfo[] = [];
const modelsDir = path.join(this.homeDirectory, AI_STUDIO_FOLDER, 'models');
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) {
Expand Down
113 changes: 77 additions & 36 deletions packages/backend/src/studio-api-impl.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,63 +18,104 @@

/* eslint-disable @typescript-eslint/no-explicit-any */

import { expect, test, vi } from 'vitest';
import { beforeEach, describe, expect, test, vi } from 'vitest';
import content from './ai-test.json';
import userContent from './ai-user-test.json';
lstocchi marked this conversation as resolved.
Show resolved Hide resolved
import type { ApplicationManager } from './managers/applicationManager';
import type { RecipeStatusRegistry } from './registries/RecipeStatusRegistry';
import { StudioApiImpl } from './studio-api-impl';
import type { PlayGroundManager } from './playground';
import type { TaskRegistry } from './registries/TaskRegistry';

import * as fs from 'node:fs';

vi.mock('./ai.json', () => {
return {
default: content,
};
});

const studioApiImpl = new StudioApiImpl(
{} as unknown as ApplicationManager,
{} as unknown as RecipeStatusRegistry,
{} as unknown as TaskRegistry,
{} as unknown as PlayGroundManager,
);
let studioApiImpl: StudioApiImpl;

test('expect correct model is returned with valid id', async () => {
const model = await studioApiImpl.getModelById('llama-2-7b-chat.Q5_K_S');
expect(model).toBeDefined();
expect(model.name).toEqual('Llama-2-7B-Chat-GGUF');
expect(model.registry).toEqual('Hugging Face');
expect(model.url).toEqual(
'https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q5_K_S.gguf',
beforeEach(async () => {
studioApiImpl = new StudioApiImpl(
{
appUserDirectory: '.',
} as unknown as ApplicationManager,
{} as unknown as RecipeStatusRegistry,
{} as unknown as TaskRegistry,
{} as unknown as PlayGroundManager,
);
vi.resetAllMocks();
vi.mock('node:fs');
});

test('expect error if id does not correspond to any model', async () => {
await expect(() => studioApiImpl.getModelById('unknown')).rejects.toThrowError('No model found having id unknown');
describe('invalid user catalog', () => {
beforeEach(async () => {
vi.spyOn(fs.promises, 'readFile').mockResolvedValue('invalid json');
await studioApiImpl.loadCatalog();
});

test('expect correct model is returned with valid id', async () => {
const model = await studioApiImpl.getModelById('llama-2-7b-chat.Q5_K_S');
expect(model).toBeDefined();
expect(model.name).toEqual('Llama-2-7B-Chat-GGUF');
expect(model.registry).toEqual('Hugging Face');
expect(model.url).toEqual(
'https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q5_K_S.gguf',
);
});

test('expect error if id does not correspond to any model', async () => {
await expect(() => studioApiImpl.getModelById('unknown')).rejects.toThrowError('No model found having id unknown');
});

test('expect array of models based on list of ids', async () => {
const models = await studioApiImpl.getModelsByIds(['llama-2-7b-chat.Q5_K_S', 'albedobase-xl-1.3']);
expect(models).toBeDefined();
expect(models.length).toBe(2);
expect(models[0].name).toEqual('Llama-2-7B-Chat-GGUF');
expect(models[0].registry).toEqual('Hugging Face');
expect(models[0].url).toEqual(
'https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q5_K_S.gguf',
);
expect(models[1].name).toEqual('AlbedoBase XL 1.3');
expect(models[1].registry).toEqual('Civital');
expect(models[1].url).toEqual('');
});

test('expect empty array if input list is empty', async () => {
const models = await studioApiImpl.getModelsByIds([]);
expect(models).toBeDefined();
expect(models.length).toBe(0);
});

test('expect empty array if input list has ids that are not in the catalog', async () => {
const models = await studioApiImpl.getModelsByIds(['1', '2']);
expect(models).toBeDefined();
expect(models.length).toBe(0);
});
});

test('expect array of models based on list of ids', async () => {
const models = await studioApiImpl.getModelsByIds(['llama-2-7b-chat.Q5_K_S', 'albedobase-xl-1.3']);
expect(models).toBeDefined();
expect(models.length).toBe(2);
expect(models[0].name).toEqual('Llama-2-7B-Chat-GGUF');
expect(models[0].registry).toEqual('Hugging Face');
expect(models[0].url).toEqual(
test('expect correct model is returned from default catalog with valid id when no user catalog exists', async () => {
vi.spyOn(fs, 'existsSync').mockReturnValue(false);
await studioApiImpl.loadCatalog();
const model = await studioApiImpl.getModelById('llama-2-7b-chat.Q5_K_S');
expect(model).toBeDefined();
expect(model.name).toEqual('Llama-2-7B-Chat-GGUF');
expect(model.registry).toEqual('Hugging Face');
expect(model.url).toEqual(
'https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q5_K_S.gguf',
);
expect(models[1].name).toEqual('AlbedoBase XL 1.3');
expect(models[1].registry).toEqual('Civital');
expect(models[1].url).toEqual('');
});

test('expect empty array if input list is empty', async () => {
const models = await studioApiImpl.getModelsByIds([]);
expect(models).toBeDefined();
expect(models.length).toBe(0);
});

test('expect empty array if input list has ids that are not in the catalog', async () => {
const models = await studioApiImpl.getModelsByIds(['1', '2']);
expect(models).toBeDefined();
expect(models.length).toBe(0);
test('expect correct model is returned with valid id when the user catalog is valid', async () => {
vi.spyOn(fs, 'existsSync').mockReturnValue(true);
vi.spyOn(fs.promises, 'readFile').mockResolvedValue(JSON.stringify(userContent));
await studioApiImpl.loadCatalog();
const model = await studioApiImpl.getModelById('model1');
expect(model).toBeDefined();
expect(model.name).toEqual('Model 1');
expect(model.registry).toEqual('Hugging Face');
expect(model.url).toEqual('https://model1.example.com');
});
61 changes: 44 additions & 17 deletions packages/backend/src/studio-api-impl.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,28 +1,61 @@
import type { StudioAPI } from '@shared/src/StudioAPI';
import type { Category } from '@shared/src/models/ICategory';
import type { Recipe } from '@shared/src/models/IRecipe';
import content from './ai.json';
import defaultCatalog from './ai.json';
import type { ApplicationManager } from './managers/applicationManager';
import { AI_STUDIO_FOLDER } from './managers/applicationManager';
import type { RecipeStatusRegistry } from './registries/RecipeStatusRegistry';
import type { RecipeStatus } from '@shared/src/models/IRecipeStatus';
import type { ModelInfo } from '@shared/src/models/IModelInfo';
import type { TaskRegistry } from './registries/TaskRegistry';
import type { Task } from '@shared/src/models/ITask';
import * as path from 'node:path';
import type { PlayGroundManager } from './playground';
import * as podmanDesktopApi from '@podman-desktop/api';
import type { QueryState } from '@shared/src/models/IPlaygroundQueryState';
import type { Catalog } from '@shared/src/models/ICatalog';

import * as path from 'node:path';
import * as fs from 'node:fs';

export const RECENT_CATEGORY_ID = 'recent-category';

export class StudioApiImpl implements StudioAPI {
private catalog: Catalog;

constructor(
private applicationManager: ApplicationManager,
private recipeStatusRegistry: RecipeStatusRegistry,
private taskRegistry: TaskRegistry,
private playgroundManager: PlayGroundManager,
) {}
) {
// We start with an empty catalog, for the methods to work before the catalog is loaded
this.catalog = {
categories: [],
models: [],
recipes: [],
};
}

async loadCatalog() {
const catalogPath = path.resolve(this.applicationManager.appUserDirectory, 'catalog.json');
try {
if (!fs.existsSync(catalogPath)) {
this.setCatalog(defaultCatalog);
return;
}
// TODO(feloy): watch catalog file and update catalog with new content
const data = await fs.promises.readFile(catalogPath, 'utf-8');
const cat = JSON.parse(data) as Catalog;
this.setCatalog(cat);
} catch (err: unknown) {
console.error('unable to read catalog file, reverting to default catalog', err);
this.setCatalog(defaultCatalog);
}
}

setCatalog(newCatalog: Catalog) {
// TODO(feloy): send message to frontend with new catalog
this.catalog = newCatalog;
}

async openURL(url: string): Promise<boolean> {
return await podmanDesktopApi.env.openExternal(podmanDesktopApi.Uri.parse(url));
Expand All @@ -41,31 +74,31 @@ export class StudioApiImpl implements StudioAPI {
}

async getCategories(): Promise<Category[]> {
return content.categories;
return this.catalog.categories;
}
Comment on lines 76 to 78
Copy link
Collaborator

Choose a reason for hiding this comment

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

not sure why it's a Promise ?

Copy link
Collaborator

Choose a reason for hiding this comment

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

(same remark for almost all methods)

as catalog is a model, I'm not sure about the usage of async/promises for reading from it

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I think they are async to respect the interface (class) defined in shared/src/StudioAPI.ts, where they are returning a Promise, and from the client side, they are really returning a Promise, due to the wrapper sending a message.

Copy link
Collaborator

Choose a reason for hiding this comment

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

ok got it.

I thought that client was getting the model object asynchronously but then it was just fields

client:

const catalog = await rpc.get(CatalogAPI.class).getCatalog();
const categories = catalog.categories;
const something = catalog.something;
const else = catalog.else;

and not a call each time we need something like

const categories = await rpc....getCategories();
const something = await rpc....getSomething();
const else = await rpc....getElse();

Copy link
Contributor Author

Choose a reason for hiding this comment

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

What I expect to do with the next PR is to have the backend (this class) post a message to the frontend with the new content of the catalog every time it changes (in newCatalog()). On the frontend, I wonder what would be the best approach: having a single catalog store, or 3 (or more) stores: recipes, models, categories.

Or, even, having 3 different messages (newRecipes, newModels, newCategories)? But inihs case we would have to check which ones have been modified, to push only the modified ones.

WDYT?

Copy link
Collaborator

Choose a reason for hiding this comment

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

I'm fine with all options (as the traffic/latency is less a problem on a local client)

I think what would matter here, is that webview is recreated each time we click on the webview icon (and webview can read/store a state)

in that case, instead of asking again the catalog, it could start from the cache/state to display faster some items

so it would be interesting to have frontend being able to refresh/restore its UI as fast as possible (so avoiding lot of computes on backend if nothing has changed)


async getRecipesByCategory(categoryId: string): Promise<Recipe[]> {
if (categoryId === RECENT_CATEGORY_ID) return this.getRecentRecipes();

return content.recipes.filter(recipe => recipe.categories.includes(categoryId));
return this.catalog.recipes.filter(recipe => recipe.categories.includes(categoryId));
}

async getRecipeById(recipeId: string): Promise<Recipe> {
const recipe = (content.recipes as Recipe[]).find(recipe => recipe.id === recipeId);
const recipe = (this.catalog.recipes as Recipe[]).find(recipe => recipe.id === recipeId);
if (recipe) return recipe;
throw new Error('Not found');
}

async getModelById(modelId: string): Promise<ModelInfo> {
const model = content.models.find(m => modelId === m.id);
const model = this.catalog.models.find(m => modelId === m.id);
if (!model) {
throw new Error(`No model found having id ${modelId}`);
}
return model;
}

async getModelsByIds(ids: string[]): Promise<ModelInfo[]> {
return content.models.filter(m => ids.includes(m.id)) ?? [];
return this.catalog.models.filter(m => ids.includes(m.id)) ?? [];
}

// eslint-disable-next-line @typescript-eslint/no-unused-vars
Expand All @@ -91,7 +124,7 @@ export class StudioApiImpl implements StudioAPI {
async getLocalModels(): Promise<ModelInfo[]> {
const local = this.applicationManager.getLocalModels();
const localIds = local.map(l => l.id);
return content.models.filter(m => localIds.includes(m.id));
return this.catalog.models.filter(m => localIds.includes(m.id));
}

async getTasksByLabel(label: string): Promise<Task[]> {
Expand All @@ -104,13 +137,7 @@ export class StudioApiImpl implements StudioAPI {
throw new Error('model not found');
}

const modelPath = path.resolve(
this.applicationManager.homeDirectory,
AI_STUDIO_FOLDER,
'models',
modelId,
localModelInfo[0].file,
);
const modelPath = path.resolve(this.applicationManager.appUserDirectory, 'models', modelId, localModelInfo[0].file);

await this.playgroundManager.startPlayground(modelId, modelPath);
}
Expand Down
1 change: 1 addition & 0 deletions packages/backend/src/studio.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ export class Studio {
const applicationManager = new ApplicationManager(gitManager, recipeStatusRegistry, this.#extensionContext);
this.playgroundManager = new PlayGroundManager(this.#panel.webview);
this.studioApi = new StudioApiImpl(applicationManager, recipeStatusRegistry, taskRegistry, this.playgroundManager);
await this.studioApi.loadCatalog();
// Register the instance
this.rpcExtension.registerInstance<StudioApiImpl>(StudioApiImpl, this.studioApi);
}
Expand Down
9 changes: 9 additions & 0 deletions packages/shared/src/models/ICatalog.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import type { Category } from './ICategory';
import type { ModelInfo } from './IModelInfo';
import type { Recipe } from './IRecipe';

export interface Catalog {
recipes: Recipe[];
models: ModelInfo[];
categories: Category[];
}