draw-things-cli

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { execSync } from "child_process";
import path from "path";
import fs from "fs";
import { z } from "zod";
import { configSchematics } from "./config";
import {
  detectCli,
  runCommand,
  buildGenerateCommand,
  parseOutputPath,
  convertPngToJpeg,
} from "./utils";

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  // ---------------------------------------------------------------------------
  // Detect CLI once at toolsProvider initialization (per chat session)
  // ---------------------------------------------------------------------------
  const config = ctl.getPluginConfig(configSchematics);
  const cliPath = config.get("cliPath");
  const cliStatus = detectCli(cliPath);

  // ---------------------------------------------------------------------------
  // Tool: generate_image
  // ---------------------------------------------------------------------------
  const generateImageTool = tool({
    name: "generate_image",
    description: [
      "Generate an image using draw-things-cli running locally on this Mac.",
      "Only 'prompt' is required — all other parameters fall back to plugin configuration.",
      "Use 'list_dt_models' first if you need to pick a specific model.",
      "Returns the generated image inline. Generation may take 30–300 seconds depending on model.",
    ].join(" "),
    parameters: {
      prompt: z
        .string()
        .describe("Image generation prompt describing what to draw."),
      model: z
        .string()
        .optional()
        .describe("Model filename (e.g. ernie_image_turbo_q8p.ckpt). Falls back to plugin default if omitted."),
      negative_prompt: z
        .string()
        .optional()
        .describe("Negative prompt — things to avoid in the image."),
      width: z
        .number()
        .optional()
        .describe("Image width in pixels. Falls back to plugin default."),
      height: z
        .number()
        .optional()
        .describe("Image height in pixels. Falls back to plugin default."),
      steps: z
        .number()
        .optional()
        .describe("Number of diffusion steps. Falls back to plugin default."),
      cfg: z
        .number()
        .optional()
        .describe("CFG scale (classifier-free guidance). Falls back to plugin default."),
      seed: z
        .number()
        .optional()
        .describe("Generation seed. Negative value = random. Falls back to plugin default."),
    },
    implementation: async (
      { prompt, model, negative_prompt, width, height, steps, cfg, seed },
      { status, warn, signal }
    ) => {
      // Always reload config — it may have changed between tool calls
      const cfg_ = ctl.getPluginConfig(configSchematics);
      const currentCliPath = cfg_.get("cliPath");

      // --- Check CLI availability ---
      if (!cliStatus.found) {
        const msg = `ERROR: draw-things-cli not found at "${currentCliPath}". ` +
          `Please install it (brew install drawthingsai/draw-things/draw-things-cli) ` +
          `and verify the path in plugin settings.`;
        warn(msg);
        return msg;
      }

      // --- Resolve model ---
      const resolvedModel = model?.trim() || cfg_.get("defaultModel").trim();
      if (!resolvedModel) {
        const msg = "ERROR: No model specified. Set a default model in plugin settings " +
          "or pass the 'model' parameter. Use 'list_dt_models' to see available models.";
        warn(msg);
        return msg;
      }

      // --- Resolve output path ---
      const timestamp = Date.now();
      const filename = `dt_${timestamp}.png`;
      const configOutputDir = cfg_.get("outputDir").trim();
      const outputDir = configOutputDir || ctl.getWorkingDirectory();

      // Ensure output directory exists
      try {
        fs.mkdirSync(outputDir, { recursive: true });
      } catch (e: any) {
        const msg = `ERROR: Cannot create output directory "${outputDir}": ${e.message}`;
        warn(msg);
        return msg;
      }

      const outputPath = path.join(outputDir, filename);

      // --- Build CLI command ---
      const command = buildGenerateCommand({
        cliPath: currentCliPath,
        model: resolvedModel,
        prompt,
        negativePrompt: negative_prompt,
        width: width ?? cfg_.get("defaultWidth"),
        height: height ?? cfg_.get("defaultHeight"),
        steps: steps ?? cfg_.get("defaultSteps"),
        cfg: cfg ?? cfg_.get("defaultCfg"),
        seed: seed ?? cfg_.get("defaultSeed"),
        outputPath,
      });

      // --- Run generation ---
      const timeoutMs = cfg_.get("generationTimeout") * 1000;
      status(`Generating: "${prompt.slice(0, 60)}${prompt.length > 60 ? "…" : ""}"`);

      const result = await runCommand(
        command,
        { shell: "/bin/bash" },
        timeoutMs,
        signal
      );

      // --- Check result ---
      if (result.timedOut) {
        const msg =
          `Image generation did not complete within ${cfg_.get("generationTimeout")} seconds. ` +
          "Please increase the 'Generation Timeout' in the plugin settings and try again. " +
          "Large models may need 3–5 minutes.";
        warn(msg);
        return msg;
      }

      if (result.error && !result.stdout.includes("Wrote:")) {
        const msg = `ERROR: draw-things-cli failed.\nSTDERR: ${result.stderr}\nSTDOUT: ${result.stdout}`;
        warn(msg);
        return msg;
      }

      // --- Parse output file path from CLI stdout ---
      const writtenPath = parseOutputPath(result.stdout) ?? outputPath;

      if (!fs.existsSync(writtenPath)) {
        const msg = `ERROR: Output file not found at "${writtenPath}". ` +
          `CLI stdout: ${result.stdout}`;
        warn(msg);
        return msg;
      }

      // --- Copy to working directory and optionally convert to JPEG ---
      // writtenPath = file in outputDir (if set) or directly in working dir
      const configOutputDir2 = cfg_.get("outputDir").trim();
      const convertToJpeg = cfg_.get("convertToJpeg");
      const workingDir = ctl.getWorkingDirectory();

      let chatFileName = path.basename(writtenPath);
      let chatFilePath: string;
      let chatMimeType = "image/png";

      if (configOutputDir2) {
        // outputDir is set — PNG lives there, copy it to working dir for inline display
        const pngInWorkingDir = path.join(workingDir, chatFileName);
        try {
          fs.mkdirSync(workingDir, { recursive: true });
          fs.copyFileSync(writtenPath, pngInWorkingDir); // copy only — original in outputDir untouched
        } catch (e: any) {
          warn(`Could not copy image to session working directory: ${e.message}`);
        }
        chatFilePath = pngInWorkingDir; // work with the copy from here on
      } else {
        // outputDir not set — PNG is already in working dir
        chatFilePath = writtenPath;
      }

      if (convertToJpeg) {
        // Convert the working dir copy — original PNG in outputDir is never touched
        status("Converting to JPEG…");
        const quality = Math.max(0, Math.min(100, cfg_.get("jpegQuality")));
        try {
          const jpegBuffer = await convertPngToJpeg(chatFilePath, quality);
          const jpegFileName = chatFileName.replace(/\.png$/i, ".jpg");
          const jpegFilePath = path.join(workingDir, jpegFileName);
          fs.writeFileSync(jpegFilePath, jpegBuffer);
          // Remove the PNG copy from working dir
          try { fs.unlinkSync(chatFilePath); } catch { /* ignore */ }
          chatFileName = jpegFileName;
          chatFilePath = jpegFilePath;
          chatMimeType = "image/jpeg";
        } catch (e: any) {
          warn(`JPEG conversion failed, using original PNG: ${e.message}`);
        }
      }

      // --- Return JSON matching the filesystem:read_media_file format ---
      // LM Studio plugin tools cannot return content[] blocks directly — they come back as raw JSON.
      // This plain JSON string is the format LM Studio recognizes for inline image rendering.
      return JSON.stringify({
        type: "image",
        fileName: chatFileName,
        mimeType: chatMimeType,
        markdown: `![Image](./${chatFileName})`,
        $hint: "This is an image file. Present the image to the user by using the markdown above.",
      });
    },
  });

  // ---------------------------------------------------------------------------
  // Tool: list_dt_models
  // ---------------------------------------------------------------------------
  const listModelsTool = tool({
    name: "list_dt_models",
    description: [
      "List all Draw Things models downloaded on this Mac.",
      "Use this to find the correct model filename for image generation.",
      "Pass the filename (e.g. 'ernie_image_turbo_q8p.ckpt') as the 'model' parameter to 'generate_image'.",
    ].join(" "),
    parameters: {},
    implementation: async (_, { status }) => {
      const cfg_ = ctl.getPluginConfig(configSchematics);
      const currentCliPath = cfg_.get("cliPath");

      if (!cliStatus.found) {
        return `ERROR: draw-things-cli not found at "${currentCliPath}". ` +
          "Please install it and verify the path in plugin settings.";
      }

      status("Loading model list…");
      try {
        const out = execSync(
          `"${currentCliPath}" models list --downloaded-only`,
          { timeout: 15000 }
        )
          .toString()
          .trim();
        return out || "No downloaded models found.";
      } catch (e: any) {
        return `ERROR: Failed to list models: ${e.message}`;
      }
    },
  });

  return [generateImageTool, listModelsTool];
}