feat: Jan Model Hub should stay updated. (#4707)
* feat: Jan Model Hub should stay updated. * chore: polish provider description
This commit is contained in:
Louis
GitHub
parent
839a00127d
commit
fddb7251fb
@ -1 +1 @@
|
||||
1.0.11-rc1
|
||||
1.0.11-rc2
|
||||
|
||||
@ -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/.*)'],
|
||||
}
|
||||
@ -7,7 +7,7 @@
|
||||
"author": "Jan <service@jan.ai>",
|
||||
"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/*",
|
||||
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
@ -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
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@ -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<T> = {
|
||||
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<T>(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<void> {
|
||||
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<void>
|
||||
}
|
||||
|
||||
/**
|
||||
@ -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<Model>): Promise<Model> {
|
||||
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<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,
|
||||
option?: OptionType
|
||||
): 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
|
||||
@ -236,7 +291,10 @@ export default class JanModelExtension extends ModelExtension {
|
||||
* @param model
|
||||
*/
|
||||
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(
|
||||
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<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
|
||||
*/
|
||||
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
|
||||
|
||||
@ -264,20 +334,38 @@ export default class JanModelExtension extends ModelExtension {
|
||||
* @param model
|
||||
*/
|
||||
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.
|
||||
*/
|
||||
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
|
||||
*/
|
||||
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<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
|
||||
}
|
||||
|
||||
@ -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')
|
||||
|
||||
@ -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)
|
||||
})
|
||||
})
|
||||
|
||||
@ -11,5 +11,5 @@
|
||||
"rootDir": "./src"
|
||||
},
|
||||
"include": ["./src"],
|
||||
"exclude": ["**/*.test.ts"]
|
||||
"exclude": ["**/*.test.ts", "vite.config.ts"]
|
||||
}
|
||||
|
||||
8
extensions/model-extension/vite.config.ts
Normal file
8
extensions/model-extension/vite.config.ts
Normal file
@ -0,0 +1,8 @@
|
||||
import { defineConfig } from "vite"
|
||||
export default defineConfig(({ mode }) => ({
|
||||
define: process.env.VITEST ? {} : { global: 'window' },
|
||||
test: {
|
||||
environment: 'jsdom',
|
||||
},
|
||||
}))
|
||||
|
||||
@ -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.`
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user