LLM-Labs/src/lib/lab2-chat.ts

import {
  DOMParser,
  XMLSerializer,
  type Element as XmlDomElement,
} from "@xmldom/xmldom";
import {
  LAB2_DEFAULT_OLLAMA_MODELS,
  LAB2_DEFAULT_OLLAMA_URL,
} from "~/lib/courseware-runtime";

export const LAB2_CHAT_STORAGE_KEY = "lab2-objective5-chat-settings";
export const LAB2_DEFAULT_ENDPOINT = LAB2_DEFAULT_OLLAMA_URL;
export const LAB2_LEGACY_DEFAULT_ENDPOINT = "https://ai.zuccaro.me/api";
export const LAB2_CUSTOM_MODEL_VALUE = "__custom__";
export const LAB2_MAX_CONTEXT_MESSAGES = 10;
export const LAB2_MAX_MESSAGE_LENGTH = 4000;
export const LAB2_MAX_SVG_LENGTH = 20000;

export const LAB2_MODEL_OPTIONS = [
  ...LAB2_DEFAULT_OLLAMA_MODELS,
  {
    label: "Custom model",
    value: LAB2_CUSTOM_MODEL_VALUE,
  },
] as const;

export type Objective5Role = "user" | "assistant";

export type Objective5Message = {
  content: string;
  role: Objective5Role;
};

export type Objective5RenderMode = "text" | "svg";

export type Objective5Metrics = {
  completionTokens?: number;
  evalDurationMs?: number;
  promptTokens?: number;
  promptEvalDurationMs?: number;
  tokensPerSecond?: number;
  totalDurationMs?: number;
};

export type Objective5StoredSettings = {
  apiKey: string;
  customModel: string;
  endpoint: string;
  selectedModel: string;
};

export type Objective5ModelOption = {
  label: string;
  value: string;
};

type AssistantMessageContentPart =
  | string
  | {
      text?: string;
      type?: string;
    };

type ChatCompletionPayload = {
  choices?: Array<{
    message?: {
      content?: AssistantMessageContentPart | AssistantMessageContentPart[];
      reasoning_content?: string;
    };
  }>;
  output?: Array<{
    content?: Array<{
      text?: string;
      type?: string;
    }>;
    type?: string;
  }>;
  output_text?: string;
  message?: {
    content?: string;
    role?: string;
    thinking?: string;
  };
  completion_tokens?: number;
  prompt_tokens?: number;
  usage?: {
    completion_tokens?: number;
    prompt_tokens?: number;
    total_tokens?: number;
  };
  eval_count?: number;
  eval_duration?: number;
  prompt_eval_count?: number;
  prompt_eval_duration?: number;
  total_duration?: number;
};

type SvgSanitizationSuccess = {
  ok: true;
  svg: string;
};

type SvgSanitizationFailure = {
  error: string;
  ok: false;
};

export type SvgSanitizationResult =
  | SvgSanitizationFailure
  | SvgSanitizationSuccess;

const predefinedModelLabels = new Map<string, string>(
  LAB2_DEFAULT_OLLAMA_MODELS.map((model) => [model.value, model.label]),
);

const SVG_NAMESPACE = "http://www.w3.org/2000/svg";
const allowedSvgElements = new Set([
  "svg",
  "g",
  "path",
  "circle",
  "ellipse",
  "rect",
  "line",
  "polyline",
  "polygon",
  "text",
  "tspan",
  "defs",
  "linearGradient",
  "radialGradient",
  "stop",
  "title",
  "desc",
]);

const allowedSvgAttributes = new Set([
  "cx",
  "cy",
  "d",
  "dominant-baseline",
  "fill",
  "fill-opacity",
  "fill-rule",
  "font-family",
  "font-size",
  "font-weight",
  "gradientTransform",
  "gradientUnits",
  "height",
  "id",
  "offset",
  "opacity",
  "points",
  "preserveAspectRatio",
  "r",
  "rect",
  "role",
  "rx",
  "ry",
  "stop-color",
  "stop-opacity",
  "stroke",
  "stroke-linecap",
  "stroke-linejoin",
  "stroke-opacity",
  "stroke-width",
  "text-anchor",
  "transform",
  "version",
  "viewBox",
  "width",
  "x",
  "x1",
  "x2",
  "xml:space",
  "xmlns",
  "y",
  "y1",
  "y2",
]);

const enumValues = {
  "dominant-baseline": new Set([
    "auto",
    "alphabetic",
    "central",
    "hanging",
    "ideographic",
    "mathematical",
    "middle",
    "text-after-edge",
    "text-before-edge",
  ]),
  "fill-rule": new Set(["evenodd", "inherit", "nonzero"]),
  "gradientUnits": new Set(["objectBoundingBox", "userSpaceOnUse"]),
  "stroke-linecap": new Set(["butt", "round", "square"]),
  "stroke-linejoin": new Set(["arcs", "bevel", "miter", "miter-clip", "round"]),
  "text-anchor": new Set(["end", "inherit", "middle", "start"]),
} as const;

const numberPattern = /^-?(?:\d+|\d*\.\d+)(?:e[-+]?\d+)?%?$/i;
const numberListPattern =
  /^[-+0-9eE.,%\s]+$/;
const pathPattern = /^[MmZzLlHhVvCcSsQqTtAa0-9eE.,\s-]+$/;
const pointsPattern = /^[-+0-9eE.,\s]+$/;
const transformPattern =
  /^[a-zA-Z0-9(),.\s-]+$/;
const viewBoxPattern = /^[-+0-9eE.,\s]+$/;
const fontFamilyPattern =
  /^[a-zA-Z0-9\s,'"_-]+$/;
const textPattern =
  /^[^\u0000-\u0008\u000b\u000c\u000e-\u001f<>]+$/;
const idPattern = /^[A-Za-z_][A-Za-z0-9_.-]*$/;

export function getLab2SystemPrompt() {
  return [
    "You are helping students compare quantized models in a lab.",
    "Answer normal questions clearly and concisely.",
    "If the user asks you to draw something, respond with only one complete standalone SVG document.",
    "Do not wrap SVG in markdown fences.",
    "Do not include any explanation before or after the SVG.",
    "Use a 512 by 512 canvas with viewBox=\"0 0 512 512\".",
    "Do not use scripts, foreignObject, animation, external references, or remote assets.",
    "Prefer simple shapes, strokes, fills, and text labels that render reliably.",
  ].join(" ");
}

export function normalizeUpstreamChatEndpoint(endpoint: string) {
  const url = new URL(endpoint);
  const trimmedPath = url.pathname.replace(/\/+$/, "");

  if (trimmedPath.endsWith("/chat/completions")) {
    url.pathname = trimmedPath;
  } else if (trimmedPath.endsWith("/v1")) {
    url.pathname = `${trimmedPath}/chat/completions`;
  } else if (trimmedPath.length === 0) {
    url.pathname = "/v1/chat/completions";
  } else {
    url.pathname = `${trimmedPath}/v1/chat/completions`;
  }

  url.hash = "";
  return url.toString();
}

export function getModelListEndpointCandidates(endpoint: string) {
  const url = new URL(endpoint);
  const trimmedPath = url.pathname.replace(/\/+$/, "");

  if (
    trimmedPath.endsWith("/models") ||
    trimmedPath.endsWith("/api/tags")
  ) {
    url.hash = "";
    return [url.toString()];
  }

  const paths = new Set<string>();

  if (trimmedPath.endsWith("/ollama/api")) {
    paths.add("/ollama/api/tags");
  } else if (trimmedPath.endsWith("/ollama")) {
    paths.add("/ollama/api/tags");
  } else if (trimmedPath.endsWith("/api")) {
    paths.add("/api/tags");
    paths.add("/api/v1/models");
    paths.add("/api/models");
  } else if (trimmedPath.endsWith("/api/v1")) {
    paths.add("/api/tags");
    paths.add("/api/v1/models");
    paths.add("/api/models");
  } else if (trimmedPath.endsWith("/v1")) {
    paths.add("/v1/models");
  } else if (trimmedPath.length === 0) {
    paths.add("/api/tags");
    paths.add("/v1/models");
  } else {
    paths.add(`${trimmedPath}/api/tags`);
    paths.add(`${trimmedPath}/v1/models`);
    paths.add(`${trimmedPath}/models`);
  }

  return Array.from(paths).map((path) => {
    const candidate = new URL(url.toString());
    candidate.pathname = path;
    candidate.hash = "";
    return candidate.toString();
  });
}

export function normalizeOllamaChatEndpoint(endpoint: string) {
  return getOllamaChatEndpointCandidates(endpoint)[0];
}

export function getOllamaChatEndpointCandidates(endpoint: string) {
  const url = new URL(endpoint);
  const trimmedPath = url.pathname.replace(/\/+$/, "");

  if (
    trimmedPath.endsWith("/api/chat") ||
    trimmedPath.endsWith("/ollama/api/chat")
  ) {
    url.pathname = trimmedPath;
    url.hash = "";
    return [url.toString()];
  }

  const paths = new Set<string>();

  if (trimmedPath.endsWith("/ollama/api")) {
    paths.add("/ollama/api/chat");
  } else if (trimmedPath.endsWith("/ollama")) {
    paths.add("/ollama/api/chat");
  } else if (trimmedPath.endsWith("/api")) {
    paths.add("/api/chat");
    paths.add("/ollama/api/chat");
  } else if (trimmedPath.endsWith("/api/v1") || trimmedPath.endsWith("/v1")) {
    paths.add("/ollama/api/chat");
    paths.add("/api/chat");
  } else if (trimmedPath.length === 0) {
    paths.add("/api/chat");
    paths.add("/ollama/api/chat");
  } else {
    paths.add(`${trimmedPath}/api/chat`);
    paths.add(`${trimmedPath}/ollama/api/chat`);
  }

  return Array.from(paths).map((path) => {
    const candidate = new URL(url.toString());
    candidate.pathname = path;
    candidate.hash = "";
    return candidate.toString();
  });
}

export function looksLikeOllamaModel(model: string) {
  return model.includes(":");
}

export function isLocalEndpoint(endpoint: string) {
  try {
    const url = new URL(endpoint);
    return (
      url.hostname === "127.0.0.1" ||
      url.hostname === "localhost" ||
      url.hostname === "::1"
    );
  } catch {
    return false;
  }
}

export function clampChatMessages(messages: Objective5Message[]) {
  return messages
    .filter((message) => {
      return (
        (message.role === "assistant" || message.role === "user") &&
        typeof message.content === "string"
      );
    })
    .map((message) => {
      return {
        content: message.content.slice(0, LAB2_MAX_MESSAGE_LENGTH),
        role: message.role,
      } satisfies Objective5Message;
    })
    .slice(-LAB2_MAX_CONTEXT_MESSAGES);
}

export function buildUpstreamMessages(messages: Objective5Message[]) {
  return [
    {
      content: getLab2SystemPrompt(),
      role: "system" as const,
    },
    ...clampChatMessages(messages),
  ];
}

export function extractAssistantTextContent(payload: ChatCompletionPayload) {
  if (
    payload.message &&
    typeof payload.message.content === "string" &&
    payload.message.content.trim()
  ) {
    return payload.message.content.trim();
  }

  if (typeof payload.output_text === "string" && payload.output_text.trim()) {
    return payload.output_text.trim();
  }

  const choiceContent = payload.choices?.[0]?.message?.content;
  const choiceText = normalizeContentParts(choiceContent);
  if (choiceText) {
    return choiceText;
  }

  const reasoningContent = payload.choices?.[0]?.message?.reasoning_content;
  if (typeof reasoningContent === "string" && reasoningContent.trim()) {
    return reasoningContent.trim();
  }

  const outputText = payload.output
    ?.flatMap((item) => item.content ?? [])
    .map((item) => item.text?.trim() ?? "")
    .filter(Boolean)
    .join("\n\n")
    .trim();

  return outputText || null;
}

export function extractModelOptions(payload: unknown): Objective5ModelOption[] {
  if (!payload || typeof payload !== "object") {
    return [];
  }

  if ("data" in payload && Array.isArray(payload.data)) {
    return payload.data
      .map((item) => {
        if (!item || typeof item !== "object") return null;

        const value =
          "id" in item && typeof item.id === "string" ? item.id.trim() : "";
        const label =
          "name" in item && typeof item.name === "string" && item.name.trim()
            ? item.name.trim()
            : getModelLabel(value);

        if (!value) return null;
        return { label, value } satisfies Objective5ModelOption;
      })
      .filter((item): item is Objective5ModelOption => item !== null);
  }

  if (!("models" in payload) || !Array.isArray(payload.models)) {
    return [];
  }

  return payload.models
    .map((item) => {
      if (!item || typeof item !== "object") return null;

      const value =
        ("model" in item && typeof item.model === "string" && item.model.trim()
          ? item.model
          : "name" in item && typeof item.name === "string"
            ? item.name
            : ""
        ).trim();

      if (!value) return null;
      return {
        label: getModelLabel(value),
        value,
      } satisfies Objective5ModelOption;
    })
    .filter((item): item is Objective5ModelOption => item !== null);
}

export function extractObjective5Metrics(
  payload: ChatCompletionPayload,
): Objective5Metrics | null {
  const promptTokens =
    payload.prompt_eval_count ??
    payload.prompt_tokens ??
    payload.usage?.prompt_tokens;
  const completionTokens =
    payload.eval_count ??
    payload.completion_tokens ??
    payload.usage?.completion_tokens;
  const promptEvalDurationMs = toMilliseconds(payload.prompt_eval_duration);
  const evalDurationMs = toMilliseconds(payload.eval_duration);
  const totalDurationMs = toMilliseconds(payload.total_duration);
  const tokensPerSecond =
    typeof completionTokens === "number" &&
    typeof evalDurationMs === "number" &&
    evalDurationMs > 0
      ? roundMetric((completionTokens / evalDurationMs) * 1000)
      : undefined;

  if (
    typeof promptTokens !== "number" &&
    typeof completionTokens !== "number" &&
    typeof promptEvalDurationMs !== "number" &&
    typeof evalDurationMs !== "number" &&
    typeof totalDurationMs !== "number" &&
    typeof tokensPerSecond !== "number"
  ) {
    return null;
  }

  return {
    completionTokens,
    evalDurationMs,
    promptEvalDurationMs,
    promptTokens,
    tokensPerSecond,
    totalDurationMs,
  };
}

function normalizeContentParts(
  content: AssistantMessageContentPart | AssistantMessageContentPart[] | undefined,
) {
  if (typeof content === "string") {
    return content.trim();
  }

  if (!Array.isArray(content)) {
    return null;
  }

  const text = content
    .map((part) => {
      if (typeof part === "string") {
        return part;
      }

      return part.text ?? "";
    })
    .join("\n\n")
    .trim();

  return text || null;
}

function toMilliseconds(durationNs?: number) {
  if (typeof durationNs !== "number" || !Number.isFinite(durationNs)) {
    return undefined;
  }

  return roundMetric(durationNs / 1_000_000);
}

function roundMetric(value: number) {
  return Math.round(value * 10) / 10;
}

export function extractSvgMarkup(content: string) {
  const trimmed = content.trim();
  if (!trimmed) return null;

  const fencedMatch =
    /^```(?:svg|xml)?\s*([\s\S]*?)\s*```$/i.exec(trimmed) ??
    /```(?:svg|xml)?\s*([\s\S]*?)\s*```/i.exec(trimmed);
  const unfenced = fencedMatch?.[1]?.trim() ?? trimmed;

  const svgMatch = unfenced.match(/<svg[\s\S]*?<\/svg>/i);
  return svgMatch?.[0]?.trim() ?? null;
}

export function sanitizeSvgDocument(svgMarkup: string): SvgSanitizationResult {
  const trimmed = svgMarkup.trim();

  if (!trimmed) {
    return {
      error: "The model returned an empty SVG response.",
      ok: false,
    };
  }

  if (trimmed.length > LAB2_MAX_SVG_LENGTH) {
    return {
      error: "The SVG response is too large to render safely.",
      ok: false,
    };
  }

  const parseErrors: string[] = [];
  const parser = new DOMParser({
    onError: (level, message) => {
      if (level !== "warning") {
        parseErrors.push(String(message));
      }
    },
  });

  let document;

  try {
    document = parser.parseFromString(trimmed, "image/svg+xml");
  } catch {
    return {
      error: "The model returned malformed SVG markup.",
      ok: false,
    };
  }
  if (parseErrors.length > 0) {
    return {
      error: "The model returned malformed SVG markup.",
      ok: false,
    };
  }

  const root = document.documentElement;
  if (!root || root.tagName !== "svg") {
    return {
      error: "The response did not contain a standalone SVG document.",
      ok: false,
    };
  }

  const validationError = validateSvgNode(root);
  if (validationError) {
    return {
      error: validationError,
      ok: false,
    };
  }

  root.setAttribute("xmlns", SVG_NAMESPACE);
  if (!root.getAttribute("viewBox")) {
    root.setAttribute("viewBox", "0 0 512 512");
  }
  if (!root.getAttribute("width")) {
    root.setAttribute("width", "512");
  }
  if (!root.getAttribute("height")) {
    root.setAttribute("height", "512");
  }

  const serialized = new XMLSerializer().serializeToString(root).trim();
  if (!serialized.startsWith("<svg") || serialized.length > LAB2_MAX_SVG_LENGTH) {
    return {
      error: "The sanitized SVG could not be rendered safely.",
      ok: false,
    };
  }

  return {
    ok: true,
    svg: serialized,
  };
}

export function svgToDataUrl(svg: string) {
  return `data:image/svg+xml;charset=utf-8,${encodeURIComponent(svg)}`;
}

export function getActiveModel(selectedModel: string, customModel: string) {
  if (selectedModel === LAB2_CUSTOM_MODEL_VALUE) {
    return customModel.trim();
  }

  return selectedModel.trim();
}

export function getDefaultObjective5Settings(): Objective5StoredSettings {
  return {
    apiKey: "",
    customModel: "",
    endpoint: LAB2_DEFAULT_ENDPOINT,
    selectedModel: LAB2_MODEL_OPTIONS[0].value,
  };
}

export function getDefaultObjective5ModelOptions(): Objective5ModelOption[] {
  return [...LAB2_MODEL_OPTIONS];
}

function getModelLabel(value: string) {
  return predefinedModelLabels.get(value) ?? value;
}

function validateSvgNode(node: XmlDomElement): string | null {
  if (!allowedSvgElements.has(node.tagName)) {
    return `The SVG used a blocked element: <${node.tagName}>.`;
  }

  const attributes = Array.from(node.attributes ?? []);
  for (const attribute of attributes) {
    const validationError = validateSvgAttribute(attribute.name, attribute.value);
    if (validationError) {
      return validationError;
    }
  }

  const children = Array.from(node.childNodes);
  for (const child of children) {
    if (child.nodeType === child.ELEMENT_NODE) {
      const childError = validateSvgNode(child as XmlDomElement);
      if (childError) {
        return childError;
      }
      continue;
    }

    if (child.nodeType === child.TEXT_NODE) {
      const textValue = child.nodeValue ?? "";
      if (textValue.trim().length === 0) {
        continue;
      }

      if (!allowsTextChildren(node.tagName)) {
        return `The SVG contained unexpected text inside <${node.tagName}>.`;
      }

      if (!textPattern.test(textValue)) {
        return "The SVG contained unsafe text content.";
      }
    }
  }

  return null;
}

function allowsTextChildren(tagName: string) {
  return tagName === "text" || tagName === "tspan" || tagName === "title" || tagName === "desc";
}

function validateSvgAttribute(name: string, value: string) {
  const normalizedName = name.trim();
  const normalizedValue = value.trim();

  if (!allowedSvgAttributes.has(normalizedName)) {
    return `The SVG used a blocked attribute: ${normalizedName}.`;
  }

  if (!normalizedValue) {
    return null;
  }

  if (/^on/i.test(normalizedName)) {
    return "The SVG contained blocked event handlers.";
  }

  if (normalizedName === "xmlns") {
    return normalizedValue === SVG_NAMESPACE
      ? null
      : "The SVG used an unexpected XML namespace.";
  }

  if (
    /(?:javascript:|data:|https?:|ftp:|file:)/i.test(normalizedValue) ||
    normalizedValue.includes("<") ||
    normalizedValue.includes(">")
  ) {
    return `The SVG used an unsafe value for ${normalizedName}.`;
  }

  if (normalizedName in enumValues) {
    return enumValues[normalizedName as keyof typeof enumValues].has(
      normalizedValue,
    )
      ? null
      : `The SVG used an unsupported value for ${normalizedName}.`;
  }

  switch (normalizedName) {
    case "cx":
    case "cy":
    case "fill-opacity":
    case "font-size":
    case "height":
    case "offset":
    case "opacity":
    case "r":
    case "rx":
    case "ry":
    case "stop-opacity":
    case "stroke-opacity":
    case "stroke-width":
    case "width":
    case "x":
    case "x1":
    case "x2":
    case "y":
    case "y1":
    case "y2":
      return numberPattern.test(normalizedValue)
        ? null
        : `The SVG used an invalid numeric value for ${normalizedName}.`;
    case "viewBox":
      return viewBoxPattern.test(normalizedValue)
        ? null
        : "The SVG used an invalid viewBox.";
    case "d":
      return pathPattern.test(normalizedValue)
        ? null
        : "The SVG used an invalid path definition.";
    case "points":
      return pointsPattern.test(normalizedValue)
        ? null
        : "The SVG used invalid polygon or polyline points.";
    case "transform":
    case "gradientTransform":
      return transformPattern.test(normalizedValue)
        ? null
        : `The SVG used an invalid ${normalizedName}.`;
    case "fill":
    case "stroke":
    case "stop-color":
      return validateSvgPaintValue(normalizedValue)
        ? null
        : `The SVG used an unsupported paint value for ${normalizedName}.`;
    case "font-family":
      return fontFamilyPattern.test(normalizedValue)
        ? null
        : "The SVG used an unsupported font family.";
    case "font-weight":
      return /^(?:normal|bold|bolder|lighter|[1-9]00)$/i.test(normalizedValue)
        ? null
        : "The SVG used an unsupported font weight.";
    case "id":
      return idPattern.test(normalizedValue)
        ? null
        : "The SVG used an invalid id attribute.";
    case "preserveAspectRatio":
      return /^[A-Za-z0-9\s]+$/.test(normalizedValue)
        ? null
        : "The SVG used an invalid preserveAspectRatio value.";
    case "role":
      return /^[a-zA-Z0-9\s-]+$/.test(normalizedValue)
        ? null
        : "The SVG used an invalid role attribute.";
    case "version":
      return /^1\.1$/.test(normalizedValue) || /^2(?:\.0)?$/.test(normalizedValue)
        ? null
        : "The SVG used an unsupported version value.";
    case "xml:space":
      return /^(?:default|preserve)$/.test(normalizedValue)
        ? null
        : "The SVG used an invalid xml:space value.";
    default:
      return numberListPattern.test(normalizedValue)
        ? null
        : `The SVG used an invalid value for ${normalizedName}.`;
  }
}

function validateSvgPaintValue(value: string) {
  if (
    /^(?:none|currentColor|transparent)$/i.test(value) ||
    /^#[0-9a-f]{3,8}$/i.test(value) ||
    /^(?:rgb|rgba|hsl|hsla)\([\d\s.,%+-]+\)$/i.test(value)
  ) {
    return true;
  }

  if (/^url\(#[-A-Za-z0-9_.]+\)$/i.test(value)) {
    return true;
  }

  return /^[a-zA-Z]+$/.test(value);
}