@jataware/beaker-client 2.0.0-b.1

package/dist/render.d.ts

@@ -0,0 +1,38 @@
+import { IRenderMime } from '@jupyterlab/rendermime';
+import { PartialJSONObject } from '@lumino/coreutils';
+export interface IBeakerRendererOptions {
+    renderers: ReadonlyArray<IMimeRenderer>;
+}
+export type MimetypeString = "text/plain" | "text/html" | string;
+/**
+ *
+ */
+export interface IMimeBundle {
+    [mimetype: MimetypeString]: PartialJSONObject;
+}
+export interface IMimeRenderer<OutputType = HTMLElement> {
+    rank: number;
+    mimetypes: MimetypeString[];
+    render: (mimeType: MimetypeString, data: PartialJSONObject, metadata: PartialJSONObject) => OutputType;
+}
+export declare class MimeRenderer implements IMimeRenderer<HTMLElement> {
+    rank: number;
+    mimetypes: string[];
+    render(mimeType: MimetypeString, data: PartialJSONObject, metadata: PartialJSONObject): HTMLElement;
+}
+export declare class JupyterMimeRenderer extends MimeRenderer {
+    constructor(factory: IRenderMime.IRendererFactory);
+    render(mimeType: MimetypeString, data: PartialJSONObject, metadata?: PartialJSONObject): HTMLElement;
+    private _factory;
+}
+export declare class BeakerRenderer {
+    constructor(options?: IBeakerRendererOptions);
+    addRenderer(renderer: IMimeRenderer): void;
+    get rankedMimetypes(): MimetypeString[];
+    render(mimeType: MimetypeString, data: PartialJSONObject, metadata?: PartialJSONObject): any;
+    renderMimeBundle(bundle: IMimeBundle, metadata?: PartialJSONObject): {
+        [key: MimetypeString]: HTMLElement;
+    };
+    rankedMimetypesInBundle(bundle: IMimeBundle): MimetypeString[];
+    private _renderers;
+}

package/dist/render.js

@@ -0,0 +1,84 @@
+import { MimeModel, standardRendererFactories } from '@jupyterlab/rendermime';
+import { Sanitizer } from '@jupyterlab/apputils';
+
+class MimeRenderer {
+  rank = 100;
+  mimetypes = [];
+  render(mimeType, data, metadata) {
+    return new HTMLElement();
+  }
+}
+class JupyterMimeRenderer extends MimeRenderer {
+  constructor(factory) {
+    super();
+    this._factory = factory;
+    this.rank = factory.defaultRank || 100;
+    this.mimetypes = [...factory.mimeTypes];
+  }
+  render(mimeType, data, metadata) {
+    const renderer = this._factory.createRenderer({
+      mimeType,
+      resolver: null,
+      sanitizer: new Sanitizer(),
+      linkHandler: null,
+      latexTypesetter: null,
+      markdownParser: null,
+      translator: void 0
+    });
+    const model = new MimeModel({
+      trusted: true,
+      data: { [mimeType]: data },
+      metadata
+    });
+    renderer.renderModel(model);
+    return renderer.node;
+  }
+  _factory;
+}
+class BeakerRenderer {
+  constructor(options) {
+    this._renderers = {};
+    for (const factory of standardRendererFactories) {
+      const renderer = new JupyterMimeRenderer(factory);
+      this.addRenderer(renderer);
+    }
+    for (const renderer of options?.renderers || []) {
+      this.addRenderer(renderer);
+    }
+  }
+  addRenderer(renderer) {
+    for (const mimetype of renderer.mimetypes) {
+      if (!Object.keys(this._renderers).includes(mimetype)) {
+        this._renderers[mimetype] = renderer;
+      } else {
+        const prev_renderer = this._renderers[mimetype];
+        if (renderer.rank <= prev_renderer.rank) {
+          this._renderers[mimetype] = renderer;
+        }
+      }
+    }
+  }
+  get rankedMimetypes() {
+    const mimetypes = Object.keys(this._renderers);
+    mimetypes.sort((a, b) => this._renderers[a].rank - this._renderers[b].rank);
+    return mimetypes;
+  }
+  render(mimeType, data, metadata) {
+    const renderer = this._renderers[mimeType];
+    if (renderer) {
+      return renderer.render(mimeType, data, metadata || {});
+    }
+  }
+  renderMimeBundle(bundle, metadata) {
+    return Object.fromEntries(Object.entries(bundle).map(([mimeType, content]) => {
+      return [mimeType, this.render(mimeType, content, metadata)];
+    }));
+  }
+  rankedMimetypesInBundle(bundle) {
+    const result = this.rankedMimetypes.filter((mime) => bundle && Object.keys(bundle).includes(mime));
+    return result;
+  }
+  _renderers;
+}
+
+export { BeakerRenderer, JupyterMimeRenderer, MimeRenderer };

package/dist/session.d.ts

@@ -0,0 +1,166 @@
+import { SessionContext } from '@jupyterlab/apputils';
+import { IKernelConnection } from '@jupyterlab/services/lib/kernel/kernel';
+import { ServiceManager } from '@jupyterlab/services';
+import { JSONObject } from '@lumino/coreutils';
+import { Slot } from '@lumino/signaling';
+import { IBeakerAvailableContexts, IBeakerFuture, IActiveContextInfo } from './util';
+import { BeakerNotebook, BeakerRawCell, BeakerCodeCell, BeakerMarkdownCell, BeakerQueryCell } from './notebook';
+import { BeakerRenderer, IBeakerRendererOptions } from './render';
+import * as messages from '@jupyterlab/services/lib/kernel/messages';
+export interface IBeakerSessionOptions {
+    settings: any;
+    name: string;
+    kernelName?: string;
+    sessionId?: string;
+    rendererOptions?: IBeakerRendererOptions;
+    messageHandler?: Slot<any, any>;
+    context?: {
+        slug: string;
+        payload: any;
+    };
+}
+export type JupyterKernelStatus = messages.Status;
+export type BeakerKernelStatus = JupyterKernelStatus | 'connected' | 'connecting' | 'reconnecting' | 'disconnected';
+/**
+ * Main class for connecting to and working with a Beaker kernel.
+ */
+export declare class BeakerSession {
+    constructor(options?: IBeakerSessionOptions);
+    /**
+     * Internal initialization logic once all the services are up and ready.
+     */
+    private initialize;
+    /**
+     * Low-level method for reconnecting to the kernel.
+     *
+     */
+    reconnect(): Promise<SessionContext>;
+    /**
+     * Low-level method for sending a message to the Beaker kernel over the "shell" channel.
+     *
+     * @param messageType - The message type, as passed in `msg.header.msg_type`
+     * @param content - Any JSON-encodable payload to be included with the message
+     * @param messageId - (Optional) Pre-defined id for the message. One will be generated if not provided.
+     */
+    sendBeakerMessage(messageType: string, content: JSONObject, messageId?: string, metadata?: {
+        [key: string]: any;
+    }): IBeakerFuture;
+    /**
+     * Handler for session-specific messages from the Beaker kernel.
+     * This handler will be evoked for all IOPub messages, but should ignore all messages that are not session-specific.
+     *
+     * @param _sessionContext - The session Context related to the incoming message
+     * @param msg - The incoming IOPub message
+     */
+    private _sessionMessageHandler;
+    private _connectionFailureHandler;
+    /**
+     * Returns a promise, that once resolved provides all Beaker contexts available in the session.
+     */
+    availableContexts(): Promise<IBeakerAvailableContexts>;
+    /**
+     * Returns a promise that once resolved provides detailed information about the active context.
+     */
+    activeContext(): Promise<IActiveContextInfo>;
+    setContext(contextPayload: any): Promise<IActiveContextInfo>;
+    /**
+     * Executes a Beaker Action, handling all of the message
+     *
+     * The usual IBeakerFuture response handlers can be applied to the returned future to do act upon the responses.
+     *
+     * @param actionName - Name of the action to execute
+     * @param payload - Payload to pass along with the action
+     * @param messageId - (Optional) Id for request message. If not provided, will be generated automatically.
+     * @returns - A future
+     */
+    executeAction(actionName: string, payload: JSONObject, messageId?: string): IBeakerFuture;
+    /**
+     * Interrupt the kernel activity, stopping execution in both the beaker LLM ReAct loop and the subkernel as needed.
+     * See
+     *
+     * @returns - A future
+     */
+    interrupt(): Promise<any>;
+    /**
+     *
+     * @param source - The `code` contents of the cell.
+     * @param metadata - (Optional) Any metadata to be associated with the cell.
+     * @param outputs - (Optional) Any outputs that should be included/displayed.
+     * @returns - A reference to the generated cell
+     */
+    addCodeCell(source: string, metadata?: {}, outputs?: never[]): BeakerCodeCell;
+    /**
+     * Convenience method for adding a MarkdownCell to the notebook
+     *
+     * @param source - The raw markdown encoded text that should be rendered upon execute
+     * @param metadata - (Optional) Any metadata to be associated with the cell.
+     * @returns - A reference to the generated cell
+     */
+    addMarkdownCell(source: string, metadata?: {}): BeakerMarkdownCell;
+    /**
+     * Convenience method for adding a RawCell to the notebook
+     *
+     * @param source - The raw contents to be included in the raw cell.
+     * @param metadata - (Optional) Any metadata to be associated with the cell.
+     * @returns - A reference to the generated cell
+     */
+    addRawCell(source: string, metadata?: {}): BeakerRawCell;
+    /**
+     * Convenience method for adding a QueryCell to the notebook
+     *
+     * @param source - The contents of the query for the LLM as a plain string
+     * @param metadata - (Optional) Any metadata to be associated with the cell.
+     * @returns - A reference to the generated cell
+     */
+    addQueryCell(source: string, metadata?: {}): BeakerQueryCell;
+    /**
+     * Populates the sessions notebook with the provided notebook json
+     *
+     * @param notebookJSONObject - The json representation of a notebook, as found inside an .ipynb file
+     */
+    loadNotebook(notebookJSONObject: object): void;
+    /**
+     * Completely resets the session, clearing the notebook and history, and restarting the fresh kernel so it is in a fresh state.
+     */
+    reset(): void;
+    /**
+     * A promise that resolves once everything the session requires are also ready.
+     */
+    get sessionReady(): Promise<void>;
+    /**
+     * A reference to the underlying Jupyter SessionContext for this Beaker Session.
+     */
+    get session(): SessionContext;
+    /**
+     * Returns the status of the Beaker kernel
+     */
+    get status(): BeakerKernelStatus;
+    /**
+     * A reference to the Jupyter KernelConnection object for this session.
+     */
+    get kernel(): IKernelConnection | null;
+    get kernelInfo(): any;
+    get lastKernelModel(): any;
+    get prevClientId(): string | undefined;
+    /**
+     * A reference to the Jupyter ServiceManager which contains all of the services for this session.
+     */
+    get services(): ServiceManager;
+    get renderer(): BeakerRenderer;
+    get sessionId(): string;
+    private _initialized;
+    private _sessionId;
+    private _sessionOptions?;
+    private _services;
+    private _serverSettings;
+    private _sessionContext;
+    private _messageHandler?;
+    private _history;
+    private _sessionInfo;
+    private _renderer;
+    private _kernelInfo;
+    private _lastKernelModel?;
+    private _prevClientId?;
+    private _hasBeenConnected;
+    notebook: BeakerNotebook;
+}

package/dist/session.js

@@ -0,0 +1,418 @@
+import { SessionContext } from '@jupyterlab/apputils';
+import { ServerConnection } from '@jupyterlab/services/lib/serverconnection';
+import { ServiceManager } from '@jupyterlab/services';
+import * as messages from '@jupyterlab/services/lib/kernel/messages';
+import { v4 } from 'uuid';
+import { createMessageId } from './util.js';
+import { BeakerNotebook, BeakerCodeCell, BeakerMarkdownCell, BeakerRawCell, BeakerQueryCell } from './notebook.js';
+import { BeakerHistory } from './history.js';
+import { BeakerRenderer } from './render.js';
+
+class BeakerSession {
+  constructor(options) {
+    this._sessionId = options?.sessionId ?? v4();
+    this._sessionOptions = options;
+    this._serverSettings = ServerConnection.makeSettings(options?.settings);
+    this._services = new ServiceManager({
+      serverSettings: this._serverSettings
+    });
+    this._hasBeenConnected = false;
+    this._services.connectionFailure.connect(this._connectionFailureHandler);
+    this._renderer = new BeakerRenderer(options?.rendererOptions);
+    this._history = new BeakerHistory(this._sessionId);
+    this.notebook = new BeakerNotebook();
+    this._initialized = new Promise(async (resolve, reject) => {
+      this._services.ready.then(async () => {
+        await this.initialize(options);
+        if (options?.context) {
+          const active_context = await this.activeContext();
+          if (active_context.slug !== options.context.slug || JSON.stringify(active_context.config) !== JSON.stringify(options.context.payload)) {
+            await this.setContext({
+              context: options.context.slug,
+              context_info: options.context.payload
+            });
+          }
+        }
+      });
+      resolve();
+    });
+  }
+  /**
+   * Internal initialization logic once all the services are up and ready.
+   */
+  async initialize(options) {
+    this._sessionContext = new SessionContext({
+      sessionManager: this._services.sessions,
+      specsManager: this._services.kernelspecs,
+      name: options?.name,
+      path: options?.sessionId,
+      kernelPreference: options?.kernelName ? {
+        name: options?.kernelName
+      } : void 0
+    });
+    this._sessionContext.kernelChanged.connect((sender, { oldValue, newValue, name }) => {
+      oldValue?.anyMessage.disconnect(this._sessionMessageHandler, this);
+      newValue?.anyMessage.disconnect(this._sessionMessageHandler, this);
+      newValue?.anyMessage.connect(this._sessionMessageHandler, this);
+      if (newValue) {
+        this._lastKernelModel = newValue?.model;
+        this._prevClientId = newValue?.clientId;
+      }
+    });
+    this._sessionContext.connectionStatusChanged.connect((sender, status) => {
+      if (!this._hasBeenConnected && status == "connected") {
+        this._hasBeenConnected = true;
+      }
+    });
+    await this._sessionContext.initialize();
+    await this._sessionContext.ready;
+    if (options?.messageHandler) {
+      this._messageHandler = options.messageHandler;
+      this._sessionContext.iopubMessage.connect(this._messageHandler);
+    } else {
+      this._messageHandler = void 0;
+    }
+  }
+  /**
+   * Low-level method for reconnecting to the kernel.
+   *
+   */
+  async reconnect() {
+    this._sessionContext.connectionStatusChanged.emit("connecting");
+    this._sessionContext.dispose;
+    await this.initialize(this._sessionOptions);
+    return this._sessionContext;
+  }
+  /**
+   * Low-level method for sending a message to the Beaker kernel over the "shell" channel.
+   *
+   * @param messageType - The message type, as passed in `msg.header.msg_type`
+   * @param content - Any JSON-encodable payload to be included with the message
+   * @param messageId - (Optional) Pre-defined id for the message. One will be generated if not provided.
+   */
+  sendBeakerMessage(messageType, content, messageId, metadata) {
+    if (!messageId) {
+      messageId = createMessageId(messageType);
+    }
+    if (!this.kernel) {
+      throw Error("Unable to send message. Not connected to kernel.");
+    }
+    const msg = messages.createMessage({
+      session: this.session.session?.id || this.kernel?.clientId,
+      channel: "shell",
+      msgType: messageType,
+      content,
+      msgId: messageId,
+      metadata
+    });
+    const future = this.kernel.sendShellMessage(
+      msg,
+      true,
+      true
+    );
+    future.msgId = messageId;
+    return future;
+  }
+  /**
+   * Handler for session-specific messages from the Beaker kernel.
+   * This handler will be evoked for all IOPub messages, but should ignore all messages that are not session-specific.
+   *
+   * @param _sessionContext - The session Context related to the incoming message
+   * @param msg - The incoming IOPub message
+   */
+  _sessionMessageHandler(kernel, { msg, direction }) {
+    if (msg.header.msg_type === "context_setup_response" || msg.header.msg_type === "context_info_response") {
+      if (msg.header.msg_type === "context_setup_response") {
+        this._sessionInfo = msg.content;
+      } else if (msg.header.msg_type === "context_info_response") {
+        this._sessionInfo = msg.content.info;
+      }
+      this.notebook.setSubkernelInfo(this._sessionInfo);
+    } else if (msg.header.msg_type === "kernel_info_reply") {
+      this._kernelInfo = msg.content;
+    } else if (msg.header.msg_type === "notebook_state_request") {
+      const outboundMsg = messages.createMessage({
+        session: kernel.clientId,
+        channel: "shell",
+        msgType: "notebook_state_response",
+        content: this.notebook.toIPynb(),
+        msgId: msg.header.msg_id + "_reply",
+        metadata: {}
+      });
+      kernel.sendShellMessage(
+        outboundMsg,
+        false,
+        true
+      );
+    }
+  }
+  _connectionFailureHandler(serviceManager, error) {
+    console.log("CONNECTION ERROR:", typeof error);
+    console.log({ cause: error.cause, message: error.message, name: error.name, stack: error.stack });
+    console.error(error);
+    console.error();
+  }
+  /**
+   * Returns a promise, that once resolved provides all Beaker contexts available in the session.
+   */
+  async availableContexts() {
+    return new Promise(async (resolve) => {
+      `${this._serverSettings.baseUrl}contexts`;
+      const response = await fetch(`${this._serverSettings.baseUrl}contexts`);
+      const data = await response.json();
+      resolve(data);
+    });
+  }
+  /**
+   * Returns a promise that once resolved provides detailed information about the active context.
+   */
+  async activeContext() {
+    return new Promise(async (resolve, reject) => {
+      await this.sessionReady;
+      const future = this.sendBeakerMessage(
+        "context_info_request",
+        {}
+      );
+      if (future !== void 0) {
+        future.onIOPub = async (msg) => {
+          if (msg.header.msg_type === "context_info_response") {
+            resolve({
+              ...msg.content,
+              kernelInfo: await this.kernel?.info
+            });
+          }
+        };
+        await future.done;
+      }
+      reject({});
+    });
+  }
+  async setContext(contextPayload) {
+    return new Promise(async (resolve, reject) => {
+      const setupResult = this.sendBeakerMessage(
+        "context_setup_request",
+        contextPayload
+      );
+      if (setupResult) {
+        await setupResult.done;
+        const contextInfo = setupResult.msg.content;
+        const kernelInfo = await this.kernel?.requestKernelInfo();
+        resolve({
+          ...contextInfo,
+          kernelInfo: kernelInfo?.content
+        });
+      }
+      reject({});
+    });
+  }
+  /**
+   * Executes a Beaker Action, handling all of the message
+   *
+   * The usual IBeakerFuture response handlers can be applied to the returned future to do act upon the responses.
+   *
+   * @param actionName - Name of the action to execute
+   * @param payload - Payload to pass along with the action
+   * @param messageId - (Optional) Id for request message. If not provided, will be generated automatically.
+   * @returns - A future
+   */
+  executeAction(actionName, payload, messageId) {
+    const requestType = `${actionName}_request`;
+    const responseType = `${actionName}_response`;
+    const messageFuture = this.sendBeakerMessage(
+      requestType,
+      payload,
+      messageId
+    );
+    if (messageFuture) {
+      const responseHandler = async (msg) => {
+        if (msg.header.msg_type === responseType && messageFuture.onResponse) {
+          await messageFuture.onResponse(msg);
+        } else if (msg.header.msg_type === "code_cell") {
+          const nb = this.notebook;
+          const codeCell = new BeakerCodeCell({
+            cell_type: "code",
+            source: msg.content.code,
+            metadata: {},
+            outputs: []
+          });
+          nb.addCell(codeCell);
+        }
+        return true;
+      };
+      messageFuture.registerMessageHook(responseHandler);
+    }
+    return messageFuture;
+  }
+  /**
+   * Interrupt the kernel activity, stopping execution in both the beaker LLM ReAct loop and the subkernel as needed.
+   * See
+   *
+   * @returns - A future
+   */
+  interrupt() {
+    if (this.session.session?.kernel) {
+      return this.session.session.kernel.interrupt();
+    } else {
+      return new Promise(() => {
+      });
+    }
+  }
+  /**
+   *
+   * @param source - The `code` contents of the cell.
+   * @param metadata - (Optional) Any metadata to be associated with the cell.
+   * @param outputs - (Optional) Any outputs that should be included/displayed.
+   * @returns - A reference to the generated cell
+   */
+  addCodeCell(source, metadata = {}, outputs = []) {
+    const cell = new BeakerCodeCell({
+      cell_type: "code",
+      source,
+      metadata,
+      outputs
+    });
+    this.notebook.addCell(cell);
+    return cell;
+  }
+  /**
+   * Convenience method for adding a MarkdownCell to the notebook
+   *
+   * @param source - The raw markdown encoded text that should be rendered upon execute
+   * @param metadata - (Optional) Any metadata to be associated with the cell.
+   * @returns - A reference to the generated cell
+   */
+  addMarkdownCell(source, metadata = {}) {
+    const cell = new BeakerMarkdownCell({
+      cell_type: "markdown",
+      source,
+      metadata
+    });
+    this.notebook.addCell(cell);
+    return cell;
+  }
+  /**
+   * Convenience method for adding a RawCell to the notebook
+   *
+   * @param source - The raw contents to be included in the raw cell.
+   * @param metadata - (Optional) Any metadata to be associated with the cell.
+   * @returns - A reference to the generated cell
+   */
+  addRawCell(source, metadata = {}) {
+    const cell = new BeakerRawCell({
+      cell_type: "raw",
+      source,
+      metadata
+    });
+    this.notebook.addCell(cell);
+    return cell;
+  }
+  /**
+   * Convenience method for adding a QueryCell to the notebook
+   *
+   * @param source - The contents of the query for the LLM as a plain string
+   * @param metadata - (Optional) Any metadata to be associated with the cell.
+   * @returns - A reference to the generated cell
+   */
+  addQueryCell(source, metadata = {}) {
+    const cell = new BeakerQueryCell({
+      cell_type: "query",
+      source,
+      metadata
+    });
+    this.notebook.addCell(cell);
+    return cell;
+  }
+  /**
+   * Populates the sessions notebook with the provided notebook json
+   *
+   * @param notebookJSONObject - The json representation of a notebook, as found inside an .ipynb file
+   */
+  loadNotebook(notebookJSONObject) {
+    this.notebook.loadFromIPynb(notebookJSONObject);
+  }
+  /**
+   * Completely resets the session, clearing the notebook and history, and restarting the fresh kernel so it is in a fresh state.
+   */
+  reset() {
+    this.notebook.cells.splice(0, this.notebook.cells.length);
+    this._history.clear();
+    this._sessionContext.restartKernel();
+  }
+  /**
+   * A promise that resolves once everything the session requires are also ready.
+   */
+  get sessionReady() {
+    return new Promise(async (resolve) => {
+      await this._services.ready;
+      await this._sessionContext.ready;
+      await this._initialized;
+      resolve();
+    });
+  }
+  /**
+   * A reference to the underlying Jupyter SessionContext for this Beaker Session.
+   */
+  get session() {
+    return this._sessionContext;
+  }
+  /**
+   * Returns the status of the Beaker kernel
+   */
+  get status() {
+    const connectionStatus = this.kernel?.connectionStatus;
+    const kernelStatus = this.kernel?.status;
+    if (connectionStatus === "connected" && kernelStatus) {
+      return kernelStatus;
+    } else if (connectionStatus === "connecting" && this._hasBeenConnected) {
+      return "reconnecting";
+    } else if (connectionStatus) {
+      return connectionStatus || "unknown";
+    } else {
+      return "disconnected";
+    }
+  }
+  /**
+   * A reference to the Jupyter KernelConnection object for this session.
+   */
+  get kernel() {
+    return this._sessionContext?.session?.kernel || null;
+  }
+  get kernelInfo() {
+    return this._kernelInfo;
+  }
+  get lastKernelModel() {
+    return this._lastKernelModel;
+  }
+  get prevClientId() {
+    return this._prevClientId;
+  }
+  /**
+   * A reference to the Jupyter ServiceManager which contains all of the services for this session.
+   */
+  get services() {
+    return this._services;
+  }
+  get renderer() {
+    return this._renderer;
+  }
+  get sessionId() {
+    return this._sessionContext.path;
+  }
+  _initialized;
+  _sessionId;
+  _sessionOptions;
+  _services;
+  _serverSettings;
+  _sessionContext;
+  _messageHandler;
+  _history;
+  _sessionInfo;
+  _renderer;
+  _kernelInfo;
+  _lastKernelModel;
+  _prevClientId;
+  _hasBeenConnected;
+  notebook;
+}
+
+export { BeakerSession };