Track tokens, retries and timing

src/pipeline.js

// The whole generator, rebuilt on LangChain chat models:
//
//   (optional) research → outline → write the deck → compile → repair what's broken
//
// It replaces generate.js. The steps are the same; what changed is that each
// step receives a `model`, so we can swap in a different model, or a fake one in tests.
//
// Every model call goes through `call()`: a timeout per attempt, retries with
// backoff for errors worth retrying, and a fallback model when one is unhealthy.

import { HumanMessage, SystemMessage } from "@langchain/core/messages";
import { contentText, createModel } from "./llm.js";
import { generateOutline } from "./outline.js";
import { researchTopic } from "./research.js";
import { buildSystemPrompt, extractTag } from "./prompt.js";
import { LIBRARY } from "./library.js";
import { compileDeck } from "./compile.js";
import { withRetry, withTimeout } from "./retry.js";
import { createModelPool } from "./fallback.js";
import { replaceSlides, splitSlides } from "./deck-edit.js";
import { createSlideStream } from "./stream.js";
import { createUsage } from "./usage.js";

const MAX_REPAIR_ROUNDS = 2;
const CALL_TIMEOUT_MS = 90_000; // one model call
const DECK_DEADLINE_MS = 240_000; // the whole deck, inside a typical 300s serverless limit

/** The models to use, best first. The second is lighter and usually less busy. */
export function defaultModels() {
  return [createModel(), createModel({ model: "gemini-flash-lite-latest" })];
}

/**
 * Outline → SlideML, following the component library. The reply is streamed, and
 * `onSlide(number, xml)` is called as soon as each slide is complete.
 */
export async function writeDeck(model, outline, { signal, onSlide = () => {}, onReply } = {}) {
  const stream = await model.stream(
    [new SystemMessage(buildSystemPrompt(LIBRARY)), new HumanMessage(`Write the deck for this outline:\n${JSON.stringify(outline, null, 2)}`)],
    { signal },
  );

  const slides = createSlideStream();
  let count = 0;
  let usage;
  for await (const chunk of stream) {
    usage = chunk.usage_metadata ?? usage; // token counts arrive with the last chunk
    for (const slideXml of slides.push(contentText(chunk))) {
      count += 1;
      onSlide(count, slideXml);
    }
  }

  onReply?.({ usage_metadata: usage });
  const xml = extractTag(slides.text, "deck");
  if (!xml) throw new Error("The model didn't reply with a <deck> element.");
  return xml;
}

/**
 * Send only the broken slides back, each with its own problems, and ask for fixed
 * versions of just those. Returns a Map of slide number → fixed XML.
 */
export async function repairSlides(model, xml, problems, { signal, onReply } = {}) {
  const slides = splitSlides(xml);
  const numbers = [...new Set(problems.map((problem) => problem.slide))].filter((number) => slides[number - 1]);

  const request = numbers
    .map((number) => {
      const list = problems.filter((problem) => problem.slide === number).map((problem) => `- ${problem.message}`).join("\n");
      return `<fix slide="${number}">\n${slides[number - 1]}\n</fix>\nProblems:\n${list}`;
    })
    .join("\n\n");

  const reply = await model.invoke(
    [
      new SystemMessage(buildSystemPrompt(LIBRARY)),
      new HumanMessage(
        "These slides from a deck have problems. Fix each one. Reply with one " +
          '<fix slide="N">…the corrected slide…</fix> per slide, and nothing else.\n\n' +
          request,
      ),
    ],
    { signal },
  );

  onReply?.(reply);
  const fixes = new Map();
  for (const match of contentText(reply).matchAll(/<fix slide="(\d+)">([\s\S]*?)<\/fix>/g)) {
    fixes.set(Number(match[1]), match[2].trim());
  }
  return fixes;
}

/** Send the deck back with the compiler's problems and ask for a corrected deck. */
export async function repairDeck(model, xml, problems, { signal, onReply } = {}) {
  const list = problems.map((problem) => `- Slide ${problem.slide}: ${problem.message}`).join("\n");
  const reply = await model.invoke(
    [
      new SystemMessage(buildSystemPrompt(LIBRARY)),
      new HumanMessage(`This deck has problems:\n${list}\n\nFix them and reply with the whole corrected <deck>.\n\n${xml}`),
    ],
    { signal },
  );
  onReply?.(reply);
  return extractTag(contentText(reply), "deck");
}

/**
 * options:
 *   models      chat models, best first (default: defaultModels()); or `model` for just one
 *   research    true to search the web first
 *   slideCount  how many slides to plan (default 6)
 *   theme       "light" or "dark"
 *   compile     the compiler to use (tests pass a fake one)
 *   onEvent     called with { type: "stage" | "slide" | "retry" | "fallback", … } so a UI can show progress
 *   wait        how retries wait (tests pass a fake one)
 */
export async function generateDeck(topic, options = {}) {
  const { research = false, slideCount = 6, theme, compile = compileDeck, onEvent = () => {}, wait } = options;
  const models = options.model ? [options.model] : (options.models ?? defaultModels());

  const usage = createUsage();
  const onReply = (message) => usage.recordReply(message);
  const pool = createModelPool(models, {
    onFallback: (info) => {
      usage.recordFallback();
      onEvent({ type: "fallback", ...info });
    },
  });
  const deadline = Date.now() + DECK_DEADLINE_MS;

  // stage timing ⟶ fallback ⟶ retries ⟶ timeout ⟶ the actual call
  const call = async (stage, fn) => {
    onEvent({ type: "stage", stage, status: "start" });
    const result = await usage.track(stage, () =>
      pool.run((model) =>
        withRetry(() => withTimeout((signal) => fn(model, signal), CALL_TIMEOUT_MS), {
          retries: 2,
          deadline,
          wait,
          onRetry: (info) => {
            usage.recordRetry();
            onEvent({ type: "retry", stage, ...info });
          },
        }),
      ),
    );
    onEvent({ type: "stage", stage, status: "end" });
    return result;
  };

  const findings = research
    ? await call("research", (model, signal) => researchTopic(model, topic, { signal, onReply }))
    : { notes: "", sources: [] };
  const outline = await call("outline", (model, signal) => generateOutline(model, topic, { slideCount, research: findings.notes, signal, onReply }));
  let xml = await call("deck", (model, signal) =>
    writeDeck(model, outline, { signal, onReply, onSlide: (number, slideXml) => onEvent({ type: "slide", number, xml: slideXml }) }),
  );
  let result = compile(xml, theme);
  let repairs = 0;

  while (result.problems.length > 0 && repairs < MAX_REPAIR_ROUNDS) {
    repairs += 1;
    // Deck-level problems (slide 0, like unreadable XML) need the whole deck; otherwise fix only the broken slides.
    const wholeDeck = result.problems.some((problem) => problem.slide === 0);
    const fixed = wholeDeck
      ? await call("repair", (model, signal) => repairDeck(model, xml, result.problems, { signal, onReply }))
      : replaceSlides(xml, await call("repair", (model, signal) => repairSlides(model, xml, result.problems, { signal, onReply })));
    if (!fixed) break;
    const next = compile(fixed, theme);
    if (next.problems.length >= result.problems.length) break;
    xml = fixed;
    result = next;
  }

  return { outline, xml, slides: result.slides, problems: result.problems, repairs, sources: findings.sources, usage: usage.summary() };
}