@jataware/beaker-client 2.0.0-b.1

package/dist/notebook.js

@@ -0,0 +1,690 @@
+import * as messages from '@jupyterlab/services/lib/kernel/messages';
+import { v4 } from 'uuid';
+import { BeakerCellFutures, BeakerCellFuture } from './util.js';
+
+class BeakerBaseCell {
+  // Override index type to allow methods to be defined on the class
+  static IPYNB_KEYS = ["cell_type", "source", "metadata", "id", "attachments", "outputs", "execution_count"];
+  id = BeakerBaseCell.generateId();
+  metadata = {};
+  source = "";
+  status = "idle";
+  children = [];
+  constructor(content) {
+    Object.assign(this, content);
+    if (content.id === void 0) {
+      this.id = BeakerBaseCell.generateId();
+    }
+  }
+  static generateId() {
+    return v4();
+  }
+  fromJSON(obj) {
+    Object.keys(obj).forEach((key) => {
+      this[key] = obj[key];
+    });
+  }
+  toJSON() {
+    return { ...this };
+  }
+  fromIPynb(obj) {
+    Object.keys(obj).forEach((key) => {
+      this[key] = obj[key];
+    });
+  }
+  toIPynb() {
+    const output = {};
+    for (const key of Object.keys(this)) {
+      if (BeakerBaseCell.IPYNB_KEYS.includes(key)) {
+        output[key] = this[key];
+      }
+    }
+    return output;
+  }
+}
+class BeakerRawCell extends BeakerBaseCell {
+  cell_type = "raw";
+  attachments;
+  constructor(content) {
+    super(content);
+    Object.assign(this, content);
+  }
+  execute(session) {
+    return null;
+  }
+}
+class BeakerCodeCell extends BeakerBaseCell {
+  cell_type = "code";
+  outputs = [];
+  execution_count = null;
+  last_execution = { status: "none" };
+  busy = false;
+  constructor(content) {
+    super({ ...content });
+    Object.assign(this, content);
+  }
+  execute(session, syntheticFuture) {
+    this.busy = true;
+    var future;
+    this.outputs.splice(0, this.outputs.length);
+    if (syntheticFuture !== void 0) {
+      future = syntheticFuture;
+    } else {
+      future = session.sendBeakerMessage(
+        "execute_request",
+        {
+          code: this.source,
+          silent: false,
+          store_history: true,
+          user_expressions: {},
+          allow_stdin: true,
+          stop_on_error: false
+        }
+      );
+      future.onIOPub = (msg) => BeakerCellFutures.handleIOPub(msg, this);
+      future.onReply = (msg) => BeakerCellFutures.handleReply(msg, this);
+    }
+    return future;
+  }
+  rollback(session) {
+    if (this?.last_execution?.status === "ok") {
+      const future = session.executeAction(
+        "rollback",
+        {
+          checkpoint_index: this.last_execution.checkpoint_index || null
+        }
+      );
+      this.busy = true;
+      future.done.then((reply_msg) => {
+        if (reply_msg.content.status === "ok") {
+          this.busy = false;
+          this.last_execution = { "status": "none" };
+          this.execution_count = null;
+          this.outputs.splice(0, this.outputs.length);
+          this.children.splice(0, this.children.length);
+        } else if (reply_msg.content.status === "error") {
+          this.last_execution = reply_msg.content;
+          this.outputs.push({ ...reply_msg.content, output_type: "error" });
+        }
+      });
+      return future;
+    }
+    return null;
+  }
+  reset_execution_state() {
+    this.last_execution = { status: "modified" };
+  }
+  toIPynb() {
+    const result = {
+      ...super.toIPynb(),
+      cell_type: "code",
+      outputs: this.outputs.map(
+        (output) => {
+          return {
+            ...output,
+            transient: void 0
+          };
+        }
+      ),
+      execution_count: this.execution_count
+    };
+    return result;
+  }
+}
+class BeakerMarkdownCell extends BeakerBaseCell {
+  static IPYNB_KEYS = ["cell_type", "source", "metadata", "id", "attachments"];
+  cell_type = "markdown";
+  attachments;
+  constructor(content) {
+    super({ ...content });
+    Object.assign(this, content);
+  }
+  execute(session) {
+    return null;
+  }
+  toIPynb() {
+    const output = {};
+    for (const key of Object.keys(this)) {
+      if (BeakerMarkdownCell.IPYNB_KEYS.includes(key)) {
+        output[key] = this[key];
+      }
+    }
+    return output;
+  }
+}
+class BeakerQueryCell extends BeakerBaseCell {
+  cell_type = "query";
+  events = [];
+  _current_input_request_message;
+  _toolCallIndex = /* @__PURE__ */ new Map();
+  constructor(content) {
+    super({ cell_type: "query", ...content });
+    Object.assign(this, content);
+  }
+  execute(session, includeNotebookState = true) {
+    this.events.splice(0, this.events.length);
+    this.children.splice(0, this.children.length);
+    this._toolCallIndex.clear();
+    const handleIOPub = async (msg) => {
+      const msg_type = msg.header.msg_type;
+      const content = msg.content;
+      if (msg_type === "status") {
+        this.status = content.execution_state;
+      } else if (msg_type === "llm_thought") {
+        const tool_calls = (content.tool_calls || []).map(
+          (tc) => ({
+            tool_call_id: tc.tool_call_id,
+            tool_name: tc.tool_name,
+            tool_input: tc.tool_input,
+            state: tc.state || "pending"
+          })
+        );
+        if (!content.thought && tool_calls.length === 0) {
+          return;
+        }
+        const event = {
+          type: "thought",
+          content: {
+            thought: content.thought || "",
+            thought_id: content.thought_id,
+            tool_calls
+          }
+        };
+        this.events.push(event);
+        const eventIndex = this.events.length - 1;
+        tool_calls.forEach((tc, slot) => {
+          if (tc.tool_call_id) {
+            this._toolCallIndex.set(tc.tool_call_id, { eventIndex, slot });
+          }
+        });
+      } else if (msg_type === "tool_call_update") {
+        const entry = this._toolCallIndex.get(content.tool_call_id);
+        if (entry === void 0) {
+          return;
+        }
+        const { eventIndex, slot } = entry;
+        const event = this.events[eventIndex];
+        if (!event || event.type !== "thought") {
+          return;
+        }
+        const existing = event.content.tool_calls[slot];
+        const { tool_call_id: _id, ...updateFields } = content;
+        event.content.tool_calls = [
+          ...event.content.tool_calls.slice(0, slot),
+          { ...existing, ...updateFields },
+          ...event.content.tool_calls.slice(slot + 1)
+        ];
+      } else if (msg_type === "llm_response" && content.name === "response_text") {
+        this.events.push({
+          type: "response",
+          content: content.text
+        });
+      } else if (msg_type === "code_cell") {
+        const nb = session.notebook;
+        const codeCell = new BeakerCodeCell({
+          cell_type: "code",
+          source: content.code,
+          metadata: {
+            parent_cell: this.id
+          },
+          outputs: []
+        });
+        const queryCellIndex = nb.cells.findIndex((cell) => cell.id === this.id);
+        if (queryCellIndex >= 0) {
+          session.notebook.cells.splice(queryCellIndex + 1, 0, codeCell);
+        } else {
+          nb.addCell(codeCell);
+        }
+      } else if (msg_type === "add_child_codecell") {
+        const autoexecute = content.autoexecute;
+        const codeCell = new BeakerCodeCell({
+          cell_type: "code",
+          source: content.code,
+          metadata: {
+            parent_cell: this.id
+          },
+          outputs: [],
+          busy: true
+        });
+        codeCell.last_execution = {
+          status: "pending",
+          checkpoint_index: content.checkpoint_index
+        };
+        this.events.push({
+          type: "code_cell",
+          content: {
+            parent_id: this.id,
+            cell_id: codeCell.id,
+            metadata: {}
+          }
+        });
+        this.children.push(codeCell);
+        const reactiveCell = this.children[this.children.length - 1];
+        if (autoexecute && content.execute_request_msg) {
+          const executeRequest = content.execute_request_msg;
+          new BeakerCellFuture(
+            executeRequest,
+            // Message that is presumably sent
+            true,
+            // expectReply
+            true,
+            // disposeOnDone
+            session.kernel,
+            // kernel
+            reactiveCell
+            // relatedCodecell
+          );
+        }
+      } else if (msg_type == "update_child_codecell") {
+        const cellId = content.id;
+        const cell = this.children.find((value) => value.id === cellId);
+        if (cell === void 0 || cell === null) {
+          console.log("Can't find cell");
+          throw Error("Cell to be updated not found ");
+        }
+        cell.execution_count = content.execution_count;
+        let status = {
+          status: content.execution_status
+        };
+        if (content.execution_status === "error") {
+          const error_details = {
+            ename: msg.content.ename,
+            evalue: msg.content.evalue,
+            traceback: msg.content.traceback
+          };
+          cell.outputs.push({
+            output_type: "error",
+            ...error_details
+          });
+          status = { ...status, ...error_details };
+        } else if (status.status === "ok") {
+          status = { ...status, checkpoint_index: content.checkpoint_index };
+          cell.outputs.push({
+            output_type: "execute_result",
+            data: {
+              "text/plain": content.execution_return
+            }
+          });
+        }
+        cell.last_execution = status;
+      }
+    };
+    const handleStdin = async (msg) => {
+      if (messages.isInputRequestMsg(msg)) {
+        this.events.push({
+          type: "user_question",
+          content: msg.content.prompt
+        });
+        this.status = "awaiting_input";
+        this._current_input_request_message = msg;
+      }
+    };
+    const handleReply = async (msg) => {
+      if (msg.content.status === "ok") {
+        this.last_execution = {
+          status: "ok"
+        };
+      } else if (msg.content.status === "error") {
+        const error_details = {
+          ename: msg.content.ename,
+          evalue: msg.content.evalue,
+          traceback: msg.content.traceback
+        };
+        this.last_execution = {
+          status: "error",
+          ...error_details
+        };
+        this.events.push({
+          type: "error",
+          content: { ...msg.content }
+        });
+      } else if (msg.content.status === "abort") {
+        this.last_execution = {
+          status: "abort"
+        };
+        this.events.push({
+          type: "abort",
+          content: "Request aborted"
+        });
+      }
+    };
+    var metadata = {};
+    if (includeNotebookState) {
+      const notebookState = session.notebook.toIPynb();
+      notebookState.cells = notebookState.cells.filter(
+        // Skip this cell and any children of this cell
+        (cell) => cell.id !== this.id && cell.metadata?.parent_cell !== this.id
+      );
+      metadata["notebook_state"] = notebookState;
+    }
+    const future = session.sendBeakerMessage(
+      "llm_request",
+      {
+        request: this.source
+      },
+      void 0,
+      metadata
+    );
+    future.onIOPub = handleIOPub;
+    future.onStdin = handleStdin;
+    future.onReply = handleReply;
+    return future;
+  }
+  respond(response, session) {
+    if (this.status !== "awaiting_input" || !this._current_input_request_message) {
+      return;
+    }
+    this.events.push({
+      type: "user_answer",
+      content: response
+    });
+    session.kernel?.sendInputReply(
+      {
+        "value": response,
+        "status": "ok"
+      },
+      this._current_input_request_message.header
+    );
+    this.status = "busy";
+    this._current_input_request_message = void 0;
+  }
+  toMultipleCells() {
+    class JoinableMarkdown extends String {
+      prefix = "";
+      suffix = "";
+    }
+    class AgentMarkdown extends JoinableMarkdown {
+      prefix = `
+### Agent:
+<div style="display: flex; width: 100%;">
+<div style="padding: 0.2rem; margin: 0.2rem; margin-left: 4rem;">
+
+`;
+      suffix = `</div></div>
+
+`;
+    }
+    class FinalResponseMarkdown extends AgentMarkdown {
+    }
+    class UserMarkdown extends JoinableMarkdown {
+      prefix = `
+### User:
+<div style="display: flex; width: 100%;">
+<div style="padding: 0.2rem; margin: 0.2rem; margin-left: 4rem;">
+
+`;
+      suffix = `</div></div>
+
+`;
+    }
+    const parentMetadata = {
+      ...this.metadata,
+      beaker_cell_type: "query",
+      prompt: this.source,
+      events: this.events
+    };
+    const eventElements = this.events.map((event) => {
+      if (event.type === "thought") {
+        const lines = [];
+        if (event.content.thought) {
+          lines.push(`${event.content.thought}  `);
+        }
+        for (const tc of event.content.tool_calls || []) {
+          const argsBlob = JSON.stringify(tc.tool_input ?? {});
+          const stateLabel = tc.state ? ` [${tc.state}]` : "";
+          lines.push(`- \`${tc.tool_name}\`${stateLabel} — \`${argsBlob}\``);
+          if (tc.output_preview) {
+            const previewSuffix = tc.output_truncated ? "…" : "";
+            lines.push(`  - output: \`${tc.output_preview}${previewSuffix}\``);
+          }
+          if (tc.error) {
+            lines.push(`  - error: \`${tc.error.ename}: ${tc.error.evalue}\``);
+          }
+        }
+        return new AgentMarkdown(`${lines.join("\n")}
+`);
+      } else if (event.type === "user_question") {
+        return new AgentMarkdown(`**Agent Question:**  
+> ${event.content}
+`);
+      } else if (event.type === "user_answer") {
+        return new UserMarkdown(`**User Response:**  
+> ${event.content}
+`);
+      } else if (event.type === "code_cell") {
+        const childCell = this.children.find((child) => child.id === event.content.cell_id);
+        if (childCell === void 0) {
+          throw `Failed to find child cell ${event.content.cell_id} in children`;
+        }
+        childCell.metadata.beaker_child_of = this.id;
+        childCell.metadata.beakerQueryCellChild = true;
+        return childCell;
+      } else if (event.type === "response") {
+        var output;
+        if (typeof event.content === "string") {
+          const lines = event.content.split("\n").map((line) => /^\s*$/.test(line) ? "" : `${line}`).map((line) => /^#+\s+/.test(line) ? `###${line}` : line);
+          output = `
+${lines.join("\n")}`;
+        } else {
+          output = `
+${event.content}`;
+        }
+        return new FinalResponseMarkdown(output);
+      } else if (event.type === "abort") {
+        return new JoinableMarkdown(event.content);
+      } else if (event.type === "error") {
+        return new BeakerMarkdownCell({
+          cell_type: "markdown",
+          id: v4(),
+          source: [JSON.stringify(event.content, void 0, 2)],
+          metadata: { ...parentMetadata, parentQueryCell: true, beakerQueryCellChild: true }
+        });
+      } else {
+        return new JoinableMarkdown(`Error: Unhandled query event type "${event.type}".`);
+      }
+    });
+    let outputElements = [
+      // new JoinableMarkdown('### User:'),
+      new UserMarkdown(`${this.source}`),
+      ...eventElements
+    ];
+    let lastSeenElement = null;
+    let cell = new BeakerMarkdownCell(
+      {
+        cell_type: "markdown",
+        id: this.id,
+        source: [],
+        // contains all parent metadata
+        metadata: { ...parentMetadata, parentQueryCell: true }
+      }
+    );
+    const outputCells = [
+      cell
+    ];
+    for (let element of outputElements) {
+      if (element instanceof JoinableMarkdown) {
+        if (lastSeenElement === null && element.prefix) {
+          if (element.prefix) {
+            cell.source.push(element.prefix);
+          }
+        }
+        if (lastSeenElement === null || lastSeenElement.constructor === element.constructor) {
+          cell.source.push(element);
+        } else if (!(lastSeenElement instanceof JoinableMarkdown)) {
+          cell = new BeakerMarkdownCell({
+            cell_type: "markdown",
+            id: v4(),
+            source: [],
+            metadata: { skipWhenLoading: true, beakerQueryCellChild: true }
+          });
+          if (element instanceof FinalResponseMarkdown) {
+            cell.metadata.finalResponse = true;
+          }
+          if (element.prefix) {
+            cell.source.push(element.prefix);
+          }
+          cell.source.push(element);
+          outputCells.push(cell);
+        } else {
+          if (lastSeenElement.suffix) {
+            cell.source.push(lastSeenElement.suffix);
+          }
+          if (element.prefix) {
+            cell.source.push(element.prefix);
+          }
+          cell.source.push(element);
+        }
+      } else {
+        if (lastSeenElement instanceof JoinableMarkdown && lastSeenElement.suffix) {
+          cell.source.push(lastSeenElement.suffix);
+        }
+        outputCells.push(element);
+      }
+      lastSeenElement = element;
+    }
+    return outputCells;
+  }
+  // public fromJSON(obj: any) {
+  //     if (this.metadata.beaker_cell_type !== "query") {
+  //         throw TypeError("Cell is trying to be parsed as a query cell, but doesn't have the required metadata.")
+  //     }
+  //     this.source = obj.metadata.prompt;
+  //     this.events = obj.metadata.events;
+  // }
+  // public toJSON() {
+  //     return this.toMarkdownCell();
+  // }
+  fromIPynb(obj) {
+    if (this.metadata.beaker_cell_type !== "query") {
+      throw TypeError("Cell is trying to be parsed as a query cell, but doesn't have the required metadata.");
+    }
+    this.source = obj.metadata.prompt;
+    this.events = obj.metadata.events;
+  }
+  toIPynb() {
+    return this.toMultipleCells().flatMap((cell) => cell.toIPynb());
+  }
+}
+class BeakerNotebookContent {
+  nbformat = 4;
+  nbformat_minor = 5;
+  metadata = {};
+  cells = [];
+}
+class BeakerNotebook {
+  constructor() {
+    this.content = {
+      nbformat: 4,
+      nbformat_minor: 5,
+      cells: [],
+      metadata: {
+        "kernelspec": {
+          "display_name": "Beaker Kernel",
+          "name": "beaker",
+          "language": "beaker"
+        },
+        "language_info": this.subkernelInfo
+      }
+    };
+    this.cells = new Proxy(this.content.cells, {});
+  }
+  toJSON() {
+    this.content.metadata.language_info = this.subkernelInfo;
+    return { ...this.content };
+  }
+  fromJSON(obj) {
+    Object.keys(obj).forEach((key) => {
+      this.content[key] = obj[key];
+    });
+    this.cells = new Proxy(this.content.cells, {});
+  }
+  loadFromIPynb(obj) {
+    this.content.nbformat = obj.nbformat;
+    this.content.nbformat_minor = obj.nbformat_minor;
+    const cellList = obj.cells.map((cell) => {
+      if (cell.cell_type === "raw") {
+        return new BeakerRawCell(cell);
+      } else if (cell.cell_type === "code") {
+        return new BeakerCodeCell(cell);
+      } else if (cell.metadata.beaker_cell_type === "query") {
+        return new BeakerQueryCell({
+          cell_type: "query",
+          id: cell.id,
+          source: cell.metadata.prompt,
+          events: cell.metadata.events,
+          metadata: cell.metadata
+        });
+      } else if (cell.cell_type === "markdown") {
+        if (cell.metadata?.skipWhenLoading === true) {
+          return;
+        }
+        return new BeakerMarkdownCell(cell);
+      } else {
+        return new BeakerRawCell(cell);
+      }
+    }).filter((cell) => cell);
+    let cellMap = {};
+    cellList.forEach((cell) => {
+      cellMap[cell.id] = cell;
+    });
+    let i = 0;
+    while (i < cellList.length) {
+      const cell = cellList[i];
+      if (typeof cell.metadata.beaker_child_of !== "undefined") {
+        cellMap[cell.metadata.beaker_child_of].children.push(cellList.splice(i, 1)[0]);
+      } else {
+        i++;
+      }
+    }
+    this.cells.splice(
+      0,
+      this.cells.length,
+      ...cellList
+    );
+    this.content.metadata = obj.metadata;
+  }
+  toIPynb() {
+    this.content.metadata.language_info = this.subkernelInfo;
+    return {
+      nbformat: this.content.nbformat,
+      nbformat_minor: this.content.nbformat_minor,
+      cells: this.content.cells.flatMap((cell) => cell.toIPynb()),
+      metadata: this.content.metadata
+    };
+  }
+  addCell(cell) {
+    this.cells.push(cell);
+  }
+  insertCell(cell, position) {
+    this.cells.splice(position, 0, cell);
+  }
+  removeCell(index) {
+    this.cells.splice(index, 1);
+  }
+  cutCell(index) {
+    const removedValues = this.cells.splice(index, 1);
+    if (removedValues.length > 0) {
+      return removedValues[0];
+    }
+    return null;
+  }
+  moveCell(fromIndex, toIndex) {
+    if (fromIndex !== toIndex) {
+      const cell = this.cutCell(fromIndex);
+      if (cell) {
+        this.insertCell(cell, toIndex);
+      }
+    }
+  }
+  setSubkernelInfo(sessionInfo) {
+    this.subkernelInfo = {
+      "name": sessionInfo.subkernel,
+      "display_name": sessionInfo.language
+    };
+  }
+  content;
+  cells;
+  subkernelInfo;
+}
+
+export { BeakerBaseCell, BeakerCodeCell, BeakerMarkdownCell, BeakerNotebook, BeakerNotebookContent, BeakerQueryCell, BeakerRawCell };