nuxt-edge-ai 0.1.0

package/dist/runtime/server/utils/edge-ai-engine.js

@@ -0,0 +1,270 @@
+import { access, readFile } from "node:fs/promises";
+import { createRequire } from "node:module";
+import { dirname, extname, join, resolve } from "node:path";
+import { pathToFileURL } from "node:url";
+const state = {
+  loading: false,
+  warmed: false
+};
+const require = createRequire(resolve(process.cwd(), "package.json"));
+function contentTypeForExtension(extension) {
+  switch (extension) {
+    case ".json":
+      return "application/json";
+    case ".wasm":
+      return "application/wasm";
+    case ".txt":
+      return "text/plain; charset=utf-8";
+    default:
+      return "application/octet-stream";
+  }
+}
+function installFileFetchShim() {
+  const globalWithShim = globalThis;
+  if (globalWithShim.__nuxtEdgeAIFileFetchShim) {
+    return;
+  }
+  const originalFetch = globalThis.fetch.bind(globalThis);
+  globalThis.fetch = async (input, init) => {
+    const url = typeof input === "string" ? input : input instanceof URL ? input.href : input.url;
+    if (url.startsWith("file://")) {
+      const fileUrl = new URL(url);
+      const buffer = await readFile(fileUrl);
+      return new Response(buffer, {
+        headers: {
+          "content-type": contentTypeForExtension(extname(fileUrl.pathname).toLowerCase())
+        }
+      });
+    }
+    return originalFetch(input, init);
+  };
+  globalWithShim.__nuxtEdgeAIFileFetchShim = true;
+}
+async function resolveVendoredOnnxRuntimeBaseUrl() {
+  const candidates = /* @__PURE__ */ new Set();
+  try {
+    const moduleEntry = require.resolve("nuxt-edge-ai");
+    candidates.add(join(dirname(moduleEntry), "runtime", "server", "vendor", "onnxruntime"));
+  } catch {
+  }
+  candidates.add(resolve(process.cwd(), "node_modules", "nuxt-edge-ai", "dist", "runtime", "server", "vendor", "onnxruntime"));
+  candidates.add(resolve(process.cwd(), "dist", "runtime", "server", "vendor", "onnxruntime"));
+  candidates.add(resolve(process.cwd(), "src", "runtime", "server", "vendor", "onnxruntime"));
+  for (const candidate of candidates) {
+    try {
+      await access(join(candidate, "ort-wasm-simd-threaded.mjs"));
+      return `${pathToFileURL(candidate).href}/`;
+    } catch {
+    }
+  }
+  throw new Error("Unable to locate vendored onnxruntime-web assets on disk.");
+}
+async function loadTransformersRuntime() {
+  if (state.runtime) {
+    return state.runtime;
+  }
+  if (!state.initPromise) {
+    state.initPromise = (async () => {
+      installFileFetchShim();
+      const ortModule = await import("../vendor/onnxruntime/ort.wasm.min.mjs");
+      const ort = (ortModule.InferenceSession ? ortModule : ortModule.default) ?? ortModule;
+      const globalWithOrt = globalThis;
+      globalWithOrt[Symbol.for("onnxruntime")] = ort;
+      const transformers = await import("../vendor/huggingface/transformers.web.mjs");
+      state.runtime = transformers;
+    })().catch((error) => {
+      state.lastError = error instanceof Error ? error.message : String(error);
+      state.initPromise = void 0;
+      throw error;
+    });
+  }
+  await state.initPromise;
+  return state.runtime;
+}
+function ensureOnnxWasmEnv(env) {
+  const envWithMutableBackends = env;
+  envWithMutableBackends.backends ??= {};
+  envWithMutableBackends.backends.onnx ??= {};
+  envWithMutableBackends.backends.onnx.wasm ??= {};
+  return envWithMutableBackends.backends.onnx.wasm;
+}
+function resolveModelSource(config, modelOverride) {
+  const requestedModel = modelOverride?.trim();
+  const source = requestedModel || config.model.localPath || config.model.id;
+  return {
+    id: requestedModel || config.model.id,
+    task: config.model.task,
+    localPath: config.model.localPath,
+    allowRemote: config.model.allowRemote,
+    dtype: config.model.dtype,
+    source
+  };
+}
+function currentEngineState(config) {
+  return {
+    active: config.runtime,
+    ready: config.runtime === "mock" || Boolean(state.pipeline),
+    warmed: config.runtime === "mock" || state.warmed,
+    loading: state.loading,
+    cacheDir: config.cacheDir,
+    lastError: state.lastError
+  };
+}
+async function ensureTransformersPipeline(config, modelOverride) {
+  const model = resolveModelSource(config, modelOverride);
+  const modelKey = `${model.source}::${model.dtype ?? "default"}`;
+  if (state.pipeline && state.modelKey === modelKey) {
+    return { pipeline: state.pipeline, loadedNow: false, model };
+  }
+  if (!state.pipelinePromise || state.modelKey !== modelKey) {
+    state.loading = true;
+    state.pipelinePromise = (async () => {
+      const transformers = await loadTransformersRuntime();
+      transformers.env.allowLocalModels = Boolean(config.model.localPath);
+      transformers.env.allowRemoteModels = config.model.allowRemote;
+      transformers.env.useFS = false;
+      transformers.env.useFSCache = false;
+      if (transformers.env.useFSCache) {
+        transformers.env.cacheDir = config.cacheDir;
+      }
+      if (config.model.localPath) {
+        transformers.env.localModelPath = config.model.localPath;
+      }
+      const wasmBaseUrl = await resolveVendoredOnnxRuntimeBaseUrl();
+      const onnxWasmEnv = ensureOnnxWasmEnv(transformers.env);
+      onnxWasmEnv.wasmPaths = wasmBaseUrl;
+      onnxWasmEnv.numThreads = 1;
+      onnxWasmEnv.proxy = false;
+      const ortRuntime = globalThis[Symbol.for("onnxruntime")];
+      if (ortRuntime) {
+        ortRuntime.env ??= {};
+        ortRuntime.env.wasm ??= {};
+        ortRuntime.env.wasm.wasmPaths = wasmBaseUrl;
+        ortRuntime.env.wasm.numThreads = 1;
+        ortRuntime.env.wasm.proxy = false;
+      }
+      return transformers.pipeline(config.model.task, model.source, {
+        device: "wasm",
+        dtype: model.dtype
+      });
+    })().then((pipeline2) => {
+      state.pipeline = pipeline2;
+      state.modelKey = modelKey;
+      state.warmed = true;
+      state.lastError = void 0;
+      return pipeline2;
+    }).catch((error) => {
+      state.lastError = error instanceof Error ? error.message : String(error);
+      state.pipeline = void 0;
+      state.modelKey = void 0;
+      throw error;
+    }).finally(() => {
+      state.loading = false;
+    });
+  }
+  const pipeline = await state.pipelinePromise;
+  return { pipeline, loadedNow: true, model };
+}
+function resolveGenerationOptions(defaults, overrides) {
+  return {
+    maxNewTokens: overrides?.maxNewTokens ?? defaults.maxNewTokens,
+    temperature: overrides?.temperature ?? defaults.temperature,
+    topP: overrides?.topP ?? defaults.topP,
+    doSample: overrides?.doSample ?? defaults.doSample,
+    repetitionPenalty: overrides?.repetitionPenalty ?? defaults.repetitionPenalty
+  };
+}
+function extractGeneratedText(prompt, output) {
+  const firstItem = Array.isArray(output) ? output[0] : output;
+  if (firstItem && typeof firstItem === "object" && "generated_text" in firstItem) {
+    const generatedText = String(firstItem.generated_text ?? "");
+    return generatedText.startsWith(prompt) ? generatedText.slice(prompt.length).trim() : generatedText.trim();
+  }
+  if (typeof firstItem === "string") {
+    return firstItem.trim();
+  }
+  return JSON.stringify(output);
+}
+function runMockInference(config, input) {
+  const generation = resolveGenerationOptions(config.model.generation, input.generation);
+  const text = [
+    `Mock runtime for "${input.model || config.model.id}".`,
+    `Prompt received: ${input.prompt}`,
+    "Switch `edgeAI.runtime` to `transformers-wasm` to run the real WASM model runtime."
+  ].join(" ");
+  return {
+    text,
+    model: input.model || config.model.id,
+    runtime: "mock",
+    provider: "mock",
+    generation,
+    metrics: {
+      latencyMs: 0,
+      promptLength: input.prompt.length,
+      completionLength: text.length
+    }
+  };
+}
+export async function getEdgeAIHealth(config) {
+  const model = resolveModelSource(config);
+  if (config.runtime === "transformers-wasm" && config.warmup && !state.pipeline && !state.loading) {
+    try {
+      await ensureTransformersPipeline(config);
+    } catch {
+    }
+  }
+  return {
+    status: "ok",
+    runtime: config.runtime,
+    model,
+    defaults: config.model.generation,
+    engine: currentEngineState(config)
+  };
+}
+export async function pullEdgeAIModel(config) {
+  if (config.runtime === "mock") {
+    return {
+      status: "ready",
+      runtime: "mock",
+      model: resolveModelSource(config),
+      engine: currentEngineState(config),
+      loadedNow: false
+    };
+  }
+  const { loadedNow, model } = await ensureTransformersPipeline(config);
+  return {
+    status: "ready",
+    runtime: "transformers-wasm",
+    model,
+    engine: currentEngineState(config),
+    loadedNow
+  };
+}
+export async function generateEdgeAIText(config, input) {
+  if (config.runtime === "mock") {
+    return runMockInference(config, input);
+  }
+  const generation = resolveGenerationOptions(config.model.generation, input.generation);
+  const start = performance.now();
+  const { pipeline, model } = await ensureTransformersPipeline(config, input.model);
+  const output = await pipeline(input.prompt, {
+    max_new_tokens: generation.maxNewTokens,
+    temperature: generation.temperature,
+    top_p: generation.topP,
+    do_sample: generation.doSample,
+    repetition_penalty: generation.repetitionPenalty
+  });
+  const text = extractGeneratedText(input.prompt, output);
+  return {
+    text,
+    model: model.id,
+    runtime: "transformers-wasm",
+    provider: "transformers.js-wasm",
+    generation,
+    metrics: {
+      latencyMs: Number((performance.now() - start).toFixed(2)),
+      promptLength: input.prompt.length,
+      completionLength: text.length
+    }
+  };
+}