@ariacompute/engine-ts 1.5.0 → 1.6.0

package/README.md

@@ -11,3 +11,16 @@ const eng = new Engine("/path/to/bundle");
 console.log(eng.complete([{ role: "user", content: "hi" }]));
 eng.close();
 ```
+
+## Auto-download by model name
+
+`Engine.open(modelRef, { token, site })` accepts a local bundle path **or** an
+Aria model name (e.g. `gemma-4-e2b-it_q4`). A value containing `/` or already on
+disk is loaded directly; otherwise the SDK downloads it from the Dashboard
+private source (requires `token`; `site` defaults to `https://ariacompute.com`)
+into `~/.ariacompute/models/{model}` and then loads it. A valid cached bundle is
+reused without re-downloading.
+
+```ts
+const eng = await Engine.open("gemma-4-e2b-it_q4", { token: "DASHBOARD_TOKEN" });
+```

package/dist/index.d.ts

@@ -1,13 +1,63 @@
-export type ChatMessage = {
-    role: string;
+export interface GenerateOptions {
+    max_tokens?: number;
+    temperature?: number;
+    top_p?: number;
+    stream?: boolean;
+    stop?: string[];
+    response_format?: {
+        type: "text" | "json_object" | "json_schema";
+        json_schema?: unknown;
+    };
+}
+export interface Turn {
+    role: "system" | "user" | "assistant";
     content: string;
+}
+export interface Generation {
+    success: boolean;
+    response: string;
+    error?: string;
+    tokens_used?: number;
+    finish_reason?: string;
+}
+export interface CompleteResult {
+    success: boolean;
+    response: string;
+    error?: string;
+    generation?: Generation;
+}
+export interface OpenOptions {
+    /** Dashboard bearer token. Required when modelRef is a model name. */
+    token?: string;
+    /** Dashboard base URL. Defaults to https://ariacompute.com. */
+    site?: string;
+    /** Explicit path to the FFI library. */
+    ffiLib?: string;
+}
+/** Parse `slug`/`quant` from a model name such as `gemma-4-e2b-it_q4`. */
+export declare function parseBundleName(model: string): {
+    slug: string;
+    quant: string;
 };
+/** Download `model` from the Dashboard private source into
+ * `~/.ariacompute/models/{model}` and return that directory.
+ * Skips the download when a valid bundle is already cached. */
+export declare function downloadModel(model: string, token: string, site?: string): Promise<string>;
+export declare function isLocalRef(modelRef: string): boolean;
 export declare class Engine {
     private lib;
     private handle;
-    constructor(bundlePath: string);
+    /** Construct from a local bundle directory. */
+    constructor(bundle: string, opts?: OpenOptions);
+    /** Auto-detect: a value containing a separator or already on disk is a local
+     * path; otherwise it is a model name that is downloaded (requires `token`)
+     * then loaded. */
+    static open(modelRef: string, opts?: OpenOptions): Promise<Engine>;
+    complete(messages: Turn[], opts?: GenerateOptions): CompleteResult;
+    embed(text: string): number[];
+    transcribe(pcm: Uint8Array): string;
     close(): void;
-    complete(messages: ChatMessage[], options?: Record<string, unknown>, tools?: unknown[]): any;
-    embed(input: string): any;
-    transcribe(pcm: Buffer): any;
 }
+export declare const _internal: {
+    encoder: TextEncoder;
+};

package/dist/index.js

@@ -1,75 +1,241 @@
 "use strict";
-var __importDefault = (this && this.__importDefault) || function (mod) {
-    return (mod && mod.__esModule) ? mod : { "default": mod };
-};
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.Engine = void 0;
-const koffi_1 = __importDefault(require("koffi"));
-function loadLib() {
-    const path = process.env.ARIA_FFI_LIB;
-    if (!path)
-        throw new Error("Set ARIA_FFI_LIB");
-    const lib = koffi_1.default.load(path);
-    return {
-        init: lib.func("aria_model_init", "void *", ["str"]),
-        destroy: lib.func("aria_model_destroy", "void", ["void *"]),
-        complete: lib.func("aria_complete", "int", [
-            "void *",
-            "str",
-            "str",
-            "str",
-            "uint8 *",
-            "size_t",
-        ]),
-        embed: lib.func("aria_embed", "int", ["void *", "str", "uint8 *", "size_t"]),
-        transcribe: lib.func("aria_transcribe", "int", [
-            "void *",
-            "uint8 *",
-            "size_t",
-            "str",
-            "uint8 *",
-            "size_t",
-        ]),
-        lastError: lib.func("aria_last_error", "str", []),
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
     };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+Object.defineProperty(exports, "__esModule", { value: true });
+exports._internal = exports.Engine = void 0;
+exports.parseBundleName = parseBundleName;
+exports.downloadModel = downloadModel;
+exports.isLocalRef = isLocalRef;
+const fs = __importStar(require("node:fs"));
+const os = __importStar(require("node:os"));
+const path = __importStar(require("node:path"));
+const fflate = __importStar(require("fflate"));
+const DEFAULT_SITE = "https://ariacompute.com";
+function ariaHome() {
+    return process.env.ARIA_COMPUTE_HOME || path.join(os.homedir(), ".ariacompute");
+}
+function cacheDir(model) {
+    return path.join(ariaHome(), "models", model);
+}
+/** Parse `slug`/`quant` from a model name such as `gemma-4-e2b-it_q4`. */
+function parseBundleName(model) {
+    if (!model || model.includes("/") || model.includes("\\")) {
+        throw new Error(`invalid model name: ${model}`);
+    }
+    const idx = model.lastIndexOf("_q");
+    if (idx !== -1) {
+        const slug = model.slice(0, idx);
+        const suffix = model.slice(idx + 2);
+        const quant = suffix === "4"
+            ? "int4"
+            : suffix === "8"
+                ? "int8"
+                : suffix === "326" || suffix === "3.26"
+                    ? "int326"
+                    : (() => {
+                        throw new Error(`unknown quant suffix _q${suffix}`);
+                    })();
+        if (!slug)
+            throw new Error(`invalid model name: ${model}`);
+        return { slug, quant };
+    }
+    return { slug: model, quant: "int4" };
 }
-function parseOut(buf) {
-    const text = buf.toString("utf8").replace(/\0.*$/s, "");
-    return JSON.parse(text);
+function isValidBundle(dir) {
+    try {
+        if (!fs.statSync(path.join(dir, "weight.bin")).isFile())
+            return false;
+        const meta = JSON.parse(fs.readFileSync(path.join(dir, "config.json"), "utf8"));
+        return meta.format === "aria-quant-bundle";
+    }
+    catch {
+        return false;
+    }
+}
+function flattenSingleSubdir(dir) {
+    const entries = fs.readdirSync(dir).filter((e) => !e.startsWith("."));
+    if (entries.length !== 1)
+        return;
+    const only = path.join(dir, entries[0]);
+    if (!fs.statSync(only).isDirectory())
+        return;
+    if (!fs.existsSync(path.join(only, "config.json")))
+        return;
+    for (const name of fs.readdirSync(only)) {
+        fs.renameSync(path.join(only, name), path.join(dir, name));
+    }
+    fs.rmdirSync(only);
+}
+async function extractZip(data, dest) {
+    if (data[0] !== 0x50 || data[1] !== 0x4b) {
+        throw new Error("downloaded archive is not a valid zip");
+    }
+    fs.mkdirSync(dest, { recursive: true });
+    const entries = fflate.unzipSync(data);
+    for (const [name, content] of Object.entries(entries)) {
+        const out = path.join(dest, name);
+        if (name.endsWith("/")) {
+            fs.mkdirSync(out, { recursive: true });
+        }
+        else {
+            fs.mkdirSync(path.dirname(out), { recursive: true });
+            fs.writeFileSync(out, content);
+        }
+    }
+    flattenSingleSubdir(dest);
+}
+async function httpGet(url, token) {
+    const resp = await fetch(url, { headers: { Authorization: `Bearer ${token}` } });
+    if (!resp.ok)
+        return { ok: false, status: resp.status, body: null };
+    const ct = resp.headers.get("content-type") || "";
+    if (ct.includes("json"))
+        return { ok: true, status: resp.status, body: await resp.json() };
+    return { ok: true, status: resp.status, body: Buffer.from(await resp.arrayBuffer()) };
+}
+const encoder = new TextEncoder();
+/** Download `model` from the Dashboard private source into
+ * `~/.ariacompute/models/{model}` and return that directory.
+ * Skips the download when a valid bundle is already cached. */
+async function downloadModel(model, token, site = DEFAULT_SITE) {
+    if (!token)
+        throw new Error("dashboard token is required to download a model");
+    const { slug, quant } = parseBundleName(model);
+    const cache = cacheDir(model);
+    if (fs.existsSync(cache) && isValidBundle(cache))
+        return cache;
+    const metaUrl = `${site.replace(/\/$/, "")}/api/dashboard/models/${encodeURIComponent(slug)}/download?quant=${encodeURIComponent(quant)}&sdk=v1.0&format=json`;
+    const metaResp = await httpGet(metaUrl, token);
+    if (!metaResp.ok)
+        throw new Error(`dashboard request failed: ${metaResp.status}`);
+    const meta = metaResp.body;
+    if (!meta.url)
+        throw new Error("dashboard meta returned empty url");
+    const zipResp = await httpGet(meta.url, token);
+    if (!zipResp.ok)
+        throw new Error(`download stream failed: ${zipResp.status}`);
+    const data = zipResp.body;
+    const staging = path.join(ariaHome(), "models", `.${model}.partial`);
+    if (fs.existsSync(staging))
+        fs.rmSync(staging, { recursive: true, force: true });
+    await extractZip(data, staging);
+    if (!isValidBundle(staging)) {
+        fs.rmSync(staging, { recursive: true, force: true });
+        throw new Error("downloaded archive did not contain a valid aria-quant-bundle");
+    }
+    if (fs.existsSync(cache))
+        fs.rmSync(cache, { recursive: true, force: true });
+    fs.renameSync(staging, cache);
+    return cache;
+}
+function loadLib(ffiLib) {
+    const libPath = ffiLib || process.env.ARIA_FFI_LIB;
+    if (!libPath)
+        throw new Error("aria FFI lib not found; set ARIA_FFI_LIB or pass ffiLib");
+    const koffiNS = require("koffi");
+    return koffiNS.load(libPath);
+}
+function isLocalRef(modelRef) {
+    return modelRef.includes("/") || modelRef.includes("\\") || fs.existsSync(modelRef);
 }
 class Engine {
-    constructor(bundlePath) {
-        this.lib = loadLib();
-        this.handle = this.lib.init(bundlePath);
-        if (!this.handle)
-            throw new Error(this.lib.lastError() || "init failed");
+    /** Construct from a local bundle directory. */
+    constructor(bundle, opts = {}) {
+        this.lib = loadLib(opts.ffiLib);
+        this.lib.func("aria_model_init", "void*", ["str"]);
+        this.lib.func("aria_model_destroy", "void", ["void*"]);
+        this.lib.func("aria_complete", "int", ["void*", "str", "str", "str", "str", "size_t"]);
+        this.lib.func("aria_embed", "int", ["void*", "str", "str", "size_t"]);
+        this.lib.func("aria_transcribe", "int", ["void*", "void*", "size_t", "str", "str", "size_t"]);
+        this.lib.func("aria_last_error", "str", []);
+        this.handle = this.lib.aria_model_init(bundle);
+        if (!this.handle) {
+            const err = this.lib.aria_last_error();
+            throw new Error(err || "init failed");
+        }
     }
-    close() {
-        if (this.handle) {
-            this.lib.destroy(this.handle);
-            this.handle = null;
+    /** Auto-detect: a value containing a separator or already on disk is a local
+     * path; otherwise it is a model name that is downloaded (requires `token`)
+     * then loaded. */
+    static async open(modelRef, opts = {}) {
+        if (isLocalRef(modelRef))
+            return new Engine(modelRef, opts);
+        if (!opts.token) {
+            throw new Error(`model name '${modelRef}' requires a dashboard token to download`);
         }
+        const bundle = await downloadModel(modelRef, opts.token, opts.site ?? DEFAULT_SITE);
+        return new Engine(bundle, opts);
     }
-    complete(messages, options = { max_tokens: 16 }, tools = []) {
-        const out = Buffer.alloc(256 * 1024);
-        const rc = this.lib.complete(this.handle, JSON.stringify(messages), JSON.stringify(options), JSON.stringify(tools), out, out.length);
-        if (rc !== 0)
-            throw new Error(this.lib.lastError() || "complete failed");
-        return parseOut(out);
+    complete(messages, opts = {}) {
+        const payload = JSON.stringify({ messages, opts });
+        const buf = Buffer.alloc(1 << 16);
+        const rc = this.lib.aria_complete(this.handle, payload, "", "", buf, buf.length);
+        if (rc !== 0) {
+            const err = this.lib.aria_last_error();
+            return { success: false, response: "", error: err || "complete failed" };
+        }
+        try {
+            const parsed = JSON.parse(buf.toString("utf8").replace(/\0+$/, ""));
+            return { success: true, response: parsed.response ?? "", generation: parsed };
+        }
+        catch {
+            return { success: true, response: buf.toString("utf8").replace(/\0+$/, "") };
+        }
     }
-    embed(input) {
-        const out = Buffer.alloc(256 * 1024);
-        const rc = this.lib.embed(this.handle, JSON.stringify({ input }), out, out.length);
-        if (rc !== 0)
-            throw new Error(this.lib.lastError() || "embed failed");
-        return parseOut(out);
+    embed(text) {
+        const buf = Buffer.alloc(1 << 20);
+        const rc = this.lib.aria_embed(this.handle, text, buf, buf.length);
+        if (rc !== 0) {
+            const err = this.lib.aria_last_error();
+            throw new Error(err || "embed failed");
+        }
+        return JSON.parse(buf.toString("utf8").replace(/\0+$/, ""));
     }
     transcribe(pcm) {
-        const out = Buffer.alloc(64 * 1024);
-        const rc = this.lib.transcribe(this.handle, pcm, pcm.length, null, out, out.length);
-        if (rc !== 0)
-            throw new Error(this.lib.lastError() || "transcribe failed");
-        return parseOut(out);
+        const buf = Buffer.alloc(1 << 16);
+        const rc = this.lib.aria_transcribe(this.handle, pcm, pcm.length, "", buf, buf.length);
+        if (rc !== 0) {
+            const err = this.lib.aria_last_error();
+            throw new Error(err || "transcribe failed");
+        }
+        return buf.toString("utf8").replace(/\0+$/, "");
+    }
+    close() {
+        if (this.handle)
+            this.lib.aria_model_destroy(this.handle);
+        this.handle = null;
     }
 }
 exports.Engine = Engine;
+exports._internal = { encoder };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ariacompute/engine-ts",
-  "version": "1.5.0",
+  "version": "1.6.0",
   "description": "TypeScript/Node SDK for Aria Engine",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -10,6 +10,7 @@
   },
   "license": "MIT",
   "dependencies": {
+    "fflate": "^0.8.2",
     "koffi": "^2.8.0"
   },
   "devDependencies": {