Local LLM (WebLLM)

export type WebLLMModelTier = "ultra-low" | "balanced" | "quality";

export type WebLLMModelProfile = {
  id: string;
  label: string;
  tier: WebLLMModelTier;
  approxDownloadMB: number;
  notes?: string;
  experimental?: boolean;
};

export function recommendWebLLMModel(): WebLLMRecommendation {
  const memory = (navigator as NavigatorWithMemory).deviceMemory;
  
  // Low memory devices (≤4GB) → SmolLM2
  if (memory && memory <= 4) {
    return {
      modelId: WEBLLM_FALLBACK_MODEL_ID,
      reason: "low-memory-device"
    };
  }
  
  // Default → Llama 3.2 1B
  return {
    modelId: WEBLLM_PRIMARY_MODEL_ID,
    reason: "default"
  };
}

class WebLLMEngineManager {
  private engine: MLCEngineInterface | null = null;
  private activeModelId: string | null = null;
  private loadingPromise: Promise<void> | null = null;

  async ensureReady(
    modelId: string,
    options: WebLLMEnsureReadyOptions
  ): Promise<void> {
    // Check WebGPU support
    if (!this.supportsWebGPU()) {
      throw new WebLLMProviderError(
        "WEBLLM_UNSUPPORTED",
        "WebGPU is not available in this browser."
      );
    }

    // Already loaded?
    if (this.engine && this.activeModelId === modelId) {
      return;
    }

    // Request download consent
    if (!options.allowDownload) {
      throw new WebLLMProviderError(
        "WEBLLM_DOWNLOAD_REQUIRED",
        "Model download consent is required."
      );
    }

    // Initialize or reload
    this.loadingPromise = (async () => {
      if (!this.engine) {
        this.engine = await CreateMLCEngine(modelId, {
          initProgressCallback: (report) => {
            const progress = Math.max(0, Math.min(1, report.progress));
            options.onProgress?.(progress, report.text || "Preparing model");
          },
          logLevel: "INFO"
        });
        this.activeModelId = modelId;
      } else {
        await this.engine.reload(modelId);
        this.activeModelId = modelId;
      }
    })();

    await this.loadingPromise;
    this.loadingPromise = null;
  }
}

async stream(
  modelId: string,
  messages: WebLLMMessage[],
  signal: AbortSignal,
  onToken: (token: string) => void
): Promise<void> {
  if (!this.engine || this.activeModelId !== modelId) {
    throw new WebLLMProviderError(
      "WEBLLM_INIT_FAILED",
      "WebLLM engine is not initialized."
    );
  }

  try {
    const chunkStream = await this.engine.chat.completions.create({
      model: modelId,
      messages,
      stream: true,
      temperature: 0.2,
      max_tokens: 700
    });

    for await (const chunk of chunkStream as AsyncIterable<ChatCompletionChunk>) {
      if (signal.aborted) {
        throw new DOMException("aborted", "AbortError");
      }
      const token = chunk.choices?.[0]?.delta?.content ?? "";
      if (token.length > 0) {
        onToken(token);
      }
    }
  } catch (error) {
    if (error instanceof DOMException && error.name === "AbortError") {
      throw error;
    }
    throw new WebLLMProviderError(
      "WEBLLM_STREAM_FAILED",
      `WebLLM stream failed: ${error.message}`
    );
  }
}

{
  temperature: 0.2,      // Low temperature for consistent, factual responses
  max_tokens: 700,       // Limit response length for speed
  stream: true           // Token-by-token streaming
}

private supportsWebGPU(): boolean {
  const nav = navigator as Navigator & { gpu?: object };
  return typeof nav.gpu !== "undefined";
}

<WebLLMDownloadDialog
  open={showWebLLMDownload}
  onOpenChange={setShowWebLLMDownload}
  onConfirm={() => {
    setWebLLMDownloadConsent(true);
    setShowWebLLMDownload(false);
  }}
  onCancel={() => {
    setShowWebLLMDownload(false);
  }}
  modelProfile={selectedWebLLMModel}
  downloadProgress={webllmDownloadProgress}
  downloadStatus={webllmDownloadStatus}
/>

export type WebLLMErrorCode =
  | "WEBLLM_UNSUPPORTED"       // WebGPU not available
  | "WEBLLM_DOWNLOAD_REQUIRED"  // Consent needed
  | "WEBLLM_INIT_FAILED"        // Initialization error
  | "WEBLLM_STREAM_FAILED";     // Inference error

export class WebLLMProviderError extends Error {
  readonly code: WebLLMErrorCode;

  constructor(code: WebLLMErrorCode, message: string) {
    super(message);
    this.name = "WebLLMProviderError";
    this.code = code;
  }
}

VITE_WEBLLM_ENABLED=true  # Enable WebLLM (default)
VITE_WEBLLM_ENABLED=false # Disable WebLLM

export function isWebLLMEnabled(): boolean {
  const raw = import.meta.env.VITE_WEBLLM_ENABLED;
  return raw === "1" || raw === "true";
}