src/adapters/llm.ts

import type {
  ConversationTurn,
  LLMAdapter,
  LLMGenerateOptions,
  RuntimeWarmupOptions,
} from "../app/types";
import { SYSTEM_PROMPT } from "../prompts/system";

type ChatMessage = {
  role: "system" | "user" | "assistant";
  content: string;
};

type WebLLMModule = typeof import("@mlc-ai/web-llm");
type EngineType = Awaited<ReturnType<WebLLMModule["CreateMLCEngine"]>>;
type InitProgress = {
  progress?: number;
  text?: string;
  timeElapsed?: number;
};

const MODEL_PREFERENCES = [
  "Qwen2.5-0.5B-Instruct-q4f16_1-MLC",
  "SmolLM2-360M-Instruct-q4f16_1-MLC",
  "Llama-3.2-1B-Instruct-q4f16_1-MLC",
  "Qwen2.5-1.5B-Instruct-q4f16_1-MLC",
];

const formatProgress = (progress: InitProgress): string => {
  const percentage =
    typeof progress.progress === "number"
      ? `${Math.round(progress.progress * 100)}%`
      : null;

  if (progress.text && percentage) {
    return `${progress.text} (${percentage})`;
  }

  if (progress.text) {
    return progress.text;
  }

  if (percentage) {
    return `Loading language model (${percentage})`;
  }

  return "Loading language model...";
};

export class WebLLMChatAdapter implements LLMAdapter {
  #engine: EngineType | null = null;
  #modelId: string | null = null;

  constructor(private readonly preferredModelId?: string) {}

  async initialize(options: RuntimeWarmupOptions = {}): Promise<void> {
    if (this.#engine) {
      options.onProgress?.(
        `Language ready (${this.#modelId ?? this.preferredModelId ?? "model loaded"}).`,
      );
      return;
    }

    const webllm = await import("@mlc-ai/web-llm");
    const modelId =
      this.preferredModelId ?? WebLLMChatAdapter.#resolveModelId(webllm);
    options.onProgress?.(`Preparing ${modelId}...`);
    this.#engine = await webllm.CreateMLCEngine(modelId, {
      initProgressCallback: (progress: InitProgress) => {
        options.onProgress?.(formatProgress(progress));
      },
    });
    this.#modelId = modelId;
    options.onProgress?.(`Language ready (${modelId}).`);
  }

  async generate(options: LLMGenerateOptions): Promise<string> {
    await this.initialize();

    if (!this.#engine) {
      throw new Error("LLM engine failed to initialize.");
    }

    const messages = WebLLMChatAdapter.#toMessages(options.turns);
    const stream = await this.#engine.chat.completions.create({
      messages,
      temperature: 0.7,
      stream: true,
    });

    let finalText = "";

    for await (const chunk of stream) {
      if (options.signal?.aborted) {
        break;
      }

      const delta = chunk.choices[0]?.delta?.content ?? "";
      if (!delta) {
        continue;
      }

      finalText += delta;
      options.onChunk(delta);
    }

    return finalText.trim();
  }

  async warmup(options: RuntimeWarmupOptions = {}): Promise<void> {
    await this.initialize(options);

    if (!this.#engine) {
      throw new Error("LLM engine failed to initialize.");
    }

    options.onProgress?.("Compiling first language response...");
    await this.#engine.chat.completions.create({
      messages: [
        {
          role: "system",
          content: SYSTEM_PROMPT,
        },
        {
          role: "user",
          content: "Reply with one short word.",
        },
      ],
      temperature: 0,
      max_tokens: 1,
    });
    options.onProgress?.("Language runtime warmed.");
  }

  static #toMessages(turns: ConversationTurn[]): ChatMessage[] {
    const messages: ChatMessage[] = [
      {
        role: "system",
        content: SYSTEM_PROMPT,
      },
    ];

    for (const turn of turns) {
      if (turn.role === "system" || !turn.transcript.trim()) {
        continue;
      }

      messages.push({
        role: turn.role,
        content: turn.transcript,
      });
    }

    return messages;
  }

  static #resolveModelId(webllm: WebLLMModule): string {
    const candidates = webllm.prebuiltAppConfig?.model_list ?? [];
    const normalized = candidates
      .map((candidate) => {
        const modelId = "model_id" in candidate ? candidate.model_id : "";
        return { candidate, modelId };
      })
      .filter((entry) => Boolean(entry.modelId));

    for (const preference of MODEL_PREFERENCES) {
      const match = normalized.find((entry) =>
        entry.modelId.includes(preference),
      );
      if (match) {
        return match.modelId;
      }
    }

    const fallback = normalized.at(0)?.modelId;
    if (!fallback) {
      throw new Error("WebLLM did not expose any prebuilt models.");
    }

    return fallback;
  }
}