draw-things

import {
  type ChatMessage,
  type FileHandle,
  type PromptPreprocessorController,
} from "@lmstudio/sdk";
import { copyFile, mkdir, stat } from "fs/promises";
import {
  basename,
  extname,
  isAbsolute,
  join,
  normalize,
  relative,
  resolve,
} from "path";
import { fileURLToPath } from "url";
import { configSchematics } from "./config";
import {
  extractGeneratedImagePathsFromToolResultContent,
  getRememberedGeneratedImages,
  rememberGeneratedImage,
} from "./mediaHistory";

const OUTPUT_SUBDIR_PATTERN = /^[A-Za-z0-9._/-]+$/;
const MAX_UPLOADED_IMAGE_OPTIONS = 5;
const MAX_GENERATED_IMAGE_OPTIONS = 5;
const IMAGE_INPUT_EXTENSIONS = new Set([
  ".avif",
  ".bmp",
  ".gif",
  ".jpeg",
  ".jpg",
  ".png",
  ".tif",
  ".tiff",
  ".webp",
]);

type ImageInputOption = {
  source: "uploaded" | "generated";
  path: string;
  label: string;
};

export async function preprocess(
  ctl: PromptPreprocessorController,
  userMessage: ChatMessage,
): Promise<ChatMessage> {
  try {
    const workingDirectory = ctl.getWorkingDirectory();
    const config = ctl.getPluginConfig(configSchematics) as { get: (key: any) => unknown };
    const outputRoot = await resolveOutputDirectory(
      workingDirectory,
      getConfigString(config, "outputSubdirectory"),
    );

    const uploadedOptions = await stageUploadedImageOptions({
      ctl,
      userMessage,
      workingDirectory,
      outputRoot: outputRoot.absolutePath,
    });
    const generatedOptions = await collectGeneratedImageOptions(ctl, workingDirectory);
    const options = dedupeImageOptions([...uploadedOptions, ...generatedOptions]);

    if (options.length === 0) {
      return userMessage;
    }

    userMessage.appendText(buildImageInputContext(options));
    return userMessage;
  } catch (error: any) {
    ctl.debug(`Draw Things prompt preprocessing skipped: ${error?.message || String(error)}`);
    return userMessage;
  }
}

async function stageUploadedImageOptions({
  ctl,
  userMessage,
  workingDirectory,
  outputRoot,
}: {
  ctl: PromptPreprocessorController;
  userMessage: ChatMessage;
  workingDirectory: string;
  outputRoot: string;
}): Promise<ImageInputOption[]> {
  const uploadedImages = userMessage
    .getFiles(ctl.client)
    .filter((file) => file.isImage())
    .slice(0, MAX_UPLOADED_IMAGE_OPTIONS);
  if (uploadedImages.length === 0) {
    return [];
  }

  const inputsDirectory = join(outputRoot, "inputs");
  await mkdir(inputsDirectory, { recursive: true });

  const options: ImageInputOption[] = [];
  for (const file of uploadedImages) {
    try {
      const stagedPath = await stageUploadedImage(file, workingDirectory, inputsDirectory);
      options.push({
        source: "uploaded",
        path: stagedPath.relativePath,
        label: `uploaded image "${file.name}"`,
      });
    } catch (error: any) {
      ctl.debug(`Could not stage uploaded image '${file.name}': ${error?.message || String(error)}`);
    }
  }
  return options;
}

async function stageUploadedImage(
  file: FileHandle,
  workingDirectory: string,
  inputsDirectory: string,
): Promise<{ relativePath: string; absolutePath: string }> {
  const sourcePath = await file.getFilePath();
  const extension = getImageExtension(file.name) || getImageExtension(sourcePath) || ".png";
  const fileBaseName = sanitizeFileBaseName(file.name) || "uploaded-image";
  const outputFileName = `input-${Date.now()}-${randomSuffix()}-${fileBaseName}${extension}`;
  const absolutePath = join(inputsDirectory, outputFileName);

  await copyFile(sourcePath, absolutePath);

  return {
    absolutePath,
    relativePath: toWorkingDirectoryRelativePath(workingDirectory, absolutePath),
  };
}

async function collectGeneratedImageOptions(
  ctl: PromptPreprocessorController,
  workingDirectory: string,
): Promise<ImageInputOption[]> {
  const candidates: string[] = [];

  for (const record of getRememberedGeneratedImages(MAX_GENERATED_IMAGE_OPTIONS * 2)) {
    candidates.push(record.path);
    if (record.absolutePath) candidates.push(record.absolutePath);
  }

  const history = await ctl.pullHistory();
  const messages = history.getMessagesArray().slice().reverse();
  for (const message of messages) {
    for (const result of message.getToolCallResults()) {
      candidates.push(...extractGeneratedImagePathsFromToolResultContent(result.content));
    }
  }

  const options: ImageInputOption[] = [];
  const seen = new Set<string>();
  for (const candidate of candidates) {
    if (options.length >= MAX_GENERATED_IMAGE_OPTIONS) break;
    const image = await resolveGeneratedImageCandidate(workingDirectory, candidate);
    if (!image || seen.has(image.relativePath)) continue;

    seen.add(image.relativePath);
    rememberGeneratedImage({
      path: image.relativePath,
      absolutePath: image.absolutePath,
    });
    options.push({
      source: "generated",
      path: image.relativePath,
      label: "previously generated image",
    });
  }

  return options;
}

async function resolveGeneratedImageCandidate(
  workingDirectory: string,
  inputPath: string,
): Promise<{ relativePath: string; absolutePath: string } | undefined> {
  const cleanedPath = normalizeMediaPathInput(inputPath);
  if (!cleanedPath) return undefined;

  const absolutePath = isAbsolute(cleanedPath)
    ? resolve(cleanedPath)
    : resolve(workingDirectory, cleanedPath);
  if (!isPathInside(workingDirectory, absolutePath)) return undefined;
  if (!IMAGE_INPUT_EXTENSIONS.has(extname(absolutePath).toLowerCase())) return undefined;

  const metadata = await stat(absolutePath).catch(() => null);
  if (!metadata?.isFile()) return undefined;

  return {
    absolutePath,
    relativePath: toWorkingDirectoryRelativePath(workingDirectory, absolutePath),
  };
}

function buildImageInputContext(options: ImageInputOption[]): string {
  const lines = options.map(
    (option) => `- ${option.label}: pass image="${option.path}"`,
  );
  return [
    "",
    "",
    "Draw Things image inputs available for this request:",
    ...lines,
    "Use one exact path above as the image parameter for image-to-image or image-to-video. Do not invent paths or omit image when the user asks to transform or animate one of these images.",
  ].join("\n");
}

function dedupeImageOptions(options: ImageInputOption[]): ImageInputOption[] {
  const seen = new Set<string>();
  const deduped: ImageInputOption[] = [];
  for (const option of options) {
    if (seen.has(option.path)) continue;
    seen.add(option.path);
    deduped.push(option);
  }
  return deduped;
}

async function resolveOutputDirectory(
  workingDirectory: string,
  configuredSubdirectory: string,
): Promise<{ absolutePath: string; relativePath: string }> {
  const rawSubdirectory = cleanString(configuredSubdirectory) || "draw-things";
  const normalizedSubdirectory = normalize(rawSubdirectory).replace(/\\/g, "/").replace(/^\.\/+/, "");

  if (isAbsolute(normalizedSubdirectory)) {
    throw new Error("Output Subdirectory must be relative to the LM Studio working directory.");
  }
  if (!OUTPUT_SUBDIR_PATTERN.test(normalizedSubdirectory)) {
    throw new Error("Output Subdirectory may only contain letters, numbers, dots, underscores, hyphens, and slashes.");
  }

  const absolutePath = resolve(workingDirectory, normalizedSubdirectory);
  if (!isPathInside(workingDirectory, absolutePath)) {
    throw new Error("Output Subdirectory must stay inside the LM Studio working directory.");
  }

  await mkdir(absolutePath, { recursive: true });
  return {
    absolutePath,
    relativePath: toWorkingDirectoryRelativePath(workingDirectory, absolutePath),
  };
}

function normalizeMediaPathInput(inputPath: string): string {
  const cleanedPath = cleanString(inputPath);
  if (!cleanedPath.toLowerCase().startsWith("file://")) {
    return cleanedPath;
  }

  try {
    return fileURLToPath(cleanedPath);
  } catch {
    return "";
  }
}

function getImageExtension(fileName: string): string | undefined {
  const extension = extname(fileName).toLowerCase();
  return IMAGE_INPUT_EXTENSIONS.has(extension) ? extension : undefined;
}

function sanitizeFileBaseName(fileName: string): string {
  const name = basename(fileName, extname(fileName))
    .replace(/[^A-Za-z0-9._-]+/g, "-")
    .replace(/^-+|-+$/g, "");
  return name.slice(0, 80);
}

function getConfigString(config: { get: (key: any) => unknown }, key: string): string {
  const value = config.get(key);
  return typeof value === "string" ? value.trim() : String(value ?? "").trim();
}

function cleanString(value?: string): string {
  return typeof value === "string" ? value.trim() : "";
}

function isPathInside(root: string, target: string): boolean {
  const resolvedRoot = resolve(root);
  const resolvedTarget = resolve(target);
  const relativePath = relative(resolvedRoot, resolvedTarget);
  return relativePath === "" || (!relativePath.startsWith("..") && !isAbsolute(relativePath));
}

function toWorkingDirectoryRelativePath(workingDirectory: string, absolutePath: string): string {
  const relativePath = relative(workingDirectory, absolutePath).replace(/\\/g, "/");
  return relativePath === "" ? basename(absolutePath) : relativePath;
}

function randomSuffix(): string {
  return Math.random().toString(36).slice(2, 8);
}