@oceanum/oceanumlab 4.6.0

package/src/index.ts

@@ -0,0 +1,327 @@
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin,
+  ILayoutRestorer,
+  ILabStatus
+} from '@jupyterlab/application';
+import { ICommandPalette, Notification } from '@jupyterlab/apputils';
+import { LabIcon } from '@jupyterlab/ui-components';
+import { find } from '@lumino/algorithm';
+import { Widget } from '@lumino/widgets';
+import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import { IStateDB } from '@jupyterlab/statedb';
+import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook';
+import { DatameshConnectWidget } from './DatameshWidget';
+import { DatameshUI } from './DatameshUI';
+import { requestAPI } from './handler';
+import { ChatRouter, ChatRouterError, ChatMessage } from './chatRouter';
+import { reporterFor } from './progress';
+import { ConversationPin } from './conversationPin';
+import { snapshotFromIpynb, snapshotOf } from './notebookContext';
+import { notebookHost } from './notebookHost';
+import { KernelHandoff } from './kernelHandoff';
+import { runChatLoop, STOPPED } from './aiLoop';
+import { MAX_OBSERVE_ROUNDS } from './constants';
+
+import '../style/index.css';
+
+import oceanumSvg from '../style/icons/oceanum.svg';
+
+// Theme-adaptive icon using currentColor (auto-adapts to light/dark themes)
+export const oceanumIcon = new LabIcon({
+  name: 'oceanum:icon',
+  svgstr: oceanumSvg
+});
+
+const PLUGIN_ID = '@oceanum/oceanumlab:datamesh-connect';
+
+/**
+ * Initialization data for the extension.
+ */
+export const datamesh_connect_extension: JupyterFrontEndPlugin<void> = {
+  id: PLUGIN_ID,
+  autoStart: true,
+  requires: [
+    ICommandPalette,
+    ILayoutRestorer,
+    ILabStatus,
+    ISettingRegistry,
+    IStateDB
+  ],
+  activate: (
+    app: JupyterFrontEnd,
+    palette: ICommandPalette,
+    restorer: ILayoutRestorer,
+    status: ILabStatus,
+    settingRegistry: ISettingRegistry
+  ) => {
+    console.log('Oceanum datamesh connect extension is loaded');
+
+    //Try to get the datamesh token from the settings
+    const updateSettings = (set: ISettingRegistry.ISettings) => {
+      const token = set.get('datameshToken');
+      if (token && token.user) {
+        window.datameshToken = token.user as string;
+        requestAPI('env/', {
+          method: 'POST',
+          body: JSON.stringify({ DATAMESH_TOKEN: token.user })
+        }).then(res => console.log(res));
+      }
+      window.injectToken = set.get('injectToken').user as boolean;
+    };
+    //Try to get the datamesh token from the envars
+
+    Promise.all([app.restored, settingRegistry.load(PLUGIN_ID)])
+      .then(([, setting]) => {
+        // Read the settings
+        updateSettings(setting);
+
+        // Listen for your plugin setting changes using Signal
+        setting.changed.connect(updateSettings);
+      })
+      .catch(reason => {
+        console.error(
+          `Something went wrong when reading the Oceanumlab settings.\n${reason}`
+        );
+      });
+
+    const getCurrentWidget = (): Widget => {
+      return app.shell!.currentWidget!;
+    };
+
+    const openDatameshUI = (event: any): void => {
+      const widgetId = 'datamesh-ui';
+      const openWidget = find(
+        app.shell.widgets('main'),
+        (widget: Widget, index: number) => {
+          return widget.id === widgetId;
+        }
+      );
+      if (openWidget) {
+        app.shell.activateById(widgetId);
+        return;
+      }
+
+      const datameshUIWidget = new DatameshUI();
+      datameshUIWidget.title.label = 'Oceanum Datamesh';
+      datameshUIWidget.id = widgetId;
+      datameshUIWidget.title.closable = true;
+      datameshUIWidget.title.icon = oceanumIcon;
+      datameshUIWidget.addClass('datamesh-ui');
+      app.shell.add(datameshUIWidget, 'main');
+    };
+
+    const datameshConnectWidget = new DatameshConnectWidget({
+      app,
+      name: 'Datamesh Connect',
+      icon: oceanumIcon,
+      openDatameshUI: openDatameshUI,
+      commands: app.commands,
+      getCurrentWidget
+    });
+    datameshConnectWidget.id = 'datamesh-connect';
+    datameshConnectWidget.title.icon = oceanumIcon;
+    datameshConnectWidget.title.caption = 'Datamesh Connect';
+
+    restorer.add(datameshConnectWidget, 'datamesh-connect');
+
+    // Rank has been chosen somewhat arbitrarily to give priority to the running
+    // sessions widget in the sidebar.
+    app.shell.add(datameshConnectWidget, 'left', { rank: 900 });
+
+    app.commands.addCommand('datamesh-ui:open', {
+      execute: (args: any) => {
+        openDatameshUI(args);
+      }
+    });
+  }
+};
+
+/**
+ * Second plugin: Oceanum AI chat integration.
+ * Wires ChatRouter → KernelHandoff and exposes the
+ * `oceanum-ai:submit-prompt` command.
+ */
+export const oceanum_ai_extension: JupyterFrontEndPlugin<void> = {
+  id: '@oceanum/oceanumlab:ai-chat',
+  autoStart: true,
+  requires: [INotebookTracker, ISettingRegistry],
+  activate: (
+    app: JupyterFrontEnd,
+    notebookTracker: INotebookTracker,
+    settingRegistry: ISettingRegistry
+  ) => {
+    console.log('Oceanum AI chat extension is loaded');
+
+    const SETTINGS_ID = '@oceanum/oceanumlab:datamesh-connect';
+
+    settingRegistry
+      .load(SETTINGS_ID)
+      .then(settings => {
+        // The notebook the current conversation lives in: what the agent is
+        // shown, and where its answers go.
+        const pin = new ConversationPin(notebookHost(app, notebookTracker));
+
+        // What a request carries from the conversation's notebook. An open
+        // notebook is read as it stands; a closed one is read from its file
+        // rather than opened again, which is left for an answer that has
+        // cells to place.
+        // The panel the current request read its selected cell from, if the
+        // notebook was open: a code answer may replace that cell only there.
+        let readFrom: NotebookPanel | null = null;
+        const router = new ChatRouter(settings, async () => {
+          const open = await pin.current();
+          readFrom = open;
+          if (open) {
+            return snapshotOf(open.content);
+          }
+          const path = pin.path();
+          if (!path) {
+            return null;
+          }
+          try {
+            const file = await app.serviceManager.contents.get(path, {
+              content: true
+            });
+            return snapshotFromIpynb(file.content);
+          } catch {
+            // Gone: an answer's cells will go into a new notebook instead.
+            return null;
+          }
+        });
+        // Asked only when an answer has blocks to place: brings the
+        // conversation's notebook to the front, opening it again if closed.
+        const handoff = new KernelHandoff(
+          () => pin.show(),
+          () => readFrom
+        );
+
+        // The run in flight, if any. Command args must be JSON, so a signal
+        // cannot be passed in; Stop is a second command that reaches it here.
+        let current: AbortController | null = null;
+
+        app.commands.addCommand('oceanum-ai:stop', {
+          label: 'Stop the current Oceanum AI response',
+          execute: () => {
+            current?.abort();
+          }
+        });
+
+        app.commands.addCommand('oceanum-ai:new-chat', {
+          label: 'Start a new Oceanum AI chat',
+          execute: () => {
+            // The run in flight belongs to the conversation being thrown away:
+            // stop it placing anything more. The panel discards its result.
+            // It stops being the current run now, not once it has finished
+            // unwinding, so it cannot report progress into the new chat.
+            const run = current;
+            current = null;
+            run?.abort();
+            return pin.start();
+          }
+        });
+
+        app.commands.addCommand('oceanum-ai:chat-context', {
+          label: 'The notebook the current Oceanum AI chat is about',
+          // Starts the conversation if nothing has -- so the first message of
+          // one nobody started with New chat pins the same way -- and reports
+          // its notebook either way.
+          execute: () => pin.ensure()
+        });
+
+        app.commands.addCommand('oceanum-ai:submit-prompt', {
+          label: 'Submit prompt to Oceanum AI',
+          execute: async (args: Record<string, unknown>) => {
+            const prompt = args['prompt'] as string | undefined;
+            const chatHistory = (args['chatHistory'] as ChatMessage[]) ?? [];
+            if (!prompt) {
+              return;
+            }
+            // A caller that skipped 'oceanum-ai:chat-context' still gets a
+            // conversation pinned the same way.
+            await pin.ensure();
+            // Read per prompt, not once at load, so a settings change applies
+            // to the next question without a reload.
+            const autoRunCode = settings.get('autoRunCode')
+              .composite as boolean;
+            const iterate = settings.get('iterate').composite as boolean;
+
+            current?.abort();
+            const controller = new AbortController();
+            current = controller;
+            // This run's progress, heard only while it is the current run.
+            const report = reporterFor(() => current === controller);
+            try {
+              // Cleared in the `finally` below however this ends, so the last
+              // phase does not sit on screen after the answer has arrived --
+              // or after Stop.
+              return await runChatLoop(
+                prompt,
+                chatHistory,
+                {
+                  route: (p, h, signal) => router.route(p, h, signal, report),
+                  observe: (p, h, runs, signal) => {
+                    // The agent's turn again: the notebook's "Running the
+                    // code…" must not stay up while the agent reads what the
+                    // code printed.
+                    report({ phase: 'interpreting' });
+                    return router.observe(p, h, runs, signal, report);
+                  },
+                  place: (response, opts) => {
+                    // The notebook's turn, not the agent's. Without this the
+                    // last phase the AGENT reported stays on screen while a
+                    // cell is running, so the user is told the agent is
+                    // reading dataset details when what is actually happening
+                    // is their own code executing. That is a worse claim than
+                    // the "Thinking…" it replaced, which was vague rather than
+                    // wrong. And "running" only when there is code that will.
+                    if (response.blocks.length > 0) {
+                      const runsCode =
+                        autoRunCode &&
+                        response.blocks.some(
+                          block => block.type === 'code' && block.content.trim()
+                        );
+                      report({ phase: runsCode ? 'running' : 'placing' });
+                    }
+                    return handoff.inject(response, opts);
+                  }
+                },
+                {
+                  autoRunCode,
+                  iterate,
+                  maxRounds: MAX_OBSERVE_ROUNDS,
+                  signal: controller.signal
+                }
+              );
+            } catch (err) {
+              // Anything that failed because the user pressed Stop is not an
+              // error to them: the aborted fetch, or its body read.
+              if (controller.signal.aborted) {
+                return STOPPED;
+              }
+              if (err instanceof ChatRouterError) {
+                Notification.error(err.message, { autoClose: 5000 });
+              } else {
+                console.error('Oceanum AI: unexpected error', err);
+              }
+            } finally {
+              // However this ended -- answered, failed, or stopped -- the
+              // agent is no longer doing anything, so the last phase must not
+              // sit on screen claiming otherwise.
+              report(null);
+              if (current === controller) {
+                current = null;
+              }
+            }
+          }
+        });
+      })
+      .catch(reason => {
+        console.error(
+          `Oceanum AI: could not load settings from ${SETTINGS_ID}.\n${reason}`
+        );
+      });
+  }
+};
+
+export default [datamesh_connect_extension, oceanum_ai_extension];

package/src/kernelHandoff.ts

@@ -0,0 +1,284 @@
+import { NotebookActions, NotebookPanel } from '@jupyterlab/notebook';
+import { CodeCell } from '@jupyterlab/cells';
+import type { ISessionContext } from '@jupyterlab/apputils';
+import type * as nbformat from '@jupyterlab/nbformat';
+import { OceanumResponse } from './chatRouter';
+import { harvestOutputs, ObservedRun, RunOutcome } from './notebookRun';
+import type { PlacedResponse } from './aiLoop';
+
+export interface InjectOptions {
+  /** If true, replace the selected code cell instead of inserting a new one */
+  replaceCodeCell?: boolean;
+  /** Run each code cell as it is placed. A user setting, not a build constant. */
+  autoRun?: boolean;
+  /**
+   * Stop. Once aborted, no further block is placed or run, and a cell that is
+   * executing is interrupted -- the kernel is the half of a round that an
+   * aborted fetch cannot reach.
+   */
+  signal?: AbortSignal;
+}
+
+/** Longest wait for a starting kernel before the cell is reported as not run. */
+const KERNEL_START_WAIT_MS = 10000;
+
+export class KernelHandoff {
+  /**
+   * @param _target The conversation's notebook, brought to the front -- and
+   *   opened again first, if it was closed. Asked only when there is
+   *   something to place, so an answer with no blocks leaves the tabs alone.
+   * @param _readFrom The panel the request read its selected cell from, if
+   *   any. A code answer replaces the selected cell only in that same panel:
+   *   one opened since has some other cell selected, and replacing it would
+   *   overwrite the user's work.
+   */
+  constructor(
+    private _target: () => Promise<NotebookPanel | null>,
+    private _readFrom: () => NotebookPanel | null
+  ) {}
+
+  /**
+   * Place every block the response carries, in order, and return the message
+   * for the chat window.
+   *
+   * This used to switch on `response.type`, because a response was exactly one
+   * of text, code or markdown. It can now carry any number of blocks in any
+   * combination (OCE-173) — a markdown table describing a dataset followed by
+   * the query that produced it is one response, not two — so the branch became
+   * a loop. No blocks is the commonest case and still touches no notebook.
+   */
+  async inject(
+    response: OceanumResponse,
+    options: InjectOptions = {}
+  ): Promise<PlacedResponse> {
+    const runs: ObservedRun[] = [];
+    let halted = false;
+    if (response.blocks.length === 0) {
+      // Nothing to place — just the message for the chat window.
+      return { message: response.message, runs };
+    }
+
+    // Every answer goes into the conversation's notebook, whichever tab is in
+    // front, and that notebook is brought to the front so the user sees it.
+    const notebookPanel = await this._target();
+    if (!notebookPanel) {
+      return {
+        message: response.message + '\n\n(Could not open a notebook)',
+        runs
+      };
+    }
+
+    const notebook = notebookPanel.content;
+    const { sessionContext } = notebookPanel;
+
+    // The user selected ONE cell, so only the first code block can replace it;
+    // everything after is inserted below in order. Captured before the loop
+    // because inserting moves `activeCell`.
+    const selected = notebook.activeCell;
+    let replaceTarget =
+      options.replaceCodeCell &&
+      notebookPanel === this._readFrom() &&
+      selected instanceof CodeCell
+        ? selected
+        : null;
+
+    for (const block of response.blocks) {
+      if (options.signal?.aborted) {
+        break;
+      }
+      if (block.type === 'code') {
+        let cell: CodeCell | null = null;
+        // The target can be gone by now: a markdown block placed before this
+        // one yielded while it rendered, and the user deleted the cell.
+        if (replaceTarget && !replaceTarget.isDisposed) {
+          replaceTarget.model.sharedModel.setSource(block.content);
+          cell = replaceTarget;
+          replaceTarget = null;
+        } else {
+          NotebookActions.insertBelow(notebook);
+          // `insertBelow` makes whatever the user's default cell type is;
+          // code goes in a code cell regardless.
+          if (
+            notebook.activeCell &&
+            !(notebook.activeCell instanceof CodeCell)
+          ) {
+            NotebookActions.changeCellType(notebook, 'code');
+          }
+          const newCell = notebook.activeCell;
+          if (!newCell) {
+            continue;
+          }
+          newCell.model.sharedModel.setSource(block.content);
+          cell = newCell instanceof CodeCell ? newCell : null;
+        }
+
+        // Run THIS cell, not the selection: the cell just written need not be
+        // the active one (a markdown block placed before it moved the
+        // selection), and the user may have several cells selected. Reading
+        // the outputs afterwards is what makes the iterate workflow possible
+        // -- it is the only place the kernel's answer can be seen. A blank
+        // block is placed but not run: `CodeCell.execute` would send nothing
+        // and the empty result would read as a kernel failure.
+        if (options.autoRun && cell && block.content.trim()) {
+          const { outcome, executed } = await this._run(
+            notebookPanel,
+            cell,
+            options.signal
+          );
+          runs.push({
+            code: block.content,
+            message: response.message,
+            ...outcome
+          });
+          if (!executed) {
+            halted = true;
+          }
+        }
+      } else {
+        NotebookActions.insertBelow(notebook);
+        if (!notebook.activeCell) {
+          continue;
+        }
+
+        // Convert to markdown cell and populate
+        NotebookActions.changeCellType(notebook, 'markdown');
+        const mdCell = notebook.activeCell;
+        if (mdCell) {
+          mdCell.model.sharedModel.setSource(block.content);
+          // Render the markdown cell
+          await NotebookActions.run(notebook, sessionContext);
+        }
+      }
+    }
+
+    return { message: response.message, runs, halted };
+  }
+
+  /**
+   * Execute one cell and report what happened, and whether it happened at all.
+   *
+   * `runCells` resolves false for a kernel error (the traceback is in the
+   * outputs) and for some of the ways a cell can fail to execute at all
+   * (pending input, interrupted); it resolves TRUE with nothing run when there
+   * is no kernel, when the kernel is terminating, or when the kernel is still
+   * starting (`CodeCell.execute` bails silently on a session without a
+   * kernel); and it rejects when the kernel is dead. None of the did-not-run
+   * paths touch the outputs, so a replaced cell would still show its previous
+   * run; they are cleared first so that whatever is there afterwards came from
+   * this run. The execution count is cleared with them and only a kernel
+   * reply sets it again, so a null count afterwards means nothing ran,
+   * whatever `runCells` resolved to; that is reported as an error rather than
+   * as a clean, silent run.
+   *
+   * Pending input is refused BEFORE clearing: clearing the outputs of a cell
+   * that is waiting on `input()` disposes the prompt widget, and the kernel
+   * then waits forever for an answer nobody can give.
+   */
+  private async _run(
+    notebookPanel: NotebookPanel,
+    cell: CodeCell,
+    signal?: AbortSignal
+  ): Promise<{ outcome: RunOutcome; executed: boolean }> {
+    const { content: notebook, sessionContext } = notebookPanel;
+    if (sessionContext.kernelDisplayStatus === 'initializing') {
+      // A kernel is starting (for a notebook this call just created, say).
+      // Bounded: `ready` never resolves if the kernel ends up needing to be
+      // selected by hand.
+      await waitForReady(sessionContext, signal);
+      if (signal?.aborted) {
+        return didNotRun('stopped before the kernel was ready');
+      }
+    }
+    // Closing the notebook, or deleting the cell, disposes it and nulls its
+    // model; every access below would throw and take the whole reply with it.
+    if (cell.isDisposed) {
+      return didNotRun('the cell was deleted or its notebook closed');
+    }
+    if (sessionContext.pendingInput) {
+      return didNotRun(
+        'the kernel is waiting for input that must be answered first'
+      );
+    }
+    const interrupt = () => {
+      void sessionContext.session?.kernel?.interrupt().catch(() => {
+        // Already dead or gone; there is nothing left to interrupt.
+      });
+    };
+    cell.model.sharedModel.transact(() => {
+      cell.model.clearExecution();
+    }, false);
+    signal?.addEventListener('abort', interrupt, { once: true });
+    let ran: boolean;
+    let failure: string | null = null;
+    try {
+      ran = await NotebookActions.runCells(notebook, [cell], sessionContext);
+    } catch (err) {
+      // A dead kernel rejects rather than resolving false. Report it like any
+      // other did-not-run so the rounds already placed keep their messages,
+      // and log it: a rejection can also be a fault inside JupyterLab's
+      // executor, and the chat line alone would hide that.
+      console.error('Oceanum AI: cell execution failed', err);
+      ran = false;
+      failure = err instanceof Error ? err.message : String(err);
+    } finally {
+      signal?.removeEventListener('abort', interrupt);
+    }
+    if (cell.isDisposed) {
+      return didNotRun('the cell was deleted or its notebook closed');
+    }
+
+    // Only stream and error outputs are wanted; serialising a rendered plot
+    // or DataFrame just to discard it is not.
+    const outputs: nbformat.IOutput[] = [];
+    const model = cell.model.outputs;
+    for (let i = 0; i < model.length; i++) {
+      const item = model.get(i);
+      if (item.type === 'stream' || item.type === 'error') {
+        outputs.push(item.toJSON());
+      }
+    }
+    const outcome = harvestOutputs(outputs);
+
+    if (
+      (!ran || cell.model.executionCount === null) &&
+      outcome.status === 'ok'
+    ) {
+      return didNotRun(
+        failure ?? 'no kernel available yet, or the execution was interrupted'
+      );
+    }
+    return { outcome, executed: true };
+  }
+}
+
+/**
+ * Resolve when the session context is ready, when Stop is pressed, or after
+ * `KERNEL_START_WAIT_MS`, whichever comes first. Stop is raced too so the
+ * button answers at once rather than after the kernel, or the timeout.
+ */
+function waitForReady(
+  sessionContext: ISessionContext,
+  signal?: AbortSignal
+): Promise<void> {
+  return new Promise<void>(resolve => {
+    const done = () => {
+      clearTimeout(timer);
+      signal?.removeEventListener('abort', done);
+      resolve();
+    };
+    const timer = setTimeout(done, KERNEL_START_WAIT_MS);
+    void sessionContext.ready.then(done);
+    signal?.addEventListener('abort', done, { once: true });
+  });
+}
+
+function didNotRun(reason: string): { outcome: RunOutcome; executed: false } {
+  return {
+    outcome: {
+      status: 'error',
+      stdout: '',
+      error: `The cell did not execute (${reason}).`
+    },
+    executed: false
+  };
+}