src/engines/onnx-transformers-engine.js

import {
  AutoConfig,
  AutoModelForImageTextToText,
  AutoProcessor,
  env,
  InterruptableStoppingCriteria,
  RawImage,
  TextStreamer,
} from '@huggingface/transformers';
import { createToolStreamFilter, displayTextFromRaw, parseToolCalls } from '../tools/tool-protocol.js';
import { prepareConversation } from './conversation-preparation.js';

const MODEL_CACHE_KEY = 'liquid-lfm-models-v4';

export async function createEngine({ manifest, telemetry }) {
  let model = null;
  let processor = null;
  let gpuDevice = null;
  const gpuHealth = {
    lost: false,
    lossReason: null,
    lossMessage: null,
    lastUncapturedError: null,
  };
  const repo = manifest.model;
  let revision = manifest.revision || 'main';

  function emit(level, message, detail = '') {
    telemetry({ level, message, detail });
  }

  function configureRuntime() {
    env.allowLocalModels = false;
    env.allowRemoteModels = true;
    env.useBrowserCache = true;
    env.cacheKey = MODEL_CACHE_KEY;
    env.backends.onnx.logLevel = 'warning';
    env.backends.onnx.webgpu.powerPreference = 'high-performance';
  }

  async function attachGpuErrorHandlers() {
    gpuDevice = await env.backends.onnx.webgpu.device;
    if (!gpuDevice) return;

    gpuDevice.addEventListener?.('uncapturederror', event => {
      const error = event.error;
      const detail = {
        type: error?.constructor?.name || error?.name || 'GPUError',
        message: error?.message || String(error || 'Unknown WebGPU error'),
      };
      gpuHealth.lastUncapturedError = { ...detail, time: new Date().toISOString() };
      emit('error', 'WebGPU reported an uncaptured error', detail);
    });

    void gpuDevice.lost.then(info => {
      gpuHealth.lost = true;
      gpuHealth.lossReason = info?.reason || 'unknown';
      gpuHealth.lossMessage = info?.message || '';
      emit('error', 'WebGPU device lost', {
        reason: gpuHealth.lossReason,
        message: gpuHealth.lossMessage,
        recovery: 'Reload the page to create a fresh GPU device. Cached model files will be reused.',
      });
    });
  }

  function progressHandler(onProgress) {
    let lastPercent = -1;
    let lastLoaded = 0;
    let lastSampleAt = performance.now();
    let lastObservedLoaded = 0;
    let smoothedBytesPerSecond = 0;
    const handler = progress => {
      if (progress.status !== 'progress_total') return;
      const percent = Math.max(0, Math.min(100, progress.progress || 0));
      if (progress.loaded > lastObservedLoaded) {
        lastObservedLoaded = progress.loaded;
        handler.lastUpdateAt = performance.now();
      }
      if (percent - lastPercent < 0.25 && percent < 100) return;
      lastPercent = percent;
      const now = performance.now();
      const elapsedSeconds = (now - lastSampleAt) / 1000;
      if (elapsedSeconds >= 0.5 && progress.loaded >= lastLoaded) {
        const currentRate = (progress.loaded - lastLoaded) / elapsedSeconds;
        smoothedBytesPerSecond = smoothedBytesPerSecond
          ? (smoothedBytesPerSecond * 0.7) + (currentRate * 0.3)
          : currentRate;
        lastLoaded = progress.loaded;
        lastSampleAt = now;
      }
      const activeFiles = Object.entries(progress.files || {}).filter(([, value]) => value.loaded < value.total);
      const activeName = [...activeFiles].sort(([, left], [, right]) => (right.total - right.loaded) - (left.total - left.loaded))[0]?.[0];
      const activeLabel = activeName?.split('/').at(-1);
      const rate = smoothedBytesPerSecond > 0 ? formatRate(smoothedBytesPerSecond) : '';
      const remainingSeconds = smoothedBytesPerSecond > 0
        ? Math.max(0, progress.total - progress.loaded) / smoothedBytesPerSecond
        : 0;
      const eta = remainingSeconds > 1 ? formatDuration(remainingSeconds) : '';
      const transferSummary = [rate, eta ? `${eta} left` : ''].filter(Boolean).join(' · ');
      const file = activeFiles.length > 1
        ? `Downloading ${activeLabel} +${activeFiles.length - 1}${transferSummary ? ` · ${transferSummary}` : ''}`
        : activeLabel || `Preparing ${repo}`;
      const snapshot = {
        status: 'loading',
        progress: percent,
        file,
        loaded: progress.loaded,
        total: progress.total,
        activeDownloads: activeFiles.length,
        activeFile: activeName || null,
        bytesPerSecond: Math.round(smoothedBytesPerSecond),
      };
      handler.latest = snapshot;
      handler.lastUpdateAt = performance.now();
      onProgress(snapshot);
    };
    handler.latest = null;
    handler.lastUpdateAt = performance.now();
    return handler;
  }

  return {
    backend: 'ONNX · Transformers.js · strict WebGPU',

    async load(onProgress) {
      configureRuntime();
      await navigator.storage?.persist?.().catch(() => false);
      if (revision === 'main') {
        revision = await resolveMainRevision(repo);
      }

      const options = {
        revision,
        device: 'webgpu',
        dtype: manifest.runtime.dtype,
        use_external_data_format: manifest.runtime.externalDataChunks,
        session_options: {
          executionProviders: ['webgpu'],
          logSeverityLevel: 2,
        },
      };

      const trackedProgress = progressHandler(onProgress);
      options.progress_callback = trackedProgress;
      const modelConfig = await AutoConfig.from_pretrained(repo, { ...options, progress_callback: null });
      modelConfig['transformers.js_config'] = {
        ...(modelConfig['transformers.js_config'] || {}),
        dtype: manifest.runtime.dtype,
        device: manifest.runtime.device,
        use_external_data_format: manifest.runtime.externalDataChunks,
      };
      options.config = modelConfig;
      let stallReported = false;
      const stallWatchdog = setInterval(() => {
        const idleSeconds = Math.floor((performance.now() - trackedProgress.lastUpdateAt) / 1000);
        if (!trackedProgress.latest || trackedProgress.latest.progress >= 100 || idleSeconds < 30) return;
        onProgress({
          ...trackedProgress.latest,
          bytesPerSecond: 0,
          file: `No transfer progress for ${idleSeconds}s · waiting on a large model shard`,
        });
        if (!stallReported && idleSeconds >= 90) {
          stallReported = true;
          emit('warn', 'Model download has not advanced for 90 seconds', {
            progress: Math.round(trackedProgress.latest.progress),
            activeDownloads: trackedProgress.latest.activeDownloads,
            activeFile: trackedProgress.latest.activeFile,
            likelyCause: 'Large browser shard, memory pressure, cache serialization, or interrupted CDN stream',
          });
        }
      }, 5000);
      try {
        [processor, model] = await Promise.all([
          AutoProcessor.from_pretrained(repo, options),
          AutoModelForImageTextToText.from_pretrained(repo, options),
        ]);
      } finally {
        clearInterval(stallWatchdog);
      }

      await attachGpuErrorHandlers();
      onProgress({ status: 'done', progress: 100, file: 'Model ready' });
    },

    async generate(messages, options = {}) {
      if (!model || !processor) throw new Error('The model is not loaded.');
      const prepared = prepareConversation(messages);
      const prompt = processor.apply_chat_template(prepared.messages, {
        add_generation_prompt: true,
        tokenize: false,
        tools: options.tools || [],
      });
      const images = await Promise.all(prepared.imageUrls.map(url => RawImage.read(url)));
      const inputs = images.length
        ? await processor(images, prompt)
        : processor.tokenizer(prompt, { add_special_tokens: false });
      const promptTokens = inputs.input_ids?.dims?.at(-1)
        ?? inputs.inputs_embeds?.dims?.at(-2)
        ?? inputs.attention_mask?.dims?.at(-1)
        ?? null;
      const stopping = new InterruptableStoppingCriteria();
      const abort = () => stopping.interrupt();
      options.signal?.addEventListener('abort', abort, { once: true });
      let streamedText = '';
      const streamFilter = createToolStreamFilter(chunk => {
        streamedText += chunk;
        options.onToken?.(chunk, null);
      }, {
        onToolCallStart: () => options.onToolCallState?.('preparing'),
        onToolCallEnd: () => options.onToolCallState?.('parsing'),
      });
      const streamer = new TextStreamer(processor.tokenizer, {
        skip_prompt: true,
        skip_special_tokens: false,
        callback_function: chunk => streamFilter.push(chunk),
      });

      let generated;
      try {
        try {
          generated = await model.generate({
            ...inputs,
            max_new_tokens: options.maxNewTokens || 384,
            do_sample: (options.temperature || 0) > 0,
            temperature: Math.max(options.temperature || 0, 0.01),
            top_p: options.topP || 0.9,
            top_k: Number.isInteger(options.topK) ? options.topK : 50,
            streamer,
            stopping_criteria: [stopping],
          });
        } catch (error) {
          const isInvalidBuffer = /Mapping WebGPU buffer failed: Invalid buffer/i.test(error.message);
          const gpuError = gpuHealth.lastUncapturedError;
          const isOutOfMemory = gpuError?.type === 'GPUOutOfMemoryError'
            || /out of memory/i.test(gpuError?.message || '')
            || /out of memory/i.test(error.message);
          emit('error', isInvalidBuffer ? 'WebGPU buffer readback failed' : 'WebGPU generation failed', {
            error: error.message,
            promptTokens,
            requestedMaxNewTokens: options.maxNewTokens || 384,
            deviceLost: gpuHealth.lost,
            deviceLossReason: gpuHealth.lossReason,
            deviceLossMessage: gpuHealth.lossMessage,
            lastUncapturedGpuError: gpuHealth.lastUncapturedError,
            interpretation: isInvalidBuffer
              ? 'ORT could not map a GPU result staging buffer. Inspect the preceding GPU error/device-loss event; this is not a context-length error by itself.'
              : null,
          });
          if (isOutOfMemory) {
            throw new Error("This browser's WebGPU session ran out of available GPU memory. This is a browser/WebGPU memory limit, not the model's context-length limit. Try a shorter conversation, fewer or smaller images, or a lower max-new-tokens setting, then reload the page.", { cause: error });
          }
          if (isInvalidBuffer) {
            throw new Error('The browser could not read a WebGPU model buffer. This often follows browser GPU-memory or resource exhaustion. Reload the page to reset the GPU session; cached model files will be reused.', { cause: error });
          }
          throw error;
        }
      } finally {
        streamFilter.finish();
        options.signal?.removeEventListener('abort', abort);
      }

      const sequences = generated?.sequences ?? generated;
      const promptLength = inputs.input_ids?.dims?.at(-1) || 0;
      const generatedIds = sequences?.tolist?.()?.[0]?.slice(promptLength) || [];
      const decodedText = generatedIds.length
        ? processor.tokenizer.decode(generatedIds, { skip_special_tokens: false }).trim()
        : '';
      const rawText = decodedText || streamedText.trim();
      let toolCalls;
      try {
        toolCalls = parseToolCalls(rawText);
      } catch (error) {
        error.rawModelOutput = rawText;
        emit('warn', 'Tool-call parsing failed', {
          error: error.message,
          rawOutputCharacters: rawText.length,
        });
        throw error;
      }
      const finalText = displayTextFromRaw(rawText) || streamedText.trim();
      if (!finalText) {
        emit('warn', 'Generation returned no displayable text', {
          sequenceTokens: generatedIds.length,
        });
      }
      return {
        text: finalText,
        toolCalls,
        rawOutput: rawText,
        finishReason: toolCalls.length ? 'tool_calls' : options.signal?.aborted ? 'stopped' : 'stop',
      };
    },

    async clearCache() {
      if (!globalThis.caches) return { cleared: false, entriesDeleted: 0 };
      const cacheName = env.cacheKey || MODEL_CACHE_KEY;
      const cache = await caches.open(cacheName);
      const entriesDeleted = (await cache.keys()).length;
      const cleared = await caches.delete(cacheName);
      if (!cleared) emit('warn', 'Browser model cache was already empty', { cacheName, entriesDeleted });
      return { cleared, entriesDeleted };
    },

    async cacheInfo() {
      if (!globalThis.caches) return { used: 0, available: 0 };
      const cache = await caches.open(env.cacheKey || MODEL_CACHE_KEY);
      const keys = await cache.keys();
      let used = 0;
      for (const key of keys) {
        const response = await cache.match(key);
        used += Number(response?.headers.get('content-length') || 0);
      }
      const estimate = await navigator.storage?.estimate?.();
      return { used, available: estimate?.quota || 0 };
    },

    clearConversationCache() {},

    async dispose() {
      await model?.dispose?.();
      model = null;
      processor = null;
    },
  };
}

async function resolveMainRevision(repo) {
  const response = await env.fetch(`https://huggingface.co/api/models/${repo}`);
  if (!response.ok) throw new Error(`Could not resolve the model main branch (${response.status}).`);
  const info = await response.json();
  if (!/^[0-9a-f]{40}$/i.test(info.sha || '')) throw new Error('Hugging Face returned an invalid model revision.');
  return info.sha;
}

function formatRate(bytesPerSecond) {
  return `${(bytesPerSecond / 1024 / 1024).toFixed(1)} MB/s`;
}

function formatDuration(seconds) {
  if (seconds < 60) return `${Math.ceil(seconds)}s`;
  const minutes = Math.ceil(seconds / 60);
  if (minutes < 60) return `${minutes}m`;
  const hours = Math.floor(minutes / 60);
  return `${hours}h ${minutes % 60}m`;
}