();
+ const [showPrompt, setShowPrompt] = useState(false);
const [isCompleterEnabled, setIsCompleterEnabled] = useState(
props.completionProvider && props.completionProvider.isEnabled()
@@ -109,6 +108,7 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element {
// whether the form is currently saving
const [saving, setSaving] = useState(false);
+ const [reseting, setReseting] = useState(false);
/**
* Effect: initialize inputs after fetching server info.
@@ -132,6 +132,8 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element {
}
setLmProvider(server.chat.lmProvider);
setClmProvider(server.completions.lmProvider);
+ setChatPromptConfig(server.config.chat_prompt);
+ setClmPromptConfig(server.config.completion_prompt)
}, [server]);
/**
@@ -190,6 +192,31 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element {
setFields(currFields);
}, [server, lmProvider]);
+ const handleReset = async () => {
+ if (server.state !== ServerInfoState.Ready) {
+ return;
+ }
+
+ setReseting(true);
+ try {
+ await apiKeysAlert.clear();
+ await AiService.deleteConfig();
+ } catch (e) {
+ console.error(e);
+ const msg =
+ e instanceof Error || typeof e === 'string'
+ ? e.toString()
+ : 'An unknown error occurred. Check the console for more details.';
+ alert.show('error', msg);
+ return;
+ } finally {
+ setSaving(false);
+ }
+ await server.refetchAll();
+ alert.show('success', 'Settings reset successfully.');
+ setReseting(false);
+ }
+
const handleSave = async () => {
// compress fields with JSON values
if (server.state !== ServerInfoState.Ready) {
@@ -226,7 +253,9 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element {
}
}),
completions_model_provider_id: clmGlobalId,
- send_with_shift_enter: sendWse
+ send_with_shift_enter: sendWse,
+ chat_prompt: chatPromptConfig,
+ completion_prompt: clmPromptConfig
};
updateRequest = minifyUpdate(server.config, updateRequest);
updateRequest.last_read = server.config.last_read;
@@ -480,68 +509,110 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element {
No Inline Completion models.
)}
- {/* API Keys section */}
- API Keys
-
- {Object.entries(apiKeys).length === 0 &&
- server.config.api_keys.length === 0 ? (
- No API keys are required by the selected models.
- ) : null}
-
- {/* API key inputs for newly-used providers */}
- {Object.entries(apiKeys).map(([apiKeyName, apiKeyValue], idx) => (
-
- setApiKeys(apiKeys => ({
- ...apiKeys,
- [apiKeyName]: e.target.value
- }))
- }
- />
- ))}
- {/* Pre-existing API keys */}
-
-
- {/* Input */}
- Input
-
-
- When writing a message, press Enter to:
-
- {
- setSendWse(e.target.value === 'newline');
- }}
- >
- }
- label="Send the message"
- />
- }
- label={
- <>
- Start a new line (use Shift+Enter to send)
- >
- }
- />
-
-
-
+
+ {/* Prompt Editor */}
+ setShowPrompt(x => !x)} sx={{ cursor: "pointer" }}>
+
+ Edit Prompts
+
+ {showPrompt ? : }
+
+ {
+ showPrompt && (
+
+ setChatPromptConfig(config => ({
+ default: "",
+ ...config,
+ system: evt.target.value
+ }))}
+ fullWidth
+ sx={{
+ scrollbarWidth: "none"
+ }}
+ />
+ setChatPromptConfig(config => ({
+ system: "",
+ ...config,
+ default: evt.target.value
+ }))}
+ fullWidth
+ sx={{
+ scrollbarWidth: "none"
+ }}
+ helperText={
+ <>
+ notebook_code (str): The code of the active nodebook in jupytext percent format
+
+ active_cell_id (str): The uuid of the active cell
+
+ selection (dict | None): The current selection of the user
+
+ variable_context (str): Description of variables used in the input using @ decorator
+ >
+ }
+ />
+ {
+ (
+ <>
+ setChatPromptConfig(config => ({
+ default: "",
+ ...config,
+ system: evt.target.value
+ }))}
+ fullWidth
+ sx={{
+ scrollbarWidth: "none"
+ }}
+ />
+ setChatPromptConfig(config => ({
+ system: "",
+ ...config,
+ default: evt.target.value
+ }))}
+ fullWidth
+ sx={{
+ scrollbarWidth: "none"
+ }}
+ helperText={
+ <>
+ prefix (str): The code of before the cursor position including previous cells
+
+ suffix (str): The code after the cursor position including previous cells
+ >
+ }
+ />
+ >
+ )
+ }
+
+ )
+ }
+
+
+
+
diff --git a/packages/jupyter-ai/src/components/chat.tsx b/packages/jupyter-ai/src/components/chat.tsx
index d55b1a5ce..d40e50d8c 100644
--- a/packages/jupyter-ai/src/components/chat.tsx
+++ b/packages/jupyter-ai/src/components/chat.tsx
@@ -1,4 +1,4 @@
-import React, { useState, useEffect } from 'react';
+import React, { useState, useEffect, useRef } from 'react';
import { Box } from '@mui/system';
import { Button, IconButton, Stack } from '@mui/material';
import SettingsIcon from '@mui/icons-material/Settings';
@@ -9,7 +9,7 @@ import type { IThemeManager } from '@jupyterlab/apputils';
import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
import type { User } from '@jupyterlab/services';
import { ISignal } from '@lumino/signaling';
-
+import { INotebookTracker } from '@jupyterlab/notebook';
import { JlThemeProvider } from './jl-theme-provider';
import { ChatMessages } from './chat-messages';
import { PendingMessages } from './pending-messages';
@@ -33,6 +33,8 @@ import { UserContextProvider, useUserContext } from '../contexts/user-context';
import { ScrollContainer } from './scroll-container';
import { TooltippedIconButton } from './mui-extras/tooltipped-icon-button';
import { TelemetryContextProvider } from '../contexts/telemetry-context';
+import { NotebookTrackerContext } from '../contexts/notebook-tracker-context';
+import { UtilsContext } from '../contexts/utils-context';
type ChatBodyProps = {
chatHandler: ChatHandler;
@@ -207,6 +209,10 @@ export type ChatProps = {
messageFooter: IJaiMessageFooter | null;
telemetryHandler: IJaiTelemetryHandler | null;
userManager: User.IManager;
+ notebookTracker: INotebookTracker | null;
+ // Helps going back to the notebook to the active cell
+ // After some action such as code addition is performed
+ goBackToNotebook: () => void;
};
enum ChatView {
@@ -217,12 +223,26 @@ enum ChatView {
export function Chat(props: ChatProps): JSX.Element {
const [view, setView] = useState(props.chatView || ChatView.Chat);
const [showWelcomeMessage, setShowWelcomeMessage] = useState(false);
+ const ref = useRef();
const openSettingsView = () => {
setShowWelcomeMessage(false);
setView(ChatView.Settings);
};
+ // TODO: Find a better way to do this
+ // This helps in keyboard accessibility where the user can simply press
+ // Escape to go back to the notebook whenever anything is active in the chat
+ // One way could be to add this to all focusable elements such as inputs and
+ // buttons that is less maintainable
+ useEffect(() => {
+ document.addEventListener("keydown", (evt) => {
+ if (evt.key === "Escape" && ref.current && ref.current.contains(document.activeElement)) {
+ props.goBackToNotebook();
+ }
+ })
+ }, []);
+
return (
@@ -231,76 +251,80 @@ export function Chat(props: ChatProps): JSX.Element {
activeCellManager={props.activeCellManager}
>
-
- =4.3.0.
- // See: https://jupyterlab.readthedocs.io/en/latest/extension/extension_migration.html#css-styling
- className="jp-ThemedContainer"
- // root box should not include padding as it offsets the vertical
- // scrollbar to the left
- sx={{
- width: '100%',
- height: '100%',
- boxSizing: 'border-box',
- background: 'var(--jp-layout-color0)',
- display: 'flex',
- flexDirection: 'column'
- }}
+
+
+
- {/* top bar */}
-
- {view !== ChatView.Chat ? (
- setView(ChatView.Chat)}>
-
-
- ) : (
-
- )}
- {view === ChatView.Chat ? (
-
- {!showWelcomeMessage && (
-
- props.chatHandler.sendMessage({ type: 'clear' })
- }
- tooltip="New chat"
- >
-
-
- )}
- openSettingsView()}>
-
+
+ {/* top bar */}
+
+ {view !== ChatView.Chat ? (
+ setView(ChatView.Chat)}>
+
-
- ) : (
-
+ ) : (
+
+ )}
+ {view === ChatView.Chat ? (
+
+ {!showWelcomeMessage && (
+
+ props.chatHandler.sendMessage({ type: 'clear' })
+ }
+ tooltip="New chat"
+ >
+
+
+ )}
+ openSettingsView()}>
+
+
+
+ ) : (
+
+ )}
+
+ {/* body */}
+ {view === ChatView.Chat && (
+
+ )}
+ {view === ChatView.Settings && (
+
)}
-
- {/* body */}
- {view === ChatView.Chat && (
-
- )}
- {view === ChatView.Settings && (
-
- )}
-
-
+
+
+
+
diff --git a/packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx b/packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx
index 315e5d4d6..bbce5bb2d 100644
--- a/packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx
+++ b/packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx
@@ -17,6 +17,7 @@ import { useCopy } from '../../hooks/use-copy';
import { AiService } from '../../handler';
import { useTelemetry } from '../../contexts/telemetry-context';
import { TelemetryEvent } from '../../tokens';
+import { useUtilsContext } from '../../contexts/utils-context';
export type CodeToolbarProps = {
/**
@@ -50,14 +51,31 @@ export function CodeToolbar(props: CodeToolbarProps): JSX.Element {
borderTop: 'none'
}}
>
+
+
-
-
);
}
+const useCellAction = (callback: () => void, command: string, props: ToolbarButtonProps) => {
+ const { goBackToNotebook } = useUtilsContext();
+ const telemetryHandler = useTelemetry();
+
+ return () => {
+ callback();
+ if (goBackToNotebook) goBackToNotebook();
+
+ try {
+ telemetryHandler.onEvent(buildTelemetryEvent(command, props));
+ } catch (e) {
+ console.error(e);
+ return;
+ }
+ }
+}
+
type ToolbarButtonProps = {
code: string;
activeCellExists: boolean;
@@ -96,24 +114,15 @@ function buildTelemetryEvent(
}
function InsertAboveButton(props: ToolbarButtonProps) {
- const telemetryHandler = useTelemetry();
const tooltip = props.activeCellExists
? 'Insert above active cell'
: 'Insert above active cell (no active cell)';
+ const cb = useCellAction(() => props.activeCellManager.insertAbove(props.code), 'insert-above', props);
return (
{
- props.activeCellManager.insertAbove(props.code);
-
- try {
- telemetryHandler.onEvent(buildTelemetryEvent('insert-above', props));
- } catch (e) {
- console.error(e);
- return;
- }
- }}
+ onClick={cb}
disabled={!props.activeCellExists}
>
@@ -122,25 +131,16 @@ function InsertAboveButton(props: ToolbarButtonProps) {
}
function InsertBelowButton(props: ToolbarButtonProps) {
- const telemetryHandler = useTelemetry();
const tooltip = props.activeCellExists
? 'Insert below active cell'
: 'Insert below active cell (no active cell)';
+ const cb = useCellAction(() => props.activeCellManager.insertBelow(props.code), 'insert-below', props);
return (
{
- props.activeCellManager.insertBelow(props.code);
-
- try {
- telemetryHandler.onEvent(buildTelemetryEvent('insert-below', props));
- } catch (e) {
- console.error(e);
- return;
- }
- }}
+ onClick={cb}
>
@@ -148,23 +148,14 @@ function InsertBelowButton(props: ToolbarButtonProps) {
}
function ReplaceButton(props: ToolbarButtonProps) {
- const telemetryHandler = useTelemetry();
const { replace, replaceDisabled, replaceLabel } = useReplace();
+ const cb = useCellAction(() => replace(props.code), 'replace', props);
return (
{
- replace(props.code);
-
- try {
- telemetryHandler.onEvent(buildTelemetryEvent('replace', props));
- } catch (e) {
- console.error(e);
- return;
- }
- }}
+ onClick={cb}
>
@@ -172,23 +163,15 @@ function ReplaceButton(props: ToolbarButtonProps) {
}
export function CopyButton(props: ToolbarButtonProps): JSX.Element {
- const telemetryHandler = useTelemetry();
const { copy, copyLabel } = useCopy();
+ const cb = useCellAction(() => copy(props.code), 'copy', props);
+
return (
{
- copy(props.code);
-
- try {
- telemetryHandler.onEvent(buildTelemetryEvent('copy', props));
- } catch (e) {
- console.error(e);
- return;
- }
- }}
+ onClick={cb}
aria-label="Copy to clipboard"
>
diff --git a/packages/jupyter-ai/src/components/rendermime-markdown.tsx b/packages/jupyter-ai/src/components/rendermime-markdown.tsx
index 9a0278517..31c996111 100644
--- a/packages/jupyter-ai/src/components/rendermime-markdown.tsx
+++ b/packages/jupyter-ai/src/components/rendermime-markdown.tsx
@@ -21,6 +21,11 @@ type RendermimeMarkdownProps = {
* `AgentStreamMessage`, in which case this should be set to `false`.
*/
complete: boolean;
+ /**
+ * Hides the code toolbar. We require it in cases when we
+ * are rendering code blocks in prompts
+ */
+ hideCodeToolbar?: boolean;
};
/**
@@ -63,7 +68,7 @@ function RendermimeMarkdownBase(props: RendermimeMarkdownProps): JSX.Element {
useEffect(() => {
const renderContent = async () => {
// initialize mime model
- const mdStr = escapeLatexDelimiters(props.markdownStr);
+ const mdStr = escapeLatexDelimiters(props.markdownStr);
const model = props.rmRegistry.createModel({
data: { [MD_MIME_TYPE]: mdStr }
});
@@ -126,7 +131,7 @@ function RendermimeMarkdownBase(props: RendermimeMarkdownProps): JSX.Element {
// Render a `CodeToolbar` element underneath each code block.
// We use ReactDOM.createPortal() so each `CodeToolbar` element is able
// to use the context in the main React tree.
- codeToolbarDefns.map(codeToolbarDefn => {
+ !props.hideCodeToolbar && codeToolbarDefns.map(codeToolbarDefn => {
const [codeToolbarRoot, codeToolbarProps] = codeToolbarDefn;
return createPortal(
,
diff --git a/packages/jupyter-ai/src/contexts/active-cell-context.tsx b/packages/jupyter-ai/src/contexts/active-cell-context.tsx
index 72e93a8ca..6979a2ffc 100644
--- a/packages/jupyter-ai/src/contexts/active-cell-context.tsx
+++ b/packages/jupyter-ai/src/contexts/active-cell-context.tsx
@@ -74,6 +74,11 @@ export class ActiveCellManager {
return this._activeCellErrorChanged;
}
+ getActiveCellId() {
+ const sharedModel = this._activeCell?.model.sharedModel;
+ return sharedModel?.id;
+ }
+
/**
* Returns an `ActiveCellContent` object that describes the current active
* cell. If no active cell exists, this method returns `null`.
diff --git a/packages/jupyter-ai/src/contexts/notebook-tracker-context.tsx b/packages/jupyter-ai/src/contexts/notebook-tracker-context.tsx
new file mode 100644
index 000000000..ace1042dc
--- /dev/null
+++ b/packages/jupyter-ai/src/contexts/notebook-tracker-context.tsx
@@ -0,0 +1,19 @@
+import { INotebookTracker } from '@jupyterlab/notebook';
+import React, { useContext } from 'react';
+
+/**
+ * NotebookTrackerContext is a React context that holds an INotebookTracker instance or null.
+ * It is used to provide access to the notebook tracker object throughout the React component tree.
+ */
+export const NotebookTrackerContext = React.createContext(null);
+
+/**
+ * useNotebookTrackerContext is a custom hook that provides access to the current value of the NotebookTrackerContext.
+ * This hook allows components to easily access the notebook tracker without directly using the context API.
+ *
+ * @returns The current value of the NotebookTrackerContext, which is either an INotebookTracker instance or null.
+ */
+export const useNotebookTrackerContext = () => {
+ const tracker = useContext(NotebookTrackerContext);
+ return tracker;
+}
diff --git a/packages/jupyter-ai/src/contexts/utils-context.tsx b/packages/jupyter-ai/src/contexts/utils-context.tsx
new file mode 100644
index 000000000..c3647b059
--- /dev/null
+++ b/packages/jupyter-ai/src/contexts/utils-context.tsx
@@ -0,0 +1,20 @@
+import React, { useContext } from 'react';
+
+/**
+ * UtilsContext is a React context that holds an object with an optional goBackToNotebook function.
+ * It is used to provide utility functions throughout the React component tree.
+ */
+export const UtilsContext = React.createContext<{
+ goBackToNotebook?: () => void;
+}>({});
+
+/**
+ * useUtilsContext is a custom hook that provides access to the current value of the UtilsContext.
+ * This hook allows components to easily access the utility functions without directly using the context API.
+ *
+ * @returns The current value of the UtilsContext, which is an object that may contain a goBackToNotebook function.
+ */
+export const useUtilsContext = () => {
+ const utils = useContext(UtilsContext);
+ return utils;
+}
diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts
index 1b58ed2af..89758bf5d 100644
--- a/packages/jupyter-ai/src/handler.ts
+++ b/packages/jupyter-ai/src/handler.ts
@@ -78,9 +78,21 @@ export namespace AiService {
| CellSelection
| CellWithErrorSelection;
+ export type NotebookCell = {
+ content?: string;
+ type: string;
+ }
+
+ export type ChatNotebookContent = {
+ notebook_code?: NotebookCell[];
+ active_cell_id?: number;
+ variable_context?: string;
+ }
+
export type ChatRequest = {
prompt: string;
- selection?: Selection;
+ selection?: Selection | null;
+ notebook?: ChatNotebookContent | null;
};
export type ClearRequest = {
@@ -201,6 +213,11 @@ export namespace AiService {
pending_messages: PendingMessage[];
};
+ export type PromptConfig = {
+ system: string;
+ default: string;
+ }
+
export type DescribeConfigResponse = {
model_provider_id: string | null;
embeddings_provider_id: string | null;
@@ -209,6 +226,8 @@ export namespace AiService {
fields: Record>;
last_read: number;
completions_model_provider_id: string | null;
+ chat_prompt: PromptConfig;
+ completion_prompt: PromptConfig
};
export type UpdateConfigRequest = {
@@ -220,6 +239,8 @@ export namespace AiService {
last_read?: number;
completions_model_provider_id?: string | null;
completions_fields?: Record>;
+ chat_prompt?: PromptConfig;
+ completion_prompt?: PromptConfig
};
export async function getConfig(): Promise {
@@ -302,6 +323,12 @@ export namespace AiService {
});
}
+ export async function deleteConfig(): Promise {
+ return requestAPI('config', {
+ method: 'DELETE'
+ });
+ }
+
export async function deleteApiKey(keyName: string): Promise {
return requestAPI(`api_keys/${keyName}`, {
method: 'DELETE'
diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts
index f24fbfa00..b6534dcf0 100644
--- a/packages/jupyter-ai/src/index.ts
+++ b/packages/jupyter-ai/src/index.ts
@@ -3,7 +3,7 @@ import {
JupyterFrontEndPlugin,
ILayoutRestorer
} from '@jupyterlab/application';
-
+import { INotebookTracker } from '@jupyterlab/notebook';
import {
IWidgetTracker,
ReactWidget,
@@ -51,7 +51,8 @@ const plugin: JupyterFrontEndPlugin = {
IThemeManager,
IJaiCompletionProvider,
IJaiMessageFooter,
- IJaiTelemetryHandler
+ IJaiTelemetryHandler,
+ INotebookTracker
],
provides: IJaiCore,
activate: async (
@@ -62,7 +63,8 @@ const plugin: JupyterFrontEndPlugin = {
themeManager: IThemeManager | null,
completionProvider: IJaiCompletionProvider | null,
messageFooter: IJaiMessageFooter | null,
- telemetryHandler: IJaiTelemetryHandler | null
+ telemetryHandler: IJaiTelemetryHandler | null,
+ notebookTracker: INotebookTracker | null
) => {
/**
* Initialize selection watcher singleton
@@ -85,6 +87,10 @@ const plugin: JupyterFrontEndPlugin = {
});
};
+ const goBackToNotebook = () => {
+ app.commands.execute('notebook:enter-edit-mode');
+ }
+
const focusInputSignal = new Signal({});
let chatWidget: ReactWidget;
@@ -102,7 +108,9 @@ const plugin: JupyterFrontEndPlugin = {
focusInputSignal,
messageFooter,
telemetryHandler,
- app.serviceManager.user
+ app.serviceManager.user,
+ notebookTracker,
+ goBackToNotebook
);
} catch (e) {
chatWidget = buildErrorWidget(themeManager);
diff --git a/packages/jupyter-ai/src/utils.ts b/packages/jupyter-ai/src/utils.ts
index e21b47c6b..55f4e088c 100644
--- a/packages/jupyter-ai/src/utils.ts
+++ b/packages/jupyter-ai/src/utils.ts
@@ -1,10 +1,13 @@
/**
* Contains various utility functions shared throughout the project.
*/
-import { Notebook } from '@jupyterlab/notebook';
+import { INotebookTracker, Notebook } from '@jupyterlab/notebook';
import { FileEditor } from '@jupyterlab/fileeditor';
import { CodeEditor } from '@jupyterlab/codeeditor';
import { Widget } from '@lumino/widgets';
+import { IKernelConnection } from '@jupyterlab/services/lib/kernel/kernel';
+import { IIOPubMessage } from '@jupyterlab/services/lib/kernel/messages';
+import { CellModel } from '@jupyterlab/cells';
/**
* Get text selection from an editor widget (DocumentWidget#content).
@@ -71,3 +74,175 @@ export function getModelLocalId(globalModelId: string): string | null {
const components = globalModelId.split(':').slice(1);
return components.join(':');
}
+
+/**
+ * Executes a given code string in the provided Jupyter Notebook kernel and retrieves the output.
+ *
+ * @param kernel - The kernel connection to execute the code.
+ * @param code - The code string to be executed.
+ * @returns A promise that resolves with the output of the executed code or null if the execution fails.
+ */
+export async function executeCode(kernel: IKernelConnection, code: string): Promise {
+ const executeRequest = kernel.requestExecute({ code });
+ let variableValue: string | null = null;
+
+ kernel.registerMessageHook(executeRequest.msg.header.msg_id, (msg: IIOPubMessage) => {
+ if (
+ msg.header.msg_type === 'stream' &&
+ // @ts-expect-error tserror
+ msg.content.name === 'stdout'
+ ) {
+ // @ts-expect-error tserror
+ variableValue = msg.content.text.trim();
+ }
+ return true;
+ });
+
+ const reply = await executeRequest.done;
+ if (reply && reply.content.status === 'ok') {
+ return variableValue;
+ } else {
+ console.error('Failed to retrieve variable value');
+ return null;
+ }
+}
+
+/**
+ * Generates the code to describe a variable using a Jupyter AI utility.
+ *
+ * @param name - The name of the variable to describe.
+ * @returns The generated code as a string.
+ */
+const getCode = (name: string) => `
+from jupyter_ai._variable_describer import describe_var
+print(describe_var('${name}', ${name}))
+`
+
+/**
+ * Retrieves the description of a variable from the current Jupyter Notebook kernel.
+ *
+ * @param variableName - The name of the variable to describe.
+ * @param notebookTracker - The notebook tracker to get the current notebook and kernel.
+ * @returns A promise that resolves with the variable description or null if retrieval fails.
+ */
+export async function getVariableDescription(
+ variableName: string,
+ notebookTracker: INotebookTracker
+): Promise {
+ const notebook = notebookTracker.currentWidget;
+ if (notebook && notebook.sessionContext.session?.kernel) {
+ const kernel = notebook.sessionContext.session.kernel;
+ try {
+ return await executeCode(kernel, getCode(variableName));
+ } catch (error) {
+ console.error('Error retrieving variable value:', error);
+ return null;
+ }
+ } else {
+ console.error('No active kernel found');
+ return null;
+ }
+}
+
+/**
+ * Processes variables in the user input by replacing them with their descriptions from the Jupyter Notebook kernel.
+ *
+ * @param userInput - The user input string containing variables to process.
+ * @param notebookTracker - The notebook tracker to get the current notebook and kernel.
+ * @returns A promise that resolves with an object containing the processed input and variable values.
+ */
+export async function processVariables(
+ userInput: string,
+ notebookTracker: INotebookTracker | null
+ ): Promise<{
+ processedInput: string,
+ varValues: string
+ }> {
+ if (!notebookTracker) {
+ return {
+ processedInput: userInput,
+ varValues: ""
+ }
+ }
+
+ const variablePattern = / @([a-zA-Z_][a-zA-Z0-9_]*(\[[^\]]+\]|(\.[a-zA-Z_][a-zA-Z0-9_]*)?)*)[\s,\-]?/g;
+ let match;
+ let variablesProcessed: string[] = [];
+ let varValues = '';
+ let processedInput = userInput;
+
+ while ((match = variablePattern.exec(userInput)) !== null) {
+ const variableName = match[1];
+ if (variablesProcessed.includes(variableName)) {
+ continue;
+ }
+ variablesProcessed.push(variableName);
+ try {
+ const varDescription = await getVariableDescription(variableName, notebookTracker);
+ varValues += `${varDescription}`;
+ processedInput = processedInput.replace(`@${variableName}`, `\`${variableName}\``);
+ } catch (error) {
+ console.error(`Error accessing variable ${variableName}:`, error);
+ }
+ }
+
+ return { processedInput, varValues}
+ }
+
+/**
+ * Code to retrieve the global variables from the Jupyter Notebook kernel.
+ */
+const GLOBALS_CODE = `
+from jupyter_ai.filter_autocomplete_globals import filter_globals
+print(filter_globals(globals()))
+`
+
+/**
+ * Retrieves completion suggestions for a given prefix from the current Jupyter Notebook kernel's global variables.
+ *
+ * @param prefix - The prefix to match global variables against.
+ * @param notebookTracker - The notebook tracker to get the current notebook and kernel.
+ * @returns A promise that resolves with an array of matching global variable names.
+ */
+export async function getCompletion(
+ prefix: string,
+ notebookTracker: INotebookTracker | null
+): Promise {
+ const notebook = notebookTracker?.currentWidget;
+ if (notebook?.sessionContext?.session?.kernel) {
+ const kernel = notebook.sessionContext.session.kernel;
+ const globalsString = await executeCode(kernel, GLOBALS_CODE);
+
+ if (!globalsString) return [];
+
+ const globals: string[] = JSON.parse(globalsString.replace(/'/g, '"'));
+ const filteredGlobals = globals.filter(x => x.startsWith(prefix))
+
+
+ return filteredGlobals;
+ }
+
+ return [];
+}
+
+/**
+ * Formats code for a Jupyter Notebook cell, adding appropriate annotations based on cell type.
+ *
+ * @param cell - The cell model containing the code to format.
+ * @param index - Optional index of the cell for annotation.
+ * @returns The formatted code as a string.
+ */
+export const formatCodeForCell = (cell: CellModel["sharedModel"], index?: number) => {
+ const cellNumber = index !== undefined ? `Cell ${index}` : '';
+
+ switch (cell.cell_type) {
+ case "code":
+ return `# %% ${cellNumber}\n${cell.source}`
+ case "markdown":
+ case "raw":
+ const source = cell.source.split("\n").map(line => `# ${line}`).join("\n")
+ return `# %% [markdown] ${cellNumber}\n${source}`
+ default:
+ return ""
+ }
+ }
\ No newline at end of file
diff --git a/packages/jupyter-ai/src/widgets/chat-sidebar.tsx b/packages/jupyter-ai/src/widgets/chat-sidebar.tsx
index 732eedd3c..76eac4e9c 100644
--- a/packages/jupyter-ai/src/widgets/chat-sidebar.tsx
+++ b/packages/jupyter-ai/src/widgets/chat-sidebar.tsx
@@ -4,7 +4,7 @@ import { ReactWidget } from '@jupyterlab/apputils';
import type { IThemeManager } from '@jupyterlab/apputils';
import type { User } from '@jupyterlab/services';
import type { Awareness } from 'y-protocols/awareness';
-
+import { INotebookTracker } from '@jupyterlab/notebook';
import { Chat } from '../components/chat';
import { chatIcon } from '../icons';
import { SelectionWatcher } from '../selection-watcher';
@@ -29,7 +29,9 @@ export function buildChatSidebar(
focusInputSignal: ISignal,
messageFooter: IJaiMessageFooter | null,
telemetryHandler: IJaiTelemetryHandler | null,
- userManager: User.IManager
+ userManager: User.IManager,
+ notebookTracker: INotebookTracker | null,
+ goBackToNotebook: () => void
): ReactWidget {
const ChatWidget = ReactWidget.create(
);
ChatWidget.id = 'jupyter-ai::chat';