storytellurx-v1/app/game-engine/models-client.ts

import type {
  ChatCompletionRequestStreaming,
  WebWorkerMLCEngine,
} from "@mlc-ai/web-llm";
import type * as webLLMType from "@mlc-ai/web-llm";
import type { ModelDefinition } from "./models-server.ts";
import {
  convertToModelMessages,
  createUIMessageStreamResponse,
  Output,
  streamText,
} from "ai";
import { addFullGameStateAsync, getGameState } from "./game-state.ts";
import { constructSystemPrompt } from "./prompt-constructor.ts";
import { AI_PERSONAS, type AIPersonaId } from "./personas.ts";
import {
  EroticNarrativeResponseSchema,
  type MessageMetadata,
  SystemNarrativeResponseSchema,
} from "./game-schemas.ts";
import {
  applyMindControlRewrites,
  includeUserHintForSystem,
  inludeSystemAsAssistant,
} from "./generation-utils.ts";
import { createOpenAI } from "@ai-sdk/openai";

interface EngineOptions {
  onProgress?: (v: unknown) => void;
}

export const AVAILABLE_CLIENT_MODELS: readonly ModelDefinition[] = [
  {
    id: "Hermes-3-Llama-3.1-8B-q4f16_1-MLC",
    name: "Hermes-3-Llama-3.1-8B-q4f16_1-MLC (WebLLM; On-Device)",
    type: "client",
    getModel: (engineOptions) =>
      createOpenAI({
        name: "webllm",
        baseURL: "http://webllm/v1",
        apiKey: "NO",
        fetch: webLLMFetch(engineOptions),
      })("Hermes-3-Llama-3.1-8B-q4f16_1-MLC"),
    parameters: {
      maxOutputTokens: 9000,
      temperature: 1,
      // temperature=0.8, repetition_penalty=1.1,
      frequencyPenalty: 1.1,
      //presencePenalty: 2,
    },
    features: {
      developerInstructionDisabled: true,
      maxPersonas: 2,
      guideStepAdditionalUserText:
        "Oh *Guide*, dirige moi vers un des personas qui satisfait mes désirs.",
      guideStepAdditionalAssistantText:
        "CRITICAL: YOU MUST FILL `personaId` NOW !!!!! NOW !!!",
    },
  },
] as const;

const getEngine = (() => {
  let engine: Promise<WebWorkerMLCEngine> | undefined;
  return (
    modelIdentifier: string,
    options?: EngineOptions,
  ) => {
    console.log("Initializing WebLLM engine...");
    const workerURL = "/public/assets/esm/worker.js?v=" + Date.now();
    engine ??= import("https://esm.run/@mlc-ai/web-llm").then(
      async (webLLM: unknown) => {
        const worker = new Worker(workerURL, {
          type: "module",
          name: "mlc-worker",
        });
        worker.addEventListener("error", (error) => {
          console.error("[MLCEngine] Worker error:", error);
          engine = undefined;
          worker.terminate();
          throw new Error(`Worker error: ${error.message}`);
        });
        worker.addEventListener("messageerror", (error) => {
          console.error("[MLCEngine] Worker message error:", error);
        });

        await new Promise<void>((resolve, reject) => {
          const timeout = setTimeout(() => {
            reject(new Error("Worker initialization timeout"));
          }, 10000);
          worker.addEventListener("message", (msg) => {
            console.log("[MLCEngine] Received worker message:", msg.data);
            if (msg.data.type === "ready") {
              clearTimeout(timeout);
              resolve();
            } else if (msg.data.type === "error") {
              clearTimeout(timeout);
              reject(new Error(msg.data.error));
            }
          });
          worker.postMessage({
            type: "init",
            moduleURL: "https://esm.run/@mlc-ai/web-llm",
          });
          console.log("[MLCEngine] Worker initialized successfully");
        });

        const mlcEngine = await (webLLM as typeof webLLMType)
          .CreateWebWorkerMLCEngine(
            worker,
            modelIdentifier,
            {
              logLevel: "TRACE",
              initProgressCallback: (progress) => {
                console.log("[MLCEngine] Loading progress:", progress);
                options?.onProgress?.(progress);
              },
            },
            { context_window_size: 8192 },
          ).catch((error: unknown) => {
            worker.terminate();
            console.error("[MLCEngine] Error loading model:", error);
            const message = error instanceof Error
              ? error.message
              : String(error);
            throw new Error(
              `Failed to load MLC model "${modelIdentifier}": ${message}`,
            );
          });
        return mlcEngine;
      },
    );
    return engine;
  };
})();

const webLLMFetch = (engineOptions: EngineOptions) =>
async (
  _url: string | URL | Request,
  options?: RequestInit,
): Promise<Response> => {
  const requestBody = JSON.parse(options?.body as string);
  const webllmMessages = requestBody.messages;
  console.log(requestBody);
  const engine = await getEngine(requestBody.model, engineOptions);
  const body: ChatCompletionRequestStreaming = {
    model: requestBody.model,
    max_tokens: requestBody.max_tokens,
    messages: webllmMessages,
    presence_penalty: requestBody.presencePenalty,
    temperature: requestBody.temperature,
    frequency_penalty: requestBody.frequencePenality,
    stream: true,
    stream_options: requestBody.stream_options,
    response_format: {
      type: "json_object",
      schema: typeof requestBody.response_format.json_schema.schema === "string"
        ? requestBody.response_format.json_schema.schema
        : JSON.stringify(requestBody.response_format.json_schema.schema),
    },
  };
  console.log(body);
  const chunks = await engine.chat.completions.create(body);
  const readableStream = new ReadableStream({
    async start(controller) {
      const encoder = new TextEncoder();
      for await (const chunk of chunks) {
        const sseLine = `data: ${JSON.stringify(chunk)}\n\n`;
        console.log(sseLine);
        controller.enqueue(encoder.encode(sseLine));
      }
      controller.enqueue(encoder.encode("data: [DONE]\n\n"));
      controller.close();
    },
  });
  return new Response(readableStream, {
    headers: { "Content-Type": "text/event-stream; charset=utf-8" },
  });
};

export async function mockFetchToLocalAPI(
  _url: string | URL | Request,
  options: RequestInit | undefined,
  engineOptions: EngineOptions,
): Promise<Response> {
  const requestBody = JSON.parse(options?.body as string);
  const { messages: originalMessages, locale, modelId } = requestBody;

  const { getModel, features, parameters } =
    AVAILABLE_CLIENT_MODELS.find((v) => v.id === modelId) ??
      AVAILABLE_CLIENT_MODELS[0];

  const webLLMModel = getModel(engineOptions);

  const enrichedMessages = await addFullGameStateAsync(originalMessages);
  const currentGameState = getGameState(enrichedMessages);
  const messagesForLLM = applyMindControlRewrites(
    originalMessages,
    enrichedMessages,
  );
  const actingPersonaId: AIPersonaId =
    currentGameState.personaId === "storytellurx_pause"
      ? "storytellurx_pause"
      : currentGameState.personaId ?? "storytellurx_gateway";
  const persona = AI_PERSONAS.find((p) => p.id === actingPersonaId) ??
    AI_PERSONAS.find((p) => p.id === "storytellurx_gateway")!;
  const schemaForThisTurn = persona.type === "system"
    ? SystemNarrativeResponseSchema
    : EroticNarrativeResponseSchema;

  const systemInstruction = await constructSystemPrompt(
    currentGameState,
    actingPersonaId,
    locale,
    features,
  );

  const result = streamText({
    model: webLLMModel,
    messages: convertToModelMessages(
      includeUserHintForSystem(
        features?.developerInstructionDisabled === false
          ? inludeSystemAsAssistant(
            messagesForLLM,
            systemInstruction,
            schemaForThisTurn,
            features,
          )
          : messagesForLLM,
        features,
      ),
    ),
    system: features?.developerInstructionDisabled
      ? undefined
      : systemInstruction,
    experimental_output: features?.structuredOutputDisabled
      ? undefined
      : Output.object({ schema: schemaForThisTurn }),
    maxRetries: 2,
    ...parameters,
  });

  // Convert the result stream into a UI-compatible stream
  const stream = result.toUIMessageStream({
    originalMessages,
    onError: (v) => {
      console.error(v);
      return "An error occurred with WebLLM.";
    },
    messageMetadata: ({ part }) => {
      if (part.type === "finish") {
        return {
          model: { id: modelId, parameters: features },
          createdAt: Date.now(),
          system: {
            tokenCount: part.totalUsage.inputTokens,
            content: systemInstruction,
          },
          totalTokenCount: part.totalUsage.totalTokens ?? 0,
          tokenCount: part.totalUsage.outputTokens,
        } as Partial<MessageMetadata>;
      }
    },
  });

  return createUIMessageStreamResponse({ stream });
}