@sema-agent/core/dist/tools/fs/notebook.js

Stateless, task-oriented AI agent core

import { imageMagicMatches } from "./safety.js";
export function isNotebookPath(p) {
    return p.toLowerCase().endsWith(".ipynb");
}
export const NOTEBOOK_CELL_OUTPUT_POINTER_CHARS = 10_000;
export const NOTEBOOK_OUTPUT_TEXT_CLIP = 10_000;
export const NOTEBOOK_IMAGE_BASE64_BUDGET = 262_144;
function joinedText(v) {
    if (v === undefined || v === null)
        return "";
    return Array.isArray(v) ? v.join("") : String(v);
}
function clipOutputText(s) {
    if (s.length <= NOTEBOOK_OUTPUT_TEXT_CLIP)
        return s;
    const droppedLines = s.slice(NOTEBOOK_OUTPUT_TEXT_CLIP).split("\n").length;
    return `${s.slice(0, NOTEBOOK_OUTPUT_TEXT_CLIP)}\n\n… [${droppedLines} lines truncated] …`;
}
function imageFromData(data) {
    if (data === undefined)
        return undefined;
    for (const mime of ["image/png", "image/jpeg"]) {
        const v = data[mime];
        if (typeof v !== "string")
            continue;
        const stripped = v.replace(/\s/g, "");
        if (!imageMagicMatches(Buffer.from(stripped, "base64"), mime))
            return undefined;
        return { image_data: stripped, media_type: mime };
    }
    return undefined;
}
function mapOutput(raw) {
    if (raw === null || typeof raw !== "object")
        return undefined;
    const o = raw;
    switch (o.output_type) {
        case "stream":
            return { output_type: "stream", text: clipOutputText(joinedText(o.text)) };
        case "execute_result":
        case "display_data": {
            const data = o.data !== null && typeof o.data === "object" ? o.data : undefined;
            const image = imageFromData(data);
            return { output_type: o.output_type, text: clipOutputText(joinedText(data?.["text/plain"])), ...(image ? { image } : {}) };
        }
        case "error": {
            const tb = Array.isArray(o.traceback) ? o.traceback.join("\n") : "";
            return { output_type: "error", text: clipOutputText(`${joinedText(o.ename)}: ${joinedText(o.evalue)}\n${tb}`) };
        }
        default:
            return undefined;
    }
}
export function parseNotebookCells(text, notebookPath) {
    let nb;
    try {
        nb = JSON.parse(text);
    }
    catch {
        return { ok: false, reason: "json" };
    }
    const top = nb;
    const rawCells = top?.cells;
    if (!Array.isArray(rawCells) || rawCells.some((c) => c === null || typeof c !== "object")) {
        return { ok: false, reason: "shape" };
    }
    const language = typeof top?.metadata?.language_info?.name === "string" ? top.metadata.language_info.name : "python";
    const cells = rawCells.map((raw, i) => {
        const c = raw;
        const cellType = typeof c.cell_type === "string" ? c.cell_type : "code";
        const cell = {
            cell_id: typeof c.id === "string" ? c.id : `cell-${i}`,
            cellType,
            source: joinedText(c.source),
        };
        if (cellType === "code") {
            cell.language = language;
            if (typeof c.execution_count === "number" && c.execution_count)
                cell.execution_count = c.execution_count;
            if (Array.isArray(c.outputs) && c.outputs.length > 0) {
                const mapped = c.outputs.map(mapOutput).filter((o) => o !== undefined);
                let sum = 0;
                let over = false;
                for (const o of mapped) {
                    sum += (o.text?.length ?? 0) + (o.image?.image_data.length ?? 0);
                    if (sum > NOTEBOOK_CELL_OUTPUT_POINTER_CHARS) {
                        over = true;
                        break;
                    }
                }
                cell.outputs = over
                    ? [{ output_type: "stream", text: `Outputs are too large to include. Use Bash with: cat "${notebookPath}" | jq '.cells[${i}].outputs'` }]
                    : mapped;
            }
        }
        return cell;
    });
    return { ok: true, cells };
}
export function renderNotebookCells(cells, imageBase64Budget) {
    const blocks = [];
    let imageChars = 0;
    const pushText = (t) => {
        const last = blocks[blocks.length - 1];
        if (last !== undefined && last.type === "text")
            last.text += `\n${t}`;
        else
            blocks.push({ type: "text", text: t });
    };
    for (const cell of cells) {
        const tags = [];
        if (cell.cellType !== "code")
            tags.push(`<cell_type>${cell.cellType}</cell_type>`);
        if (cell.language !== undefined && cell.language !== "python" && cell.cellType === "code")
            tags.push(`<language>${cell.language}</language>`);
        pushText(`<cell id="${cell.cell_id}">${tags.join("")}${cell.source}</cell id="${cell.cell_id}">`);
        for (const o of cell.outputs ?? []) {
            if (o.text !== undefined && o.text !== "")
                pushText(`\n${o.text}`);
            if (o.image !== undefined) {
                if (imageChars + o.image.image_data.length > imageBase64Budget) {
                    pushText("[image output omitted: notebook image budget exceeded]");
                }
                else {
                    imageChars += o.image.image_data.length;
                    blocks.push({ type: "image", data: o.image.image_data, mimeType: o.image.media_type });
                }
            }
        }
    }
    const textChars = blocks.reduce((n, b) => n + (b.type === "text" ? b.text.length : 0), 0);
    return { blocks, textChars, imageChars };
}
export function stripNotebookImageData(cell) {
    if (cell.outputs === undefined)
        return cell;
    return {
        ...cell,
        outputs: cell.outputs.map((o) => o.image === undefined ? o : { ...o, image: { ...o.image, image_data: `[${o.image.image_data.length} base64 chars omitted]` } }),
    };
}