refactor(nitro-node): move checkMagicBytes functions out to utils

Author: InNoobWeTrust

nitro-node/jest.config.ts

@@ -1,9 +1,14 @@
-import type { JestConfigWithTsJest } from 'ts-jest'
+import type { JestConfigWithTsJest } from "ts-jest";
 
 const jestConfig: JestConfigWithTsJest = {
-  preset: 'ts-jest',
-  testEnvironment: 'node',
-  transformIgnorePatterns: ['/node_modules/']
-}
+  preset: "ts-jest",
+  testEnvironment: "node",
+  transformIgnorePatterns: ["/node_modules/"],
+  globals: {
+    RELEASE_URL_PREFIX: "https://api.github.com/repos/janhq/nitro/releases/",
+    TAGGED_RELEASE_URL_PREFIX:
+      "https://api.github.com/repos/janhq/nitro/releases/tags",
+  },
+};
 
-export default jestConfig
+export default jestConfig;

nitro-node/src/index.ts

@@ -1,2 +1,3 @@
 export * from "./types";
 export * from "./nitro";
+export { setLogger } from "./logger";

nitro-node/src/logger.ts

@@ -0,0 +1,14 @@
+import os from "node:os";
+import { NitroLogger } from "./types";
+
+// The logger to use, default to stdout
+export let log: NitroLogger = (message, ..._) =>
+  process.stdout.write(message + os.EOL);
+
+/**
+ * Set logger before running nitro
+ * @param {NitroLogger} logger The logger to use
+ */
+export async function setLogger(logger: NitroLogger): Promise<void> {
+  log = logger;
+}

nitro-node/src/nitro.ts

@@ -1,4 +1,3 @@
-import os from "node:os";
 import fs from "node:fs";
 import path from "node:path";
 import { ChildProcessWithoutNullStreams, spawn } from "node:child_process";
@@ -14,12 +13,13 @@ import {
   NitroNvidiaConfig,
   NitroModelSetting,
   NitroPromptSetting,
-  NitroLogger,
   NitroModelOperationResponse,
   NitroModelInitOptions,
   ResourcesInfo,
 } from "./types";
 import { downloadNitro } from "./scripts";
+import { checkMagicBytes } from "./utils";
+import { log } from "./logger";
 // Polyfill fetch with retry
 const fetchRetry = fetchRT(fetch);
 
@@ -58,7 +58,7 @@ const NVIDIA_DEFAULT_CONFIG: NitroNvidiaConfig = {
 const SUPPORTED_MODEL_FORMATS = [".gguf"];
 
 // The supported model magic number
-const SUPPORTED_MODEL_MAGIC_NUMBERS = ["GGUF"];
+const SUPPORTED_MODEL_MAGIC_BYTES = ["GGUF"];
 
 // The subprocess instance for Nitro
 let subprocess: ChildProcessWithoutNullStreams | undefined = undefined;
@@ -68,9 +68,6 @@ let currentModelFile: string = "";
 let currentSettings: NitroModelSetting | undefined = undefined;
 // The Nvidia info file for checking for CUDA support on the system
 let nvidiaConfig: NitroNvidiaConfig = NVIDIA_DEFAULT_CONFIG;
-// The logger to use, default to stdout
-let log: NitroLogger = (message, ..._) =>
-  process.stdout.write(message + os.EOL);
 // The absolute path to bin directory
 let binPath: string = path.join(__dirname, "..", "bin");
 
@@ -118,14 +115,6 @@ export async function setNvidiaConfig(
   nvidiaConfig = config;
 }
 
-/**
- * Set logger before running nitro
- * @param {NitroLogger} logger The logger to use
- */
-export async function setLogger(logger: NitroLogger): Promise<void> {
-  log = logger;
-}
-
 /**
  * Stops a Nitro subprocess.
  * @param wrapper - The model wrapper.
@@ -135,55 +124,32 @@ export function stopModel(): Promise<NitroModelOperationResponse> {
   return killSubprocess();
 }
 
-/**
- * Read the magic bytes from a file and check if they match the provided magic bytes
- */
-export async function checkMagicBytes(
-  filePath: string,
-  magicBytes: string,
-): Promise<boolean> {
-  const desired = Buffer.from(magicBytes);
-  const nBytes = desired.byteLength;
-  const chunks = [];
-  for await (let chunk of fs.createReadStream(filePath, {
-    start: 0,
-    end: nBytes - 1,
-  })) {
-    chunks.push(chunk);
-  }
-  const actual = Buffer.concat(chunks);
-  log(
-    `Comparing file's magic bytes <${actual.toString()}> and desired <${desired.toString()}>`,
-  );
-  return Buffer.compare(actual, desired) === 0;
-}
-
 /**
  * Initializes a Nitro subprocess to load a machine learning model.
  * @param modelFullPath - The absolute full path to model directory.
  * @param promptTemplate - The template to use for generating prompts.
  * @returns A Promise that resolves when the model is loaded successfully, or rejects with an error message if the model is not found or fails to load.
  */
 export async function runModel({
-  modelFullPath,
+  modelPath,
   promptTemplate,
 }: NitroModelInitOptions): Promise<NitroModelOperationResponse> {
   // Download nitro binaries if it's not already downloaded
   await downloadNitro(binPath);
-  const files: string[] = fs.readdirSync(modelFullPath);
+  const files: string[] = fs.readdirSync(modelPath);
 
   // Look for model file with supported format
   let ggufBinFile = files.find(
     (file) =>
-      file === path.basename(modelFullPath) ||
+      file === path.basename(modelPath) ||
       SUPPORTED_MODEL_FORMATS.some((ext) => file.toLowerCase().endsWith(ext)),
   );
 
   // If not found from path and extension, try from magic number
   if (!ggufBinFile) {
     for (const f of files) {
-      for (const magicNum of SUPPORTED_MODEL_MAGIC_NUMBERS) {
-        if (await checkMagicBytes(path.join(modelFullPath, f), magicNum)) {
+      for (const magicBytes of SUPPORTED_MODEL_MAGIC_BYTES) {
+        if (await checkMagicBytes(path.join(modelPath, f), magicBytes)) {
           ggufBinFile = f;
           break;
         }
@@ -194,7 +160,7 @@ export async function runModel({
 
   if (!ggufBinFile) throw new Error("No GGUF model file found");
 
-  currentModelFile = path.join(modelFullPath, ggufBinFile);
+  currentModelFile = path.join(modelPath, ggufBinFile);
 
   const nitroResourceProbe = await getResourcesInfo();
   // Convert promptTemplate to system_prompt, user_prompt, ai_prompt

nitro-node/src/types/index.ts

@@ -33,7 +33,7 @@ export interface NitroModelSetting extends NitroPromptSetting {
  * The response object for model init operation.
  */
 export interface NitroModelInitOptions {
-  modelFullPath: string;
+  modelPath: string;
   promptTemplate?: string;
 }
 

nitro-node/src/utils/index.ts

@@ -0,0 +1,25 @@
+import fs from "node:fs";
+import { log } from "../logger";
+
+/**
+ * Read the magic bytes from a file and check if they match the provided magic bytes
+ */
+export async function checkMagicBytes(
+  filePath: string,
+  magicBytes: string,
+): Promise<boolean> {
+  const desired = Buffer.from(magicBytes);
+  const nBytes = desired.byteLength;
+  const chunks = [];
+  for await (let chunk of fs.createReadStream(filePath, {
+    start: 0,
+    end: nBytes - 1,
+  })) {
+    chunks.push(chunk);
+  }
+  const actual = Buffer.concat(chunks);
+  log(
+    `Comparing file's magic bytes <${actual.toString()}> and desired <${desired.toString()}>`,
+  );
+  return Buffer.compare(actual, desired) === 0;
+}

nitro-node/test/nitro-process.test.ts

@@ -14,8 +14,8 @@ import {
   loadLLMModel,
   validateModelStatus,
   chatCompletion,
-  checkMagicBytes,
 } from "../src";
+import { checkMagicBytes } from "../src/utils";
 
 // FIXME: Shorthand only possible for es6 targets and up
 //import * as model from './model.json' assert {type: 'json'}
@@ -103,24 +103,22 @@ const sleep = async (ms: number): Promise<NodeJS.Timeout> =>
  */
 describe("Manage nitro process", () => {
   /// BEGIN SUITE CONFIG
-  const modelFullPath = fs.mkdtempSync(
-    path.join(os.tmpdir(), "nitro-node-test"),
-  );
+  const modelPath = fs.mkdtempSync(path.join(os.tmpdir(), "nitro-node-test"));
   let modelCfg: Record<string, any> = {};
 
   // Setup steps before running the suite
   const setupHooks = [
     // Get model config from json
     getModelConfigHook((cfg) => Object.assign(modelCfg, cfg)),
     // Download model before starting tests
-    downloadModelHook(modelCfg, modelFullPath),
+    downloadModelHook(modelCfg, modelPath),
   ];
   // Teardown steps after running the suite
   const teardownHooks = [
     // Stop nitro after running, regardless of error or not
     () => stopModel(),
     // On teardown, cleanup tmp directory that was created earlier
-    cleanupTargetDirHook(modelFullPath),
+    cleanupTargetDirHook(modelPath),
   ];
   /// END SUITE CONFIG
 
@@ -145,7 +143,7 @@ describe("Manage nitro process", () => {
     async () => {
       // Start nitro
       await runModel({
-        modelFullPath,
+        modelPath,
         promptTemplate: modelCfg.settings.prompt_template,
       });
       // Wait 5s for nitro to start
@@ -161,14 +159,14 @@ describe("Manage nitro process", () => {
     async () => {
       // Start nitro
       await runModel({
-        modelFullPath,
+        modelPath,
         promptTemplate: modelCfg.settings.prompt_template,
       });
       // Wait 5s for nitro to start
       await sleep(5 * 1000);
       // Load LLM model
       await loadLLMModel({
-        llama_model_path: modelFullPath,
+        llama_model_path: modelPath,
         ctx_len: modelCfg.settings.ctx_len,
         ngl: modelCfg.settings.ngl,
         embedding: false,
@@ -240,27 +238,25 @@ describe("Manage nitro process", () => {
       const fileName = modelCfg.source_url.split("/")?.pop() ?? "model.gguf";
       // Rename the extension of model file
       fs.renameSync(
-        path.join(modelFullPath, fileName),
-        path.join(modelFullPath, `${fileName.replace(/\.gguf$/gi, ".bak")}`),
+        path.join(modelPath, fileName),
+        path.join(modelPath, `${fileName.replace(/\.gguf$/gi, ".bak")}`),
       );
     });
     afterEach(async () => {
       const fileName = modelCfg.source_url.split("/")?.pop() ?? "model.gguf";
       // Restore the extension of model file
       fs.renameSync(
-        path.join(modelFullPath, `${fileName.replace(/\.gguf$/gi, ".bak")}`),
-        path.join(modelFullPath, fileName),
+        path.join(modelPath, `${fileName.replace(/\.gguf$/gi, ".bak")}`),
+        path.join(modelPath, fileName),
       );
     });
     test(
       "should be able to detect model file by magic number",
       async () => {
-        const files = fs.readdirSync(modelFullPath) as string[];
+        const files = fs.readdirSync(modelPath) as string[];
         // Test checking magic bytes
         const res = await Promise.all(
-          files.map((f) =>
-            checkMagicBytes(path.join(modelFullPath, f), "GGUF"),
-          ),
+          files.map((f) => checkMagicBytes(path.join(modelPath, f), "GGUF")),
         );
         expect(res).toContain(true);
       },