draw-things-chat

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,
  formatToolMetaBlock,
  localTimestamp,
  readState,
  writeStateAtomic,
  getLogsDir,
} from "../core-bundle.mjs";

const ReviewSequenceParamsShape = {
  variant: z.coerce
    .string()
    .describe(
      "Reference to the video to review. Use standard media notation: " +
      "'vN' for variants, 'iN' for images, 'pN' for pictures (e.g. 'v1', 'i3'). " +
      "A bare integer N is also accepted and treated as vN. " +
      "Must correspond to a generated video.",
    ),

  fps: z.coerce
    .number()
    .min(0.1)
    .max(30)
    .optional()
    .default(2)
    .describe(
      "Frame sampling rate in fps. Default: 2. Higher values send more frames to the model.",
    ),
} satisfies Record<string, z.ZodTypeAny>;

function appendReviewSequenceToPluginLog(line: string): void {
  try {
    const logsDir = getLogsDir();
    if (!fs.existsSync(logsDir)) fs.mkdirSync(logsDir, { recursive: true });
    fs.appendFileSync(
      path.join(logsDir, "generate-image-plugin.log"),
      `${localTimestamp()} - ${line}\n`,
    );
  } catch {
    // best-effort
  }
}

export function createReviewSequenceTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "review_sequence",
    description: `Send a frame sequence of a generated video to the model for visual review.

Use this when the user asks you to watch, review or analyse the content of a generated video.
LM Studio does not support native video objects; this tool compensates by sampling frames at a configurable fps and transmitting them as a labelled image sequence.

Parameters:
- variant: Media reference for the video — 'vN', 'iN', or 'pN' notation (e.g. 'v1', 'i3').
- fps: Frame sampling rate (default: 2). Higher → more frames → larger context.

${formatToolMetaBlock()}`,
    parameters: ReviewSequenceParamsShape,
    implementation: async (args: any, ctx: ToolCallContext) => {
      try {
        // Refresh active chat context (mirrors review_image behaviour).
        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 workingDir = ctl.getWorkingDirectory();
        if (typeof workingDir !== "string" || !workingDir.trim()) {
          return "review_sequence failed: working directory not available.";
        }
        const chatWd = workingDir;

        // ── Parse reference (vN / iN / pN / bare integer) ───────────────────────────

        const rawRef = typeof args?.variant === "string" ? args.variant.trim()
          : String(args?.variant ?? "").trim();

        let refType: "v" | "i" | "p" = "v";
        let refNum = 0;
        const refMatch = /^([vip])(\d+)$/i.exec(rawRef);
        if (refMatch) {
          refType = refMatch[1].toLowerCase() as "v" | "i" | "p";
          refNum = parseInt(refMatch[2], 10);
        } else if (/^\d+$/.test(rawRef)) {
          refType = "v";
          refNum = parseInt(rawRef, 10);
        }

        if (refNum <= 0) {
          return (
            "review_sequence failed: variant must be a media reference like 'v1', 'i3', 'p2', " +
            "or a positive integer."
          );
        }

        const displayRef = `${refType}${refNum}`;

        const fps = typeof args?.fps === "number" && args.fps > 0 ? args.fps : 2;

        // ── Locate the MOV file via chat_media_state.json ────────────────────────────

        const state = await readState(chatWd);

        // movAbs resolution:
        //   vN / iN  → filename in chatWd, .png → .mov (file lives alongside the PNG)
        //   pN       → sourceUrl is already the absolute path to the .mov
        let movAbs: string | undefined;

        if (refType === "v") {
          const rec = (state.variants || []).find(
            (x: any) => typeof x?.v === "number" && x.v === refNum,
          );
          if (!rec) {
            const available = [
              ...(state.variants || []).map((x: any) => typeof x?.v === "number" ? `v${x.v}` : null),
              ...(state.images  || []).map((x: any) => typeof x?.i === "number" ? `i${x.i}` : null),
            ].filter(Boolean).join(", ");
            return (
              `review_sequence failed: ${displayRef} not found in state.` +
              (available ? ` Available: ${available}.` : " No variants or images recorded yet.")
            );
          }
          const fn = typeof (rec as any).filename === "string" ? String((rec as any).filename) : "";
          if (!fn) return `review_sequence failed: ${displayRef} has no filename recorded.`;
          movAbs = path.join(chatWd, fn.replace(/\.png$/i, ".mov"));

        } else if (refType === "i") {
          const rec = (state.images || []).find(
            (x: any) => typeof x?.i === "number" && x.i === refNum,
          );
          if (!rec) {
            const available = [
              ...(state.variants || []).map((x: any) => typeof x?.v === "number" ? `v${x.v}` : null),
              ...(state.images  || []).map((x: any) => typeof x?.i === "number" ? `i${x.i}` : null),
            ].filter(Boolean).join(", ");
            return (
              `review_sequence failed: ${displayRef} not found in state.` +
              (available ? ` Available: ${available}.` : " No variants or images recorded yet.")
            );
          }
          const fn = typeof (rec as any).filename === "string" ? String((rec as any).filename) : "";
          if (!fn) return `review_sequence failed: ${displayRef} has no filename recorded.`;
          movAbs = path.join(chatWd, fn.replace(/\.png$/i, ".mov"));

        } else {
          // pN: sourceUrl is the absolute path to the .mov (cross-chat reference)
          const rec = (state.pictures || []).find(
            (x: any) => typeof x?.p === "number" && x.p === refNum,
          );
          if (!rec) {
            const available = (state.pictures || [])
              .map((x: any) => typeof x?.p === "number" ? `p${x.p}` : null)
              .filter(Boolean).join(", ");
            return (
              `review_sequence failed: ${displayRef} not found in state.` +
              (available ? ` Available: ${available}.` : " No pictures recorded yet.")
            );
          }
          const src = typeof (rec as any).sourceUrl === "string" ? String((rec as any).sourceUrl) : "";
          if (!src) return `review_sequence failed: ${displayRef} has no sourceUrl recorded.`;
          movAbs = src;
        }

        let movExists = false;
        try {
          await fs.promises.access(movAbs, fs.constants.F_OK);
          movExists = true;
        } catch {
          movExists = false;
        }

        if (!movExists) {
          return (
            `review_sequence failed: no video file found for ${displayRef}. ` +
            `Expected: ${movAbs}. ` +
            `${displayRef} may be a still image, not a video.`
          );
        }

        // ── Schedule sequence review ─────────────────────────────────────────────────

        (state as any).pendingSequenceReview = {
          requestedAt: localTimestamp(),
          movAbs,
          variant: refNum,
          variantLabel: displayRef,
          fps,
          ttlMs: 5 * 60 * 1000,
        };

        await writeStateAtomic(chatWd, state);

        const chatId = (() => {
          try {
            return path.basename(chatWd);
          } catch {
            return "(unknown)";
          }
        })();

        try {
          appendReviewSequenceToPluginLog(
            `Sequence review scheduled: chatId=${chatId} ref=${displayRef} fps=${fps}`,
          );
        } catch {
          // best-effort
        }

        try {
          ctx.status(`Extracting frames for ${displayRef} at ${fps} fps…`);
        } catch {
          // best-effort
        }

        return `Scheduled sequence review for ${displayRef} at ${fps} fps.`;
      } catch (e) {
        return `review_sequence failed: ${String((e as any)?.message || e)}`;
      }
    },
  });
}