freezmo/lib/ai.ts

/**
 * Server-side AI pipeline: vision extraction, activity generation and
 * short-answer grading. Every step asks for JSON, validates with zod and
 * retries once with the validation error fed back to the model.
 */

import { OllamaError, ollamaChat, type OllamaMessage } from "./ollama";
import { needsLanguageRepair } from "./language";
import {
  buildGenerationUserPrompt,
  buildGradingUserPrompt,
  buildHomeworkGuideUserPrompt,
  buildHomeworkMakeUserPrompt,
  buildRepairPrompt,
  GENERATION_SYSTEM_PROMPT,
  GRADING_SYSTEM_PROMPT,
  HOMEWORK_GUIDE_SYSTEM_PROMPT,
  HOMEWORK_MAKE_SYSTEM_PROMPT,
  VISION_SYSTEM_PROMPT,
  VISION_USER_INSTRUCTION,
} from "./prompts";
import {
  extractedNotesSchema,
  generatedActivitySchema,
  gradeResultSchema,
  homeworkMakeSchema,
  homeworkStepSchema,
  type Difficulty,
  type ExtractedNotes,
  type GeneratedActivity,
  type GradeResult,
  type HomeworkGuide,
  type HomeworkMake,
} from "./schemas";
import { parseJsonLoose } from "./utils";
import type { OllamaAuth } from "./config";

export type PipelineStage = "vision" | "generation" | "done";

type JsonStepOptions<T> = {
  model: string;
  auth: OllamaAuth;
  system: string;
  user: string;
  images?: string[];
  /** Validates and can also "tighten" the parsed value. */
  validate: (value: unknown) => { ok: true; value: T } | { ok: false; error: string };
  shapeHint: string;
  label: string;
  signal?: AbortSignal;
  temperature?: number;
};

type JsonStepResult<T> = {
  value: T;
  /** true when the first attempt failed and the repair round was used. */
  repaired: boolean;
  attempts: number;
};

async function runJsonStep<T>({
  model,
  auth,
  system,
  user,
  images,
  validate,
  shapeHint,
  label,
  signal,
  temperature,
}: JsonStepOptions<T>): Promise<JsonStepResult<T>> {
  const messages: OllamaMessage[] = [
    { role: "system", content: system },
    { role: "user", content: user, ...(images && images.length ? { images } : {}) },
  ];

  let lastRaw = "";
  let lastError = "";

  for (let attempt = 1; attempt <= 2; attempt++) {
    const result = await ollamaChat({
      model,
      auth,
      messages,
      json: true,
      signal,
      temperature: temperature ?? 0.4,
    });
    lastRaw = result.content;

    let parsed: unknown;
    try {
      parsed = parseJsonLoose(result.content);
    } catch (error) {
      lastError = `Response was not valid JSON: ${error instanceof Error ? error.message : String(error)}`;
      if (attempt === 1) {
        messages.push({ role: "assistant", content: result.content });
        messages.push({
          role: "user",
          content: buildRepairPrompt(lastRaw, lastError, shapeHint),
        });
        continue;
      }
      throw new OllamaError("bad_response", `${label} failed: ${lastError}`);
    }

    const checked = validate(parsed);
    if (checked.ok) {
      return { value: checked.value, repaired: attempt > 1, attempts: attempt };
    }

    lastError = checked.error;
    if (attempt === 1) {
      messages.push({ role: "assistant", content: result.content });
      messages.push({
        role: "user",
        content: buildRepairPrompt(lastRaw, lastError, shapeHint),
      });
    }
  }

  throw new OllamaError(
    "bad_response",
    `${label} returned data that did not match the expected format: ${lastError}`,
  );
}

/* ------------------------------------------------------------------ */
/* Step 1 - vision                                                     */
/* ------------------------------------------------------------------ */

const VISION_SHAPE_HINT = `{"title":"...","subject":"...","language":"English name of the notes' language","summary":"...","notes":"plain text transcription","keyConcepts":["..."],"definitions":[{"term":"...","definition":"..."}],"formulas":[{"name":"...","formula":"..."}],"facts":["..."]}`;

/**
 * Read the uploaded images (and any pasted text) into structured notes.
 * Any pasted text is passed through verbatim as additional note content.
 *
 * A small vision model can accidentally translate on hard handwriting, so the
 * result is checked and — when it clearly came back in another language —
 * restored by the (stronger, multilingual) text model.
 */
export async function extractNotes(input: {
  model: string;
  /** Model used to restore the language if the transcription was translated. */
  repairModel: string;
  auth: OllamaAuth;
  images: string[];
  text?: string;
  signal?: AbortSignal;
  onProgress?: (message: string) => void;
}): Promise<{ notes: ExtractedNotes; repaired: boolean; languageRepaired: boolean }> {
  const user = input.text?.trim()
    ? `${VISION_USER_INSTRUCTION}\n\nADDITIONAL TEXT THE STUDENT PASTED (treat as notes too):\n${input.text.slice(0, 12_000)}`
    : VISION_USER_INSTRUCTION;

  const { value, repaired } = await runJsonStep<ExtractedNotes>({
    model: input.model,
    auth: input.auth,
    system: VISION_SYSTEM_PROMPT,
    user,
    images: input.images,
    validate: (parsed) => {
      const result = extractedNotesSchema.safeParse(parsed);
      if (!result.success) return { ok: false, error: summariseZodError(result.error) };
      if (!result.data.notes.trim()) return { ok: false, error: "The notes field is empty." };
      return { ok: true, value: result.data };
    },
    shapeHint: VISION_SHAPE_HINT,
    label: "Reading your notes",
    signal: input.signal,
    temperature: 0.2,
  });

  const check = needsLanguageRepair({ language: value.language, text: value.notes });
  if (!check.suspicious) return { notes: value, repaired, languageRepaired: false };

  console.warn(`[freezmo] restoring the notes' language: ${check.reason}`);
  input.onProgress?.(`Restoring the notes' language (${value.language ?? "original"})…`);
  try {
    const restored = await restoreLanguage({
      model: input.repairModel,
      auth: input.auth,
      language: value.language ?? "the original language of the notes",
      extraction: value,
      signal: input.signal,
    });
    return { notes: restored, repaired, languageRepaired: true };
  } catch (error) {
    // The guard is best-effort: a failure must never break generation.
    console.warn("[freezmo] language restore failed, using the original transcription", error);
    return { notes: value, repaired, languageRepaired: false };
  }
}

const RESTORE_SYSTEM_PROMPT = `You are Freezmo's language restorer. A student's study notes were wrongly translated out of their original language (usually into English) before reaching you.

Your job: rewrite the JSON you receive back into the ORIGINAL language, keeping every fact, number, formula, name and definition exactly as it is.

Rules:
- Translate ONLY what was wrongly translated. Do not add, remove, summarise or explain anything.
- Keep the same JSON keys and the same array structure. Keep the same number of items in every array.
- Keep technical terms, units and formulas as the original notes used them.
- Keep proper nouns as they are unless the original language would inflect them.
- Output JSON only, no markdown fences, no commentary.`;

/**
 * Ask the text model to put a translated transcription back into the language
 * the vision model reported. Returns the original value if the repair is
 * unusable.
 */
export async function restoreLanguage(input: {
  model: string;
  auth: OllamaAuth;
  language: string;
  extraction: ExtractedNotes;
  signal?: AbortSignal;
}): Promise<ExtractedNotes> {
  const value = await runJsonStep<ExtractedNotes>({
    model: input.model,
    auth: input.auth,
    system: RESTORE_SYSTEM_PROMPT,
    user: [
      `ORIGINAL LANGUAGE: ${input.language}`,
      "",
      "JSON TO RESTORE:",
      JSON.stringify(input.extraction).slice(0, 40_000),
      "",
      "Return the same JSON with the text fields written in the original language.",
    ].join("\n"),
    validate: (parsed) => {
      const result = extractedNotesSchema.safeParse(parsed);
      if (!result.success) return { ok: false, error: summariseZodError(result.error) };
      if (!result.data.notes.trim()) return { ok: false, error: "The notes field is empty." };
      return { ok: true, value: result.data };
    },
    shapeHint: VISION_SHAPE_HINT,
    label: "Restoring the notes' language",
    signal: input.signal,
    temperature: 0.1,
  });
  return value.value;
}

/* ------------------------------------------------------------------ */
/* Step 2 - generation                                                 */
/* ------------------------------------------------------------------ */

const GENERATION_SHAPE_HINT = `{"title":"...","subject":"...","language":"English name of the language used","summary":"...","activities":[{"type":"multiple_choice","topic":"...","question":"...","options":["...","...","...","..."],"correctIndex":0,"explanation":"..."}, ...]}`;

const VALID_TYPES = [
  "flashcard",
  "multiple_choice",
  "true_false",
  "fill_blank",
  "matching",
  "ordering",
  "short_answer",
] as const;

/**
 * Giveaways that make an ordering card solvable without knowing the material:
 * years, clock times, "first/next/finally", "step 2", "1.".
 */
const ORDERING_GIVEAWAY =
  /\b(1[0-9]{3}|20[0-9]{2})\b|\b\d{1,2}[:.]\d{2}\b|\b(first|second|third|fourth|fifth|sixth|seventh|eighth|1st|2nd|3rd|4th|5th|6th|7th|8th)\b|\b(step|stage|phase)\s*\d+\b|^\s*\d+[.)]\s/i;

/**
 * Be forgiving: validate each activity on its own and keep the good ones.
 * Only if too few survive does the step get retried.
 */
function validateDeck(parsed: unknown, minimum: number, requested: number) {
  if (!parsed || typeof parsed !== "object") {
    return { ok: false as const, error: "Expected a JSON object at the top level." };
  }
  const record = parsed as Record<string, unknown>;
  const title = typeof record.title === "string" && record.title.trim() ? record.title.trim() : "Untitled deck";
  const subject = typeof record.subject === "string" ? record.subject.trim() : undefined;
  const summary = typeof record.summary === "string" ? record.summary.trim() : undefined;
  const language = typeof record.language === "string" && record.language.trim() ? record.language.trim() : undefined;

  const rawActivities = Array.isArray(record.activities)
    ? record.activities
    : Array.isArray(record.items)
      ? record.items
      : null;

  if (!rawActivities) {
    return { ok: false as const, error: `Missing "activities" array. Expected an array of ${requested} activity objects.` };
  }

  const activities: GeneratedActivity[] = [];
  const rejected: string[] = [];

  rawActivities.slice(0, 24).forEach((item, index) => {
    const result = generatedActivitySchema.safeParse(item);
    if (result.success) {
      const activity = sanitiseActivity(result.data);
      // An ordering card with dates/ordinals on it gives the answer away —
      // drop it rather than ship a card the learner can solve without knowing.
      if (activity.type === "ordering" && activity.items.some((entry) => ORDERING_GIVEAWAY.test(entry))) {
        rejected.push(
          `#${index + 1} (ordering): items must not contain dates, years, times or ordinals — the order has to come from the content.`,
        );
        return;
      }
      activities.push(activity);
    } else {
      const type = (item as { type?: unknown } | null)?.type;
      rejected.push(`#${index + 1} (${typeof type === "string" ? type : "no type"}): ${summariseZodError(result.error)}`);
    }
  });

  if (activities.length < minimum) {
    return {
      ok: false as const,
      error: `Only ${activities.length} valid activities were found, at least ${minimum} are needed. Rejected: ${rejected.slice(0, 3).join("; ") || "none"}. Valid types are: ${VALID_TYPES.join(", ")}.`,
    };
  }

  return { ok: true as const, value: { title, subject, summary, language, activities } };
}

export async function generateActivities(input: {
  model: string;
  auth: OllamaAuth;
  notes: string;
  difficulty: Difficulty;
  count: number;
  language?: string;
  existing?: { title: string; types: string[]; count: number };
  focus?: string;
  signal?: AbortSignal;
}): Promise<{
  deck: { title: string; subject?: string; summary?: string; language?: string; activities: GeneratedActivity[] };
  requested: number;
}> {
  const requested = Math.min(Math.max(input.count, 4), 18);

  const { value } = await runJsonStep<{
    title: string;
    subject?: string;
    summary?: string;
    language?: string;
    activities: GeneratedActivity[];
  }>({
    model: input.model,
    auth: input.auth,
    system: GENERATION_SYSTEM_PROMPT,
    user: buildGenerationUserPrompt({
      notes: input.notes,
      difficulty: input.difficulty,
      count: requested,
      language: input.language,
      existing: input.existing,
      focus: input.focus,
    }),
    validate: (parsed) => validateDeck(parsed, Math.max(4, Math.ceil(requested * 0.5)), requested),
    shapeHint: GENERATION_SHAPE_HINT,
    label: "Building your activities",
    signal: input.signal,
    temperature: 0.7,
  });

  return { deck: value, requested };
}

/* ------------------------------------------------------------------ */
/* Grading                                                             */
/* ------------------------------------------------------------------ */

const GRADE_SHAPE_HINT = `{"correct":true,"score":85,"feedback":"short feedback","missing":["key point missed"]}`;

export async function gradeShortAnswer(input: {
  model: string;
  auth: OllamaAuth;
  question: string;
  referenceAnswer: string;
  keyPoints: string[];
  userAnswer: string;
  language?: string;
  signal?: AbortSignal;
}): Promise<GradeResult> {
  const { value } = await runJsonStep<GradeResult>({
    model: input.model,
    auth: input.auth,
    system: GRADING_SYSTEM_PROMPT,
    user: buildGradingUserPrompt(input),
    validate: (parsed) => {
      const result = gradeResultSchema.safeParse(parsed);
      if (!result.success) return { ok: false, error: summariseZodError(result.error) };
      return { ok: true, value: result.data };
    },
    shapeHint: GRADE_SHAPE_HINT,
    label: "Grading your answer",
    signal: input.signal,
    temperature: 0.1,
  });
  // An accepted answer never has "missing" points — anything the question did
  // not ask about is a bonus, and the UI must not show it as a gap.
  return value.correct ? { ...value, missing: [] } : value;
}

/* ------------------------------------------------------------------ */
/* Make my homework                                                    */
/* ------------------------------------------------------------------ */

const HOMEWORK_MAKE_SHAPE_HINT = `{"title":"...","language":"English name of the language used","answer":"the finished homework as plain text"}`;

/** Mode 1: the finished homework, with no commentary at all. */
export async function makeHomework(input: {
  model: string;
  auth: OllamaAuth;
  assignment: string;
  language?: string;
  signal?: AbortSignal;
}): Promise<HomeworkMake> {
  const { value } = await runJsonStep<HomeworkMake>({
    model: input.model,
    auth: input.auth,
    system: HOMEWORK_MAKE_SYSTEM_PROMPT,
    user: buildHomeworkMakeUserPrompt({ assignment: input.assignment, language: input.language }),
    validate: (parsed) => {
      const result = homeworkMakeSchema.safeParse(parsed);
      if (!result.success) return { ok: false, error: summariseZodError(result.error) };
      if (!result.data.answer.trim()) return { ok: false, error: "The answer field is empty." };
      return { ok: true, value: result.data };
    },
    shapeHint: HOMEWORK_MAKE_SHAPE_HINT,
    label: "Writing the homework",
    signal: input.signal,
    temperature: 0.4,
  });
  return value;
}

const HOMEWORK_GUIDE_SHAPE_HINT = `{"title":"...","language":"English name of the language used","steps":[{"title":"...","explanation":"...","hints":["nudge","hint","answer"],"answer":"...","checkpoint":{"type":"multiple_choice","topic":"...","question":"...","options":["...","...","...","..."],"correctIndex":0,"explanation":"..."}}]}`;

/**
 * Mode 2: the interactive guide. Be forgiving with checkpoints — a step with a
 * broken checkpoint still ships (the checkpoint is simply dropped), because the
 * explanation, hints and answer are the valuable part.
 */
export async function buildHomeworkGuide(input: {
  model: string;
  auth: OllamaAuth;
  assignment: string;
  language?: string;
  signal?: AbortSignal;
}): Promise<HomeworkGuide> {
  const { value } = await runJsonStep<HomeworkGuide>({
    model: input.model,
    auth: input.auth,
    system: HOMEWORK_GUIDE_SYSTEM_PROMPT,
    user: buildHomeworkGuideUserPrompt({ assignment: input.assignment, language: input.language }),
    validate: (parsed) => {
      if (!parsed || typeof parsed !== "object") {
        return { ok: false, error: "Expected a JSON object at the top level." };
      }
      const record = parsed as Record<string, unknown>;
      if (!Array.isArray(record.steps)) {
        return { ok: false, error: `Missing "steps" array. Expected 3-8 step objects.` };
      }

      const steps: HomeworkGuide["steps"] = [];
      const rejected: string[] = [];

      record.steps.slice(0, 12).forEach((raw, index) => {
        const result = homeworkStepSchema.safeParse(raw);
        if (!result.success) {
          rejected.push(`#${index + 1}: ${summariseZodError(result.error)}`);
          return;
        }
        const step = result.data;
        if (!step.answer.trim() || !step.explanation.trim()) {
          rejected.push(`#${index + 1}: explanation or answer is empty.`);
          return;
        }
        steps.push({
          ...step,
          hints: step.hints.map((hint) => hint.trim()).filter(Boolean).slice(0, 4),
          checkpoint: step.checkpoint ? sanitiseActivity(step.checkpoint) : undefined,
        });
      });

      if (steps.length < 2) {
        return {
          ok: false,
          error: `Only ${steps.length} usable steps were found, at least 2 are needed. Rejected: ${rejected.slice(0, 3).join("; ") || "none"}.`,
        };
      }

      const title = typeof record.title === "string" && record.title.trim() ? record.title.trim() : "Homework guide";
      const language = typeof record.language === "string" && record.language.trim() ? record.language.trim() : undefined;
      return { ok: true, value: { title, language, steps } };
    },
    shapeHint: HOMEWORK_GUIDE_SHAPE_HINT,
    label: "Building your guide",
    signal: input.signal,
    temperature: 0.5,
  });
  return value;
}

/* ------------------------------------------------------------------ */
/* Helpers                                                             */
/* ------------------------------------------------------------------ */

/**
 * Small fixes so the UI can trust the data: ordering items must be unique
 * (they are used as React keys and are ambiguous otherwise), fill_blank
 * sentences must actually contain a blank, and word banks must contain
 * the answer.
 */
function sanitiseActivity(activity: GeneratedActivity): GeneratedActivity {
  switch (activity.type) {
    case "ordering": {
      const seen = new Set<string>();
      const items = activity.items.filter((item) => {
        const key = item.trim().toLowerCase();
        if (seen.has(key)) return false;
        seen.add(key);
        return true;
      });
      return items.length === activity.items.length ? activity : { ...activity, items };
    }
    case "fill_blank": {
      if (/_{2,}/.test(activity.sentence)) return activity;
      return { ...activity, sentence: `${activity.sentence.trim()} ____` };
    }
    case "matching": {
      const seen = new Set<string>();
      const pairs = activity.pairs.filter((pair) => {
        const key = pair.term.trim().toLowerCase();
        if (seen.has(key)) return false;
        seen.add(key);
        return true;
      });
      return pairs.length === activity.pairs.length ? activity : { ...activity, pairs };
    }
    default:
      return activity;
  }
}

/** Compact, human-readable zod error for the repair prompt. */
function summariseZodError(error: { issues: Array<{ path: PropertyKey[]; message: string }> }): string {
  return error.issues
    .slice(0, 6)
    .map((issue) => `${issue.path.map(String).join(".") || "(root)"}: ${issue.message}`)
    .join("; ");
}

/** Turn the structured extraction into the plain text the generator reads. */
export function notesToText(notes: ExtractedNotes): string {
  const parts: string[] = [];
  parts.push(`TITLE: ${notes.title}`);
  if (notes.subject) parts.push(`SUBJECT: ${notes.subject}`);
  if (notes.language) parts.push(`LANGUAGE: ${notes.language}`);
  if (notes.summary) parts.push(`SUMMARY: ${notes.summary}`);
  parts.push("", "TRANSCRIPTION:", notes.notes);
  if (notes.keyConcepts.length) parts.push("", "KEY CONCEPTS:", ...notes.keyConcepts.map((c) => `- ${c}`));
  if (notes.definitions.length)
    parts.push("", "DEFINITIONS:", ...notes.definitions.map((d) => `- ${d.term}: ${d.definition}`));
  if (notes.formulas.length)
    parts.push("", "FORMULAS:", ...notes.formulas.map((f) => `- ${f.name}: ${f.formula}`));
  if (notes.facts.length) parts.push("", "FACTS:", ...notes.facts.map((f) => `- ${f}`));
  return parts.join("\n").slice(0, 48_000);
}