feat: Add nitro engine json reader and writer

This commit is contained in:
hiro 2023-12-03 14:36:01 +07:00
parent 8ab36d7cb2
commit 56b778675f

View File

@ -19,7 +19,7 @@ import {
events, events,
executeOnMain, executeOnMain,
getUserSpace, getUserSpace,
fs fs,
} from "@janhq/core"; } from "@janhq/core";
import { InferenceExtension } from "@janhq/core"; import { InferenceExtension } from "@janhq/core";
import { requestInference } from "./helpers/sse"; import { requestInference } from "./helpers/sse";
@ -31,7 +31,7 @@ import { join } from "path";
* The class provides methods for initializing and stopping a model, and for making inference requests. * The class provides methods for initializing and stopping a model, and for making inference requests.
* It also subscribes to events emitted by the @janhq/core package and handles new message requests. * It also subscribes to events emitted by the @janhq/core package and handles new message requests.
*/ */
export default class JanInferenceExtension implements InferenceExtension { export default class JanInferenceNitroExtension implements InferenceExtension {
private static readonly _homeDir = 'engines' private static readonly _homeDir = 'engines'
private static readonly _engineMetadataFileName = 'nitro.json' private static readonly _engineMetadataFileName = 'nitro.json'
@ -49,10 +49,10 @@ export default class JanInferenceExtension implements InferenceExtension {
* Subscribes to events emitted by the @janhq/core package. * Subscribes to events emitted by the @janhq/core package.
*/ */
onLoad(): void { onLoad(): void {
fs.mkdir(JanInferenceExtension._homeDir) fs.mkdir(JanInferenceNitroExtension._homeDir)
this.writeDefaultEngineSettings()
events.on(EventName.OnMessageSent, (data) => events.on(EventName.OnMessageSent, (data) =>
JanInferenceExtension.handleMessageRequest(data, this) JanInferenceNitroExtension.handleMessageRequest(data, this)
); );
} }
@ -74,14 +74,10 @@ export default class JanInferenceExtension implements InferenceExtension {
): Promise<void> { ): Promise<void> {
const userSpacePath = await getUserSpace(); const userSpacePath = await getUserSpace();
const modelFullPath = join(userSpacePath, "models", modelId, modelId); const modelFullPath = join(userSpacePath, "models", modelId, modelId);
let engine_settings = JSON.parse(await fs.readFile(join(JanInferenceExtension._homeDir, JanInferenceExtension._engineMetadataFileName)))
engine_settings = {
engine_settings
...settings,
};
return executeOnMain(MODULE, "initModel", { return executeOnMain(MODULE, "initModel", {
modelFullPath, modelFullPath,
engine_settings, settings,
}); });
} }
@ -102,6 +98,28 @@ export default class JanInferenceExtension implements InferenceExtension {
this.controller?.abort(); this.controller?.abort();
} }
private async writeDefaultEngineSettings() {
try {
const destPath = join(JanInferenceNitroExtension._homeDir, JanInferenceNitroExtension._engineMetadataFileName)
// TODO: Check with @louis for adding new binding
// if (await fs.checkFileExists(destPath)) {
const default_engine_settings = {
"ctx_len": 2048,
"ngl": 100,
"cont_batching": false,
"embedding": false
}
console.log(`Writing nitro engine settings to ${destPath}`)
await fs.writeFile(destPath, JSON.stringify(default_engine_settings, null, 2))
// }
// else {
// console.log(`Using existing nitro engine settings at ${destPath}`)
// }
} catch (err) {
console.error(err)
}
}
/** /**
* Makes a single response inference request. * Makes a single response inference request.
* @param {MessageRequest} data - The data for the inference request. * @param {MessageRequest} data - The data for the inference request.
@ -141,7 +159,7 @@ export default class JanInferenceExtension implements InferenceExtension {
*/ */
private static async handleMessageRequest( private static async handleMessageRequest(
data: MessageRequest, data: MessageRequest,
instance: JanInferenceExtension instance: JanInferenceNitroExtension
) { ) {
const timestamp = Date.now(); const timestamp = Date.now();
const message: ThreadMessage = { const message: ThreadMessage = {