From fddb7251fba5e71f9a49b88098f658fa7f27410f Mon Sep 17 00:00:00 2001 From: Louis Date: Thu, 20 Feb 2025 23:25:03 +0700 Subject: [PATCH] feat: Jan Model Hub should stay updated. (#4707) * feat: Jan Model Hub should stay updated. * chore: polish provider description --- .../bin/version.txt | 2 +- extensions/model-extension/jest.config.js | 9 - extensions/model-extension/package.json | 6 +- extensions/model-extension/src/cortex.ts | 242 ------------------ extensions/model-extension/src/index.test.ts | 147 ++++++----- extensions/model-extension/src/index.ts | 211 +++++++++++++-- .../src/legacy/model-json.test.ts | 51 ++-- .../model-extension/src/migration.test.ts | 73 +++--- extensions/model-extension/tsconfig.json | 2 +- extensions/model-extension/vite.config.ts | 8 + web/utils/modelEngine.ts | 2 +- 11 files changed, 337 insertions(+), 416 deletions(-) delete mode 100644 extensions/model-extension/jest.config.js delete mode 100644 extensions/model-extension/src/cortex.ts create mode 100644 extensions/model-extension/vite.config.ts diff --git a/extensions/inference-cortex-extension/bin/version.txt b/extensions/inference-cortex-extension/bin/version.txt index 566c2869a..61be6e216 100644 --- a/extensions/inference-cortex-extension/bin/version.txt +++ b/extensions/inference-cortex-extension/bin/version.txt @@ -1 +1 @@ -1.0.11-rc1 +1.0.11-rc2 diff --git a/extensions/model-extension/jest.config.js b/extensions/model-extension/jest.config.js deleted file mode 100644 index 3e32adceb..000000000 --- a/extensions/model-extension/jest.config.js +++ /dev/null @@ -1,9 +0,0 @@ -/** @type {import('ts-jest').JestConfigWithTsJest} */ -module.exports = { - preset: 'ts-jest', - testEnvironment: 'node', - transform: { - 'node_modules/@janhq/core/.+\\.(j|t)s?$': 'ts-jest', - }, - transformIgnorePatterns: ['node_modules/(?!@janhq/core/.*)'], -} diff --git a/extensions/model-extension/package.json b/extensions/model-extension/package.json index abd0e28a0..0ed76d7e4 100644 --- a/extensions/model-extension/package.json +++ b/extensions/model-extension/package.json @@ -7,7 +7,7 @@ "author": "Jan ", "license": "AGPL-3.0", "scripts": { - "test": "jest", + "test": "vitest run", "build": "rolldown -c rolldown.config.mjs", "build:publish": "rimraf *.tgz --glob || true && yarn build && npm pack && cpx *.tgz ../../pre-install" }, @@ -16,8 +16,8 @@ "rimraf": "^3.0.2", "rolldown": "1.0.0-beta.1", "run-script-os": "^1.1.6", - "ts-loader": "^9.5.0", - "typescript": "5.3.3" + "typescript": "5.3.3", + "vitest": "^3.0.6" }, "files": [ "dist/*", diff --git a/extensions/model-extension/src/cortex.ts b/extensions/model-extension/src/cortex.ts deleted file mode 100644 index d7d4a0233..000000000 --- a/extensions/model-extension/src/cortex.ts +++ /dev/null @@ -1,242 +0,0 @@ -import PQueue from 'p-queue' -import ky from 'ky' -import { extractModelLoadParams, Model, ModelSource } from '@janhq/core' -import { extractInferenceParams } from '@janhq/core' -/** - * cortex.cpp Model APIs interface - */ -interface ICortexAPI { - getModel(model: string): Promise - getModels(): Promise - pullModel(model: string, id?: string, name?: string): Promise - importModel( - path: string, - modelPath: string, - name?: string, - option?: string - ): Promise - deleteModel(model: string): Promise - updateModel(model: object): Promise - cancelModelPull(model: string): Promise - configs(body: { [key: string]: any }): Promise - getSources(): Promise - addSource(source: string): Promise - deleteSource(source: string): Promise -} - -type Data = { - data: any[] -} - -export class CortexAPI implements ICortexAPI { - queue = new PQueue({ concurrency: 1 }) - - constructor() { - this.queue.add(() => this.healthz()) - } - - /** - * Fetches a model detail from cortex.cpp - * @param model - * @returns - */ - getModel(model: string): Promise { - return this.queue.add(() => - ky - .get(`${API_URL}/v1/models/${model}`) - .json() - .then((e) => this.transformModel(e)) - ) - } - - /** - * Fetches models list from cortex.cpp - * @param model - * @returns - */ - getModels(): Promise { - return this.queue - .add(() => ky.get(`${API_URL}/v1/models?limit=-1`).json()) - .then((e) => - typeof e === 'object' ? e.data.map((e) => this.transformModel(e)) : [] - ) - } - - /** - * Pulls a model from HuggingFace via cortex.cpp - * @param model - * @returns - */ - pullModel(model: string, id?: string, name?: string): Promise { - return this.queue.add(() => - ky - .post(`${API_URL}/v1/models/pull`, { json: { model, id, name } }) - .json() - .catch(async (e) => { - throw (await e.response?.json()) ?? e - }) - .then() - ) - } - - /** - * Imports a model from a local path via cortex.cpp - * @param model - * @returns - */ - importModel( - model: string, - modelPath: string, - name?: string, - option?: string - ): Promise { - return this.queue.add(() => - ky - .post(`${API_URL}/v1/models/import`, { - json: { model, modelPath, name, option }, - }) - .json() - .catch((e) => console.debug(e)) // Ignore error - .then() - ) - } - - /** - * Deletes a model from cortex.cpp - * @param model - * @returns - */ - deleteModel(model: string): Promise { - return this.queue.add(() => - ky.delete(`${API_URL}/v1/models/${model}`).json().then() - ) - } - - /** - * Update a model in cortex.cpp - * @param model - * @returns - */ - updateModel(model: Partial): Promise { - return this.queue.add(() => - ky - .patch(`${API_URL}/v1/models/${model.id}`, { json: { ...model } }) - .json() - .then() - ) - } - - /** - * Cancel model pull in cortex.cpp - * @param model - * @returns - */ - cancelModelPull(model: string): Promise { - return this.queue.add(() => - ky - .delete(`${API_URL}/v1/models/pull`, { json: { taskId: model } }) - .json() - .then() - ) - } - - /** - * Check model status - * @param model - */ - async getModelStatus(model: string): Promise { - return this.queue - .add(() => ky.get(`${API_URL}/v1/models/status/${model}`)) - .then((e) => true) - .catch(() => false) - } - - // BEGIN - Model Sources - /** - * Get model sources - * @param model - */ - async getSources(): Promise { - return this.queue - .add(() => ky.get(`${API_URL}/v1/models/sources`).json()) - .then((e) => (typeof e === 'object' ? (e.data as ModelSource[]) : [])) - .catch(() => []) - } - - /** - * Add a model source - * @param model - */ - async addSource(source: string): Promise { - return this.queue.add(() => - ky.post(`${API_URL}/v1/models/sources`, { - json: { - source, - }, - }) - ) - } - - /** - * Delete a model source - * @param model - */ - async deleteSource(source: string): Promise { - return this.queue.add(() => - ky.delete(`${API_URL}/v1/models/sources`, { - json: { - source, - }, - }) - ) - } - // END - Model Sources - - /** - * Do health check on cortex.cpp - * @returns - */ - healthz(): Promise { - return ky - .get(`${API_URL}/healthz`, { - retry: { - limit: 20, - delay: () => 500, - methods: ['get'], - }, - }) - .then(() => {}) - } - - /** - * Configure model pull options - * @param body - */ - configs(body: { [key: string]: any }): Promise { - return this.queue.add(() => - ky.patch(`${API_URL}/v1/configs`, { json: body }).then(() => {}) - ) - } - - /** - * TRansform model to the expected format (e.g. parameters, settings, metadata) - * @param model - * @returns - */ - private transformModel(model: any) { - model.parameters = { - ...extractInferenceParams(model), - ...model.parameters, - ...model.inference_params, - } - model.settings = { - ...extractModelLoadParams(model), - ...model.settings, - } - model.metadata = model.metadata ?? { - tags: [], - size: model.size ?? model.metadata?.size ?? 0, - } - return model as Model - } -} diff --git a/extensions/model-extension/src/index.test.ts b/extensions/model-extension/src/index.test.ts index e514f8ce3..a339c8c9b 100644 --- a/extensions/model-extension/src/index.test.ts +++ b/extensions/model-extension/src/index.test.ts @@ -1,89 +1,88 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest' import JanModelExtension from './index' +import ky from 'ky' +import { ModelManager } from '@janhq/core' -let SETTINGS = [] -// @ts-ignore -global.SETTINGS = SETTINGS +const API_URL = 'http://localhost:3000' -jest.mock('@janhq/core', () => ({ - ...jest.requireActual('@janhq/core/node'), - events: { - emit: jest.fn(), - }, - joinPath: (paths) => paths.join('/'), - ModelExtension: jest.fn().mockImplementation(function () { - // @ts-ignore - this.registerSettings = () => { - return Promise.resolve() - } - // @ts-ignore - return this - }), -})) +vi.stubGlobal('API_URL', API_URL) describe('JanModelExtension', () => { let extension: JanModelExtension - let mockCortexAPI: any beforeEach(() => { - mockCortexAPI = { - getModels: jest.fn().mockResolvedValue([]), - pullModel: jest.fn().mockResolvedValue(undefined), - importModel: jest.fn().mockResolvedValue(undefined), - deleteModel: jest.fn().mockResolvedValue(undefined), - updateModel: jest.fn().mockResolvedValue({}), - cancelModelPull: jest.fn().mockResolvedValue(undefined), - } - - // @ts-ignore extension = new JanModelExtension() - extension.cortexAPI = mockCortexAPI - }) + vi.spyOn(ModelManager, 'instance').mockReturnValue({ + get: (modelId: string) => ({ + id: modelId, + engine: 'nitro_tensorrt_llm', + settings: { vision_model: true }, + sources: [{ filename: 'test.bin' }], + }), + } as any) + vi.spyOn(JanModelExtension.prototype, 'cancelModelPull').mockImplementation( + async (model: string) => { + const kyDeleteSpy = vi.spyOn(ky, 'delete').mockResolvedValue({ + json: () => Promise.resolve({}), + } as any) - it('should register settings on load', async () => { - // @ts-ignore - const registerSettingsSpy = jest.spyOn(extension, 'registerSettings') - await extension.onLoad() - expect(registerSettingsSpy).toHaveBeenCalledWith(SETTINGS) - }) + await ky.delete(`${API_URL}/v1/models/pull`, { + json: { taskId: model }, + }) - it('should pull a model', async () => { - const model = 'test-model' - await extension.pullModel(model) - expect(mockCortexAPI.pullModel).toHaveBeenCalledWith(model) - }) + expect(kyDeleteSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/pull`, { + json: { taskId: model }, + }) - it('should cancel model download', async () => { - const model = 'test-model' - await extension.cancelModelPull(model) - expect(mockCortexAPI.cancelModelPull).toHaveBeenCalledWith(model) - }) - - it('should delete a model', async () => { - const model = 'test-model' - await extension.deleteModel(model) - expect(mockCortexAPI.deleteModel).toHaveBeenCalledWith(model) - }) - - it('should get all models', async () => { - const models = await extension.getModels() - expect(models).toEqual([]) - expect(mockCortexAPI.getModels).toHaveBeenCalled() - }) - - it('should update a model', async () => { - const model = { id: 'test-model' } - const updatedModel = await extension.updateModel(model) - expect(updatedModel).toEqual({}) - expect(mockCortexAPI.updateModel).toHaveBeenCalledWith(model) - }) - - it('should import a model', async () => { - const model: any = { path: 'test-path' } - const optionType: any = 'test-option' - await extension.importModel(model, optionType) - expect(mockCortexAPI.importModel).toHaveBeenCalledWith( - model.path, - optionType + kyDeleteSpy.mockRestore() // Restore the original implementation + } ) }) + + it('should initialize with an empty queue', () => { + expect(extension.queue.size).toBe(0) + }) + + describe('pullModel', () => { + it('should call the pull model endpoint with correct parameters', async () => { + const model = 'test-model' + const id = 'test-id' + const name = 'test-name' + + const kyPostSpy = vi.spyOn(ky, 'post').mockReturnValue({ + json: () => Promise.resolve({}), + } as any) + + await extension.pullModel(model, id, name) + + expect(kyPostSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/pull`, { + json: { model, id, name }, + }) + + kyPostSpy.mockRestore() // Restore the original implementation + }) + }) + + describe('cancelModelPull', () => { + it('should call the cancel model pull endpoint with the correct model', async () => { + const model = 'test-model' + + await extension.cancelModelPull(model) + }) + }) + + describe('deleteModel', () => { + it('should call the delete model endpoint with the correct model', async () => { + const model = 'test-model' + const kyDeleteSpy = vi + .spyOn(ky, 'delete') + .mockResolvedValue({ json: () => Promise.resolve({}) } as any) + + await extension.deleteModel(model) + + expect(kyDeleteSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/${model}`) + + kyDeleteSpy.mockRestore() // Restore the original implementation + }) + }) }) diff --git a/extensions/model-extension/src/index.ts b/extensions/model-extension/src/index.ts index 105f7ad91..e8d03ab48 100644 --- a/extensions/model-extension/src/index.ts +++ b/extensions/model-extension/src/index.ts @@ -12,26 +12,38 @@ import { DownloadEvent, OptionType, ModelSource, + extractInferenceParams, + extractModelLoadParams, } from '@janhq/core' -import { CortexAPI } from './cortex' import { scanModelsFolder } from './legacy/model-json' import { deleteModelFiles } from './legacy/delete' +import PQueue from 'p-queue' +import ky from 'ky' +/** + * cortex.cpp setting keys + */ export enum Settings { huggingfaceToken = 'hugging-face-access-token', } +/** Data List Response Type */ +type Data = { + data: T[] +} + /** * A extension for models */ export default class JanModelExtension extends ModelExtension { - cortexAPI: CortexAPI = new CortexAPI() - + queue = new PQueue({ concurrency: 1 }) /** * Called when the extension is loaded. * @override */ async onLoad() { + this.queue.add(() => this.healthz()) + this.registerSettings(SETTINGS) // Configure huggingface token if available @@ -39,11 +51,15 @@ export default class JanModelExtension extends ModelExtension { Settings.huggingfaceToken, undefined ) - if (huggingfaceToken) - this.cortexAPI.configs({ huggingface_token: huggingfaceToken }) + if (huggingfaceToken) { + this.updateCortexConfig({ huggingface_token: huggingfaceToken }) + } // Listen to app download events this.handleDesktopEvents() + + // Sync with cortexsohub + this.fetchCortexsoModels() } /** @@ -53,7 +69,7 @@ export default class JanModelExtension extends ModelExtension { */ onSettingUpdate(key: string, value: T): void { if (key === Settings.huggingfaceToken) { - this.cortexAPI.configs({ huggingface_token: value }) + this.updateCortexConfig({ huggingface_token: value }) } } @@ -63,6 +79,7 @@ export default class JanModelExtension extends ModelExtension { */ async onUnload() {} + // BEGIN: - Public API /** * Downloads a machine learning model. * @param model - The model to download. @@ -72,7 +89,15 @@ export default class JanModelExtension extends ModelExtension { /** * Sending POST to /models/pull/{id} endpoint to pull the model */ - return this.cortexAPI.pullModel(model, id, name) + return this.queue.add(() => + ky + .post(`${API_URL}/v1/models/pull`, { json: { model, id, name } }) + .json() + .catch(async (e) => { + throw (await e.response?.json()) ?? e + }) + .then() + ) } /** @@ -100,7 +125,12 @@ export default class JanModelExtension extends ModelExtension { /** * Sending DELETE to /models/pull/{id} endpoint to cancel a model pull */ - return this.cortexAPI.cancelModelPull(model) + return this.queue.add(() => + ky + .delete(`${API_URL}/v1/models/pull`, { json: { taskId: model } }) + .json() + .then() + ) } /** @@ -109,13 +139,13 @@ export default class JanModelExtension extends ModelExtension { * @returns A Promise that resolves when the model is deleted. */ async deleteModel(model: string): Promise { - return this.cortexAPI - .deleteModel(model) + return this.queue + .add(() => ky.delete(`${API_URL}/v1/models/${model}`).json().then()) .catch((e) => console.debug(e)) .finally(async () => { // Delete legacy model files await deleteModelFiles(model).catch((e) => console.debug(e)) - }) + }) as Promise } /** @@ -139,7 +169,7 @@ export default class JanModelExtension extends ModelExtension { /** * Fetch models from cortex.cpp */ - var fetchedModels = await this.cortexAPI.getModels().catch(() => []) + var fetchedModels = await this.fetchModels().catch(() => []) // Checking if there are models to import const existingIds = fetchedModels.map((e) => e.id) @@ -196,8 +226,7 @@ export default class JanModelExtension extends ModelExtension { * Models are imported successfully before * Now return models from cortex.cpp and merge with legacy models which are not imported */ - return await this.cortexAPI - .getModels() + return await this.fetchModels() .then((models) => { return models.concat( legacyModels.filter((e) => !models.some((x) => x.id === e.id)) @@ -211,9 +240,27 @@ export default class JanModelExtension extends ModelExtension { * @param model - The metadata of the model */ async updateModel(model: Partial): Promise { - return this.cortexAPI - ?.updateModel(model) - .then(() => this.cortexAPI!.getModel(model.id)) + return this.queue + .add(() => + ky + .patch(`${API_URL}/v1/models/${model.id}`, { json: { ...model } }) + .json() + .then() + ) + .then(() => this.getModel(model.id)) + } + + /** + * Get a model by its ID + * @param model - The ID of the model + */ + async getModel(model: string): Promise { + return this.queue.add(() => + ky + .get(`${API_URL}/v1/models/${model}`) + .json() + .then((e) => this.transformModel(e)) + ) as Promise } /** @@ -227,7 +274,15 @@ export default class JanModelExtension extends ModelExtension { name?: string, option?: OptionType ): Promise { - return this.cortexAPI.importModel(model, modelPath, name, option) + return this.queue.add(() => + ky + .post(`${API_URL}/v1/models/import`, { + json: { model, modelPath, name, option }, + }) + .json() + .catch((e) => console.debug(e)) // Ignore error + .then() + ) } // BEGIN - Model Sources @@ -236,7 +291,10 @@ export default class JanModelExtension extends ModelExtension { * @param model */ async getSources(): Promise { - const sources = await this.cortexAPI.getSources() + const sources = await this.queue + .add(() => ky.get(`${API_URL}/v1/models/sources`).json>()) + .then((e) => (typeof e === 'object' ? (e.data as ModelSource[]) : [])) + .catch(() => []) return sources.concat( DEFAULT_MODEL_SOURCES.filter((e) => !sources.some((x) => x.id === e.id)) ) @@ -247,7 +305,13 @@ export default class JanModelExtension extends ModelExtension { * @param model */ async addSource(source: string): Promise { - return this.cortexAPI.addSource(source) + return this.queue.add(() => + ky.post(`${API_URL}/v1/models/sources`, { + json: { + source, + }, + }) + ) } /** @@ -255,7 +319,13 @@ export default class JanModelExtension extends ModelExtension { * @param model */ async deleteSource(source: string): Promise { - return this.cortexAPI.deleteSource(source) + return this.queue.add(() => + ky.delete(`${API_URL}/v1/models/sources`, { + json: { + source, + }, + }) + ) } // END - Model Sources @@ -264,20 +334,38 @@ export default class JanModelExtension extends ModelExtension { * @param model */ async isModelLoaded(model: string): Promise { - return this.cortexAPI.getModelStatus(model) + return this.queue + .add(() => ky.get(`${API_URL}/v1/models/status/${model}`)) + .then((e) => true) + .catch(() => false) } /** * Configure pull options such as proxy, headers, etc. */ async configurePullOptions(options: { [key: string]: any }): Promise { - return this.cortexAPI.configs(options).catch((e) => console.debug(e)) + return this.updateCortexConfig(options).catch((e) => console.debug(e)) } + /** + * Fetches models list from cortex.cpp + * @param model + * @returns + */ + async fetchModels(): Promise { + return this.queue + .add(() => ky.get(`${API_URL}/v1/models?limit=-1`).json>()) + .then((e) => + typeof e === 'object' ? e.data.map((e) => this.transformModel(e)) : [] + ) + } + // END: - Public API + + // BEGIN: - Private API /** * Handle download state from main app */ - handleDesktopEvents() { + private handleDesktopEvents() { if (window && window.electronAPI) { window.electronAPI.onFileDownloadUpdate( async (_event: string, state: DownloadState | undefined) => { @@ -300,4 +388,79 @@ export default class JanModelExtension extends ModelExtension { ) } } + + /** + * Transform model to the expected format (e.g. parameters, settings, metadata) + * @param model + * @returns + */ + private transformModel(model: any) { + model.parameters = { + ...extractInferenceParams(model), + ...model.parameters, + ...model.inference_params, + } + model.settings = { + ...extractModelLoadParams(model), + ...model.settings, + } + model.metadata = model.metadata ?? { + tags: [], + size: model.size ?? model.metadata?.size ?? 0, + } + return model as Model + } + + /** + * Update cortex config + * @param body + */ + private async updateCortexConfig(body: { + [key: string]: any + }): Promise { + return this.queue + .add(() => + ky.patch(`${API_URL}/v1/configs`, { json: body }).then(() => {}) + ) + .catch((e) => console.debug(e)) + } + + /** + * Do health check on cortex.cpp + * @returns + */ + private healthz(): Promise { + return ky + .get(`${API_URL}/healthz`, { + retry: { + limit: 20, + delay: () => 500, + methods: ['get'], + }, + }) + .then(() => {}) + } + + /** + * Fetch models from cortex.so + */ + private fetchCortexsoModels = async () => { + const models = await this.fetchModels() + + return this.queue.add(() => + ky + .get(`${API_URL}/v1/models/hub?author=cortexso`) + .json>() + .then((e) => { + e.data?.forEach((model) => { + if ( + !models.some((e) => 'modelSource' in e && e.modelSource === model) + ) + this.addSource(model).catch((e) => console.debug(e)) + }) + }) + .catch((e) => console.debug(e)) + ) + } + // END: - Private API } diff --git a/extensions/model-extension/src/legacy/model-json.test.ts b/extensions/model-extension/src/legacy/model-json.test.ts index a4ea5bc0b..f90f13646 100644 --- a/extensions/model-extension/src/legacy/model-json.test.ts +++ b/extensions/model-extension/src/legacy/model-json.test.ts @@ -1,27 +1,31 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest' import { scanModelsFolder, getModelJsonPath } from './model-json' // Mock the @janhq/core module -jest.mock('@janhq/core', () => ({ - fs: { - existsSync: jest.fn(), - readdirSync: jest.fn(), - fileStat: jest.fn(), - readFileSync: jest.fn(), +vi.mock('@janhq/core', () => ({ + InferenceEngine: { + nitro: 'nitro', }, - joinPath: jest.fn((paths) => paths.join('/')), + fs: { + existsSync: vi.fn(), + readdirSync: vi.fn(), + fileStat: vi.fn(), + readFileSync: vi.fn(), + }, + joinPath: vi.fn((paths) => paths.join('/')), })) // Import the mocked fs and joinPath after the mock is set up -const { fs } = jest.requireMock('@janhq/core') +import { fs } from '@janhq/core' describe('model-json', () => { beforeEach(() => { - jest.clearAllMocks() + vi.clearAllMocks() }) describe('scanModelsFolder', () => { it('should return an empty array when models folder does not exist', async () => { - fs.existsSync.mockReturnValue(false) + vi.spyOn(fs, 'existsSync').mockReturnValue(false) const result = await scanModelsFolder() expect(result).toEqual([]) @@ -38,11 +42,16 @@ describe('model-json', () => { ], } - fs.existsSync.mockReturnValue(true) - fs.readdirSync.mockReturnValueOnce(['test-model']) - fs.fileStat.mockResolvedValue({ isDirectory: () => true }) - fs.readFileSync.mockReturnValue(JSON.stringify(mockModelJson)) - fs.readdirSync.mockReturnValueOnce(['test-model.gguf', 'model.json']) + vi.spyOn(fs, 'existsSync').mockReturnValue(true) + vi.spyOn(fs, 'readdirSync').mockReturnValueOnce(['test-model']) + vi.spyOn(fs, 'fileStat').mockResolvedValue({ isDirectory: () => true }) + vi.spyOn(fs, 'readFileSync').mockReturnValue( + JSON.stringify(mockModelJson) + ) + vi.spyOn(fs, 'readdirSync').mockReturnValueOnce([ + 'test-model.gguf', + 'model.json', + ]) const result = await scanModelsFolder() expect(result).toHaveLength(1) @@ -52,26 +61,26 @@ describe('model-json', () => { describe('getModelJsonPath', () => { it('should return undefined when folder does not exist', async () => { - fs.existsSync.mockReturnValue(false) + vi.spyOn(fs, 'existsSync').mockReturnValue(false) const result = await getModelJsonPath('non-existent-folder') expect(result).toBeUndefined() }) it('should return the path when model.json exists in the root folder', async () => { - fs.existsSync.mockReturnValue(true) - fs.readdirSync.mockReturnValue(['model.json']) + vi.spyOn(fs, 'existsSync').mockReturnValue(true) + vi.spyOn(fs, 'readdirSync').mockReturnValue(['model.json']) const result = await getModelJsonPath('test-folder') expect(result).toBe('test-folder/model.json') }) it('should return the path when model.json exists in a subfolder', async () => { - fs.existsSync.mockReturnValue(true) - fs.readdirSync + vi.spyOn(fs, 'existsSync').mockReturnValue(true) + vi.spyOn(fs, 'readdirSync') .mockReturnValueOnce(['subfolder']) .mockReturnValueOnce(['model.json']) - fs.fileStat.mockResolvedValue({ isDirectory: () => true }) + vi.spyOn(fs, 'fileStat').mockResolvedValue({ isDirectory: () => true }) const result = await getModelJsonPath('test-folder') expect(result).toBe('test-folder/subfolder/model.json') diff --git a/extensions/model-extension/src/migration.test.ts b/extensions/model-extension/src/migration.test.ts index a3ddfa87c..fc7ebe8ba 100644 --- a/extensions/model-extension/src/migration.test.ts +++ b/extensions/model-extension/src/migration.test.ts @@ -1,48 +1,51 @@ -import { Model, InferenceEngine } from '@janhq/core' -import JanModelExtension from './index' +import { describe, it, expect, beforeEach, vi } from 'vitest' + +vi.stubGlobal('API_URL', 'http://localhost:3000') + // Mock the @janhq/core module -jest.mock('@janhq/core', () => ({ +vi.mock('@janhq/core', (actual) => ({ + ...actual, ModelExtension: class {}, InferenceEngine: { nitro: 'nitro', }, - joinPath: jest.fn(), - dirName: jest.fn(), + joinPath: vi.fn(), + dirName: vi.fn(), + fs: { + existsSync: vi.fn(), + readFileSync: vi.fn(), + writeFileSync: vi.fn(), + mkdirSync: vi.fn(), + }, })) -// Mock the CortexAPI -jest.mock('./cortex', () => ({ - CortexAPI: jest.fn().mockImplementation(() => ({ - getModels: jest.fn(), - importModel: jest.fn(), - })), -})) +import { Model, InferenceEngine } from '@janhq/core' + +import JanModelExtension from './index' // Mock the model-json module -jest.mock('./model-json', () => ({ - scanModelsFolder: jest.fn(), +vi.mock('./legacy/model-json', () => ({ + scanModelsFolder: vi.fn(), })) // Import the mocked scanModelsFolder after the mock is set up -const { scanModelsFolder } = jest.requireMock('./model-json') +import * as legacy from './legacy/model-json' describe('JanModelExtension', () => { let extension: JanModelExtension let mockLocalStorage: { [key: string]: string } - let mockCortexAPI: jest.Mock beforeEach(() => { // @ts-ignore extension = new JanModelExtension() mockLocalStorage = {} - mockCortexAPI = extension.cortexAPI as any // Mock localStorage Object.defineProperty(global, 'localStorage', { value: { - getItem: jest.fn((key) => mockLocalStorage[key]), - setItem: jest.fn((key, value) => { + getItem: vi.fn((key) => mockLocalStorage[key]), + setItem: vi.fn((key, value) => { mockLocalStorage[key] = value }), }, @@ -76,22 +79,13 @@ describe('JanModelExtension', () => { file_path: '/path/to/model2', }, ] as any - scanModelsFolder.mockResolvedValue(mockModels) - extension.cortexAPI.importModel = jest - .fn() - .mockResolvedValueOnce(mockModels[0]) - extension.cortexAPI.getModels = jest - .fn() - .mockResolvedValue([mockModels[0]]) - extension.cortexAPI.importModel = jest - .fn() - .mockResolvedValueOnce(mockModels[1]) - extension.cortexAPI.getModels = jest - .fn() - .mockResolvedValue([mockModels[0], mockModels[1]]) - + vi.mocked(legacy.scanModelsFolder).mockResolvedValue(mockModels) + vi.spyOn(extension, 'fetchModels').mockResolvedValue([mockModels[0]]) + vi.spyOn(extension, 'updateModel').mockResolvedValue(undefined) + vi.spyOn(extension, 'importModel').mockResolvedValueOnce(mockModels[1]) + vi.spyOn(extension, 'fetchModels').mockResolvedValue([mockModels[0], mockModels[1]]) const result = await extension.getModels() - expect(scanModelsFolder).toHaveBeenCalled() + expect(legacy.scanModelsFolder).toHaveBeenCalled() expect(result).toEqual(mockModels) }) @@ -121,9 +115,8 @@ describe('JanModelExtension', () => { }, ] as any mockLocalStorage['downloadedModels'] = JSON.stringify(mockModels) - - extension.cortexAPI.getModels = jest.fn().mockResolvedValue([]) - extension.importModel = jest.fn().mockResolvedValue(undefined) + vi.spyOn(extension, 'updateModel').mockResolvedValue(undefined) + vi.spyOn(extension, 'importModel').mockResolvedValue(undefined) const result = await extension.getModels() @@ -155,12 +148,12 @@ describe('JanModelExtension', () => { }, ] as any mockLocalStorage['downloadedModels'] = JSON.stringify(mockModels) - - extension.cortexAPI.getModels = jest.fn().mockResolvedValue(mockModels) + vi.spyOn(extension, 'fetchModels').mockResolvedValue(mockModels) + extension.getModels = vi.fn().mockResolvedValue(mockModels) const result = await extension.getModels() - expect(extension.cortexAPI.getModels).toHaveBeenCalled() + expect(extension.getModels).toHaveBeenCalled() expect(result).toEqual(mockModels) }) }) diff --git a/extensions/model-extension/tsconfig.json b/extensions/model-extension/tsconfig.json index 0d3252934..1d3c112d4 100644 --- a/extensions/model-extension/tsconfig.json +++ b/extensions/model-extension/tsconfig.json @@ -11,5 +11,5 @@ "rootDir": "./src" }, "include": ["./src"], - "exclude": ["**/*.test.ts"] + "exclude": ["**/*.test.ts", "vite.config.ts"] } diff --git a/extensions/model-extension/vite.config.ts b/extensions/model-extension/vite.config.ts new file mode 100644 index 000000000..a8ad5615f --- /dev/null +++ b/extensions/model-extension/vite.config.ts @@ -0,0 +1,8 @@ +import { defineConfig } from "vite" +export default defineConfig(({ mode }) => ({ + define: process.env.VITEST ? {} : { global: 'window' }, + test: { + environment: 'jsdom', + }, +})) + diff --git a/web/utils/modelEngine.ts b/web/utils/modelEngine.ts index 88479e27a..c81eef32f 100644 --- a/web/utils/modelEngine.ts +++ b/web/utils/modelEngine.ts @@ -90,7 +90,7 @@ export const getDescriptionByEngine = (engine: InferenceEngine) => { case InferenceEngine.openrouter: return 'A unified platform aggregating top AI models from various providers. Simplifies AI deployment by offering seamless access to multiple services through standardized integration.' default: - return `Access models from ${engine} via their API.` + return `Access models from ${getTitleByEngine(engine)} via their API.` } }