pixlstash-lmstudio

import { type PredictionLoopHandlerController } from "@lmstudio/sdk";
import { Agent, fetch } from "undici";
import { Jimp } from "jimp";
import { writeFile } from "node:fs/promises";
import { join } from "node:path";
import { configSchematics, globalConfigSchematics } from "./config.js";

/**
 * Image formats LM Studio can render inline as a markdown image. The picture
 * search may return other things (mp4, mov, webm, heif, avif, …) when the
 * format filter is empty — those are skipped because they can't render.
 */
const RENDERABLE = new Set(["png", "jpg", "jpeg", "webp", "gif"]);

type ListPath = "characters" | "picture_sets" | "projects";

interface PictureHit {
  id: number;
  format: string;
}

/**
 * Search PixlStash for `query`, pick one of the top-N renderable results that
 * hasn't been used in this chat, download it into the chat working directory,
 * downscale if needed, and return a local-path markdown image string. Returns
 * null when no fresh match exists.
 *
 * The picked id is added to `usedIds` so the caller can keep deduping across
 * calls within the same response.
 */
export async function fetchImageMarkdown(
  ctl: PredictionLoopHandlerController,
  query: string,
  signal: AbortSignal,
  usedIds: Set<number>,
): Promise<string | null> {
  const cfg = ctl.getPluginConfig(configSchematics);
  const conn = ctl.getGlobalPluginConfig(globalConfigSchematics);

  const baseUrl = conn.get("baseUrl").replace(/\/$/, "");
  const token = conn.get("apiToken");
  const topN = cfg.get("topN");
  const useThumbnail = cfg.get("imageSource") === "thumbnail";

  // Self-signed HTTPS: skip TLS verification only when the user opted in.
  const dispatcher =
    baseUrl.startsWith("https:") && conn.get("ignoreCertErrors")
      ? new Agent({ connect: { rejectUnauthorized: false } })
      : undefined;

  const headers: Record<string, string> = token ? { Authorization: `Bearer ${token}` } : {};
  const fetchOpts = { headers, signal, ...(dispatcher ? { dispatcher } : {}) };

  const params = await buildSearchParams(ctl, baseUrl, fetchOpts, query);

  const searchRes = await fetch(`${baseUrl}/api/v1/pictures/search?${params}`, fetchOpts);
  if (!searchRes.ok) throw new Error(`search HTTP ${searchRes.status}`);
  const results = (await searchRes.json()) as PictureHit[];
  const fresh = results.filter(
    (r) => RENDERABLE.has(r.format.toLowerCase()) && !usedIds.has(r.id),
  );
  console.info(
    `PixlStash: search returned ${results.length} result(s), ${fresh.length} fresh renderable`,
  );
  if (fresh.length === 0) return null;

  const pool = fresh.slice(0, topN);
  const pick = pool[Math.floor(Math.random() * pool.length)];
  usedIds.add(pick.id);

  const ext = useThumbnail ? "webp" : pick.format.toLowerCase();
  const imgUrl = useThumbnail
    ? `${baseUrl}/api/v1/pictures/thumbnails/${pick.id}.webp`
    : `${baseUrl}/api/v1/pictures/${pick.id}.${ext}`;
  const imgRes = await fetch(imgUrl, fetchOpts);
  if (!imgRes.ok) throw new Error(`image fetch HTTP ${imgRes.status} — ${imgUrl}`);
  const bytes = Buffer.from(await imgRes.arrayBuffer());

  const resized = await downscale(bytes, ext, cfg.get("maxImageHeight"));
  const fileName = `pixlstash-${pick.id}.${ext}`;
  await writeFile(join(ctl.getWorkingDirectory(), fileName), resized);
  console.info(`PixlStash: injected picture #${pick.id} (${fileName})`);

  // Relative local path resolves against the working directory and renders inline.
  return `![pixlstash ${pick.id}](${fileName})`;
}

/** Format an unknown error as a one-line string, surfacing `cause` when present. */
export function describeError(err: unknown): string {
  return err instanceof Error && err.cause
    ? `${err.message} — ${String(err.cause)}`
    : String(err);
}

// ---------------------------------------------------------------- internals --

/** Build the query string for `/pictures/search` from the per-chat config. */
async function buildSearchParams(
  ctl: PredictionLoopHandlerController,
  baseUrl: string,
  fetchOpts: Parameters<typeof fetch>[1],
  query: string,
): Promise<URLSearchParams> {
  const cfg = ctl.getPluginConfig(configSchematics);
  const params = new URLSearchParams({
    query,
    limit: String(cfg.get("topN")),
    threshold: String(cfg.get("threshold")),
  });
  if (cfg.get("minScore") > 0) params.set("min_score", String(cfg.get("minScore")));

  // PixlStash's `format` filter is a case-sensitive `IN (...)` over an
  // upper-cased column (e.g. "PNG"); send each configured value in both cases
  // so a lowercase config still matches.
  const formatVariants = new Set<string>();
  for (const f of cfg.get("formats")) {
    formatVariants.add(f);
    formatVariants.add(f.toUpperCase());
  }
  for (const f of formatVariants) params.append("format", f);
  for (const t of cfg.get("requiredTags")) params.append("tag", t);
  for (const t of cfg.get("rejectedTags")) params.append("rejected_tag", t);

  const characterIds = await resolveScopeIds(
    baseUrl,
    fetchOpts,
    "characters",
    cfg.get("scopeCharacters"),
  );
  for (const id of characterIds) params.append("character_ids", String(id));
  const setIds = await resolveScopeIds(baseUrl, fetchOpts, "picture_sets", cfg.get("scopeSets"));
  for (const id of setIds) params.append("set_ids", String(id));
  const projectId = await resolveScopeId(baseUrl, fetchOpts, "projects", cfg.get("scopeProject"));
  if (projectId !== null) params.set("project_id", String(projectId));

  return params;
}

/**
 * Downscale a PNG/JPEG taller than `maxHeight`, preserving aspect ratio. Other
 * formats, `maxHeight <= 0`, and decode errors fall back to the original bytes.
 * Exported for testing.
 */
export async function downscale(bytes: Buffer, ext: string, maxHeight: number): Promise<Buffer> {
  if (maxHeight <= 0 || !/^(png|jpe?g)$/i.test(ext)) return bytes;
  try {
    const image = await Jimp.read(bytes);
    if (image.height <= maxHeight) return bytes;
    const width = Math.max(1, Math.round((image.width * maxHeight) / image.height));
    image.resize({ w: width, h: maxHeight });
    const out = /png/i.test(ext)
      ? await image.getBuffer("image/png")
      : await image.getBuffer("image/jpeg");
    return Buffer.from(out);
  } catch (err) {
    console.info(`PixlStash: resize skipped — ${describeError(err)}`);
    return bytes;
  }
}

/**
 * Resolve a configured character/set/project name to its PixlStash id by
 * case-insensitive match against the list endpoint. Cached because ids are
 * stable. Returns null for a blank name or any failed lookup, in which case
 * the search just runs unscoped rather than failing the whole turn.
 */
const scopeIdCache = new Map<string, number>();
async function resolveScopeId(
  baseUrl: string,
  fetchOpts: Parameters<typeof fetch>[1],
  listPath: ListPath,
  rawName: string,
): Promise<number | null> {
  const wanted = rawName.trim().toLowerCase();
  if (!wanted) return null;
  const cacheKey = `${baseUrl}|${listPath}|${wanted}`;
  const cached = scopeIdCache.get(cacheKey);
  if (cached !== undefined) return cached;
  try {
    const res = await fetch(`${baseUrl}/api/v1/${listPath}`, fetchOpts);
    if (!res.ok) throw new Error(`HTTP ${res.status}`);
    const items = (await res.json()) as Array<{ id: number; name?: string }>;
    const match = items.find((it) => (it.name ?? "").trim().toLowerCase() === wanted);
    if (match) {
      scopeIdCache.set(cacheKey, match.id);
      return match.id;
    }
    console.info(`PixlStash: scope ${listPath} "${rawName}" not found`);
    return null;
  } catch (err) {
    console.info(
      `PixlStash: scope ${listPath} "${rawName}" lookup failed — ${describeError(err)}`,
    );
    return null;
  }
}

/** Resolve a list of names to ids in parallel, dropping blanks and misses. */
async function resolveScopeIds(
  baseUrl: string,
  fetchOpts: Parameters<typeof fetch>[1],
  listPath: ListPath,
  rawNames: string[],
): Promise<number[]> {
  const resolved = await Promise.all(
    rawNames.map((name) => resolveScopeId(baseUrl, fetchOpts, listPath, name)),
  );
  return resolved.filter((id): id is number => id !== null);
}