lmexif-sniffer

import * as fsp from "node:fs/promises";
import * as path from "node:path";

import { text, tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";

import { configSchematics } from "./config";
import { extractMetadataList } from "./extract";
import { flattenToMetadataList } from "./flatten";
import { downloadUrlToPath } from "./fetch";
import {
  fsCreateDirectory,
  fsExtractMetadata,
  fsListFiles,
  fsReadFile,
  fsUpdateLocalMediaPngCharacterCard,
  fsWriteFile,
  fsWriteMetadata,
  fsWritePngCharacterCard,
} from "./filesystem";
import { updateImageExif } from "./exifEdit";
import { listImageRelativePaths, parseLocalMediaRoot } from "./localMedia";
import { policyFromConfig } from "./security";
import { resolveUnderRoot } from "./paths";

const optionalRel = z.union([z.string(), z.null()]).optional();

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const pluginConfig = ctl.getPluginConfig(configSchematics);

  const listFiles = tool({
    name: "list_files",
    description: text`
      List files and subdirectories in the top level of the configured Base Directory
      (plugin settings), matching the LM Studio filesystem-access style. Returns JSON rows
      {path, value}.
    `,
    parameters: {},
    implementation: async () => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsListFiles(base);
    },
  });

  const readFile = tool({
    name: "read_file",
    description: text`
      Read a UTF-8 text file under the configured Base Directory. Parameter file_name is relative
      with the same naming rules as the LM Studio filesystem-access plugin (letters, digits,
      underscore, hyphen, dot, slash; no '..').
    `,
    parameters: { file_name: z.string().min(1) },
    implementation: async ({ file_name }) => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsReadFile(base, file_name);
    },
  });

  const writeFile = tool({
    name: "write_file",
    description: text`
      Write or overwrite a UTF-8 text file under the configured Base Directory, creating parent
      subdirectories as needed.
    `,
    parameters: { file_name: z.string().min(1), content: z.string() },
    implementation: async ({ file_name, content }) => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsWriteFile(base, file_name, content);
    },
  });

  const createDirectory = tool({
    name: "create_directory",
    description: text`Create a subdirectory under the configured Base Directory (mkdir -p semantics).`,
    parameters: { directory_name: z.string().min(1) },
    implementation: async ({ directory_name }) => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsCreateDirectory(base, directory_name);
    },
  });

  const extractMetadata = tool({
    name: "extract_metadata",
    description: text`
      Read an image or video under the configured Base Directory. Extract EXIF (images), PNG tEXt/zTXt
      chunks, decoded SillyTavern character card JSON when present (image.character_card), or a
      placeholder for video (no ffprobe in this plugin build). Returns a flat JSON list of {path, value}
      rows. Optional output_json_file_name writes the same array under the base.
    `,
    parameters: {
      source_file_name: z.string().min(1),
      output_json_file_name: optionalRel,
      include_piexif: z.boolean().optional().default(false),
    },
    implementation: async ({ source_file_name, output_json_file_name, include_piexif }) => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsExtractMetadata(base, source_file_name, output_json_file_name ?? null, Boolean(include_piexif));
    },
  });

  const writePngCharacterCard = tool({
    name: "write_png_character_card",
    description: text`
      Embed or replace SillyTavern-style character card data on a .png under the configured Base Directory
      (in-place). Strips existing chara/ccv3/character text chunks, then writes Base64(JSON) in a new tEXt
      or zTXt chunk after IHDR. card_json is either a JSON string or an object (stringified).
    `,
    parameters: {
      file_name: z.string().min(1),
      card_json: z.union([z.string().min(1), z.record(z.unknown())]),
      keyword: z.enum(["chara", "ccv3"]).optional().default("chara"),
      compress: z.boolean().optional().default(false),
    },
    implementation: async ({ file_name, card_json, keyword, compress }) => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsWritePngCharacterCard(
        base,
        file_name,
        card_json as string | Record<string, unknown>,
        keyword ?? "chara",
        Boolean(compress),
      );
    },
  });

  const writeMetadata = tool({
    name: "write_metadata",
    description: text`
      Update or remove EXIF on a .jpg, .jpeg, or .webp file under the configured Base Directory.
      set_tags maps IFD names to {tag_name: value}; remove_tags maps IFD names to tag name lists.
    `,
    parameters: {
      file_name: z.string().min(1),
      set_tags: z.record(z.record(z.unknown())).nullable().optional(),
      remove_tags: z.record(z.array(z.string())).nullable().optional(),
    },
    implementation: async ({ file_name, set_tags, remove_tags }) => {
      const base = String(pluginConfig.get("folderName") ?? "");
      return fsWriteMetadata(
        base,
        file_name,
        (set_tags ?? undefined) as Record<string, Record<string, unknown>> | undefined,
        (remove_tags ?? undefined) as Record<string, string[]> | undefined,
      );
    },
  });

  const fetchRemoteMedia = tool({
    name: "fetch_remote_media",
    description: text`
      Download an image or video from an HTTP(S) URL into the plugin working directory (DATA_DIR
      equivalent). Validates URL, blocks SSRF to private networks unless allowed in settings,
      follows redirects up to the configured max, and enforces max download bytes.
    `,
    parameters: {
      url: z.string().min(8),
      destination_relative_path: z.string().min(1),
    },
    implementation: async ({ url, destination_relative_path }) => {
      const dataRoot = ctl.getWorkingDirectory();
      const dest = resolveUnderRoot(dataRoot, destination_relative_path);
      const policy = policyFromConfig({
        fetchAllowPrivateHosts: Boolean(pluginConfig.get("fetchAllowPrivateHosts")),
        fetchAllowedHostSuffixes: String(pluginConfig.get("fetchAllowedHostSuffixes") ?? ""),
        fetchBlockedHostSuffixes: String(pluginConfig.get("fetchBlockedHostSuffixes") ?? ""),
      });
      const limits = {
        maxDownloadBytes: Number(pluginConfig.get("maxDownloadBytes") ?? 100_000_000),
        connectTimeoutMs: Number(pluginConfig.get("fetchConnectTimeoutSeconds") ?? 10) * 1000,
        readTimeoutMs: Number(pluginConfig.get("fetchReadTimeoutSeconds") ?? 120) * 1000,
        maxRedirects: Number(pluginConfig.get("fetchMaxRedirects") ?? 8),
      };
      const meta = await downloadUrlToPath(url, dest, policy, limits);
      const summary = { saved_path: dest, ...meta };
      return flattenToMetadataList(summary, "fetch");
    },
  });

  const extractMetadataToJson = tool({
    name: "extract_metadata_to_json",
    description: text`
      Read a local image or video file under the plugin working directory, extract metadata (including
      PNG SillyTavern character_card when applicable), and write a JSON array file under the same
      directory. Returns the same array as tool output.
    `,
    parameters: {
      source_relative_path: z.string().min(1),
      output_json_relative_path: z.string().min(1),
      include_piexif: z.boolean().optional().default(false),
    },
    implementation: async ({ source_relative_path, output_json_relative_path, include_piexif }) => {
      const dataRoot = ctl.getWorkingDirectory();
      const src = resolveUnderRoot(dataRoot, source_relative_path);
      const out = resolveUnderRoot(dataRoot, output_json_relative_path);
      const rows = await extractMetadataList(src, Boolean(include_piexif));
      await fsp.mkdir(path.dirname(out), { recursive: true });
      await fsp.writeFile(out, JSON.stringify(rows, null, 2), "utf8");
      return rows;
    },
  });

  const validateLocalMediaRoot = tool({
    name: "validate_local_media_root",
    description: text`
      Check whether local_media_root is usable: must be an absolute path to an existing directory.
    `,
    parameters: { local_media_root: z.string().min(1) },
    implementation: async ({ local_media_root }) => {
      try {
        const root = parseLocalMediaRoot(local_media_root);
        const summary = {
          usable: true,
          resolved_root: root,
          exists_and_is_dir: true,
          error: null,
        };
        return flattenToMetadataList(summary, "local_media");
      } catch (e) {
        const summary = {
          usable: false,
          resolved_root: null,
          exists_and_is_dir: false,
          error: String(e),
        };
        return flattenToMetadataList(summary, "local_media");
      }
    },
  });

  const listLocalMediaImages = tool({
    name: "list_local_media_images",
    description: text`
      List image files under the given local_media_root (absolute directory path). relative_directory
      is relative to that root (empty string means the root). recursive includes subfolders. max_files
      caps results (default 500, max 5000).
    `,
    parameters: {
      local_media_root: z.string().min(1),
      relative_directory: z.string().optional().default(""),
      recursive: z.boolean().optional().default(false),
      max_files: z.number().int().min(1).max(5000).optional().default(500),
    },
    implementation: async ({ local_media_root, relative_directory, recursive, max_files }) => {
      const root = parseLocalMediaRoot(local_media_root);
      const rels = await listImageRelativePaths(root, relative_directory ?? "", Boolean(recursive), max_files ?? 500);
      const rows: Array<{ path: string; value: unknown }> = [];
      for (let i = 0; i < rels.length; i++) {
        rows.push({ path: `local_media.images[${i}]`, value: rels[i] });
      }
      return rows;
    },
  });

  const extractLocalMediaMetadataToJson = tool({
    name: "extract_local_media_metadata_to_json",
    description: text`
      Like extract_metadata_to_json but the source file lives under local_media_root. JSON output is
      written under the plugin working directory via output_json_relative_path.
    `,
    parameters: {
      local_media_root: z.string().min(1),
      source_relative_path: z.string().min(1),
      output_json_relative_path: z.string().min(1),
      include_piexif: z.boolean().optional().default(false),
    },
    implementation: async ({
      local_media_root,
      source_relative_path,
      output_json_relative_path,
      include_piexif,
    }) => {
      const localRoot = parseLocalMediaRoot(local_media_root);
      const dataRoot = ctl.getWorkingDirectory();
      const src = resolveUnderRoot(localRoot, source_relative_path);
      const out = resolveUnderRoot(dataRoot, output_json_relative_path);
      const rows = await extractMetadataList(src, Boolean(include_piexif));
      await fsp.mkdir(path.dirname(out), { recursive: true });
      await fsp.writeFile(out, JSON.stringify(rows, null, 2), "utf8");
      return rows;
    },
  });

  const updateLocalMediaPngCharacterCard = tool({
    name: "update_local_media_png_character_card",
    description: text`
      Same as write_png_character_card but the PNG path is resolved under local_media_root (absolute
      directory) plus image_relative_path. Writes in-place on the local file.
    `,
    parameters: {
      local_media_root: z.string().min(1),
      image_relative_path: z.string().min(1),
      card_json: z.union([z.string().min(1), z.record(z.unknown())]),
      keyword: z.enum(["chara", "ccv3"]).optional().default("chara"),
      compress: z.boolean().optional().default(false),
    },
    implementation: async ({ local_media_root, image_relative_path, card_json, keyword, compress }) => {
      return fsUpdateLocalMediaPngCharacterCard(
        local_media_root,
        image_relative_path,
        card_json as string | Record<string, unknown>,
        keyword ?? "chara",
        Boolean(compress),
      );
    },
  });

  const updateLocalMediaExif = tool({
    name: "update_local_media_exif",
    description: text`
      Update or remove EXIF tags on a JPEG or WebP file under local_media_root (in-place), using the
      same IFD/tag semantics as the Python exif-sniffer server.
    `,
    parameters: {
      local_media_root: z.string().min(1),
      image_relative_path: z.string().min(1),
      set_tags: z.record(z.record(z.unknown())).nullable().optional(),
      remove_tags: z.record(z.array(z.string())).nullable().optional(),
    },
    implementation: async ({ local_media_root, image_relative_path, set_tags, remove_tags }) => {
      const localRoot = parseLocalMediaRoot(local_media_root);
      const src = resolveUnderRoot(localRoot, image_relative_path);
      const summary = updateImageExif(
        src,
        (set_tags ?? {}) as Record<string, Record<string, unknown>>,
        (remove_tags ?? {}) as Record<string, string[]>,
      );
      return flattenToMetadataList(summary, "local_media.exif_update");
    },
  });

  return [
    listFiles,
    readFile,
    writeFile,
    createDirectory,
    extractMetadata,
    writePngCharacterCard,
    writeMetadata,
    fetchRemoteMedia,
    extractMetadataToJson,
    validateLocalMediaRoot,
    listLocalMediaImages,
    extractLocalMediaMetadataToJson,
    updateLocalMediaPngCharacterCard,
    updateLocalMediaExif,
  ];
}