user-docs

import {
  tool,
  type Tool,
  type ToolsProviderController,
  type ToolCallContext,
} from "@lmstudio/sdk";
import { z } from "zod";
import path from "path";
import fs from "fs";
import {
  setActiveChatContext,
  readState,
  writeStateAtomic,
  generatePreviewFromBuffer,
  appendImages,
  getSelfPluginIdentifier,
  formatToolMetaBlock,
} from "../core-bundle.mjs";
import { executePythonScript } from "../utils/pythonRunner.js";
import { globalConfigSchematics } from "../config.js";

export function createExtractImageTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "extract_image",
    description: `Render a page of a PDF document as a PNG image and register it for visual inspection.

Use this tool when:
- find_doc returned a result from a PDF and you want to view a specific page visually
- The user asks to look at, analyse, or annotate a page in a PDF document

The rendered page is registered as an iN image entry. Use show_image, analyse_image, or annotate_image to inspect it.

Parameters:
- source: Absolute path to the PDF file. Use the documentPath from a find_doc result.
- page: Page number to render (1-based).
- dpi: Render resolution in dots per inch (72–300). Default: 150. Use 200–300 for text-heavy pages.

${formatToolMetaBlock()}`,
    parameters: {
      source: z
        .string()
        .describe("Absolute path or bare filename of the PDF (e.g. 'AV10_EU_DE.pdf'). Bare filenames are resolved against notesDirectory and all contentDirectories."),
      page: z
        .number()
        .int()
        .min(1)
        .describe("Page number to render (1-based)."),
      dpi: z
        .number()
        .int()
        .min(72)
        .max(300)
        .optional()
        .default(150)
        .describe("Render resolution in DPI (72–300). Default: 150."),
    },
    implementation: async (args: any, ctx: ToolCallContext) => {
      // ── 1. Resolve working directory ─────────────────────────────────────
      try {
        const workingDir = ctl.getWorkingDirectory();
        if (typeof workingDir === "string" && workingDir.trim().length > 0) {
          const chatId = path.basename(workingDir);
          if (/^\d+$/.test(chatId)) {
            setActiveChatContext({ chatId, workingDir, requestId: `tool-${Date.now()}` });
          }
        }
      } catch {
        // best-effort
      }

      const chatWd = ctl.getWorkingDirectory();
      if (typeof chatWd !== "string" || !chatWd.trim()) {
        return "extract_image failed: could not resolve LM Studio chat working directory.";
      }

      // ── 2. Validate args — resolve bare filename against all indexed dirs ──
      const rawSource: string = typeof args?.source === "string" ? args.source.trim() : "";
      if (!rawSource) return "extract_image: source must not be empty.";

      let pdfPath: string;
      if (fs.existsSync(rawSource)) {
        // Already an absolute (or resolvable) path.
        pdfPath = rawSource;
      } else {
        // Bare filename — search notesDirectory + every contentDirectory.
        const getter: any =
          (ctl as any).getGlobalPluginConfig || (ctl as any).getGlobalConfig;
        const gcfg = getter ? getter.call(ctl, globalConfigSchematics) : null;
        const notesDir: string = gcfg?.get("notesDirectory") ?? "";
        const contentDirs: string[] = gcfg?.get("contentDirectories") ?? [];
        const searchDirs = [notesDir, ...contentDirs].filter(Boolean);

        const basename = path.basename(rawSource);
        let found: string | null = null;
        for (const dir of searchDirs) {
          const candidate = path.join(dir, basename);
          if (fs.existsSync(candidate)) { found = candidate; break; }
        }
        if (!found) {
          return `extract_image: file not found — "${rawSource}". Searched: ${searchDirs.join(", ")}`;
        }
        pdfPath = found;
      }

      const page: number = typeof args?.page === "number" ? Math.max(1, Math.round(args.page)) : 1;
      const dpi: number = typeof args?.dpi === "number" ? Math.min(300, Math.max(72, Math.round(args.dpi))) : 150;

      // ── 3. Determine output path ─────────────────────────────────────────
      await fs.promises.mkdir(chatWd, { recursive: true }).catch(() => {});
      const pdfStem = path.basename(pdfPath, path.extname(pdfPath)).replace(/[^a-z0-9_-]/gi, "_");
      const outputFilename = `pdf-${pdfStem}-p${page}.png`;
      const outputPath = path.join(chatWd, outputFilename);

      // ── 4. Run Python renderer ───────────────────────────────────────────
      try { ctx.status(`Rendering page ${page} of ${path.basename(pdfPath)}…`); } catch {}

      const scriptPath = path.join(__dirname, "..", "python", "extract_image_page.py");
      const result = await executePythonScript<{
        success: boolean;
        output_path?: string;
        width?: number;
        height?: number;
        page?: number;
        page_count?: number;
        dpi?: number;
        error?: string;
      }>(scriptPath, [pdfPath, String(page), outputPath, String(dpi)], { timeout: 60000 });

      if (!result.success || !result.data?.success) {
        const errMsg = result.error ?? result.data?.error ?? "Page render failed";
        return `extract_image: ${errMsg}`;
      }

      const data = result.data!;

      // ── 5. Generate preview ──────────────────────────────────────────────
      try { ctx.status("Generating preview…"); } catch {}

      const srcBuf = await fs.promises.readFile(outputPath);
      const previewSpec = { maxDim: 2560, maxSum: 2560, mode: "sum" as const, quality: 85 };
      const previewStem = `preview-pdf-${pdfStem}-p${page}`;

      let previewFilename: string | undefined;
      try {
        const preview = await generatePreviewFromBuffer(srcBuf, chatWd, outputFilename, previewSpec, {
          customFilename: `${previewStem}.jpg`,
        });
        previewFilename = preview.previewFilename;
      } catch (e) {
        console.warn("[extract_image] preview generation failed:", String(e));
      }

      // ── 6. Register in state ─────────────────────────────────────────────
      const pluginId = getSelfPluginIdentifier() ?? "unknown";
      const imageRecord: any = {
        filename: outputFilename,
        preview: previewFilename,
        sourceTool: `${pluginId}/extract_image`,
        sourceUrl: `file://${pdfPath}`,
        sourcePage: page,
      };

      let assignedKey: string | null = null;
      try {
        const state = await readState(chatWd);
        const appendResult = appendImages(state, [imageRecord]);
        if (appendResult.changed) {
          await writeStateAtomic(chatWd, state);
        }
        const rec = (appendResult.records as any[])[0];
        if (typeof rec?.i === "number") assignedKey = `i${rec.i}`;
      } catch (err) {
        console.warn("[extract_image] state update failed:", String(err));
      }

      // ── 7. Tool result ───────────────────────────────────────────────────
      const keyLabel = assignedKey ?? "(key unavailable)";
      const pageInfo = data.page_count ? `page ${page} of ${data.page_count}` : `page ${page}`;
      const dimInfo = data.width && data.height ? ` (${data.width}×${data.height}px, ${dpi} DPI)` : "";

      try {
        ctx.status(`Done — ${pageInfo} rendered as ${keyLabel}`);
      } catch {}

      const envPreviewRaw = process.env["PREVIEW_IN_CHAT"];
      const previewInChat =
        envPreviewRaw === undefined
          ? true
          : envPreviewRaw === "1" || envPreviewRaw.toLowerCase() === "true";

      const header = `Rendered ${pageInfo} of "${path.basename(pdfPath)}"${dimInfo}.\nRegistered as ${keyLabel}.`;

      if (previewInChat && previewFilename) {
        return (
          `${header} Use ![${keyLabel}](./${previewFilename}) to view, ` +
          `analyse_image({"targets":["${keyLabel}"]}) to inspect it, ` +
          `or annotate_image({"targets":["${keyLabel}"]}) to highlight requested elements.`
        );
      }

      return (
        `${header} Use show_image({"targets":["${keyLabel}"]}) to view, ` +
        `or annotate_image({"targets":["${keyLabel}"]}) to highlight requested elements.`
      );
    },
  });
}