export-to-pdf-plugin

import { text, tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { createWriteStream } from "fs";
import { join } from "path";
import { z } from "zod";
import PDFDocument from "pdfkit";

const chatMessageSchema = z.object({
  role: z.enum(["user", "assistant", "system"]),
  content: z.string(),
  timestamp: z.number().optional(),
});

function writePdf(outputPath: string, messages: z.infer<typeof chatMessageSchema>[]): Promise<void> {
  return new Promise((resolve, reject) => {
    const doc = new PDFDocument({ margin: 40 });
    const stream = createWriteStream(outputPath);
    doc.pipe(stream);

    const now = new Date();

    doc.fontSize(18).text("EXPORTACION DE CHAT - LM STUDIO", { align: "center" });
    doc.moveDown(0.5);
    doc.fontSize(10).text(`Fecha: ${now.toLocaleString("es-ES")}`, { align: "center" });
    doc.fontSize(10).text(`Mensajes: ${messages.length}`, { align: "center" });
    doc.moveDown(1);

    for (let i = 0; i < messages.length; i++) {
      const msg = messages[i];
      const role =
        msg.role === "user" ? "USUARIO" :
        msg.role === "assistant" ? "ASISTENTE" : "SISTEMA";
      const time = msg.timestamp
        ? new Date(msg.timestamp).toLocaleTimeString("es-ES")
        : "--:--";

      doc.fontSize(11).font("Helvetica-Bold").text(`[${i + 1}] ${role} (${time})`);
      doc.moveDown(0.3);
      doc.fontSize(10).font("Helvetica").text(msg.content.replace(/\r/g, "").trim());
      doc.moveDown(1);
    }

    doc.end();
    stream.on("finish", resolve);
    stream.on("error", reject);
  });
}

export async function toolsProvider(ctl: ToolsProviderController) {
  const tools: Tool[] = [];

  const exportPdfTool = tool({
    name: "export_chat_to_pdf",
    description: text`
      Exports the current chat conversation to a PDF file.

      You must receive the full list of messages from the user to export.
      Each message should include the role (user/assistant/system) and content.

      The file will be saved in the current working directory with a timestamped name.
      Returns the path where the file was saved.
    `,
    parameters: {
      messages: z.array(chatMessageSchema).describe("Array of chat messages to export"),
      filename: z.string().optional().describe("Optional custom filename (without extension)"),
    },
    implementation: async ({ messages, filename }) => {
      const workingDirectory = ctl.getWorkingDirectory();
      const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
      const outputFilename = (filename ?? `chat-export-${timestamp}`) + ".pdf";
      const outputPath = join(workingDirectory, outputFilename);

      await writePdf(outputPath, messages);
      return { success: true, path: outputPath };
    },
  });

  tools.push(exportPdfTool);
  return tools;
}