116 lines · 3.8 KB
import { z } from "zod";
import { renderMarkdown } from "../markdown";
import type { MarkdownContext } from "../markdown";
import type { Notebook, NotebookCell, NotebookOutput } from "~/lib/code-intel";

const MAX_CELLS = 500;
const MAX_TEXT_CHARS = 20_000;
const MAX_IMAGE_BYTES = 4 * 1024 * 1024;
const IMAGE_MIMES = ["image/png", "image/jpeg", "image/gif"];
const ANSI = /\u001b\[[0-9;]*[A-Za-z]/g;

const multiline = z
  .union([z.string(), z.array(z.string())])
  .transform((value) => (Array.isArray(value) ? value.join("") : value));

const outputSchema = z
  .object({
    output_type: z.string(),
    name: z.string().optional(),
    text: multiline.optional(),
    data: z.record(z.string(), z.unknown()).optional(),
    traceback: z.array(z.string()).optional(),
    ename: z.string().optional(),
    evalue: z.string().optional(),
  })
  .loose();

const cellSchema = z
  .object({
    cell_type: z.string(),
    source: multiline.default(""),
    execution_count: z.number().nullable().optional(),
    outputs: z.array(outputSchema).optional(),
  })
  .loose();

const notebookSchema = z
  .object({
    cells: z.array(cellSchema),
    metadata: z