Nicholai/jan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Jan Model Hub should stay updated. (#4707)

Browse Source 
* feat: Jan Model Hub should stay updated.

* chore: polish provider description
This commit is contained in:
Louis 2025-02-20 23:25:03 +07:00 committed by GitHub
parent 839a00127d
commit fddb7251fb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 337 additions and 416 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
extensions/inference-cortex-extension/bin/version.txt
View File

@ -1 +1 @@
1.0.11-rc1 1.0.11-rc2

9
extensions/model-extension/jest.config.js
View File

@ -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/.*)'],
}

6
extensions/model-extension/package.json
View File

@ -7,7 +7,7 @@
"author": "Jan <service@jan.ai>", "author": "Jan <service@jan.ai>",
"license": "AGPL-3.0", "license": "AGPL-3.0",
"scripts": { "scripts": {
"test": "jest", "test": "vitest run",
"build": "rolldown -c rolldown.config.mjs", "build": "rolldown -c rolldown.config.mjs",
"build:publish": "rimraf *.tgz --glob || true && yarn build && npm pack && cpx *.tgz ../../pre-install" "build:publish": "rimraf *.tgz --glob || true && yarn build && npm pack && cpx *.tgz ../../pre-install"
}, },
@ -16,8 +16,8 @@
"rimraf": "^3.0.2", "rimraf": "^3.0.2",
"rolldown": "1.0.0-beta.1", "rolldown": "1.0.0-beta.1",
"run-script-os": "^1.1.6", "run-script-os": "^1.1.6",
"ts-loader": "^9.5.0", "typescript": "5.3.3",
"typescript": "5.3.3" "vitest": "^3.0.6"
}, },
"files": [ "files": [
"dist/*", "dist/*",

242
extensions/model-extension/src/cortex.ts
View File

@ -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<Model>
getModels(): Promise<Model[]>
pullModel(model: string, id?: string, name?: string): Promise<void>
importModel(
path: string,
modelPath: string,
name?: string,
option?: string
): Promise<void>
deleteModel(model: string): Promise<void>
updateModel(model: object): Promise<void>
cancelModelPull(model: string): Promise<void>
configs(body: { [key: string]: any }): Promise<void>
getSources(): Promise<ModelSource[]>
addSource(source: string): Promise<void>
deleteSource(source: string): Promise<void>
}
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<any> {
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<Model[]> {
return this.queue
.add(() => ky.get(`${API_URL}/v1/models?limit=-1`).json<Data>())
.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<void> {
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<void> {
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<void> {
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<Model>): Promise<void> {
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<void> {
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<boolean> {
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<ModelSource[]> {
return this.queue
.add(() => ky.get(`${API_URL}/v1/models/sources`).json<Data>())
.then((e) => (typeof e === 'object' ? (e.data as ModelSource[]) : []))
.catch(() => [])
}
/**
* Add a model source
* @param model
*/
async addSource(source: string): Promise<any> {
return this.queue.add(() =>
ky.post(`${API_URL}/v1/models/sources`, {
json: {
source,
},
})
)
}
/**
* Delete a model source
* @param model
*/
async deleteSource(source: string): Promise<any> {
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<void> {
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<void> {
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
}
}

145
extensions/model-extension/src/index.test.ts
View File

@ -1,89 +1,88 @@
import { describe, it, expect, beforeEach, vi } from 'vitest'
import JanModelExtension from './index' import JanModelExtension from './index'
import ky from 'ky'
import { ModelManager } from '@janhq/core'
let SETTINGS = [] const API_URL = 'http://localhost:3000'
// @ts-ignore
global.SETTINGS = SETTINGS
jest.mock('@janhq/core', () => ({ vi.stubGlobal('API_URL', API_URL)
...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
}),
}))
describe('JanModelExtension', () => { describe('JanModelExtension', () => {
let extension: JanModelExtension let extension: JanModelExtension
let mockCortexAPI: any
beforeEach(() => { 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 = 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)
await ky.delete(`${API_URL}/v1/models/pull`, {
json: { taskId: model },
}) })
it('should register settings on load', async () => { expect(kyDeleteSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/pull`, {
// @ts-ignore json: { taskId: model },
const registerSettingsSpy = jest.spyOn(extension, 'registerSettings')
await extension.onLoad()
expect(registerSettingsSpy).toHaveBeenCalledWith(SETTINGS)
}) })
it('should pull a model', async () => { kyDeleteSpy.mockRestore() // Restore the original implementation
const model = 'test-model' }
await extension.pullModel(model)
expect(mockCortexAPI.pullModel).toHaveBeenCalledWith(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
) )
}) })
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
})
})
}) })

211
extensions/model-extension/src/index.ts
View File

@ -12,26 +12,38 @@ import {
DownloadEvent, DownloadEvent,
OptionType, OptionType,
ModelSource, ModelSource,
extractInferenceParams,
extractModelLoadParams,
} from '@janhq/core' } from '@janhq/core'
import { CortexAPI } from './cortex'
import { scanModelsFolder } from './legacy/model-json' import { scanModelsFolder } from './legacy/model-json'
import { deleteModelFiles } from './legacy/delete' import { deleteModelFiles } from './legacy/delete'
import PQueue from 'p-queue'
import ky from 'ky'
/**
* cortex.cpp setting keys
*/
export enum Settings { export enum Settings {
huggingfaceToken = 'hugging-face-access-token', huggingfaceToken = 'hugging-face-access-token',
} }
/** Data List Response Type */
type Data<T> = {
data: T[]
}
/** /**
* A extension for models * A extension for models
*/ */
export default class JanModelExtension extends ModelExtension { export default class JanModelExtension extends ModelExtension {
cortexAPI: CortexAPI = new CortexAPI() queue = new PQueue({ concurrency: 1 })
/** /**
* Called when the extension is loaded. * Called when the extension is loaded.
* @override * @override
*/ */
async onLoad() { async onLoad() {
this.queue.add(() => this.healthz())
this.registerSettings(SETTINGS) this.registerSettings(SETTINGS)
// Configure huggingface token if available // Configure huggingface token if available
@ -39,11 +51,15 @@ export default class JanModelExtension extends ModelExtension {
Settings.huggingfaceToken, Settings.huggingfaceToken,
undefined undefined
) )
if (huggingfaceToken) if (huggingfaceToken) {
this.cortexAPI.configs({ huggingface_token: huggingfaceToken }) this.updateCortexConfig({ huggingface_token: huggingfaceToken })
}
// Listen to app download events // Listen to app download events
this.handleDesktopEvents() this.handleDesktopEvents()
// Sync with cortexsohub
this.fetchCortexsoModels()
} }
/** /**
@ -53,7 +69,7 @@ export default class JanModelExtension extends ModelExtension {
*/ */
onSettingUpdate<T>(key: string, value: T): void { onSettingUpdate<T>(key: string, value: T): void {
if (key === Settings.huggingfaceToken) { 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() {} async onUnload() {}
// BEGIN: - Public API
/** /**
* Downloads a machine learning model. * Downloads a machine learning model.
* @param model - The model to download. * @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 * 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 * 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. * @returns A Promise that resolves when the model is deleted.
*/ */
async deleteModel(model: string): Promise<void> { async deleteModel(model: string): Promise<void> {
return this.cortexAPI return this.queue
.deleteModel(model) .add(() => ky.delete(`${API_URL}/v1/models/${model}`).json().then())
.catch((e) => console.debug(e)) .catch((e) => console.debug(e))
.finally(async () => { .finally(async () => {
// Delete legacy model files // Delete legacy model files
await deleteModelFiles(model).catch((e) => console.debug(e)) await deleteModelFiles(model).catch((e) => console.debug(e))
}) }) as Promise<void>
} }
/** /**
@ -139,7 +169,7 @@ export default class JanModelExtension extends ModelExtension {
/** /**
* Fetch models from cortex.cpp * 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 // Checking if there are models to import
const existingIds = fetchedModels.map((e) => e.id) const existingIds = fetchedModels.map((e) => e.id)
@ -196,8 +226,7 @@ export default class JanModelExtension extends ModelExtension {
* Models are imported successfully before * Models are imported successfully before
* Now return models from cortex.cpp and merge with legacy models which are not imported * Now return models from cortex.cpp and merge with legacy models which are not imported
*/ */
return await this.cortexAPI return await this.fetchModels()
.getModels()
.then((models) => { .then((models) => {
return models.concat( return models.concat(
legacyModels.filter((e) => !models.some((x) => x.id === e.id)) 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 * @param model - The metadata of the model
*/ */
async updateModel(model: Partial<Model>): Promise<Model> { async updateModel(model: Partial<Model>): Promise<Model> {
return this.cortexAPI return this.queue
?.updateModel(model) .add(() =>
.then(() => this.cortexAPI!.getModel(model.id)) 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<Model> {
return this.queue.add(() =>
ky
.get(`${API_URL}/v1/models/${model}`)
.json()
.then((e) => this.transformModel(e))
) as Promise<Model>
} }
/** /**
@ -227,7 +274,15 @@ export default class JanModelExtension extends ModelExtension {
name?: string, name?: string,
option?: OptionType option?: OptionType
): Promise<void> { ): Promise<void> {
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 // BEGIN - Model Sources
@ -236,7 +291,10 @@ export default class JanModelExtension extends ModelExtension {
* @param model * @param model
*/ */
async getSources(): Promise<ModelSource[]> { async getSources(): Promise<ModelSource[]> {
const sources = await this.cortexAPI.getSources() const sources = await this.queue
.add(() => ky.get(`${API_URL}/v1/models/sources`).json<Data<ModelSource>>())
.then((e) => (typeof e === 'object' ? (e.data as ModelSource[]) : []))
.catch(() => [])
return sources.concat( return sources.concat(
DEFAULT_MODEL_SOURCES.filter((e) => !sources.some((x) => x.id === e.id)) DEFAULT_MODEL_SOURCES.filter((e) => !sources.some((x) => x.id === e.id))
) )
@ -247,7 +305,13 @@ export default class JanModelExtension extends ModelExtension {
* @param model * @param model
*/ */
async addSource(source: string): Promise<any> { async addSource(source: string): Promise<any> {
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 * @param model
*/ */
async deleteSource(source: string): Promise<any> { async deleteSource(source: string): Promise<any> {
return this.cortexAPI.deleteSource(source) return this.queue.add(() =>
ky.delete(`${API_URL}/v1/models/sources`, {
json: {
source,
},
})
)
} }
// END - Model Sources // END - Model Sources
@ -264,20 +334,38 @@ export default class JanModelExtension extends ModelExtension {
* @param model * @param model
*/ */
async isModelLoaded(model: string): Promise<boolean> { async isModelLoaded(model: string): Promise<boolean> {
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. * Configure pull options such as proxy, headers, etc.
*/ */
async configurePullOptions(options: { [key: string]: any }): Promise<any> { async configurePullOptions(options: { [key: string]: any }): Promise<any> {
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<Model[]> {
return this.queue
.add(() => ky.get(`${API_URL}/v1/models?limit=-1`).json<Data<Model>>())
.then((e) =>
typeof e === 'object' ? e.data.map((e) => this.transformModel(e)) : []
)
}
// END: - Public API
// BEGIN: - Private API
/** /**
* Handle download state from main app * Handle download state from main app
*/ */
handleDesktopEvents() { private handleDesktopEvents() {
if (window && window.electronAPI) { if (window && window.electronAPI) {
window.electronAPI.onFileDownloadUpdate( window.electronAPI.onFileDownloadUpdate(
async (_event: string, state: DownloadState | undefined) => { 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<void> {
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<void> {
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<Data<string>>()
.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
} }

51
extensions/model-extension/src/legacy/model-json.test.ts
View File

@ -1,27 +1,31 @@
import { describe, it, expect, beforeEach, vi } from 'vitest'
import { scanModelsFolder, getModelJsonPath } from './model-json' import { scanModelsFolder, getModelJsonPath } from './model-json'
// Mock the @janhq/core module // Mock the @janhq/core module
jest.mock('@janhq/core', () => ({ vi.mock('@janhq/core', () => ({
fs: { InferenceEngine: {
existsSync: jest.fn(), nitro: 'nitro',
readdirSync: jest.fn(),
fileStat: jest.fn(),
readFileSync: jest.fn(),
}, },
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 // 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', () => { describe('model-json', () => {
beforeEach(() => { beforeEach(() => {
jest.clearAllMocks() vi.clearAllMocks()
}) })
describe('scanModelsFolder', () => { describe('scanModelsFolder', () => {
it('should return an empty array when models folder does not exist', async () => { 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() const result = await scanModelsFolder()
expect(result).toEqual([]) expect(result).toEqual([])
@ -38,11 +42,16 @@ describe('model-json', () => {
], ],
} }
fs.existsSync.mockReturnValue(true) vi.spyOn(fs, 'existsSync').mockReturnValue(true)
fs.readdirSync.mockReturnValueOnce(['test-model']) vi.spyOn(fs, 'readdirSync').mockReturnValueOnce(['test-model'])
fs.fileStat.mockResolvedValue({ isDirectory: () => true }) vi.spyOn(fs, 'fileStat').mockResolvedValue({ isDirectory: () => true })
fs.readFileSync.mockReturnValue(JSON.stringify(mockModelJson)) vi.spyOn(fs, 'readFileSync').mockReturnValue(
fs.readdirSync.mockReturnValueOnce(['test-model.gguf', 'model.json']) JSON.stringify(mockModelJson)
)
vi.spyOn(fs, 'readdirSync').mockReturnValueOnce([
'test-model.gguf',
'model.json',
])
const result = await scanModelsFolder() const result = await scanModelsFolder()
expect(result).toHaveLength(1) expect(result).toHaveLength(1)
@ -52,26 +61,26 @@ describe('model-json', () => {
describe('getModelJsonPath', () => { describe('getModelJsonPath', () => {
it('should return undefined when folder does not exist', async () => { 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') const result = await getModelJsonPath('non-existent-folder')
expect(result).toBeUndefined() expect(result).toBeUndefined()
}) })
it('should return the path when model.json exists in the root folder', async () => { it('should return the path when model.json exists in the root folder', async () => {
fs.existsSync.mockReturnValue(true) vi.spyOn(fs, 'existsSync').mockReturnValue(true)
fs.readdirSync.mockReturnValue(['model.json']) vi.spyOn(fs, 'readdirSync').mockReturnValue(['model.json'])
const result = await getModelJsonPath('test-folder') const result = await getModelJsonPath('test-folder')
expect(result).toBe('test-folder/model.json') expect(result).toBe('test-folder/model.json')
}) })
it('should return the path when model.json exists in a subfolder', async () => { it('should return the path when model.json exists in a subfolder', async () => {
fs.existsSync.mockReturnValue(true) vi.spyOn(fs, 'existsSync').mockReturnValue(true)
fs.readdirSync vi.spyOn(fs, 'readdirSync')
.mockReturnValueOnce(['subfolder']) .mockReturnValueOnce(['subfolder'])
.mockReturnValueOnce(['model.json']) .mockReturnValueOnce(['model.json'])
fs.fileStat.mockResolvedValue({ isDirectory: () => true }) vi.spyOn(fs, 'fileStat').mockResolvedValue({ isDirectory: () => true })
const result = await getModelJsonPath('test-folder') const result = await getModelJsonPath('test-folder')
expect(result).toBe('test-folder/subfolder/model.json') expect(result).toBe('test-folder/subfolder/model.json')

73
extensions/model-extension/src/migration.test.ts
View File

@ -1,48 +1,51 @@
import { Model, InferenceEngine } from '@janhq/core' import { describe, it, expect, beforeEach, vi } from 'vitest'
import JanModelExtension from './index'
vi.stubGlobal('API_URL', 'http://localhost:3000')
// Mock the @janhq/core module // Mock the @janhq/core module
jest.mock('@janhq/core', () => ({ vi.mock('@janhq/core', (actual) => ({
...actual,
ModelExtension: class {}, ModelExtension: class {},
InferenceEngine: { InferenceEngine: {
nitro: 'nitro', nitro: 'nitro',
}, },
joinPath: jest.fn(), joinPath: vi.fn(),
dirName: jest.fn(), dirName: vi.fn(),
fs: {
existsSync: vi.fn(),
readFileSync: vi.fn(),
writeFileSync: vi.fn(),
mkdirSync: vi.fn(),
},
})) }))
// Mock the CortexAPI import { Model, InferenceEngine } from '@janhq/core'
jest.mock('./cortex', () => ({
CortexAPI: jest.fn().mockImplementation(() => ({ import JanModelExtension from './index'
getModels: jest.fn(),
importModel: jest.fn(),
})),
}))
// Mock the model-json module // Mock the model-json module
jest.mock('./model-json', () => ({ vi.mock('./legacy/model-json', () => ({
scanModelsFolder: jest.fn(), scanModelsFolder: vi.fn(),
})) }))
// Import the mocked scanModelsFolder after the mock is set up // 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', () => { describe('JanModelExtension', () => {
let extension: JanModelExtension let extension: JanModelExtension
let mockLocalStorage: { [key: string]: string } let mockLocalStorage: { [key: string]: string }
let mockCortexAPI: jest.Mock
beforeEach(() => { beforeEach(() => {
// @ts-ignore // @ts-ignore
extension = new JanModelExtension() extension = new JanModelExtension()
mockLocalStorage = {} mockLocalStorage = {}
mockCortexAPI = extension.cortexAPI as any
// Mock localStorage // Mock localStorage
Object.defineProperty(global, 'localStorage', { Object.defineProperty(global, 'localStorage', {
value: { value: {
getItem: jest.fn((key) => mockLocalStorage[key]), getItem: vi.fn((key) => mockLocalStorage[key]),
setItem: jest.fn((key, value) => { setItem: vi.fn((key, value) => {
mockLocalStorage[key] = value mockLocalStorage[key] = value
}), }),
}, },
@ -76,22 +79,13 @@ describe('JanModelExtension', () => {
file_path: '/path/to/model2', file_path: '/path/to/model2',
}, },
] as any ] as any
scanModelsFolder.mockResolvedValue(mockModels) vi.mocked(legacy.scanModelsFolder).mockResolvedValue(mockModels)
extension.cortexAPI.importModel = jest vi.spyOn(extension, 'fetchModels').mockResolvedValue([mockModels[0]])
.fn() vi.spyOn(extension, 'updateModel').mockResolvedValue(undefined)
.mockResolvedValueOnce(mockModels[0]) vi.spyOn(extension, 'importModel').mockResolvedValueOnce(mockModels[1])
extension.cortexAPI.getModels = jest vi.spyOn(extension, 'fetchModels').mockResolvedValue([mockModels[0], mockModels[1]])
.fn()
.mockResolvedValue([mockModels[0]])
extension.cortexAPI.importModel = jest
.fn()
.mockResolvedValueOnce(mockModels[1])
extension.cortexAPI.getModels = jest
.fn()
.mockResolvedValue([mockModels[0], mockModels[1]])
const result = await extension.getModels() const result = await extension.getModels()
expect(scanModelsFolder).toHaveBeenCalled() expect(legacy.scanModelsFolder).toHaveBeenCalled()
expect(result).toEqual(mockModels) expect(result).toEqual(mockModels)
}) })
@ -121,9 +115,8 @@ describe('JanModelExtension', () => {
}, },
] as any ] as any
mockLocalStorage['downloadedModels'] = JSON.stringify(mockModels) mockLocalStorage['downloadedModels'] = JSON.stringify(mockModels)
vi.spyOn(extension, 'updateModel').mockResolvedValue(undefined)
extension.cortexAPI.getModels = jest.fn().mockResolvedValue([]) vi.spyOn(extension, 'importModel').mockResolvedValue(undefined)
extension.importModel = jest.fn().mockResolvedValue(undefined)
const result = await extension.getModels() const result = await extension.getModels()
@ -155,12 +148,12 @@ describe('JanModelExtension', () => {
}, },
] as any ] as any
mockLocalStorage['downloadedModels'] = JSON.stringify(mockModels) mockLocalStorage['downloadedModels'] = JSON.stringify(mockModels)
vi.spyOn(extension, 'fetchModels').mockResolvedValue(mockModels)
extension.cortexAPI.getModels = jest.fn().mockResolvedValue(mockModels) extension.getModels = vi.fn().mockResolvedValue(mockModels)
const result = await extension.getModels() const result = await extension.getModels()
expect(extension.cortexAPI.getModels).toHaveBeenCalled() expect(extension.getModels).toHaveBeenCalled()
expect(result).toEqual(mockModels) expect(result).toEqual(mockModels)
}) })
}) })

2
extensions/model-extension/tsconfig.json
View File

@ -11,5 +11,5 @@
"rootDir": "./src" "rootDir": "./src"
}, },
"include": ["./src"], "include": ["./src"],
"exclude": ["**/*.test.ts"] "exclude": ["**/*.test.ts", "vite.config.ts"]
} }

8
extensions/model-extension/vite.config.ts Normal file
View File

@ -0,0 +1,8 @@
import { defineConfig } from "vite"
export default defineConfig(({ mode }) => ({
define: process.env.VITEST ? {} : { global: 'window' },
test: {
environment: 'jsdom',
},
}))

2
web/utils/modelEngine.ts
View File

@ -90,7 +90,7 @@ export const getDescriptionByEngine = (engine: InferenceEngine) => {
case InferenceEngine.openrouter: 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.' 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: default:
return `Access models from ${engine} via their API.` return `Access models from ${getTitleByEngine(engine)} via their API.`
} }
} }