847 lines
24 KiB
TypeScript

/**
* @jest-environment jsdom
*/
const readDirSyncMock = jest.fn()
const existMock = jest.fn()
const readFileSyncMock = jest.fn()
const downloadMock = jest.fn()
const mkdirMock = jest.fn()
const writeFileSyncMock = jest.fn()
const copyFileMock = jest.fn()
const dirNameMock = jest.fn()
const executeMock = jest.fn()
jest.mock('@janhq/core', () => ({
...jest.requireActual('@janhq/core/node'),
events: {
emit: jest.fn(),
},
fs: {
existsSync: existMock,
readdirSync: readDirSyncMock,
readFileSync: readFileSyncMock,
writeFileSync: writeFileSyncMock,
mkdir: mkdirMock,
copyFile: copyFileMock,
fileStat: () => ({
isDirectory: false,
}),
},
dirName: dirNameMock,
joinPath: (paths) => paths.join('/'),
ModelExtension: jest.fn(),
downloadFile: downloadMock,
executeOnMain: executeMock,
}))
jest.mock('@huggingface/gguf')
global.fetch = jest.fn(() =>
Promise.resolve({
json: () => Promise.resolve({ test: 100 }),
arrayBuffer: jest.fn(),
})
) as jest.Mock
import JanModelExtension from '.'
import { fs, dirName } from '@janhq/core'
import { gguf } from '@huggingface/gguf'
describe('JanModelExtension', () => {
let sut: JanModelExtension
beforeAll(() => {
// @ts-ignore
sut = new JanModelExtension()
})
beforeEach(() => {
jest.clearAllMocks()
})
describe('getConfiguredModels', () => {
describe("when there's no models are pre-populated", () => {
it('should return empty array', async () => {
// Mock configured models data
const configuredModels = []
existMock.mockReturnValue(true)
readDirSyncMock.mockReturnValue([])
const result = await sut.getConfiguredModels()
expect(result).toEqual([])
})
})
describe("when there's are pre-populated models - all flattened", () => {
it('returns configured models data - flatten folder - with correct file_path and model id', async () => {
// Mock configured models data
const configuredModels = [
{
id: '1',
name: 'Model 1',
version: '1.0.0',
description: 'Model 1 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model1',
},
format: 'onnx',
sources: [],
created: new Date(),
updated: new Date(),
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
{
id: '2',
name: 'Model 2',
version: '2.0.0',
description: 'Model 2 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model2',
},
format: 'onnx',
sources: [],
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
]
existMock.mockReturnValue(true)
readDirSyncMock.mockImplementation((path) => {
if (path === 'file://models') return ['model1', 'model2']
else return ['model.json']
})
readFileSyncMock.mockImplementation((path) => {
if (path.includes('model1'))
return JSON.stringify(configuredModels[0])
else return JSON.stringify(configuredModels[1])
})
const result = await sut.getConfiguredModels()
expect(result).toEqual(
expect.arrayContaining([
expect.objectContaining({
file_path: 'file://models/model1/model.json',
id: '1',
}),
expect.objectContaining({
file_path: 'file://models/model2/model.json',
id: '2',
}),
])
)
})
})
describe("when there's are pre-populated models - there are nested folders", () => {
it('returns configured models data - flatten folder - with correct file_path and model id', async () => {
// Mock configured models data
const configuredModels = [
{
id: '1',
name: 'Model 1',
version: '1.0.0',
description: 'Model 1 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model1',
},
format: 'onnx',
sources: [],
created: new Date(),
updated: new Date(),
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
{
id: '2',
name: 'Model 2',
version: '2.0.0',
description: 'Model 2 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model2',
},
format: 'onnx',
sources: [],
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
]
existMock.mockReturnValue(true)
readDirSyncMock.mockImplementation((path) => {
if (path === 'file://models') return ['model1', 'model2/model2-1']
else return ['model.json']
})
readFileSyncMock.mockImplementation((path) => {
if (path.includes('model1'))
return JSON.stringify(configuredModels[0])
else if (path.includes('model2/model2-1'))
return JSON.stringify(configuredModels[1])
})
const result = await sut.getConfiguredModels()
expect(result).toEqual(
expect.arrayContaining([
expect.objectContaining({
file_path: 'file://models/model1/model.json',
id: '1',
}),
expect.objectContaining({
file_path: 'file://models/model2/model2-1/model.json',
id: '2',
}),
])
)
})
})
})
describe('getDownloadedModels', () => {
describe('no models downloaded', () => {
it('should return empty array', async () => {
// Mock downloaded models data
existMock.mockReturnValue(true)
readDirSyncMock.mockReturnValue([])
const result = await sut.getDownloadedModels()
expect(result).toEqual([])
})
})
describe('only one model is downloaded', () => {
describe('flatten folder', () => {
it('returns downloaded models - with correct file_path and model id', async () => {
// Mock configured models data
const configuredModels = [
{
id: '1',
name: 'Model 1',
version: '1.0.0',
description: 'Model 1 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model1',
},
format: 'onnx',
sources: [],
created: new Date(),
updated: new Date(),
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
{
id: '2',
name: 'Model 2',
version: '2.0.0',
description: 'Model 2 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model2',
},
format: 'onnx',
sources: [],
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
]
existMock.mockReturnValue(true)
readDirSyncMock.mockImplementation((path) => {
if (path === 'file://models') return ['model1', 'model2']
else if (path === 'file://models/model1')
return ['model.json', 'test.gguf']
else return ['model.json']
})
readFileSyncMock.mockImplementation((path) => {
if (path.includes('model1'))
return JSON.stringify(configuredModels[0])
else return JSON.stringify(configuredModels[1])
})
const result = await sut.getDownloadedModels()
expect(result).toEqual(
expect.arrayContaining([
expect.objectContaining({
file_path: 'file://models/model1/model.json',
id: '1',
}),
])
)
})
})
})
describe('all models are downloaded', () => {
describe('nested folders', () => {
it('returns downloaded models - with correct file_path and model id', async () => {
// Mock configured models data
const configuredModels = [
{
id: '1',
name: 'Model 1',
version: '1.0.0',
description: 'Model 1 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model1',
},
format: 'onnx',
sources: [],
created: new Date(),
updated: new Date(),
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
{
id: '2',
name: 'Model 2',
version: '2.0.0',
description: 'Model 2 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model2',
},
format: 'onnx',
sources: [],
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
]
existMock.mockReturnValue(true)
readDirSyncMock.mockImplementation((path) => {
if (path === 'file://models') return ['model1', 'model2/model2-1']
else return ['model.json', 'test.gguf']
})
readFileSyncMock.mockImplementation((path) => {
if (path.includes('model1'))
return JSON.stringify(configuredModels[0])
else return JSON.stringify(configuredModels[1])
})
const result = await sut.getDownloadedModels()
expect(result).toEqual(
expect.arrayContaining([
expect.objectContaining({
file_path: 'file://models/model1/model.json',
id: '1',
}),
expect.objectContaining({
file_path: 'file://models/model2/model2-1/model.json',
id: '2',
}),
])
)
})
})
})
describe('all models are downloaded with uppercased GGUF files', () => {
it('returns downloaded models - with correct file_path and model id', async () => {
// Mock configured models data
const configuredModels = [
{
id: '1',
name: 'Model 1',
version: '1.0.0',
description: 'Model 1 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model1',
},
format: 'onnx',
sources: [],
created: new Date(),
updated: new Date(),
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
{
id: '2',
name: 'Model 2',
version: '2.0.0',
description: 'Model 2 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model2',
},
format: 'onnx',
sources: [],
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
]
existMock.mockReturnValue(true)
readDirSyncMock.mockImplementation((path) => {
if (path === 'file://models') return ['model1', 'model2/model2-1']
else if (path === 'file://models/model1')
return ['model.json', 'test.GGUF']
else return ['model.json', 'test.gguf']
})
readFileSyncMock.mockImplementation((path) => {
if (path.includes('model1'))
return JSON.stringify(configuredModels[0])
else return JSON.stringify(configuredModels[1])
})
const result = await sut.getDownloadedModels()
expect(result).toEqual(
expect.arrayContaining([
expect.objectContaining({
file_path: 'file://models/model1/model.json',
id: '1',
}),
expect.objectContaining({
file_path: 'file://models/model2/model2-1/model.json',
id: '2',
}),
])
)
})
})
describe('all models are downloaded - GGUF & Tensort RT', () => {
it('returns downloaded models - with correct file_path and model id', async () => {
// Mock configured models data
const configuredModels = [
{
id: '1',
name: 'Model 1',
version: '1.0.0',
description: 'Model 1 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model1',
},
format: 'onnx',
sources: [],
created: new Date(),
updated: new Date(),
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
{
id: '2',
name: 'Model 2',
version: '2.0.0',
description: 'Model 2 description',
object: {
type: 'model',
uri: 'http://localhost:5000/models/model2',
},
format: 'onnx',
sources: [],
parameters: {},
settings: {},
metadata: {},
engine: 'test',
} as any,
]
existMock.mockReturnValue(true)
readDirSyncMock.mockImplementation((path) => {
if (path === 'file://models') return ['model1', 'model2/model2-1']
else if (path === 'file://models/model1')
return ['model.json', 'test.gguf']
else return ['model.json', 'test.engine']
})
readFileSyncMock.mockImplementation((path) => {
if (path.includes('model1'))
return JSON.stringify(configuredModels[0])
else return JSON.stringify(configuredModels[1])
})
const result = await sut.getDownloadedModels()
expect(result).toEqual(
expect.arrayContaining([
expect.objectContaining({
file_path: 'file://models/model1/model.json',
id: '1',
}),
expect.objectContaining({
file_path: 'file://models/model2/model2-1/model.json',
id: '2',
}),
])
)
})
})
})
describe('deleteModel', () => {
describe('model is a GGUF model', () => {
it('should delete the GGUF file', async () => {
fs.unlinkSync = jest.fn()
const dirMock = dirName as jest.Mock
dirMock.mockReturnValue('file://models/model1')
fs.readFileSync = jest.fn().mockReturnValue(JSON.stringify({}))
readDirSyncMock.mockImplementation((path) => {
return ['model.json', 'test.gguf']
})
existMock.mockReturnValue(true)
await sut.deleteModel({
file_path: 'file://models/model1/model.json',
} as any)
expect(fs.unlinkSync).toHaveBeenCalledWith(
'file://models/model1/test.gguf'
)
})
it('no gguf file presented', async () => {
fs.unlinkSync = jest.fn()
const dirMock = dirName as jest.Mock
dirMock.mockReturnValue('file://models/model1')
fs.readFileSync = jest.fn().mockReturnValue(JSON.stringify({}))
readDirSyncMock.mockReturnValue(['model.json'])
existMock.mockReturnValue(true)
await sut.deleteModel({
file_path: 'file://models/model1/model.json',
} as any)
expect(fs.unlinkSync).toHaveBeenCalledTimes(0)
})
it('delete an imported model', async () => {
fs.rm = jest.fn()
const dirMock = dirName as jest.Mock
dirMock.mockReturnValue('file://models/model1')
readDirSyncMock.mockReturnValue(['model.json', 'test.gguf'])
// MARK: This is a tricky logic implement?
// I will just add test for now but will align on the legacy implementation
fs.readFileSync = jest.fn().mockReturnValue(
JSON.stringify({
metadata: {
author: 'user',
},
})
)
existMock.mockReturnValue(true)
await sut.deleteModel({
file_path: 'file://models/model1/model.json',
} as any)
expect(fs.rm).toHaveBeenCalledWith('file://models/model1')
})
it('delete tensorrt-models', async () => {
fs.rm = jest.fn()
const dirMock = dirName as jest.Mock
dirMock.mockReturnValue('file://models/model1')
readDirSyncMock.mockReturnValue(['model.json', 'test.engine'])
fs.readFileSync = jest.fn().mockReturnValue(JSON.stringify({}))
existMock.mockReturnValue(true)
await sut.deleteModel({
file_path: 'file://models/model1/model.json',
} as any)
expect(fs.unlinkSync).toHaveBeenCalledWith(
'file://models/model1/test.engine'
)
})
})
})
describe('downloadModel', () => {
const model: any = {
id: 'model-id',
name: 'Test Model',
sources: [
{ url: 'http://example.com/model.gguf', filename: 'model.gguf' },
],
engine: 'test-engine',
}
const network = {
ignoreSSL: true,
proxy: 'http://proxy.example.com',
}
const gpuSettings: any = {
gpus: [{ name: 'nvidia-rtx-3080', arch: 'ampere' }],
}
it('should reject with invalid gguf metadata', async () => {
existMock.mockImplementation(() => false)
expect(
sut.downloadModel(model, gpuSettings, network)
).rejects.toBeTruthy()
})
it('should download corresponding ID', async () => {
existMock.mockImplementation(() => true)
dirNameMock.mockImplementation(() => 'file://models/model1')
downloadMock.mockImplementation(() => {
return Promise.resolve({})
})
expect(
await sut.downloadModel(
{ ...model, file_path: 'file://models/model1/model.json' },
gpuSettings,
network
)
).toBeUndefined()
expect(downloadMock).toHaveBeenCalledWith(
{
localPath: 'file://models/model1/model.gguf',
modelId: 'model-id',
url: 'http://example.com/model.gguf',
},
{ ignoreSSL: true, proxy: 'http://proxy.example.com' }
)
})
it('should handle invalid model file', async () => {
executeMock.mockResolvedValue({})
fs.readFileSync = jest.fn(() => {
return JSON.stringify({ metadata: { author: 'user' } })
})
expect(
sut.downloadModel(
{ ...model, file_path: 'file://models/model1/model.json' },
gpuSettings,
network
)
).resolves.not.toThrow()
expect(downloadMock).not.toHaveBeenCalled()
})
it('should handle model file with no sources', async () => {
executeMock.mockResolvedValue({})
const modelWithoutSources = { ...model, sources: [] }
expect(
sut.downloadModel(
{
...modelWithoutSources,
file_path: 'file://models/model1/model.json',
},
gpuSettings,
network
)
).resolves.toBe(undefined)
expect(downloadMock).not.toHaveBeenCalled()
})
it('should handle model file with multiple sources', async () => {
const modelWithMultipleSources = {
...model,
sources: [
{ url: 'http://example.com/model1.gguf', filename: 'model1.gguf' },
{ url: 'http://example.com/model2.gguf', filename: 'model2.gguf' },
],
}
executeMock.mockResolvedValue({
metadata: { 'tokenizer.ggml.eos_token_id': 0 },
})
;(gguf as jest.Mock).mockResolvedValue({
metadata: { 'tokenizer.ggml.eos_token_id': 0 },
})
// @ts-ignore
global.NODE = 'node'
// @ts-ignore
global.DEFAULT_MODEL = {
parameters: { stop: [] },
}
downloadMock.mockImplementation(() => {
return Promise.resolve({})
})
expect(
await sut.downloadModel(
{
...modelWithMultipleSources,
file_path: 'file://models/model1/model.json',
},
gpuSettings,
network
)
).toBeUndefined()
expect(downloadMock).toHaveBeenCalledWith(
{
localPath: 'file://models/model1/model1.gguf',
modelId: 'model-id',
url: 'http://example.com/model1.gguf',
},
{ ignoreSSL: true, proxy: 'http://proxy.example.com' }
)
expect(downloadMock).toHaveBeenCalledWith(
{
localPath: 'file://models/model1/model2.gguf',
modelId: 'model-id',
url: 'http://example.com/model2.gguf',
},
{ ignoreSSL: true, proxy: 'http://proxy.example.com' }
)
})
it('should handle model file with no file_path', async () => {
executeMock.mockResolvedValue({
metadata: { 'tokenizer.ggml.eos_token_id': 0 },
})
;(gguf as jest.Mock).mockResolvedValue({
metadata: { 'tokenizer.ggml.eos_token_id': 0 },
})
// @ts-ignore
global.NODE = 'node'
// @ts-ignore
global.DEFAULT_MODEL = {
parameters: { stop: [] },
}
const modelWithoutFilepath = { ...model, file_path: undefined }
await sut.downloadModel(modelWithoutFilepath, gpuSettings, network)
expect(downloadMock).toHaveBeenCalledWith(
expect.objectContaining({
localPath: 'file://models/model-id/model.gguf',
}),
expect.anything()
)
})
it('should handle model file with invalid file_path', async () => {
executeMock.mockResolvedValue({
metadata: { 'tokenizer.ggml.eos_token_id': 0 },
})
;(gguf as jest.Mock).mockResolvedValue({
metadata: { 'tokenizer.ggml.eos_token_id': 0 },
})
// @ts-ignore
global.NODE = 'node'
// @ts-ignore
global.DEFAULT_MODEL = {
parameters: { stop: [] },
}
const modelWithInvalidFilepath = {
...model,
file_path: 'file://models/invalid-model.json',
}
await sut.downloadModel(modelWithInvalidFilepath, gpuSettings, network)
expect(downloadMock).toHaveBeenCalledWith(
expect.objectContaining({
localPath: 'file://models/model1/model.gguf',
}),
expect.anything()
)
})
it('should handle model with valid chat_template', async () => {
executeMock.mockResolvedValue('{prompt}')
;(gguf as jest.Mock).mockResolvedValue({
metadata: {},
})
// @ts-ignore
global.NODE = 'node'
// @ts-ignore
global.DEFAULT_MODEL = {
parameters: { stop: [] },
settings: {
prompt_template: '<|im-start|>{prompt}<|im-end|>',
},
}
const result = await sut.retrieveGGUFMetadata({})
expect(result).toEqual({
parameters: {
stop: [],
},
settings: {
ctx_len: 4096,
ngl: 33,
prompt_template: '{prompt}',
},
})
})
it('should handle model without chat_template', async () => {
executeMock.mockRejectedValue({})
;(gguf as jest.Mock).mockResolvedValue({
metadata: {},
})
// @ts-ignore
global.NODE = 'node'
// @ts-ignore
global.DEFAULT_MODEL = {
parameters: { stop: [] },
settings: {
prompt_template: '<|im-start|>{prompt}<|im-end|>',
},
}
const result = await sut.retrieveGGUFMetadata({})
expect(result).toEqual({
parameters: {
stop: [],
},
settings: {
ctx_len: 4096,
ngl: 33,
prompt_template: '<|im-start|>{prompt}<|im-end|>',
},
})
})
})
})