From 358ab3287edd5b2f9088f9cca390bd1286b729b4 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Sat, 12 Oct 2024 10:24:17 -0400 Subject: [PATCH 01/42] mitosheet: fix convert to datetime with mixed dtype column --- .../column_steps/change_column_dtype_code_chunk.py | 5 ++++- .../public/v1/sheet_functions/types/utils.py | 12 ++++++++++-- .../column_steps/test_change_column_dtype.py | 11 ++++++++++- 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/mitosheet/mitosheet/code_chunks/step_performers/column_steps/change_column_dtype_code_chunk.py b/mitosheet/mitosheet/code_chunks/step_performers/column_steps/change_column_dtype_code_chunk.py index e02ffa1ad..c0ede77ab 100644 --- a/mitosheet/mitosheet/code_chunks/step_performers/column_steps/change_column_dtype_code_chunk.py +++ b/mitosheet/mitosheet/code_chunks/step_performers/column_steps/change_column_dtype_code_chunk.py @@ -89,7 +89,10 @@ def get_conversion_code(state: State, sheet_index: int, column_id: ColumnID, old datetime_params_string = get_param_dict_as_code(to_datetime_params, as_single_line=True) - return f'{df_name}[{transpiled_column_header}] = pd.to_datetime({df_name}[{transpiled_column_header}], {datetime_params_string}, errors=\'coerce\')' + if datetime_params_string: + return f'{df_name}[{transpiled_column_header}] = pd.to_datetime({df_name}[{transpiled_column_header}], {datetime_params_string}, errors=\'coerce\')' + else: + return f'{df_name}[{transpiled_column_header}] = pd.to_datetime({df_name}[{transpiled_column_header}], errors=\'coerce\')' elif is_timedelta_dtype(new_dtype): return f'{df_name}[{transpiled_column_header}] = pd.to_timedelta({df_name}[{transpiled_column_header}], errors=\'coerce\')' elif is_datetime_dtype(old_dtype): diff --git a/mitosheet/mitosheet/public/v1/sheet_functions/types/utils.py b/mitosheet/mitosheet/public/v1/sheet_functions/types/utils.py index 802947841..6409bf07a 100644 --- a/mitosheet/mitosheet/public/v1/sheet_functions/types/utils.py +++ b/mitosheet/mitosheet/public/v1/sheet_functions/types/utils.py @@ -60,6 +60,12 @@ def put_nan_indexes_back(series: pd.Series, original_index: pd.Index) -> pd.Seri def get_to_datetime_params(string_series: pd.Series) -> Dict[str, Any]: + # If the series is already full of datetime objects, then we don't need to + # guess a datetime format. + if is_series_full_of_datetimes(string_series): + return {} + + detected_format = get_datetime_format(string_series) # If we detect a format, we return that. This works for all pandas versions @@ -80,6 +86,9 @@ def get_to_datetime_params(string_series: pd.Series) -> Dict[str, Any]: 'format': 'mixed' } +def is_series_full_of_datetimes(string_series: pd.Series) -> bool: + return all(isinstance(x, datetime.datetime) for x in string_series) + def get_datetime_format(string_series: pd.Series) -> Optional[str]: """ @@ -87,7 +96,7 @@ def get_datetime_format(string_series: pd.Series) -> Optional[str]: """ # Import log function here to avoid circular import from mitosheet.telemetry.telemetry_utils import log - + # Filter to only the strings since that is all we're convertin string_series = string_series[string_series.apply(lambda x: isinstance(x, str))] @@ -102,7 +111,6 @@ def get_datetime_format(string_series: pd.Series) -> Optional[str]: # TODO: Add the most popular formats to here and check them first before # trying all of the formats below for performance. - sample_string_datetime = string_series[string_series.first_valid_index()] FORMATS = [ '%m{s}%d{s}%Y', diff --git a/mitosheet/mitosheet/tests/step_performers/column_steps/test_change_column_dtype.py b/mitosheet/mitosheet/tests/step_performers/column_steps/test_change_column_dtype.py index 00f4dd0b6..290d6c10b 100644 --- a/mitosheet/mitosheet/tests/step_performers/column_steps/test_change_column_dtype.py +++ b/mitosheet/mitosheet/tests/step_performers/column_steps/test_change_column_dtype.py @@ -387,4 +387,13 @@ def test_change_dtype_to_datetime_finds_first_string(): pd.to_datetime('12/22/2023') ] }) - ) \ No newline at end of file + ) + +def test_change_dtype_to_datetime_with_datetime_objects_in_string_series(): + df = pd.DataFrame({'A': [pd.to_datetime('2020-01-01'), pd.to_datetime('2020-01-02'), pd.to_datetime('2020-01-03')]}, dtype=object) + mito = create_mito_wrapper(df) + + mito.change_column_dtype(0, ['A'], 'datetime') + + expected_df = pd.DataFrame({'A': [pd.to_datetime('2020-01-01'), pd.to_datetime('2020-01-02'), pd.to_datetime('2020-01-03')]}) + assert mito.dfs[0].equals(expected_df) \ No newline at end of file From 2f293d68dc1c98e3f34923918476103f81625a51 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Tue, 15 Oct 2024 12:07:12 -0400 Subject: [PATCH 02/42] mito-ai: teach mito-ai about operating inside of jupyter --- mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index ed62e6cf2..a5414b965 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -40,7 +40,7 @@ export class ChatHistoryManager { return this.history.displayOptimizedChatHistory; } - addUserMessage(input: string, activeCellCode?: string, variables?: Variable[]): void { + addUserMessage(input: string, activeCellCode?: string, variables?: Variable[], error: boolean = false): void { const displayMessage: OpenAI.Chat.ChatCompletionMessageParam = { role: 'user', @@ -74,6 +74,14 @@ Do: - Keep as much of the original code as possible - Ask for more context if you need it. +Important: Remember that you are executing code inside a Jupyter notebook. That means you will have persistent state issues where variables from previous cells or previous code executions might still affect current code. When those errors occur, here are a few possible solutions: +1. Restarting the kernel to reset the environment if a function or variable has been unintentionally overwritten. +2. Identify which cell might need to be rerun to properly initialize the function or variable that is causing the issue. + +For example, if an error occurs because the built-in function 'print' is overwritten by an integer, you should return the code cell with the modification to the print function removed and also return an explanation that tell the user to restart their kernel. Do not add new comments to the code cell, just return the code cell with the modification removed. + +When a user hits an error because of a persistent state issue, tell them how to resolve it. + Your task: ${input}`}; this.history.displayOptimizedChatHistory.push({message: displayMessage, error: false}); From 5264c242b45ac223f3d5cbd584791194ddd93a8e Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Tue, 15 Oct 2024 12:11:04 -0400 Subject: [PATCH 03/42] mito-ai: cleanup --- mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index a5414b965..63cd6ecf2 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -40,7 +40,7 @@ export class ChatHistoryManager { return this.history.displayOptimizedChatHistory; } - addUserMessage(input: string, activeCellCode?: string, variables?: Variable[], error: boolean = false): void { + addUserMessage(input: string, activeCellCode?: string, variables?: Variable[]): void { const displayMessage: OpenAI.Chat.ChatCompletionMessageParam = { role: 'user', From 339b8d989adef8bb9a4ec2c3d72fbdb86d6b57ac Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 17 Oct 2024 12:35:51 -0400 Subject: [PATCH 04/42] mito-ai: render cell toolbar button from example --- mito-ai/package.json | 6 +- mito-ai/schema/plugin.json | 20 +++++++ .../CellToolbarButtonsPlugin.tsx | 57 +++++++++++++++++++ .../ErrorMimeRendererPlugin.tsx | 2 +- mito-ai/src/index.ts | 8 ++- 5 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 mito-ai/schema/plugin.json create mode 100644 mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx diff --git a/mito-ai/package.json b/mito-ai/package.json index c29f80bd3..ee66ea1c1 100644 --- a/mito-ai/package.json +++ b/mito-ai/package.json @@ -23,7 +23,8 @@ "files": [ "lib/**/*.{d.ts,eot,gif,html,jpg,js,js.map,json,png,svg,woff2,ttf}", "style/**/*.{css,js,eot,gif,html,jpg,json,png,svg,woff2,ttf}", - "src/**/*.{ts,tsx}" + "src/**/*.{ts,tsx}", + "schema/*.json" ], "main": "lib/index.js", "types": "lib/index.d.ts", @@ -102,7 +103,8 @@ }, "jupyterlab": { "extension": true, - "outputDir": "mito-ai/labextension" + "outputDir": "mito-ai/labextension", + "schemaDir": "schema" }, "eslintIgnore": [ "node_modules", diff --git a/mito-ai/schema/plugin.json b/mito-ai/schema/plugin.json new file mode 100644 index 000000000..6b40fcd7d --- /dev/null +++ b/mito-ai/schema/plugin.json @@ -0,0 +1,20 @@ +{ + "jupyter.lab.shortcuts": [], + "title": "mito-ai", + "description": "mito-ai settings.", + "type": "object", + "properties": {}, + "additionalProperties": false, + "jupyter.lab.toolbars": { + "Cell": [ + { + "name": "run-code-cell", + "command": "toolbar-button:run-code-cell" + }, + { + "name": "render-markdows-cell", + "command": "toolbar-button:render-markdown-cell" + } + ] + } +} \ No newline at end of file diff --git a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx new file mode 100644 index 000000000..ba1bbe0ff --- /dev/null +++ b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx @@ -0,0 +1,57 @@ +import { + JupyterFrontEnd, + JupyterFrontEndPlugin + } from '@jupyterlab/application'; + import { INotebookTracker } from '@jupyterlab/notebook'; + import { markdownIcon, runIcon } from '@jupyterlab/ui-components'; + + const CommandIds = { + /** + * Command to render a markdown cell. + */ + renderMarkdownCell: 'toolbar-button:render-markdown-cell', + /** + * Command to run a code cell. + */ + runCodeCell: 'toolbar-button:run-code-cell' + }; + + /** + * Initialization data for the @jupyterlab-examples/cell-toolbar extension. + */ + const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { + // Important: The Cell Toolbar Buttons are added to the toolbar registry via the schema/plugin.json file. + // The id here must be mito-ai:plugin otherwise the buttons are not successfull added. My understanding is that + // the id must match the name of the package and `plugin` must be used when working with the schema/plugin.json file. + id: 'mito-ai:plugin', + description: 'A JupyterLab extension to add cell toolbar buttons.', + autoStart: true, + requires: [INotebookTracker], + activate: (app: JupyterFrontEnd, tracker: INotebookTracker) => { + const { commands } = app; + + console.log("HERE NOW") + + /* Adds a command enabled only on code cell */ + commands.addCommand(CommandIds.runCodeCell, { + icon: runIcon, + caption: 'Run a code cell', + execute: () => { + commands.execute('notebook:run-cell'); + }, + isVisible: () => tracker.activeCell?.model.type === 'code' + }); + + /* Adds a command enabled only on markdown cell */ + commands.addCommand(CommandIds.renderMarkdownCell, { + icon: markdownIcon, + caption: 'Render a markdown cell', + execute: () => { + commands.execute('notebook:run-cell'); + }, + isVisible: () => tracker.activeCell?.model.type === 'markdown' + }); + } + }; + + export default CellToolbarButtonsPlugin; \ No newline at end of file diff --git a/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx b/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx index a4f3cbe0e..d61e095c0 100644 --- a/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx +++ b/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx @@ -29,7 +29,7 @@ const ErrorMessage: React.FC = ({ onDebugClick }) => { * This plugin augments the standard error output with a prompt to debug the error in the chat interface. */ const ErrorMimeRendererPlugin: JupyterFrontEndPlugin = { - id: 'jupyterlab-debug-prompt', + id: 'mito-ai:debug-error-with-ai', autoStart: true, requires: [IRenderMimeRegistry], activate: (app: JupyterFrontEnd, rendermime: IRenderMimeRegistry) => { diff --git a/mito-ai/src/index.ts b/mito-ai/src/index.ts index e47e02bdb..bfae4e464 100644 --- a/mito-ai/src/index.ts +++ b/mito-ai/src/index.ts @@ -2,7 +2,13 @@ import AiChatPlugin from './Extensions/AiChat/AiChatPlugin'; import VariableManagerPlugin from './Extensions/VariableManager/VariableManagerPlugin'; import ErrorMimeRendererPlugin from './Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin'; +import CellToolbarButtonsPlugin from './Extensions/CellToolbarButtons/CellToolbarButtonsPlugin'; // This is the main entry point to the mito-ai extension. It must export all of the top level // extensions that we want to load. -export default [AiChatPlugin, ErrorMimeRendererPlugin, VariableManagerPlugin]; +export default [ + AiChatPlugin, + ErrorMimeRendererPlugin, + VariableManagerPlugin, + CellToolbarButtonsPlugin +]; From cac1398235cfb78b700fe9a69b87d06037950c9e Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 17 Oct 2024 14:46:13 -0400 Subject: [PATCH 05/42] mito-ai: add explain code button --- mito-ai/schema/plugin.json | 10 +-- .../CellToolbarButtonsPlugin.tsx | 74 ++++++++----------- mito-ai/src/icons/LightbulbIcon.svg | 7 ++ 3 files changed, 39 insertions(+), 52 deletions(-) create mode 100644 mito-ai/src/icons/LightbulbIcon.svg diff --git a/mito-ai/schema/plugin.json b/mito-ai/schema/plugin.json index 6b40fcd7d..0aac6a326 100644 --- a/mito-ai/schema/plugin.json +++ b/mito-ai/schema/plugin.json @@ -8,13 +8,9 @@ "jupyter.lab.toolbars": { "Cell": [ { - "name": "run-code-cell", - "command": "toolbar-button:run-code-cell" - }, - { - "name": "render-markdows-cell", - "command": "toolbar-button:render-markdown-cell" + "name": "explain-code", + "command": "toolbar-button:explain-code" } ] } -} \ No newline at end of file +} diff --git a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx index ba1bbe0ff..cc3ffff19 100644 --- a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx +++ b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx @@ -1,25 +1,18 @@ import { JupyterFrontEnd, JupyterFrontEndPlugin - } from '@jupyterlab/application'; - import { INotebookTracker } from '@jupyterlab/notebook'; - import { markdownIcon, runIcon } from '@jupyterlab/ui-components'; - - const CommandIds = { - /** - * Command to render a markdown cell. - */ - renderMarkdownCell: 'toolbar-button:render-markdown-cell', - /** - * Command to run a code cell. - */ - runCodeCell: 'toolbar-button:run-code-cell' - }; - - /** - * Initialization data for the @jupyterlab-examples/cell-toolbar extension. - */ - const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { +} from '@jupyterlab/application'; +import { INotebookTracker } from '@jupyterlab/notebook'; +import { COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; +import { LabIcon } from '@jupyterlab/ui-components'; +import LightbulbIcon from '../../../src/icons/LightbulbIcon.svg' + +export const lightBulbIcon = new LabIcon({ + name: 'mito_ai', + svgstr: LightbulbIcon +}); + +const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { // Important: The Cell Toolbar Buttons are added to the toolbar registry via the schema/plugin.json file. // The id here must be mito-ai:plugin otherwise the buttons are not successfull added. My understanding is that // the id must match the name of the package and `plugin` must be used when working with the schema/plugin.json file. @@ -27,31 +20,22 @@ import { description: 'A JupyterLab extension to add cell toolbar buttons.', autoStart: true, requires: [INotebookTracker], - activate: (app: JupyterFrontEnd, tracker: INotebookTracker) => { - const { commands } = app; + activate: (app: JupyterFrontEnd, notebookTracker: INotebookTracker) => { + const { commands } = app; - console.log("HERE NOW") - - /* Adds a command enabled only on code cell */ - commands.addCommand(CommandIds.runCodeCell, { - icon: runIcon, - caption: 'Run a code cell', - execute: () => { - commands.execute('notebook:run-cell'); - }, - isVisible: () => tracker.activeCell?.model.type === 'code' - }); - - /* Adds a command enabled only on markdown cell */ - commands.addCommand(CommandIds.renderMarkdownCell, { - icon: markdownIcon, - caption: 'Render a markdown cell', - execute: () => { - commands.execute('notebook:run-cell'); - }, - isVisible: () => tracker.activeCell?.model.type === 'markdown' - }); + // Important: To add a button to the cell toolbar, the command must start with "toolbar-button:" + // and the command must match the command in the schema/plugin.json file. + commands.addCommand('toolbar-button:explain-code', { + icon: lightBulbIcon, + caption: 'Explain code', + execute: () => { + // In order to click on the cell toolbar button, that cell must be the active cell, + // so the Ai Chat taskpane will take care of providing the cell context. + app.commands.execute(COMMAND_MITO_AI_SEND_MESSAGE, { input: `Explain this code` }); + }, + isVisible: () => notebookTracker.activeCell?.model.type === 'code' + }); } - }; - - export default CellToolbarButtonsPlugin; \ No newline at end of file +}; + +export default CellToolbarButtonsPlugin; \ No newline at end of file diff --git a/mito-ai/src/icons/LightbulbIcon.svg b/mito-ai/src/icons/LightbulbIcon.svg new file mode 100644 index 000000000..59f57eaee --- /dev/null +++ b/mito-ai/src/icons/LightbulbIcon.svg @@ -0,0 +1,7 @@ + + + + + + + From 6b135ecadcb0ffc55beb1567f2c69fa0af7e6f27 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 17 Oct 2024 15:01:57 -0400 Subject: [PATCH 06/42] mito-ai: cleanup and document --- .../Extensions/AiChat/ChatHistoryManager.tsx | 13 +++++++++++++ .../CellToolbarButtonsPlugin.tsx | 18 ++++++++++++------ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index 63cd6ecf2..cb11ee4a6 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -17,7 +17,20 @@ export interface IChatHistory { displayOptimizedChatHistory: IDisplayOptimizedChatHistory[] } +/* + The ChatHistoryManager is responsible for managing the AI chat history. + It keeps track of two types of messages: + 1. aiOptimizedChatHistory: Messages sent to the AI that include things like: instructions on how to respond, the code context, etc. + 2. displayOptimizedChatHistory: Messages displayed in the chat interface that only display info the user wants to see, + like their original input. + + TODO: In the future, we should make this its own extension that provides an interface for adding new messages to the chat history, + creating new chats, etc. Doing so would allow us to easily append new messages from other extensions without having to do so + by calling commands with untyped arguments. + + Whenever, the chatHistoryManager is updated, it should automatically send a message to the AI. +*/ export class ChatHistoryManager { private history: IChatHistory; diff --git a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx index cc3ffff19..ef6bb877d 100644 --- a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx +++ b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx @@ -8,7 +8,7 @@ import { LabIcon } from '@jupyterlab/ui-components'; import LightbulbIcon from '../../../src/icons/LightbulbIcon.svg' export const lightBulbIcon = new LabIcon({ - name: 'mito_ai', + name: 'lightbulb-icon', svgstr: LightbulbIcon }); @@ -17,7 +17,7 @@ const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { // The id here must be mito-ai:plugin otherwise the buttons are not successfull added. My understanding is that // the id must match the name of the package and `plugin` must be used when working with the schema/plugin.json file. id: 'mito-ai:plugin', - description: 'A JupyterLab extension to add cell toolbar buttons.', + description: 'Adds an "explain code cell with AI" button to the cell toolbar', autoStart: true, requires: [INotebookTracker], activate: (app: JupyterFrontEnd, notebookTracker: INotebookTracker) => { @@ -27,13 +27,19 @@ const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { // and the command must match the command in the schema/plugin.json file. commands.addCommand('toolbar-button:explain-code', { icon: lightBulbIcon, - caption: 'Explain code', + caption: 'Explain code in AI Chat', execute: () => { - // In order to click on the cell toolbar button, that cell must be the active cell, - // so the Ai Chat taskpane will take care of providing the cell context. + /* + In order to click on the cell toolbar button, that cell must be the active cell, + so the Ai Chat taskpane will take care of providing the cell context. + + TODO: In the future, instead of directly calling COMMAND_MITO_AI_SEND_MESSAGE, we should + update the AI Chat History Manager so that we can generate an AI optimzied message and a + display optimized message. + */ app.commands.execute(COMMAND_MITO_AI_SEND_MESSAGE, { input: `Explain this code` }); }, - isVisible: () => notebookTracker.activeCell?.model.type === 'code' + isVisible: () => notebookTracker.activeCell?.model.type === 'code' && app.commands.hasCommand(COMMAND_MITO_AI_SEND_MESSAGE) }); } }; From 5ccea5dabc54412e3b7aabfd160604ffad47f496 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 18 Oct 2024 11:36:49 -0400 Subject: [PATCH 07/42] mito-ai: fix bug with registering command in UseEffect --- .../src/Extensions/AiChat/ChatTaskpane.tsx | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index f3d596f75..11e67fb0a 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -44,6 +44,7 @@ const getDefaultChatHistoryManager = (): ChatHistoryManager => { return new ChatHistoryManager(chatHistory) } else { + console.log("HERE") const chatHistoryManager = new ChatHistoryManager() chatHistoryManager.addSystemMessage('You are an expert Python programmer.') return chatHistoryManager @@ -129,13 +130,29 @@ const ChatTaskpane: React.FC = ({ _sendMessage(input) } + const getChatHistoryManager = () => { + return chatHistoryManagerRef.current + } + const _sendMessage = async (input: string) => { const variables = variableManager.variables const activeCellCode = getActiveCellCode(notebookTracker) - // Create a new chat history manager so we can trigger a re-render of the chat - const updatedManager = new ChatHistoryManager(chatHistoryManager.getHistory()); + /* + 1. Access ChatHistoryManager via a function: + We use getChatHistoryManager() instead of directly accessing the state variable because + the COMMAND_MITO_AI_SEND_MESSAGE is registered in a useEffect on initial render, which + would otherwise always use the initial state values. By using a function, we ensure we always + get the most recent chat history, even when the command is executed later. + + 2. Create a new ChatHistoryManager instance: + We create a copy of the current chat history and use it to initialize a new ChatHistoryManager to + trigger a re-render in React, as simply appending to the existing ChatHistoryManager + (an immutable object) wouldn't be detected as a state change. + */ + const currentChatHistory = getChatHistoryManager().getHistory() + const updatedManager = new ChatHistoryManager(currentChatHistory); updatedManager.addUserMessage(input, activeCellCode, variables) setInput(''); From 6d4cd9a61bf05064f58dec63d0365e8d6f1f6284 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 18 Oct 2024 11:39:51 -0400 Subject: [PATCH 08/42] Update mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx --- mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 11e67fb0a..e4988d407 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -44,7 +44,6 @@ const getDefaultChatHistoryManager = (): ChatHistoryManager => { return new ChatHistoryManager(chatHistory) } else { - console.log("HERE") const chatHistoryManager = new ChatHistoryManager() chatHistoryManager.addSystemMessage('You are an expert Python programmer.') return chatHistoryManager From 7fd0cc44bb65018d2027257fcc1e94aa9f8753c8 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 21 Oct 2024 09:47:06 -0400 Subject: [PATCH 09/42] mito-ai: style code diffs exploration --- mito-ai/mito-ai/OpenAICompletionHandler.py | 3 ++ .../AiChat/ChatMessage/PythonCode.tsx | 40 ++++++++++++--- .../src/Extensions/AiChat/ChatTaskpane.tsx | 10 +++- mito-ai/src/utils/codeDiff.tsx | 50 +++++++++++++++++++ mito-ai/style/CodeMessagePart.css | 13 ++++- 5 files changed, 106 insertions(+), 10 deletions(-) create mode 100644 mito-ai/src/utils/codeDiff.tsx diff --git a/mito-ai/mito-ai/OpenAICompletionHandler.py b/mito-ai/mito-ai/OpenAICompletionHandler.py index 42d922e1e..3141b8b77 100644 --- a/mito-ai/mito-ai/OpenAICompletionHandler.py +++ b/mito-ai/mito-ai/OpenAICompletionHandler.py @@ -42,6 +42,9 @@ def post(self): # TODO: In the future, instead of returning the raw response, # return a cleaned up version of the response so we can support # multiple models + + print(response_dict) + self.finish(json.dumps(response_dict)) except Exception as e: self.set_status(500) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx index db57bc0c9..e96fda249 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx @@ -1,8 +1,8 @@ import React, { useEffect, useState } from 'react'; import { IRenderMimeRegistry, MimeModel } from '@jupyterlab/rendermime'; -import { addMarkdownCodeFormatting } from '../../../utils/strings'; import '../../../../style/PythonCode.css'; +// import { addMarkdownCodeFormatting } from '../../../utils/strings'; interface IPythonCodeProps { code: string; @@ -10,21 +10,45 @@ interface IPythonCodeProps { } const PythonCode: React.FC = ({ code, rendermime }) => { - const [node, setNode] = useState(null) + // const newCode = '' + useEffect(() => { - const model = new MimeModel({ - data: { ['text/markdown']: addMarkdownCodeFormatting(code) }, + const deletedLines = [0, 1, 2]; + + const newCode = 'print("hello world")\n# This is a comment\nprint("foobar")' + + const wrappedCode = ` +\`\`\`python +${newCode} +\`\`\` + `; + + const model = new MimeModel({ + data: { ['text/markdown']: wrappedCode}, }); const renderer = rendermime.createRenderer('text/markdown'); - renderer.renderModel(model) + + // After rendering, add the background to specific lines + const codeElement = renderer.node.querySelector('pre'); + if (codeElement) { + const codeLines = codeElement.innerHTML.split('\n'); + const highlightedCode = codeLines.map((line, index) => { + if (deletedLines.includes(index + 1)) { + return `${line}`; + } + return line; + }).join('\n'); + codeElement.innerHTML = highlightedCode; + } + + const node = renderer.node setNode(node) - }, []) - + }, [code, rendermime]) // Add dependencies to useEffect if (node) { return
el && el.appendChild(node)} /> @@ -33,4 +57,4 @@ const PythonCode: React.FC = ({ code, rendermime }) => { } }; -export default PythonCode; \ No newline at end of file +export default PythonCode; diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index f3d596f75..a2615785d 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -11,12 +11,13 @@ import { requestAPI } from '../../utils/handler'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import LoadingDots from '../../components/LoadingDots'; import { JupyterFrontEnd } from '@jupyterlab/application'; -import { getCodeBlockFromMessage } from '../../utils/strings'; +import { getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; import { ReadonlyPartialJSONObject } from '@lumino/coreutils'; import ResetIcon from '../../icons/ResetIcon'; import IconButton from '../../components/IconButton'; import { OperatingSystem } from '../../utils/user'; +import { getCodeWithDiffsMarked } from '../../utils/codeDiff'; // IMPORTANT: In order to improve the development experience, we allow you dispaly a @@ -152,6 +153,13 @@ const ChatTaskpane: React.FC = ({ if (apiResponse.type === 'success') { const response = apiResponse.response; const aiMessage = response.choices[0].message; + + + const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); + const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); + const aiGeneratedCodeWithDiffs = getCodeWithDiffsMarked(activeCellCode, aiGeneratedCodeCleaned) + aiMessage.content = aiGeneratedCodeWithDiffs + updatedManager.addAIMessageFromResponse(aiMessage); setChatHistoryManager(updatedManager); } else { diff --git a/mito-ai/src/utils/codeDiff.tsx b/mito-ai/src/utils/codeDiff.tsx new file mode 100644 index 000000000..e37d71841 --- /dev/null +++ b/mito-ai/src/utils/codeDiff.tsx @@ -0,0 +1,50 @@ +import { DiffComputer, IDiffComputerOpts, ILineChange } from "vscode-diff"; + +export const getCodeDiffLineRanges = (originalLines: string | undefined | null, modifiedLines: string | undefined | null): ILineChange[] => { + if (originalLines === undefined || originalLines === null) { + originalLines = '' + } + + if (modifiedLines === undefined || modifiedLines === null) { + modifiedLines = '' + } + + const originalLinesArray = originalLines.split('\n') + const modifiedLinesArray = modifiedLines.split('\n') + + console.log("originalLinesArray", originalLinesArray) + console.log("modifiedLinesArray", modifiedLinesArray) + + let options: IDiffComputerOpts = { + shouldPostProcessCharChanges: true, + shouldIgnoreTrimWhitespace: true, + shouldMakePrettyDiff: true, + shouldComputeCharChanges: true, + maxComputationTime: 0 // time in milliseconds, 0 => no computation limit. + } + + let diffComputer = new DiffComputer(originalLinesArray, modifiedLinesArray, options); + let lineChanges: ILineChange[] = diffComputer.computeDiff().changes; + + return lineChanges || [] +} + +export const getCodeWithDiffsMarked = (originalLines: string | undefined | null, modifiedLines: string | undefined | null): string => { + + let originalLinesTest: string = "hello\noriginal\nworld"; + let modifiedLinesTest: string = "hello\nnew\nworld\nfoobar"; + + const lineChanges = getCodeDiffLineRanges(originalLinesTest, modifiedLinesTest); + + const diffedLines = originalLinesTest.split('\n') + + let numNewLinesAdded = 0 + for (const lineChange of lineChanges) { + diffedLines[lineChange.originalStartLineNumber] = '' + diffedLines[lineChange.originalStartLineNumber] + '' + numNewLinesAdded = numNewLinesAdded + 1 + } + + console.log("diffedLines", diffedLines) + + return "```python\n" + diffedLines.join('\n') + "\n```" +} \ No newline at end of file diff --git a/mito-ai/style/CodeMessagePart.css b/mito-ai/style/CodeMessagePart.css index 131f76fc8..603b5f269 100644 --- a/mito-ai/style/CodeMessagePart.css +++ b/mito-ai/style/CodeMessagePart.css @@ -39,4 +39,15 @@ .code-message-part-toolbar button:hover { background-color: var(--chat-background-color); color: var(--chat-assistant-message-font-color); -} \ No newline at end of file +} + + + .deleted-line .jp-RenderedMarkdown { + background-color: rgba(255, 0, 0, 0.2) !important; + } + + .deleted-line .jp-RenderedMarkdown pre { + background-color: transparent !important; + margin: 0; + padding: 0; + } \ No newline at end of file From a278e68e167dfd0e1bf7206da0e53b9c296aa524 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 21 Oct 2024 16:26:09 -0400 Subject: [PATCH 10/42] mito-ai: apply colors to code cell --- mito-ai/src/Extensions/AiChat/AiChatPlugin.ts | 112 +++++++++++++++++- 1 file changed, 111 insertions(+), 1 deletion(-) diff --git a/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts b/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts index 17bc9c588..1896ce7ca 100644 --- a/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts +++ b/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts @@ -3,14 +3,101 @@ import { JupyterFrontEnd, JupyterFrontEndPlugin, } from '@jupyterlab/application'; +import { Extension, Facet, RangeSetBuilder } from '@codemirror/state'; +import { + Decoration, + DecorationSet, + EditorView, + ViewPlugin, + ViewUpdate +} from '@codemirror/view'; import { ICommandPalette, WidgetTracker } from '@jupyterlab/apputils'; import { INotebookTracker } from '@jupyterlab/notebook'; import { buildChatWidget } from './ChatWidget'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; +import { + EditorExtensionRegistry, + IEditorExtensionRegistry +} from '@jupyterlab/codemirror'; + import { COMMAND_MITO_AI_OPEN_CHAT } from '../../commands'; +// Defines new styles for this extension +const baseTheme = EditorView.baseTheme({ + // We need to set some transparency because the stripe are above + // the selection layer + '&light .cm-zebraStripe': { backgroundColor: '#d4fafaaa' }, + '&dark .cm-zebraStripe': { backgroundColor: '#1a2727aa' } +}); + +// Resolve step to use in the editor +const stepSize = Facet.define({ + combine: values => (values.length ? Math.min(...values) : 2) +}); + +// Add decoration to editor lines +const stripe = Decoration.line({ + attributes: { class: 'cm-zebraStripe' } +}); + +// Create the range of lines requiring decorations +function stripeDeco(view: EditorView) { + const step = view.state.facet(stepSize) as number; + const builder = new RangeSetBuilder(); + for (const { from, to } of view.visibleRanges) { + for (let pos = from; pos <= to; ) { + const line = view.state.doc.lineAt(pos); + if (line.number % step === 0) { + builder.add(line.from, line.from, stripe); + } + pos = line.to + 1; + } + } + return builder.finish(); +} + +// Update the decoration status of the editor view +const showStripes = ViewPlugin.fromClass( + class { + decorations: DecorationSet; + + constructor(view: EditorView) { + this.decorations = stripeDeco(view); + } + + update(update: ViewUpdate) { + // Update the stripes if the document changed, + // the viewport changed or the stripes step changed. + const oldStep = update.startState.facet(stepSize); + if ( + update.docChanged || + update.viewportChanged || + oldStep !== update.view.state.facet(stepSize) + ) { + this.decorations = stripeDeco(update.view); + } + } + }, + { + decorations: v => v.decorations + } +); + +// Full extension composed of elemental extensions +export function zebraStripes(options: { step?: number, on?: boolean} = {}): Extension { + if (options.on) { + return [ + baseTheme, + typeof options.step !== 'number' ? [] : stepSize.of(options.step), + showStripes + ]; + } + return []; +} + + /** * Initialization data for the mito-ai extension. @@ -19,7 +106,7 @@ const AiChatPlugin: JupyterFrontEndPlugin = { id: 'mito_ai:plugin', description: 'AI chat for JupyterLab', autoStart: true, - requires: [INotebookTracker, ICommandPalette, IRenderMimeRegistry, IVariableManager], + requires: [INotebookTracker, ICommandPalette, IRenderMimeRegistry, IVariableManager, IEditorExtensionRegistry], optional: [ILayoutRestorer], activate: ( app: JupyterFrontEnd, @@ -27,6 +114,7 @@ const AiChatPlugin: JupyterFrontEndPlugin = { palette: ICommandPalette, rendermime: IRenderMimeRegistry, variableManager: IVariableManager, + editorExtensionRegistry: IEditorExtensionRegistry, restorer: ILayoutRestorer | null ) => { @@ -96,6 +184,28 @@ const AiChatPlugin: JupyterFrontEndPlugin = { if (restorer) { restorer.add(widget, 'mito_ai'); } + + + + editorExtensionRegistry.addExtension( + Object.freeze({ + name: '@jupyterlab-examples/codemirror:zebra-stripes', + // Default CodeMirror extension parameters + default: 2, + factory: () => + // The factory will be called for every new CodeMirror editor + EditorExtensionRegistry.createConfigurableExtension((step: number) => + zebraStripes({ step, on: true }) + ), + // JSON schema defining the CodeMirror extension parameters + schema: { + type: 'number', + title: 'Show stripes', + description: + 'Display zebra stripes every "step" in CodeMirror editors.' + } + }) + ); } }; From 811c4b7bb5b763c906c86774b2f341849e256f9f Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 21 Oct 2024 17:32:59 -0400 Subject: [PATCH 11/42] mito-ai: conditionally render background color --- mito-ai/src/Extensions/AiChat/AiChatPlugin.ts | 113 +---------------- .../AiChat/ChatMessage/ChatMessage.tsx | 5 +- .../AiChat/ChatMessage/CodeBlock.tsx | 11 +- .../src/Extensions/AiChat/ChatTaskpane.tsx | 120 +++++++++++++++--- mito-ai/src/Extensions/AiChat/ChatWidget.tsx | 3 + .../src/Extensions/AiChat/CodeDiffDisplay.tsx | 80 ++++++++++++ 6 files changed, 201 insertions(+), 131 deletions(-) create mode 100644 mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx diff --git a/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts b/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts index 1896ce7ca..c48fad276 100644 --- a/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts +++ b/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts @@ -3,101 +3,14 @@ import { JupyterFrontEnd, JupyterFrontEndPlugin, } from '@jupyterlab/application'; -import { Extension, Facet, RangeSetBuilder } from '@codemirror/state'; -import { - Decoration, - DecorationSet, - EditorView, - ViewPlugin, - ViewUpdate -} from '@codemirror/view'; - import { ICommandPalette, WidgetTracker } from '@jupyterlab/apputils'; import { INotebookTracker } from '@jupyterlab/notebook'; import { buildChatWidget } from './ChatWidget'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; -import { - EditorExtensionRegistry, - IEditorExtensionRegistry -} from '@jupyterlab/codemirror'; - +import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; import { COMMAND_MITO_AI_OPEN_CHAT } from '../../commands'; -// Defines new styles for this extension -const baseTheme = EditorView.baseTheme({ - // We need to set some transparency because the stripe are above - // the selection layer - '&light .cm-zebraStripe': { backgroundColor: '#d4fafaaa' }, - '&dark .cm-zebraStripe': { backgroundColor: '#1a2727aa' } -}); - -// Resolve step to use in the editor -const stepSize = Facet.define({ - combine: values => (values.length ? Math.min(...values) : 2) -}); - -// Add decoration to editor lines -const stripe = Decoration.line({ - attributes: { class: 'cm-zebraStripe' } -}); - -// Create the range of lines requiring decorations -function stripeDeco(view: EditorView) { - const step = view.state.facet(stepSize) as number; - const builder = new RangeSetBuilder(); - for (const { from, to } of view.visibleRanges) { - for (let pos = from; pos <= to; ) { - const line = view.state.doc.lineAt(pos); - if (line.number % step === 0) { - builder.add(line.from, line.from, stripe); - } - pos = line.to + 1; - } - } - return builder.finish(); -} - -// Update the decoration status of the editor view -const showStripes = ViewPlugin.fromClass( - class { - decorations: DecorationSet; - - constructor(view: EditorView) { - this.decorations = stripeDeco(view); - } - - update(update: ViewUpdate) { - // Update the stripes if the document changed, - // the viewport changed or the stripes step changed. - const oldStep = update.startState.facet(stepSize); - if ( - update.docChanged || - update.viewportChanged || - oldStep !== update.view.state.facet(stepSize) - ) { - this.decorations = stripeDeco(update.view); - } - } - }, - { - decorations: v => v.decorations - } -); - -// Full extension composed of elemental extensions -export function zebraStripes(options: { step?: number, on?: boolean} = {}): Extension { - if (options.on) { - return [ - baseTheme, - typeof options.step !== 'number' ? [] : stepSize.of(options.step), - showStripes - ]; - } - return []; -} - - /** * Initialization data for the mito-ai extension. @@ -122,7 +35,7 @@ const AiChatPlugin: JupyterFrontEndPlugin = { // then call it to make a new widget const newWidget = () => { // Create a blank content widget inside of a MainAreaWidget - const chatWidget = buildChatWidget(app, notebookTracker, rendermime, variableManager) + const chatWidget = buildChatWidget(app, notebookTracker, rendermime, variableManager, editorExtensionRegistry) return chatWidget } @@ -184,28 +97,6 @@ const AiChatPlugin: JupyterFrontEndPlugin = { if (restorer) { restorer.add(widget, 'mito_ai'); } - - - - editorExtensionRegistry.addExtension( - Object.freeze({ - name: '@jupyterlab-examples/codemirror:zebra-stripes', - // Default CodeMirror extension parameters - default: 2, - factory: () => - // The factory will be called for every new CodeMirror editor - EditorExtensionRegistry.createConfigurableExtension((step: number) => - zebraStripes({ step, on: true }) - ), - // JSON schema defining the CodeMirror extension parameters - schema: { - type: 'number', - title: 'Show stripes', - description: - 'Display zebra stripes every "step" in CodeMirror editors.' - } - }) - ); } }; diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index b0cbfa881..325142f05 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -19,6 +19,7 @@ interface IChatMessageProps { app: JupyterFrontEnd isLastAiMessage: boolean operatingSystem: OperatingSystem + setDisplayCodeDiff: React.Dispatch>; } const ChatMessage: React.FC = ({ @@ -29,7 +30,8 @@ const ChatMessage: React.FC = ({ rendermime, app, isLastAiMessage, - operatingSystem + operatingSystem, + setDisplayCodeDiff }): JSX.Element | null => { if (message.role !== 'user' && message.role !== 'assistant') { // Filter out other types of messages, like system messages @@ -60,6 +62,7 @@ const ChatMessage: React.FC = ({ app={app} isLastAiMessage={isLastAiMessage} operatingSystem={operatingSystem} + setDisplayCodeDiff={setDisplayCodeDiff} /> ) } diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index 6d819a5d1..f8620cf1d 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -16,7 +16,8 @@ interface ICodeBlockProps { notebookTracker: INotebookTracker, app: JupyterFrontEnd, isLastAiMessage: boolean, - operatingSystem: OperatingSystem + operatingSystem: OperatingSystem, + setDisplayCodeDiff: React.Dispatch>; } const CodeBlock: React.FC = ({ @@ -26,7 +27,8 @@ const CodeBlock: React.FC = ({ notebookTracker, app, isLastAiMessage, - operatingSystem + operatingSystem, + setDisplayCodeDiff }): JSX.Element => { const notebookName = getNotebookName(notebookTracker) @@ -47,6 +49,8 @@ const CodeBlock: React.FC = ({ ) } + + if (role === 'assistant') { return (
@@ -55,6 +59,9 @@ const CodeBlock: React.FC = ({ {notebookName}
+
{ if (USE_DEV_AI_CONVERSATION) { const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ - {role: 'system', content: 'You are an expert Python programmer.'}, - {role: 'user', content: "```python x = 5\ny=10\nx+y``` update x to 10"}, - {role: 'assistant', content: "```python x = 10\ny=10\nx+y```"}, - {role: 'user', content: "```python x = 5\ny=10\nx+y``` Explain what this code does to me"}, - {role: 'assistant', content: "This code defines two variables, x and y. Variables are named buckets that store a value. ```python x = 5\ny=10``` It then adds them together ```python x+y``` Let me know if you want me to further explain any of those concepts"} + { role: 'system', content: 'You are an expert Python programmer.' }, + { role: 'user', content: "```python x = 5\ny=10\nx+y``` update x to 10" }, + { role: 'assistant', content: "```python x = 10\ny=10\nx+y```" }, + { role: 'user', content: "```python x = 5\ny=10\nx+y``` Explain what this code does to me" }, + { role: 'assistant', content: "This code defines two variables, x and y. Variables are named buckets that store a value. ```python x = 5\ny=10``` It then adds them together ```python x+y``` Let me know if you want me to further explain any of those concepts" } ] const chatHistory: IChatHistory = { aiOptimizedChatHistory: [...messages], - displayOptimizedChatHistory: [...messages].map(message => ({message: message, error: false})) + displayOptimizedChatHistory: [...messages].map(message => ({ message: message, error: false })) } return new ChatHistoryManager(chatHistory) @@ -55,14 +62,16 @@ interface IChatTaskpaneProps { notebookTracker: INotebookTracker rendermime: IRenderMimeRegistry variableManager: IVariableManager + editorExtensionRegistry: IEditorExtensionRegistry app: JupyterFrontEnd operatingSystem: OperatingSystem } const ChatTaskpane: React.FC = ({ - notebookTracker, - rendermime, - variableManager, + notebookTracker, + rendermime, + variableManager, + editorExtensionRegistry, app, operatingSystem }) => { @@ -70,6 +79,7 @@ const ChatTaskpane: React.FC = ({ const [chatHistoryManager, setChatHistoryManager] = useState(() => getDefaultChatHistoryManager()); const [input, setInput] = useState(''); const [loadingAIResponse, setLoadingAIResponse] = useState(false) + const [displayCodeDiff, setDisplayCodeDiff] = useState(false) const chatHistoryManagerRef = useRef(chatHistoryManager); useEffect(() => { @@ -179,7 +189,7 @@ const ChatTaskpane: React.FC = ({ const applyLatestCode = () => { const latestChatHistoryManager = chatHistoryManagerRef.current; const lastAIMessage = latestChatHistoryManager.getLastAIMessage() - + if (!lastAIMessage) { return } @@ -188,7 +198,7 @@ const ChatTaskpane: React.FC = ({ writeCodeToActiveCell(notebookTracker, code, true) } - useEffect(() => { + useEffect(() => { /* Add a new command to the JupyterLab command registry that applies the latest AI generated code to the active code cell. Do this inside of the useEffect so that we only register the command @@ -218,15 +228,90 @@ const ChatTaskpane: React.FC = ({ } } }) + + // editorExtensionRegistry.addExtension( + // Object.freeze({ + // name: '@jupyterlab-examples/codemirror:zebra-stripes', + // // Default CodeMirror extension parameters + // default: 2, + // factory: () => { + + // return EditorExtensionRegistry.createConfigurableExtension((step: number) => + // zebraStripes({ step }) + // ) + // }, + // // JSON schema defining the CodeMirror extension parameters + // schema: { + // type: 'number', + // title: 'Show stripes', + // description: 'Display zebra stripes every "step" in CodeMirror editors.', + + // } + // }) + // ); }, []) + // Create a WeakMap to store compartments per code cell + const zebraStripesCompartments = React.useRef(new WeakMap()); + + // Function to update the extensions of code cells + const updateCodeCellsExtensions = useCallback(() => { + const notebook = notebookTracker.currentWidget?.content; + if (!notebook) { + return; + } + + console.log(1) + notebook.widgets.forEach(cell => { + if (cell.model.type === 'code') { + console.log(2) + const codeCell = cell as CodeCell; + const cmEditor = codeCell.editor as CodeMirrorEditor; + const editorView = cmEditor?.editor; + + if (editorView) { + let compartment = zebraStripesCompartments.current.get(codeCell); + + if (!compartment) { + // Create a new compartment and store it + compartment = new Compartment(); + zebraStripesCompartments.current.set(codeCell, compartment); + + // Apply the initial configuration + editorView.dispatch({ + effects: StateEffect.appendConfig.of( + compartment.of(displayCodeDiff ? zebraStripes() : []) + ), + }); + } else { + // Reconfigure the compartment + editorView.dispatch({ + effects: compartment.reconfigure( + displayCodeDiff ? zebraStripes() : [] + ), + }); + } + } else { + console.log('editor view not found') + } + } + }); + }, [displayCodeDiff, notebookTracker]); + + // Update code cells when displayCodeDiff changes + useEffect(() => { + console.log("Calling updateCodeCellsExtensions") + updateCodeCellsExtensions(); + }, [displayCodeDiff, updateCodeCellsExtensions]); + + const lastAIMessagesIndex = chatHistoryManager.getLastAIMessageIndex() return (

- } title="Clear the chat history" onClick={() => { @@ -237,7 +322,7 @@ const ChatTaskpane: React.FC = ({
{displayOptimizedChatHistory.map((displayOptimizedChat, index) => { return ( - = ({ app={app} isLastAiMessage={index === lastAIMessagesIndex} operatingSystem={operatingSystem} + setDisplayCodeDiff={setDisplayCodeDiff} /> ) }).filter(message => message !== null)}
- {loadingAIResponse && + {loadingAIResponse &&
Loading AI Response
@@ -260,14 +346,14 @@ const ChatTaskpane: React.FC = ({ className={classNames("message", "message-user", 'chat-input')} placeholder={displayOptimizedChatHistory.length < 2 ? "Ask your personal Python expert anything!" : "Follow up on the conversation"} value={input} - onChange={(e) => {setInput(e.target.value)}} + onChange={(e) => { setInput(e.target.value) }} onKeyDown={(e) => { // Enter key sends the message, but we still want to allow // shift + enter to add a new line. if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); sendMessageFromChat(); - } + } }} />
diff --git a/mito-ai/src/Extensions/AiChat/ChatWidget.tsx b/mito-ai/src/Extensions/AiChat/ChatWidget.tsx index 129903e70..1674508ea 100644 --- a/mito-ai/src/Extensions/AiChat/ChatWidget.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatWidget.tsx @@ -8,6 +8,7 @@ import chatIconSvg from '../../../src/icons/ChatIcon.svg' import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { getOperatingSystem } from '../../utils/user'; +import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; export const chatIcon = new LabIcon({ name: 'mito_ai', @@ -19,6 +20,7 @@ export function buildChatWidget( notebookTracker: INotebookTracker, rendermime: IRenderMimeRegistry, variableManager: IVariableManager, + editorExtensionRegistry: IEditorExtensionRegistry ) { // Get the operating system here so we don't have to do it each time the chat changes. @@ -32,6 +34,7 @@ export function buildChatWidget( rendermime={rendermime} variableManager={variableManager} operatingSystem={operatingSystem} + editorExtensionRegistry={editorExtensionRegistry} /> ) chatWidget.id = 'mito_ai'; diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx new file mode 100644 index 000000000..c432e05aa --- /dev/null +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -0,0 +1,80 @@ +import { Extension, Facet, RangeSetBuilder } from '@codemirror/state'; +import { + Decoration, + DecorationSet, + EditorView, + ViewPlugin, + ViewUpdate +} from '@codemirror/view'; + + +// Defines new styles for this extension +const baseTheme = EditorView.baseTheme({ + // We need to set some transparency because the stripe are above + // the selection layer + '&light .cm-zebraStripe': { backgroundColor: '#d4fafaaa' }, + '&dark .cm-zebraStripe': { backgroundColor: '#1a2727aa' } + }); + + // Resolve step to use in the editor + const stepSize = Facet.define({ + combine: values => (values.length ? Math.min(...values) : 2) + }); + + // Add decoration to editor lines + const stripe = Decoration.line({ + attributes: { class: 'cm-zebraStripe' } + }); + + // Create the range of lines requiring decorations + function stripeDeco(view: EditorView) { + const step = view.state.facet(stepSize) as number; + const builder = new RangeSetBuilder(); + for (const { from, to } of view.visibleRanges) { + for (let pos = from; pos <= to; ) { + const line = view.state.doc.lineAt(pos); + if (line.number % step === 0) { + builder.add(line.from, line.from, stripe); + } + pos = line.to + 1; + } + } + return builder.finish(); + } + + // Update the decoration status of the editor view + const showStripes = ViewPlugin.fromClass( + class { + decorations: DecorationSet; + + constructor(view: EditorView) { + this.decorations = stripeDeco(view); + } + + update(update: ViewUpdate) { + // Update the stripes if the document changed, + // the viewport changed or the stripes step changed. + const oldStep = update.startState.facet(stepSize); + if ( + update.docChanged || + update.viewportChanged || + oldStep !== update.view.state.facet(stepSize) + ) { + this.decorations = stripeDeco(update.view); + } + } + }, + { + decorations: v => v.decorations + } + ); + + // Full extension composed of elemental extensions + export function zebraStripes(options: { step?: number} = {}): Extension { + return [ + baseTheme, + typeof options.step !== 'number' ? [] : stepSize.of(options.step), + showStripes + ]; + } + \ No newline at end of file From 5f6268fb8d74c8a717d7dbe8cbe8633326c53c69 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 21 Oct 2024 21:32:17 -0400 Subject: [PATCH 12/42] mito-ai: create unified diffs algorithm --- .../src/Extensions/AiChat/ChatTaskpane.tsx | 56 +++--- mito-ai/src/utils/codeDiff.tsx | 173 +++++++++++++++++- 2 files changed, 199 insertions(+), 30 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 67c956d3c..01eac61b8 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -17,7 +17,7 @@ import { ReadonlyPartialJSONObject } from '@lumino/coreutils'; import ResetIcon from '../../icons/ResetIcon'; import IconButton from '../../components/IconButton'; import { OperatingSystem } from '../../utils/user'; -import { getCodeWithDiffsMarked } from '../../utils/codeDiff'; +import { createUnifiedDiff, getCodeDiffLineRanges } from '../../utils/codeDiff'; import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; import { CodeMirrorEditor } from '@jupyterlab/codemirror'; import { CodeCell } from '@jupyterlab/cells'; @@ -67,6 +67,11 @@ interface IChatTaskpaneProps { operatingSystem: OperatingSystem } +// interface IDisplayCodeDiffInfo { +// // In order to display the code diff, I need to merge the old code with the new code +// // then keep track of which lines are deleted and which lines are added/modified +// } + const ChatTaskpane: React.FC = ({ notebookTracker, rendermime, @@ -161,14 +166,38 @@ const ChatTaskpane: React.FC = ({ }); if (apiResponse.type === 'success') { + + let originalLinesTest: string = "hello\noriginal\nworld"; + let modifiedLinesTest: string = "hello\nmodified\nworld\nfoobar"; + const response = apiResponse.response; const aiMessage = response.choices[0].message; const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); - const aiGeneratedCodeWithDiffs = getCodeWithDiffsMarked(activeCellCode, aiGeneratedCodeCleaned) - aiMessage.content = aiGeneratedCodeWithDiffs + //const aiGeneratedCodeWithDiffs = getCodeWithDiffsMarked(activeCellCode, aiGeneratedCodeCleaned) + + const lineChanges = getCodeDiffLineRanges(activeCellCode, aiGeneratedCodeCleaned) + + const unifiedDiffs = createUnifiedDiff(originalLinesTest, modifiedLinesTest, lineChanges) + + // Display the unified diff + unifiedDiffs.forEach(line => { + let prefix = ''; + if (line.type === 'added') { + prefix = '+ '; + } else if (line.type === 'removed') { + prefix = '- '; + } else { + prefix = ' '; + } + console.log(`${prefix}${line.content}`); + }); + + + console.log("unifiedDiffs", unifiedDiffs) + aiMessage.content = aiGeneratedCode || '' updatedManager.addAIMessageFromResponse(aiMessage); setChatHistoryManager(updatedManager); @@ -228,27 +257,6 @@ const ChatTaskpane: React.FC = ({ } } }) - - // editorExtensionRegistry.addExtension( - // Object.freeze({ - // name: '@jupyterlab-examples/codemirror:zebra-stripes', - // // Default CodeMirror extension parameters - // default: 2, - // factory: () => { - - // return EditorExtensionRegistry.createConfigurableExtension((step: number) => - // zebraStripes({ step }) - // ) - // }, - // // JSON schema defining the CodeMirror extension parameters - // schema: { - // type: 'number', - // title: 'Show stripes', - // description: 'Display zebra stripes every "step" in CodeMirror editors.', - - // } - // }) - // ); }, []) // Create a WeakMap to store compartments per code cell diff --git a/mito-ai/src/utils/codeDiff.tsx b/mito-ai/src/utils/codeDiff.tsx index e37d71841..0aba18960 100644 --- a/mito-ai/src/utils/codeDiff.tsx +++ b/mito-ai/src/utils/codeDiff.tsx @@ -31,12 +31,22 @@ export const getCodeDiffLineRanges = (originalLines: string | undefined | null, export const getCodeWithDiffsMarked = (originalLines: string | undefined | null, modifiedLines: string | undefined | null): string => { - let originalLinesTest: string = "hello\noriginal\nworld"; - let modifiedLinesTest: string = "hello\nnew\nworld\nfoobar"; + /* + originalCodeLines: string + newCodeLines: string + allCodeLines: string + - Ordered by line number with: + - Original code lines + - New Code Lines + deletedLineIndexes: List[int] + modifiedLineIndexes: List[int] + */ + + - const lineChanges = getCodeDiffLineRanges(originalLinesTest, modifiedLinesTest); + const lineChanges = getCodeDiffLineRanges(originalLines, modifiedLines); - const diffedLines = originalLinesTest.split('\n') + const diffedLines = originalLines?.split('\n') || [] let numNewLinesAdded = 0 for (const lineChange of lineChanges) { @@ -44,7 +54,158 @@ export const getCodeWithDiffsMarked = (originalLines: string | undefined | null, numNewLinesAdded = numNewLinesAdded + 1 } - console.log("diffedLines", diffedLines) - return "```python\n" + diffedLines.join('\n') + "\n```" +} + +interface UnifiedDiffLine { + content: string; // The content of the line + type: 'unchanged' | 'added' | 'removed'; // The type of change + originalLineNumber: number | null; // Line number in the original code + modifiedLineNumber: number | null; // Line number in the modified code +} + +export const createUnifiedDiff = ( + originalCode: string | undefined | null, + modifiedCode: string | undefined | null, + lineChanges: ILineChange[] +): UnifiedDiffLine[] => { + if (originalCode === undefined || originalCode === null) { + originalCode = '' + } + + if (modifiedCode === undefined || modifiedCode === null) { + modifiedCode = '' + } + + const originalLines = originalCode.split('\n') + const modifiedLines = modifiedCode.split('\n') + + const result: UnifiedDiffLine[] = []; + let originalLineNum = 1; + let modifiedLineNum = 1; + let changeIndex = 0; + + while ( + originalLineNum <= originalLines.length || + modifiedLineNum <= modifiedLines.length + ) { + if (changeIndex < lineChanges.length) { + const change = lineChanges[changeIndex]; + + // Process unchanged lines before the next change + while ( + (originalLineNum < change.originalStartLineNumber || + modifiedLineNum < change.modifiedStartLineNumber) && + originalLineNum <= originalLines.length && + modifiedLineNum <= modifiedLines.length + ) { + result.push({ + content: originalLines[originalLineNum - 1], + type: 'unchanged', + originalLineNumber: originalLineNum, + modifiedLineNumber: modifiedLineNum, + }); + originalLineNum++; + modifiedLineNum++; + } + + // Process the change + if ( + change.originalEndLineNumber > 0 && + change.modifiedEndLineNumber > 0 + ) { + // Modification + for ( + ; + originalLineNum <= change.originalEndLineNumber; + originalLineNum++ + ) { + result.push({ + content: originalLines[originalLineNum - 1], + type: 'removed', + originalLineNumber: originalLineNum, + modifiedLineNumber: null, + }); + } + for ( + ; + modifiedLineNum <= change.modifiedEndLineNumber; + modifiedLineNum++ + ) { + result.push({ + content: modifiedLines[modifiedLineNum - 1], + type: 'added', + originalLineNumber: null, + modifiedLineNumber: modifiedLineNum, + }); + } + } else if (change.originalEndLineNumber === 0) { + // Addition + for ( + ; + modifiedLineNum <= change.modifiedEndLineNumber; + modifiedLineNum++ + ) { + result.push({ + content: modifiedLines[modifiedLineNum - 1], + type: 'added', + originalLineNumber: null, + modifiedLineNumber: modifiedLineNum, + }); + } + } else if (change.modifiedEndLineNumber === 0) { + // Deletion + for ( + ; + originalLineNum <= change.originalEndLineNumber; + originalLineNum++ + ) { + result.push({ + content: originalLines[originalLineNum - 1], + type: 'removed', + originalLineNumber: originalLineNum, + modifiedLineNumber: null, + }); + } + } + changeIndex++; + } else { + // Process any remaining unchanged lines + if ( + originalLineNum <= originalLines.length && + modifiedLineNum <= modifiedLines.length + ) { + result.push({ + content: originalLines[originalLineNum - 1], + type: 'unchanged', + originalLineNumber: originalLineNum, + modifiedLineNumber: modifiedLineNum, + }); + originalLineNum++; + modifiedLineNum++; + } else if (originalLineNum <= originalLines.length) { + // Remaining lines were removed + result.push({ + content: originalLines[originalLineNum - 1], + type: 'removed', + originalLineNumber: originalLineNum, + modifiedLineNumber: null, + }); + originalLineNum++; + } else if (modifiedLineNum <= modifiedLines.length) { + // Remaining lines were added + result.push({ + content: modifiedLines[modifiedLineNum - 1], + type: 'added', + originalLineNumber: null, + modifiedLineNumber: modifiedLineNum, + }); + modifiedLineNum++; + } else { + break; + } + } + } + + return result; } \ No newline at end of file From e96983379e0be6c6b02ff72dfbe52501458af627 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 21 Oct 2024 22:20:38 -0400 Subject: [PATCH 13/42] mito-ai: color diffs in cm code cells --- .../AiChat/ChatMessage/ChatMessage.tsx | 3 +- .../AiChat/ChatMessage/CodeBlock.tsx | 7 +-- .../AiChat/ChatMessage/PythonCode.tsx | 31 ++----------- .../src/Extensions/AiChat/ChatTaskpane.tsx | 17 ++++--- .../src/Extensions/AiChat/CodeDiffDisplay.tsx | 44 +++++++++++++------ mito-ai/src/utils/codeDiff.tsx | 14 +++--- mito-ai/src/utils/strings.tsx | 10 ++++- 7 files changed, 67 insertions(+), 59 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index 325142f05..b5092590d 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -8,6 +8,7 @@ import { splitStringWithCodeBlocks } from '../../../utils/strings'; import ErrorIcon from '../../../icons/ErrorIcon'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { OperatingSystem } from '../../../utils/user'; +import { UnifiedDiffLine } from '../../../utils/codeDiff'; interface IChatMessageProps { @@ -19,7 +20,7 @@ interface IChatMessageProps { app: JupyterFrontEnd isLastAiMessage: boolean operatingSystem: OperatingSystem - setDisplayCodeDiff: React.Dispatch>; + setDisplayCodeDiff: React.Dispatch>; } const ChatMessage: React.FC = ({ diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index f8620cf1d..28b51ea5f 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -7,6 +7,7 @@ import { removeMarkdownCodeFormatting } from '../../../utils/strings'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { OperatingSystem } from '../../../utils/user'; import '../../../../style/CodeMessagePart.css' +import { UnifiedDiffLine } from '../../../utils/codeDiff'; interface ICodeBlockProps { @@ -17,7 +18,7 @@ interface ICodeBlockProps { app: JupyterFrontEnd, isLastAiMessage: boolean, operatingSystem: OperatingSystem, - setDisplayCodeDiff: React.Dispatch>; + setDisplayCodeDiff: React.Dispatch>; } const CodeBlock: React.FC = ({ @@ -59,8 +60,8 @@ const CodeBlock: React.FC = ({ {notebookName}
- diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx index e96fda249..dccfd8d2b 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx @@ -2,7 +2,7 @@ import React, { useEffect, useState } from 'react'; import { IRenderMimeRegistry, MimeModel } from '@jupyterlab/rendermime'; import '../../../../style/PythonCode.css'; -// import { addMarkdownCodeFormatting } from '../../../utils/strings'; +import { addMarkdownCodeFormatting } from '../../../utils/strings'; interface IPythonCodeProps { code: string; @@ -12,40 +12,15 @@ interface IPythonCodeProps { const PythonCode: React.FC = ({ code, rendermime }) => { const [node, setNode] = useState(null) - // const newCode = '' - useEffect(() => { - const deletedLines = [0, 1, 2]; - - const newCode = 'print("hello world")\n# This is a comment\nprint("foobar")' - const wrappedCode = ` -\`\`\`python -${newCode} -\`\`\` - `; - const model = new MimeModel({ - data: { ['text/markdown']: wrappedCode}, + data: { ['text/markdown']: addMarkdownCodeFormatting(code) }, }); - + const renderer = rendermime.createRenderer('text/markdown'); renderer.renderModel(model) - // After rendering, add the background to specific lines - const codeElement = renderer.node.querySelector('pre'); - if (codeElement) { - const codeLines = codeElement.innerHTML.split('\n'); - const highlightedCode = codeLines.map((line, index) => { - if (deletedLines.includes(index + 1)) { - return `${line}`; - } - return line; - }).join('\n'); - codeElement.innerHTML = highlightedCode; - } - - const node = renderer.node setNode(node) }, [code, rendermime]) // Add dependencies to useEffect diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 01eac61b8..d25f1add5 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -11,13 +11,13 @@ import { requestAPI } from '../../utils/handler'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import LoadingDots from '../../components/LoadingDots'; import { JupyterFrontEnd } from '@jupyterlab/application'; -import { getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; +import { addMarkdownCodeFormatting, getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; import { ReadonlyPartialJSONObject } from '@lumino/coreutils'; import ResetIcon from '../../icons/ResetIcon'; import IconButton from '../../components/IconButton'; import { OperatingSystem } from '../../utils/user'; -import { createUnifiedDiff, getCodeDiffLineRanges } from '../../utils/codeDiff'; +import { createUnifiedDiff, getCodeDiffLineRanges, UnifiedDiffLine } from '../../utils/codeDiff'; import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; import { CodeMirrorEditor } from '@jupyterlab/codemirror'; import { CodeCell } from '@jupyterlab/cells'; @@ -84,7 +84,7 @@ const ChatTaskpane: React.FC = ({ const [chatHistoryManager, setChatHistoryManager] = useState(() => getDefaultChatHistoryManager()); const [input, setInput] = useState(''); const [loadingAIResponse, setLoadingAIResponse] = useState(false) - const [displayCodeDiff, setDisplayCodeDiff] = useState(false) + const [displayCodeDiff, setDisplayCodeDiff] = useState(undefined) const chatHistoryManagerRef = useRef(chatHistoryManager); useEffect(() => { @@ -181,6 +181,13 @@ const ChatTaskpane: React.FC = ({ const lineChanges = getCodeDiffLineRanges(activeCellCode, aiGeneratedCodeCleaned) const unifiedDiffs = createUnifiedDiff(originalLinesTest, modifiedLinesTest, lineChanges) + const unifiedCodeString = addMarkdownCodeFormatting(unifiedDiffs.map(line => { + return line.content !== undefined ? line.content : '' + }).join('\n')) + + console.log("unified code string", unifiedCodeString) + writeCodeToActiveCell(notebookTracker, unifiedCodeString) + setDisplayCodeDiff(unifiedDiffs) // Display the unified diff unifiedDiffs.forEach(line => { @@ -288,14 +295,14 @@ const ChatTaskpane: React.FC = ({ // Apply the initial configuration editorView.dispatch({ effects: StateEffect.appendConfig.of( - compartment.of(displayCodeDiff ? zebraStripes() : []) + compartment.of(displayCodeDiff ? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : []) ), }); } else { // Reconfigure the compartment editorView.dispatch({ effects: compartment.reconfigure( - displayCodeDiff ? zebraStripes() : [] + displayCodeDiff ? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : [] ), }); } diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index c432e05aa..29701417d 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -6,35 +6,51 @@ import { ViewPlugin, ViewUpdate } from '@codemirror/view'; +import { UnifiedDiffLine } from '../../utils/codeDiff'; // Defines new styles for this extension const baseTheme = EditorView.baseTheme({ // We need to set some transparency because the stripe are above // the selection layer - '&light .cm-zebraStripe': { backgroundColor: '#d4fafaaa' }, - '&dark .cm-zebraStripe': { backgroundColor: '#1a2727aa' } + '&light .cm-codeDiffRemovedStripe': { backgroundColor: '#fad4d4aa' }, + '&dark .cm-codeDiffRemovedStripe': { backgroundColor: '#3b0101aa' }, + + '&light .cm-codeDiffInsertedStripe': { backgroundColor: '#009e08aa' }, + '&dark .cm-codeDiffInsertedStripe': { backgroundColor: '#013b12aa' } + + }); // Resolve step to use in the editor - const stepSize = Facet.define({ - combine: values => (values.length ? Math.min(...values) : 2) + const unifiedDiffLines = Facet.define({ + // TODO: Do I need to provide a combine step? }); // Add decoration to editor lines - const stripe = Decoration.line({ - attributes: { class: 'cm-zebraStripe' } + const removedStripe = Decoration.line({ + attributes: { class: 'cm-codeDiffRemovedStripe' } + }); + + const insertedStripe = Decoration.line({ + attributes: { class: 'cm-codeDiffInsertedStripe' } }); // Create the range of lines requiring decorations function stripeDeco(view: EditorView) { - const step = view.state.facet(stepSize) as number; + console.log(view.state.facet(unifiedDiffLines)) + const unifiedDiffLinesFacet = view.state.facet(unifiedDiffLines)[0]; const builder = new RangeSetBuilder(); for (const { from, to } of view.visibleRanges) { for (let pos = from; pos <= to; ) { const line = view.state.doc.lineAt(pos); - if (line.number % step === 0) { - builder.add(line.from, line.from, stripe); + console.log('LINE NUMBER', line.number, unifiedDiffLinesFacet[line.number]) + // The code mirror line numbers are 1-indexed, but our diff lines are 0-indexed + if (unifiedDiffLinesFacet[line.number - 1].type === 'removed') { + builder.add(line.from, line.from, removedStripe); + } + if (unifiedDiffLinesFacet[line.number - 1].type === 'added') { + builder.add(line.from, line.from, insertedStripe); } pos = line.to + 1; } @@ -54,11 +70,11 @@ const baseTheme = EditorView.baseTheme({ update(update: ViewUpdate) { // Update the stripes if the document changed, // the viewport changed or the stripes step changed. - const oldStep = update.startState.facet(stepSize); + const oldUnifiedDiffLines = update.startState.facet(unifiedDiffLines); if ( update.docChanged || update.viewportChanged || - oldStep !== update.view.state.facet(stepSize) + oldUnifiedDiffLines !== update.view.state.facet(unifiedDiffLines) ) { this.decorations = stripeDeco(update.view); } @@ -70,11 +86,11 @@ const baseTheme = EditorView.baseTheme({ ); // Full extension composed of elemental extensions - export function zebraStripes(options: { step?: number} = {}): Extension { + export function zebraStripes(options: { unifiedDiffLines?: UnifiedDiffLine[]} = {}): Extension { return [ baseTheme, - typeof options.step !== 'number' ? [] : stepSize.of(options.step), - showStripes + options.unifiedDiffLines ? unifiedDiffLines.of(options.unifiedDiffLines) : [], + showStripes ]; } \ No newline at end of file diff --git a/mito-ai/src/utils/codeDiff.tsx b/mito-ai/src/utils/codeDiff.tsx index 0aba18960..39a241525 100644 --- a/mito-ai/src/utils/codeDiff.tsx +++ b/mito-ai/src/utils/codeDiff.tsx @@ -1,5 +1,12 @@ import { DiffComputer, IDiffComputerOpts, ILineChange } from "vscode-diff"; +export interface UnifiedDiffLine { + content: string; // The content of the line + type: 'unchanged' | 'added' | 'removed'; // The type of change + originalLineNumber: number | null; // Line number in the original code + modifiedLineNumber: number | null; // Line number in the modified code +} + export const getCodeDiffLineRanges = (originalLines: string | undefined | null, modifiedLines: string | undefined | null): ILineChange[] => { if (originalLines === undefined || originalLines === null) { originalLines = '' @@ -57,13 +64,6 @@ export const getCodeWithDiffsMarked = (originalLines: string | undefined | null, return "```python\n" + diffedLines.join('\n') + "\n```" } -interface UnifiedDiffLine { - content: string; // The content of the line - type: 'unchanged' | 'added' | 'removed'; // The type of change - originalLineNumber: number | null; // Line number in the original code - modifiedLineNumber: number | null; // Line number in the modified code -} - export const createUnifiedDiff = ( originalCode: string | undefined | null, modifiedCode: string | undefined | null, diff --git a/mito-ai/src/utils/strings.tsx b/mito-ai/src/utils/strings.tsx index bee7ea178..3fd00c233 100644 --- a/mito-ai/src/utils/strings.tsx +++ b/mito-ai/src/utils/strings.tsx @@ -59,7 +59,15 @@ export const getCodeBlockFromMessage = (message: OpenAI.Chat.ChatCompletionMessa */ export const addMarkdownCodeFormatting = (code: string) => { - const codeWithoutBackticks = code.split('```python')[1].split('```')[0].trim() + let codeWithoutBackticks = code + + // If the code already has the code formatting backticks, remove them + // so we can add them back in the correct format + if (code.split('```python').length > 1) { + codeWithoutBackticks = code.split('```python')[1].split('```')[0].trim() + } else { + codeWithoutBackticks = code.trim() + } // Note: We add a space after the code because for some unknown reason, the markdown // renderer is cutting off the last character in the code block. From f74776250f66bed0679c9bacbd4baf2fdee84d43 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 21 Oct 2024 22:34:39 -0400 Subject: [PATCH 14/42] mito-ai: only highlight the active code cell --- mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx | 10 ++++++---- mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx | 1 + 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index d25f1add5..69aaedfbb 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -276,10 +276,12 @@ const ChatTaskpane: React.FC = ({ return; } + const activeCellIndex = notebook.activeCellIndex + console.log(1) - notebook.widgets.forEach(cell => { + notebook.widgets.forEach((cell, index) => { if (cell.model.type === 'code') { - console.log(2) + const isActiveCodeCell = activeCellIndex === index const codeCell = cell as CodeCell; const cmEditor = codeCell.editor as CodeMirrorEditor; const editorView = cmEditor?.editor; @@ -295,14 +297,14 @@ const ChatTaskpane: React.FC = ({ // Apply the initial configuration editorView.dispatch({ effects: StateEffect.appendConfig.of( - compartment.of(displayCodeDiff ? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : []) + compartment.of(displayCodeDiff && isActiveCodeCell? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : []) ), }); } else { // Reconfigure the compartment editorView.dispatch({ effects: compartment.reconfigure( - displayCodeDiff ? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : [] + displayCodeDiff && isActiveCodeCell? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : [] ), }); } diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index 29701417d..692e11da1 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -39,6 +39,7 @@ const baseTheme = EditorView.baseTheme({ // Create the range of lines requiring decorations function stripeDeco(view: EditorView) { console.log(view.state.facet(unifiedDiffLines)) + console.log(view) const unifiedDiffLinesFacet = view.state.facet(unifiedDiffLines)[0]; const builder = new RangeSetBuilder(); for (const { from, to } of view.visibleRanges) { From f24b52be74b8c7f8499a4d45a6253ff85f6e1605 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Tue, 22 Oct 2024 17:50:04 -0400 Subject: [PATCH 15/42] mito-ai: clear code diffs on apply --- .../AiChat/ChatMessage/CodeBlock.tsx | 9 ++++-- .../src/Extensions/AiChat/ChatTaskpane.tsx | 30 ++++--------------- .../src/Extensions/AiChat/CodeDiffDisplay.tsx | 5 +--- mito-ai/src/utils/codeDiff.tsx | 4 +++ 4 files changed, 18 insertions(+), 30 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index 28b51ea5f..d87199570 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -59,9 +59,14 @@ const CodeBlock: React.FC = ({
{notebookName}
- + diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 69aaedfbb..5ac5be049 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -167,43 +167,23 @@ const ChatTaskpane: React.FC = ({ if (apiResponse.type === 'success') { - let originalLinesTest: string = "hello\noriginal\nworld"; - let modifiedLinesTest: string = "hello\nmodified\nworld\nfoobar"; - const response = apiResponse.response; const aiMessage = response.choices[0].message; - const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); - //const aiGeneratedCodeWithDiffs = getCodeWithDiffsMarked(activeCellCode, aiGeneratedCodeCleaned) const lineChanges = getCodeDiffLineRanges(activeCellCode, aiGeneratedCodeCleaned) - const unifiedDiffs = createUnifiedDiff(originalLinesTest, modifiedLinesTest, lineChanges) + const unifiedDiffs = createUnifiedDiff(activeCellCode, aiGeneratedCodeCleaned, lineChanges) + const unifiedCodeString = addMarkdownCodeFormatting(unifiedDiffs.map(line => { return line.content !== undefined ? line.content : '' }).join('\n')) - console.log("unified code string", unifiedCodeString) writeCodeToActiveCell(notebookTracker, unifiedCodeString) setDisplayCodeDiff(unifiedDiffs) - // Display the unified diff - unifiedDiffs.forEach(line => { - let prefix = ''; - if (line.type === 'added') { - prefix = '+ '; - } else if (line.type === 'removed') { - prefix = '- '; - } else { - prefix = ' '; - } - console.log(`${prefix}${line.content}`); - }); - - - console.log("unifiedDiffs", unifiedDiffs) aiMessage.content = aiGeneratedCode || '' updatedManager.addAIMessageFromResponse(aiMessage); @@ -244,6 +224,7 @@ const ChatTaskpane: React.FC = ({ app.commands.addCommand(COMMAND_MITO_AI_APPLY_LATEST_CODE, { execute: () => { applyLatestCode() + setDisplayCodeDiff(undefined) } }) @@ -271,6 +252,7 @@ const ChatTaskpane: React.FC = ({ // Function to update the extensions of code cells const updateCodeCellsExtensions = useCallback(() => { + console.log("Calling update") const notebook = notebookTracker.currentWidget?.content; if (!notebook) { return; @@ -297,14 +279,14 @@ const ChatTaskpane: React.FC = ({ // Apply the initial configuration editorView.dispatch({ effects: StateEffect.appendConfig.of( - compartment.of(displayCodeDiff && isActiveCodeCell? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : []) + compartment.of(displayCodeDiff !== undefined && isActiveCodeCell? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : []) ), }); } else { // Reconfigure the compartment editorView.dispatch({ effects: compartment.reconfigure( - displayCodeDiff && isActiveCodeCell? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : [] + displayCodeDiff !== undefined && isActiveCodeCell ? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : [] ), }); } diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index 692e11da1..4912929ce 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -16,10 +16,8 @@ const baseTheme = EditorView.baseTheme({ '&light .cm-codeDiffRemovedStripe': { backgroundColor: '#fad4d4aa' }, '&dark .cm-codeDiffRemovedStripe': { backgroundColor: '#3b0101aa' }, - '&light .cm-codeDiffInsertedStripe': { backgroundColor: '#009e08aa' }, + '&light .cm-codeDiffInsertedStripe': { backgroundColor: 'rgba(79, 255, 105, 0.38)' }, '&dark .cm-codeDiffInsertedStripe': { backgroundColor: '#013b12aa' } - - }); // Resolve step to use in the editor @@ -45,7 +43,6 @@ const baseTheme = EditorView.baseTheme({ for (const { from, to } of view.visibleRanges) { for (let pos = from; pos <= to; ) { const line = view.state.doc.lineAt(pos); - console.log('LINE NUMBER', line.number, unifiedDiffLinesFacet[line.number]) // The code mirror line numbers are 1-indexed, but our diff lines are 0-indexed if (unifiedDiffLinesFacet[line.number - 1].type === 'removed') { builder.add(line.from, line.from, removedStripe); diff --git a/mito-ai/src/utils/codeDiff.tsx b/mito-ai/src/utils/codeDiff.tsx index 39a241525..517d3f3e0 100644 --- a/mito-ai/src/utils/codeDiff.tsx +++ b/mito-ai/src/utils/codeDiff.tsx @@ -85,6 +85,10 @@ export const createUnifiedDiff = ( let modifiedLineNum = 1; let changeIndex = 0; + // Create unified lines + // Loop through the lineChages + // If unmodified + while ( originalLineNum <= originalLines.length || modifiedLineNum <= modifiedLines.length From 5d9aba43a64542155813e7910d02d65e1863d6d7 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Wed, 23 Oct 2024 10:57:50 -0400 Subject: [PATCH 16/42] mito-ai: fix bug with diffing empty cells --- .../AiChat/ChatMessage/ChatMessage.tsx | 4 +-- .../AiChat/ChatMessage/PythonCode.tsx | 2 +- .../src/Extensions/AiChat/ChatTaskpane.tsx | 10 ++++-- mito-ai/src/utils/codeDiff.tsx | 34 +++---------------- mito-ai/src/utils/strings.tsx | 33 ++++++++++++++---- 5 files changed, 42 insertions(+), 41 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index b5092590d..77cf1aea4 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -4,7 +4,7 @@ import { classNames } from '../../../utils/classNames'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import CodeBlock from './CodeBlock'; import { INotebookTracker } from '@jupyterlab/notebook'; -import { splitStringWithCodeBlocks } from '../../../utils/strings'; +import { PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE, splitStringWithCodeBlocks } from '../../../utils/strings'; import ErrorIcon from '../../../icons/ErrorIcon'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { OperatingSystem } from '../../../utils/user'; @@ -49,7 +49,7 @@ const ChatMessage: React.FC = ({ {'message-error': error} )}> {messageContentParts.map(messagePart => { - if (messagePart.startsWith('```python')) { + if (messagePart.startsWith(PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE)) { // Make sure that there is actually code in the message. // An empty code will look like this '```python ```' // TODO: Add a test for this since its broke a few times now. diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx index dccfd8d2b..4dee86b4e 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/PythonCode.tsx @@ -15,7 +15,7 @@ const PythonCode: React.FC = ({ code, rendermime }) => { useEffect(() => { const model = new MimeModel({ - data: { ['text/markdown']: addMarkdownCodeFormatting(code) }, + data: { ['text/markdown']: addMarkdownCodeFormatting(code, true) }, }); const renderer = rendermime.createRenderer('text/markdown'); diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 5ac5be049..82d055af8 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -11,7 +11,7 @@ import { requestAPI } from '../../utils/handler'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import LoadingDots from '../../components/LoadingDots'; import { JupyterFrontEnd } from '@jupyterlab/application'; -import { addMarkdownCodeFormatting, getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; +import { getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; import { ReadonlyPartialJSONObject } from '@lumino/coreutils'; import ResetIcon from '../../icons/ResetIcon'; @@ -177,10 +177,16 @@ const ChatTaskpane: React.FC = ({ const unifiedDiffs = createUnifiedDiff(activeCellCode, aiGeneratedCodeCleaned, lineChanges) - const unifiedCodeString = addMarkdownCodeFormatting(unifiedDiffs.map(line => { + console.log("unifiedDiffs") + console.log(unifiedDiffs) + + const unifiedCodeString = (unifiedDiffs.map(line => { return line.content !== undefined ? line.content : '' }).join('\n')) + console.log("unifiedCodeString1") + console.log(unifiedCodeString) + writeCodeToActiveCell(notebookTracker, unifiedCodeString) setDisplayCodeDiff(unifiedDiffs) diff --git a/mito-ai/src/utils/codeDiff.tsx b/mito-ai/src/utils/codeDiff.tsx index 517d3f3e0..a24214e6a 100644 --- a/mito-ai/src/utils/codeDiff.tsx +++ b/mito-ai/src/utils/codeDiff.tsx @@ -36,33 +36,6 @@ export const getCodeDiffLineRanges = (originalLines: string | undefined | null, return lineChanges || [] } -export const getCodeWithDiffsMarked = (originalLines: string | undefined | null, modifiedLines: string | undefined | null): string => { - - /* - originalCodeLines: string - newCodeLines: string - allCodeLines: string - - Ordered by line number with: - - Original code lines - - New Code Lines - deletedLineIndexes: List[int] - modifiedLineIndexes: List[int] - */ - - - - const lineChanges = getCodeDiffLineRanges(originalLines, modifiedLines); - - const diffedLines = originalLines?.split('\n') || [] - - let numNewLinesAdded = 0 - for (const lineChange of lineChanges) { - diffedLines[lineChange.originalStartLineNumber] = '' + diffedLines[lineChange.originalStartLineNumber] + '' - numNewLinesAdded = numNewLinesAdded + 1 - } - - return "```python\n" + diffedLines.join('\n') + "\n```" -} export const createUnifiedDiff = ( originalCode: string | undefined | null, @@ -85,9 +58,10 @@ export const createUnifiedDiff = ( let modifiedLineNum = 1; let changeIndex = 0; - // Create unified lines - // Loop through the lineChages - // If unmodified + console.log("originalLines") + console.log(originalLines) + console.log(modifiedLines) + console.log(lineChanges) while ( originalLineNum <= originalLines.length || diff --git a/mito-ai/src/utils/strings.tsx b/mito-ai/src/utils/strings.tsx index 3fd00c233..96ebadc54 100644 --- a/mito-ai/src/utils/strings.tsx +++ b/mito-ai/src/utils/strings.tsx @@ -1,5 +1,11 @@ import OpenAI from "openai"; +export const PYTHON_CODE_BLOCK_START_WITH_NEW_LINE = '```python\n' +export const PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE = '```python' +export const PYTHON_CODE_BLOCK_END_WITH_NEW_LINE = '\n```' +export const PYTHON_CODE_BLOCK_END_WITHOUT_NEW_LINE = '```' + + /* Given a message from the OpenAI API, returns the content as a string. If the content is not a string, returns undefined. @@ -56,22 +62,32 @@ export const getCodeBlockFromMessage = (message: OpenAI.Chat.ChatCompletionMessa ```python x + 1 ``` + + Sometimes, we also want to trim the code to remove any leading or trailing whitespace. For example, + when we're displaying the code in the chat history this is useful. Othertimes we don't want to trim. + For example, when we're displaying the code in the active cell, we want to keep the users's whitespace. + This is important for showing diffs. If the code cell contains no code, the first line will be marked as + removed in the code diff. To ensure the diff lines up with the code, we need to leave this whitespace line. */ -export const addMarkdownCodeFormatting = (code: string) => { +export const addMarkdownCodeFormatting = (code: string, trim?: boolean) => { let codeWithoutBackticks = code // If the code already has the code formatting backticks, remove them // so we can add them back in the correct format - if (code.split('```python').length > 1) { - codeWithoutBackticks = code.split('```python')[1].split('```')[0].trim() + if (code.split(PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE).length > 1) { + codeWithoutBackticks = code.split(PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE)[1].split(PYTHON_CODE_BLOCK_END_WITHOUT_NEW_LINE)[0] } else { - codeWithoutBackticks = code.trim() + codeWithoutBackticks = code + } + + if (trim) { + codeWithoutBackticks = codeWithoutBackticks.trim() } // Note: We add a space after the code because for some unknown reason, the markdown // renderer is cutting off the last character in the code block. - return "```python\n" + codeWithoutBackticks + " " + "\n```" + return `${PYTHON_CODE_BLOCK_START_WITH_NEW_LINE}${codeWithoutBackticks} ${PYTHON_CODE_BLOCK_END_WITH_NEW_LINE}` } /* @@ -88,5 +104,10 @@ export const addMarkdownCodeFormatting = (code: string) => { Jupyter does not need the backticks. */ export const removeMarkdownCodeFormatting = (code: string) => { - return code.split('```python')[1].split('```')[0].trim() + + if (code.split(PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE).length > 1) { + return code.split(PYTHON_CODE_BLOCK_START_WITH_NEW_LINE)[1].split(PYTHON_CODE_BLOCK_END_WITH_NEW_LINE)[0] + } + + return code } \ No newline at end of file From 916782f3db39419cf83867e612b041fe2cd52af7 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Wed, 23 Oct 2024 15:49:04 -0400 Subject: [PATCH 17/42] mito-ai: cleanup --- mito-ai/mito-ai/OpenAICompletionHandler.py | 2 - .../AiChat/ChatMessage/CodeBlock.tsx | 2 +- .../src/Extensions/AiChat/ChatTaskpane.tsx | 56 ++++--- .../src/Extensions/AiChat/CodeDiffDisplay.tsx | 137 ++++++++---------- mito-ai/src/commands.tsx | 1 + mito-ai/src/utils/codeDiff.tsx | 84 ++++++----- mito-ai/style/CodeMessagePart.css | 13 +- 7 files changed, 144 insertions(+), 151 deletions(-) diff --git a/mito-ai/mito-ai/OpenAICompletionHandler.py b/mito-ai/mito-ai/OpenAICompletionHandler.py index 3141b8b77..7dce283af 100644 --- a/mito-ai/mito-ai/OpenAICompletionHandler.py +++ b/mito-ai/mito-ai/OpenAICompletionHandler.py @@ -43,8 +43,6 @@ def post(self): # return a cleaned up version of the response so we can support # multiple models - print(response_dict) - self.finish(json.dumps(response_dict)) except Exception as e: self.set_status(500) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index d87199570..1f0bfb11a 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -66,7 +66,7 @@ const CodeBlock: React.FC = ({ Apply {isLastAiMessage ? (operatingSystem === 'mac' ? 'CMD+Y' : 'CTRL+Y') : ''} diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 82d055af8..9207578fc 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -12,17 +12,17 @@ import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import LoadingDots from '../../components/LoadingDots'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; -import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; +import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_REJECT_LATEST_CODE, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; import { ReadonlyPartialJSONObject } from '@lumino/coreutils'; import ResetIcon from '../../icons/ResetIcon'; import IconButton from '../../components/IconButton'; import { OperatingSystem } from '../../utils/user'; -import { createUnifiedDiff, getCodeDiffLineRanges, UnifiedDiffLine } from '../../utils/codeDiff'; +import { getCodeDiffsAndUnifiedCodeString, UnifiedDiffLine } from '../../utils/codeDiff'; import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; import { CodeMirrorEditor } from '@jupyterlab/codemirror'; import { CodeCell } from '@jupyterlab/cells'; import { StateEffect, Compartment } from '@codemirror/state'; -import { zebraStripes } from './CodeDiffDisplay'; +import { codeDiffStripesExtension } from './CodeDiffDisplay'; @@ -170,28 +170,17 @@ const ChatTaskpane: React.FC = ({ const response = apiResponse.response; const aiMessage = response.choices[0].message; + // Extract the code from the AI's message and then calculate the code diffs const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); + const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) - const lineChanges = getCodeDiffLineRanges(activeCellCode, aiGeneratedCodeCleaned) - - const unifiedDiffs = createUnifiedDiff(activeCellCode, aiGeneratedCodeCleaned, lineChanges) - - console.log("unifiedDiffs") - console.log(unifiedDiffs) - - const unifiedCodeString = (unifiedDiffs.map(line => { - return line.content !== undefined ? line.content : '' - }).join('\n')) - - console.log("unifiedCodeString1") - console.log(unifiedCodeString) - + // Temporarily write the unified code string to the active cell so we can display + // the code diffs to the user. Once the user accepts or rejects the code, we'll + // apply the correct version of the code. writeCodeToActiveCell(notebookTracker, unifiedCodeString) setDisplayCodeDiff(unifiedDiffs) - aiMessage.content = aiGeneratedCode || '' - updatedManager.addAIMessageFromResponse(aiMessage); setChatHistoryManager(updatedManager); } else { @@ -234,12 +223,24 @@ const ChatTaskpane: React.FC = ({ } }) + app.commands.addCommand(COMMAND_MITO_AI_REJECT_LATEST_CODE, { + execute: () => { + setDisplayCodeDiff(undefined) + } + }) + app.commands.addKeyBinding({ command: COMMAND_MITO_AI_APPLY_LATEST_CODE, keys: ['Accel Y'], selector: 'body', }); + app.commands.addKeyBinding({ + command: COMMAND_MITO_AI_REJECT_LATEST_CODE, + keys: ['Accel D'], + selector: 'body', + }); + /* Add a new command to the JupyterLab command registry that sends the current chat message. We use this to automatically send the message when the user adds an error to the chat. @@ -254,11 +255,10 @@ const ChatTaskpane: React.FC = ({ }, []) // Create a WeakMap to store compartments per code cell - const zebraStripesCompartments = React.useRef(new WeakMap()); + const codeDiffStripesCompartments = React.useRef(new WeakMap()); // Function to update the extensions of code cells const updateCodeCellsExtensions = useCallback(() => { - console.log("Calling update") const notebook = notebookTracker.currentWidget?.content; if (!notebook) { return; @@ -266,7 +266,6 @@ const ChatTaskpane: React.FC = ({ const activeCellIndex = notebook.activeCellIndex - console.log(1) notebook.widgets.forEach((cell, index) => { if (cell.model.type === 'code') { const isActiveCodeCell = activeCellIndex === index @@ -275,37 +274,36 @@ const ChatTaskpane: React.FC = ({ const editorView = cmEditor?.editor; if (editorView) { - let compartment = zebraStripesCompartments.current.get(codeCell); + let compartment = codeDiffStripesCompartments.current.get(codeCell); if (!compartment) { // Create a new compartment and store it compartment = new Compartment(); - zebraStripesCompartments.current.set(codeCell, compartment); + codeDiffStripesCompartments.current.set(codeCell, compartment); // Apply the initial configuration editorView.dispatch({ effects: StateEffect.appendConfig.of( - compartment.of(displayCodeDiff !== undefined && isActiveCodeCell? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : []) + compartment.of(displayCodeDiff !== undefined && isActiveCodeCell? codeDiffStripesExtension({ unifiedDiffLines: displayCodeDiff }) : []) ), }); } else { // Reconfigure the compartment editorView.dispatch({ effects: compartment.reconfigure( - displayCodeDiff !== undefined && isActiveCodeCell ? zebraStripes({ unifiedDiffLines: displayCodeDiff }) : [] + displayCodeDiff !== undefined && isActiveCodeCell ? codeDiffStripesExtension({ unifiedDiffLines: displayCodeDiff }) : [] ), }); } } else { - console.log('editor view not found') + console.log('Mito AI: editor view not found when applying code diff stripes') } } }); }, [displayCodeDiff, notebookTracker]); - // Update code cells when displayCodeDiff changes + useEffect(() => { - console.log("Calling updateCodeCellsExtensions") updateCodeCellsExtensions(); }, [displayCodeDiff, updateCodeCellsExtensions]); diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index 4912929ce..99c64efd0 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -11,84 +11,75 @@ import { UnifiedDiffLine } from '../../utils/codeDiff'; // Defines new styles for this extension const baseTheme = EditorView.baseTheme({ - // We need to set some transparency because the stripe are above - // the selection layer - '&light .cm-codeDiffRemovedStripe': { backgroundColor: '#fad4d4aa' }, - '&dark .cm-codeDiffRemovedStripe': { backgroundColor: '#3b0101aa' }, + // We need to set some transparency because the stripes are above the selection layer + '.cm-codeDiffRemovedStripe': { backgroundColor: 'rgba(250, 212, 212, 0.62)' }, + '.cm-codeDiffInsertedStripe': { backgroundColor: 'rgba(79, 255, 105, 0.38)' }, +}); - '&light .cm-codeDiffInsertedStripe': { backgroundColor: 'rgba(79, 255, 105, 0.38)' }, - '&dark .cm-codeDiffInsertedStripe': { backgroundColor: '#013b12aa' } - }); - - // Resolve step to use in the editor - const unifiedDiffLines = Facet.define({ - // TODO: Do I need to provide a combine step? - }); - - // Add decoration to editor lines - const removedStripe = Decoration.line({ - attributes: { class: 'cm-codeDiffRemovedStripe' } - }); +// Resolve step to use in the editor +const unifiedDiffLines = Facet.define({ + // TODO: Do I need to provide a combine step? +}); - const insertedStripe = Decoration.line({ - attributes: { class: 'cm-codeDiffInsertedStripe' } - }); - - // Create the range of lines requiring decorations - function stripeDeco(view: EditorView) { - console.log(view.state.facet(unifiedDiffLines)) - console.log(view) - const unifiedDiffLinesFacet = view.state.facet(unifiedDiffLines)[0]; - const builder = new RangeSetBuilder(); - for (const { from, to } of view.visibleRanges) { - for (let pos = from; pos <= to; ) { - const line = view.state.doc.lineAt(pos); - // The code mirror line numbers are 1-indexed, but our diff lines are 0-indexed - if (unifiedDiffLinesFacet[line.number - 1].type === 'removed') { - builder.add(line.from, line.from, removedStripe); - } - if (unifiedDiffLinesFacet[line.number - 1].type === 'added') { - builder.add(line.from, line.from, insertedStripe); - } - pos = line.to + 1; +// Add decoration to editor lines +const removedStripe = Decoration.line({ + attributes: { class: 'cm-codeDiffRemovedStripe' } +}); + +const insertedStripe = Decoration.line({ + attributes: { class: 'cm-codeDiffInsertedStripe' } +}); + +// Create the range of lines requiring decorations +const getCodeDiffStripesDecoration = (view: EditorView): DecorationSet => { + const unifiedDiffLinesFacet = view.state.facet(unifiedDiffLines)[0]; + const builder = new RangeSetBuilder(); + for (const { from, to } of view.visibleRanges) { + for (let pos = from; pos <= to;) { + const line = view.state.doc.lineAt(pos); + // The code mirror line numbers are 1-indexed, but our diff lines are 0-indexed + if (unifiedDiffLinesFacet[line.number - 1].type === 'removed') { + builder.add(line.from, line.from, removedStripe); + } + if (unifiedDiffLinesFacet[line.number - 1].type === 'inserted') { + builder.add(line.from, line.from, insertedStripe); } + pos = line.to + 1; } - return builder.finish(); } - - // Update the decoration status of the editor view - const showStripes = ViewPlugin.fromClass( - class { - decorations: DecorationSet; - - constructor(view: EditorView) { - this.decorations = stripeDeco(view); - } - - update(update: ViewUpdate) { - // Update the stripes if the document changed, - // the viewport changed or the stripes step changed. - const oldUnifiedDiffLines = update.startState.facet(unifiedDiffLines); - if ( - update.docChanged || - update.viewportChanged || - oldUnifiedDiffLines !== update.view.state.facet(unifiedDiffLines) - ) { - this.decorations = stripeDeco(update.view); - } + return builder.finish(); +} + +// Update the decoration status of the editor view +const showStripes = ViewPlugin.fromClass( + class { + decorations: DecorationSet; + + constructor(view: EditorView) { + this.decorations = getCodeDiffStripesDecoration(view); + } + + update(update: ViewUpdate) { + const oldUnifiedDiffLines = update.startState.facet(unifiedDiffLines); + if ( + update.docChanged || + update.viewportChanged || + oldUnifiedDiffLines !== update.view.state.facet(unifiedDiffLines) + ) { + this.decorations = getCodeDiffStripesDecoration(update.view); } - }, - { - decorations: v => v.decorations } - ); - - // Full extension composed of elemental extensions - export function zebraStripes(options: { unifiedDiffLines?: UnifiedDiffLine[]} = {}): Extension { - return [ - baseTheme, - options.unifiedDiffLines ? unifiedDiffLines.of(options.unifiedDiffLines) : [], - showStripes - ]; + }, + { + decorations: v => v.decorations } - \ No newline at end of file +); + +// Create the Code Mirror Extension to apply the code diff stripes to the code mirror editor +export function codeDiffStripesExtension(options: { unifiedDiffLines?: UnifiedDiffLine[] } = {}): Extension { + return [ + baseTheme, + options.unifiedDiffLines ? unifiedDiffLines.of(options.unifiedDiffLines) : [], + showStripes + ]; +} diff --git a/mito-ai/src/commands.tsx b/mito-ai/src/commands.tsx index cfaf94e82..4fcc20e70 100644 --- a/mito-ai/src/commands.tsx +++ b/mito-ai/src/commands.tsx @@ -2,4 +2,5 @@ const MITO_AI = 'mito_ai' export const COMMAND_MITO_AI_OPEN_CHAT = `${MITO_AI}:open-chat` export const COMMAND_MITO_AI_APPLY_LATEST_CODE = `${MITO_AI}:apply-latest-code` +export const COMMAND_MITO_AI_REJECT_LATEST_CODE = `${MITO_AI}:reject-latest-code` export const COMMAND_MITO_AI_SEND_MESSAGE = `${MITO_AI}:send-message` \ No newline at end of file diff --git a/mito-ai/src/utils/codeDiff.tsx b/mito-ai/src/utils/codeDiff.tsx index a24214e6a..2276fe711 100644 --- a/mito-ai/src/utils/codeDiff.tsx +++ b/mito-ai/src/utils/codeDiff.tsx @@ -2,7 +2,7 @@ import { DiffComputer, IDiffComputerOpts, ILineChange } from "vscode-diff"; export interface UnifiedDiffLine { content: string; // The content of the line - type: 'unchanged' | 'added' | 'removed'; // The type of change + type: 'unchanged' | 'inserted' | 'removed'; // The type of change originalLineNumber: number | null; // Line number in the original code modifiedLineNumber: number | null; // Line number in the modified code } @@ -19,9 +19,6 @@ export const getCodeDiffLineRanges = (originalLines: string | undefined | null, const originalLinesArray = originalLines.split('\n') const modifiedLinesArray = modifiedLines.split('\n') - console.log("originalLinesArray", originalLinesArray) - console.log("modifiedLinesArray", modifiedLinesArray) - let options: IDiffComputerOpts = { shouldPostProcessCharChanges: true, shouldIgnoreTrimWhitespace: true, @@ -42,6 +39,7 @@ export const createUnifiedDiff = ( modifiedCode: string | undefined | null, lineChanges: ILineChange[] ): UnifiedDiffLine[] => { + if (originalCode === undefined || originalCode === null) { originalCode = '' } @@ -53,16 +51,29 @@ export const createUnifiedDiff = ( const originalLines = originalCode.split('\n') const modifiedLines = modifiedCode.split('\n') + /* + Algorithm explanation: + + This function creates a unified diff by comparing the original and modified code. + It iterates through both versions of the code simultaneously, creating a new representation + of the code called result that is UnifiedDiffLine[]. Each time the algorithm sees a new line + of code, it adds it to the result, marking it as unchanged, removed, or inserted. + + The algorithm works as follows: + 1. Process unchanged lines until a change is encountered. + 2. When a change is found, handle it based on its type: + a. Modification: Mark original lines as removed, mark modified lines as inserted. + b. Inserted: Add new lines from the modified code and mark as Inserted. + c. Removed: Add removed lines from the original code to the result and mark as Removed. + 3. After processing all changes, handle any remaining lines. + The result is a unified diff that shows all changes in context. + */ + const result: UnifiedDiffLine[] = []; let originalLineNum = 1; let modifiedLineNum = 1; let changeIndex = 0; - console.log("originalLines") - console.log(originalLines) - console.log(modifiedLines) - console.log(lineChanges) - while ( originalLineNum <= originalLines.length || modifiedLineNum <= modifiedLines.length @@ -93,11 +104,8 @@ export const createUnifiedDiff = ( change.modifiedEndLineNumber > 0 ) { // Modification - for ( - ; - originalLineNum <= change.originalEndLineNumber; - originalLineNum++ - ) { + // First add removed lines + for (;originalLineNum <= change.originalEndLineNumber; originalLineNum++) { result.push({ content: originalLines[originalLineNum - 1], type: 'removed', @@ -105,39 +113,28 @@ export const createUnifiedDiff = ( modifiedLineNumber: null, }); } - for ( - ; - modifiedLineNum <= change.modifiedEndLineNumber; - modifiedLineNum++ - ) { + // Then add inserted lines + for (;modifiedLineNum <= change.modifiedEndLineNumber; modifiedLineNum++) { result.push({ content: modifiedLines[modifiedLineNum - 1], - type: 'added', + type: 'inserted', originalLineNumber: null, modifiedLineNumber: modifiedLineNum, }); } } else if (change.originalEndLineNumber === 0) { - // Addition - for ( - ; - modifiedLineNum <= change.modifiedEndLineNumber; - modifiedLineNum++ - ) { + // Inserted Lines + for (;modifiedLineNum <= change.modifiedEndLineNumber;modifiedLineNum++) { result.push({ content: modifiedLines[modifiedLineNum - 1], - type: 'added', + type: 'inserted', originalLineNumber: null, modifiedLineNumber: modifiedLineNum, }); } } else if (change.modifiedEndLineNumber === 0) { - // Deletion - for ( - ; - originalLineNum <= change.originalEndLineNumber; - originalLineNum++ - ) { + // Removed lines + for (;originalLineNum <= change.originalEndLineNumber; originalLineNum++) { result.push({ content: originalLines[originalLineNum - 1], type: 'removed', @@ -174,7 +171,7 @@ export const createUnifiedDiff = ( // Remaining lines were added result.push({ content: modifiedLines[modifiedLineNum - 1], - type: 'added', + type: 'inserted', originalLineNumber: null, modifiedLineNumber: modifiedLineNum, }); @@ -186,4 +183,23 @@ export const createUnifiedDiff = ( } return result; +} + + +export const getCodeDiffsAndUnifiedCodeString = (originalCode: string | undefined | null, modifiedCode: string | undefined | null): { + unifiedCodeString: string; + unifiedDiffs: UnifiedDiffLine[]; +} => { + + const lineChanges = getCodeDiffLineRanges(originalCode, modifiedCode) + const unifiedDiffs = createUnifiedDiff(originalCode, modifiedCode, lineChanges) + + const unifiedCodeString = (unifiedDiffs.map(line => { + return line.content !== undefined ? line.content : '' + }).join('\n')) + + return { + unifiedCodeString, + unifiedDiffs + } } \ No newline at end of file diff --git a/mito-ai/style/CodeMessagePart.css b/mito-ai/style/CodeMessagePart.css index 603b5f269..131f76fc8 100644 --- a/mito-ai/style/CodeMessagePart.css +++ b/mito-ai/style/CodeMessagePart.css @@ -39,15 +39,4 @@ .code-message-part-toolbar button:hover { background-color: var(--chat-background-color); color: var(--chat-assistant-message-font-color); -} - - - .deleted-line .jp-RenderedMarkdown { - background-color: rgba(255, 0, 0, 0.2) !important; - } - - .deleted-line .jp-RenderedMarkdown pre { - background-color: transparent !important; - margin: 0; - padding: 0; - } \ No newline at end of file +} \ No newline at end of file From 59116b56069f4616122042bfbd4f54fa23dea2c1 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 24 Oct 2024 15:32:11 -0400 Subject: [PATCH 18/42] mito-ai: update prompt --- .../Extensions/AiChat/ChatHistoryManager.tsx | 43 ++++++++++++++++--- .../AiChat/ChatMessage/CodeBlock.tsx | 2 - .../src/Extensions/AiChat/ChatTaskpane.tsx | 9 ++-- mito-ai/src/utils/notebook.tsx | 2 +- 4 files changed, 42 insertions(+), 14 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index cb11ee4a6..ca7f6f238 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -68,15 +68,10 @@ ${input}`}; content: `You have access to the following variables: ${variables?.map(variable => `${JSON.stringify(variable, null, 2)}\n`).join('')} - -Code in the active code cell: - -\`\`\`python -${activeCellCode} -\`\`\` Complete the task below. Decide what variables to use and what changes you need to make to the active code cell. Only return the full new active code cell and a concise explanation of the changes you made. + Do not: - Use the word "I" - Include multiple approaches in your response @@ -87,7 +82,33 @@ Do: - Keep as much of the original code as possible - Ask for more context if you need it. -Important: Remember that you are executing code inside a Jupyter notebook. That means you will have persistent state issues where variables from previous cells or previous code executions might still affect current code. When those errors occur, here are a few possible solutions: + + + + +Code in the active code cell: + +\`\`\`python +import pandas as pd +loans_df = pd.read_csv('./loans.csv') +\`\`\` + +Your task: convert the issue_date column to datetime. + +Output: + +\`\`\`python +import pandas as pd +loans_df = pd.read_csv('./loans.csv') +loans_df['issue_date'] = pd.to_datetime(loans_df['issue_date']) +\`\`\` + +Use the pd.to_datetime function to convert the issue_date column to datetime. + + + + +Remember that you are executing code inside a Jupyter notebook. That means you will have persistent state issues where variables from previous cells or previous code executions might still affect current code. When those errors occur, here are a few possible solutions: 1. Restarting the kernel to reset the environment if a function or variable has been unintentionally overwritten. 2. Identify which cell might need to be rerun to properly initialize the function or variable that is causing the issue. @@ -95,6 +116,14 @@ For example, if an error occurs because the built-in function 'print' is overwri When a user hits an error because of a persistent state issue, tell them how to resolve it. + + +Code in the active code cell: + +\`\`\`python +${activeCellCode} +\`\`\` + Your task: ${input}`}; this.history.displayOptimizedChatHistory.push({message: displayMessage, error: false}); diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index 1f0bfb11a..a34182832 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -50,8 +50,6 @@ const CodeBlock: React.FC = ({ ) } - - if (role === 'assistant') { return (
diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 9207578fc..a915456f2 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -170,6 +170,9 @@ const ChatTaskpane: React.FC = ({ const response = apiResponse.response; const aiMessage = response.choices[0].message; + updatedManager.addAIMessageFromResponse(aiMessage); + setChatHistoryManager(updatedManager); + // Extract the code from the AI's message and then calculate the code diffs const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); @@ -180,9 +183,6 @@ const ChatTaskpane: React.FC = ({ // apply the correct version of the code. writeCodeToActiveCell(notebookTracker, unifiedCodeString) setDisplayCodeDiff(unifiedDiffs) - - updatedManager.addAIMessageFromResponse(aiMessage); - setChatHistoryManager(updatedManager); } else { updatedManager.addAIMessageFromMessageContent(apiResponse.errorMessage, true) setChatHistoryManager(updatedManager); @@ -191,8 +191,9 @@ const ChatTaskpane: React.FC = ({ setLoadingAIResponse(false) } catch (error) { console.error('Error calling OpenAI API:', error); + } finally { + setLoadingAIResponse(false) } - }; const displayOptimizedChatHistory = chatHistoryManager.getDisplayOptimizedHistory() diff --git a/mito-ai/src/utils/notebook.tsx b/mito-ai/src/utils/notebook.tsx index 97fc10454..acf2c128d 100644 --- a/mito-ai/src/utils/notebook.tsx +++ b/mito-ai/src/utils/notebook.tsx @@ -26,7 +26,7 @@ export const writeCodeToActiveCell = (notebookTracker: INotebookTracker, code: s const codeMirrorValidCode = removeMarkdownCodeFormatting(code) const activeCell = getActiveCell(notebookTracker) if (activeCell !== undefined) { - activeCell.model.sharedModel.source = codeMirrorValidCode + activeCell.model.sharedModel.source = codeMirrorValidCode if (focus) { activeCell.node.focus() From 972c469a846e68e8aea140da3f8395ce5ac20ae4 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 24 Oct 2024 16:24:58 -0400 Subject: [PATCH 19/42] mito-ai: fix code mirror extension crash bug --- .../src/Extensions/AiChat/CodeDiffDisplay.tsx | 23 ++++++++++++++++++- mito-ai/src/utils/arrays.tsx | 11 +++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 mito-ai/src/utils/arrays.tsx diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index 99c64efd0..77cb7984d 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -7,6 +7,7 @@ import { ViewUpdate } from '@codemirror/view'; import { UnifiedDiffLine } from '../../utils/codeDiff'; +import { deepEqualArrays } from '../../utils/arrays'; // Defines new styles for this extension @@ -18,6 +19,10 @@ const baseTheme = EditorView.baseTheme({ // Resolve step to use in the editor const unifiedDiffLines = Facet.define({ + combine: (unifiedDiffLines) => { + console.log('combining unified diff lines') + return unifiedDiffLines + } // TODO: Do I need to provide a combine step? }); @@ -37,7 +42,21 @@ const getCodeDiffStripesDecoration = (view: EditorView): DecorationSet => { for (const { from, to } of view.visibleRanges) { for (let pos = from; pos <= to;) { const line = view.state.doc.lineAt(pos); + // console.log('unifiedDiffLinesFacet[line.number - 1]', unifiedDiffLinesFacet[line.number - 1]) + // console.log('line', line.number) // The code mirror line numbers are 1-indexed, but our diff lines are 0-indexed + if (line.number - 1 >= unifiedDiffLinesFacet.length) { + /* + Because we need to rerender the decorations each time the doc changes or viewport updates + (maybe we don't need to, but the code mirror examples does this so we will to for now) there + is a race condition where sometimes the content of the code cell updates before the unified diff lines + are updated. As a result, we need to break out of the loop before we get a null pointer error. + + This isn't a problem because right afterwards, the code mirror updates again due to the unified diff lines + being updated. In that render, we get the correct results. + */ + break + } if (unifiedDiffLinesFacet[line.number - 1].type === 'removed') { builder.add(line.from, line.from, removedStripe); } @@ -61,10 +80,12 @@ const showStripes = ViewPlugin.fromClass( update(update: ViewUpdate) { const oldUnifiedDiffLines = update.startState.facet(unifiedDiffLines); + const newUnifiedDiffLines = update.view.state.facet(unifiedDiffLines); + if ( update.docChanged || update.viewportChanged || - oldUnifiedDiffLines !== update.view.state.facet(unifiedDiffLines) + !deepEqualArrays(oldUnifiedDiffLines[0], newUnifiedDiffLines[0]) ) { this.decorations = getCodeDiffStripesDecoration(update.view); } diff --git a/mito-ai/src/utils/arrays.tsx b/mito-ai/src/utils/arrays.tsx new file mode 100644 index 000000000..2ecb5a806 --- /dev/null +++ b/mito-ai/src/utils/arrays.tsx @@ -0,0 +1,11 @@ +export const deepEqualArrays = (arr1: any[], arr2: any[]): boolean => { + if (arr1.length !== arr2.length) return false; + for (let i = 0; i < arr1.length; i++) { + if (typeof arr1[i] === 'object' && typeof arr2[i] === 'object') { + if (!deepEqualArrays(arr1[i], arr2[i])) return false; + } else if (arr1[i] !== arr2[i]) { + return false; + } + } + return true; +} \ No newline at end of file From a1d09d3908b525c01ec28f3fce4f5c9ff2053383 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 24 Oct 2024 16:43:49 -0400 Subject: [PATCH 20/42] mito-ai: cleanup --- mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx | 2 -- 1 file changed, 2 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index 77cb7984d..3d90e54d8 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -20,10 +20,8 @@ const baseTheme = EditorView.baseTheme({ // Resolve step to use in the editor const unifiedDiffLines = Facet.define({ combine: (unifiedDiffLines) => { - console.log('combining unified diff lines') return unifiedDiffLines } - // TODO: Do I need to provide a combine step? }); // Add decoration to editor lines From bb177a47149f9ca5e653579ac409c4251a0fe4e7 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 24 Oct 2024 16:55:41 -0400 Subject: [PATCH 21/42] mito-ai: fix dependencies in package.json and yarn.lock --- mito-ai/package.json | 3 ++- mito-ai/yarn.lock | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/mito-ai/package.json b/mito-ai/package.json index ee66ea1c1..d2f6b9006 100644 --- a/mito-ai/package.json +++ b/mito-ai/package.json @@ -64,7 +64,8 @@ "@types/react-dom": "^18.0.10", "openai": "^4.1.0", "react": "^18.0.0", - "react-dom": "^18.0.0" + "react-dom": "^18.0.0", + "vscode-diff": "^2.1.1" }, "devDependencies": { "@jupyterlab/builder": "^4.0.0", diff --git a/mito-ai/yarn.lock b/mito-ai/yarn.lock index d62259887..bfcbee7d8 100644 --- a/mito-ai/yarn.lock +++ b/mito-ai/yarn.lock @@ -7926,6 +7926,7 @@ __metadata: stylelint-csstree-validator: ^3.0.0 stylelint-prettier: ^4.0.0 typescript: ~5.0.2 + vscode-diff: ^2.1.1 yjs: ^13.5.0 languageName: unknown linkType: soft @@ -10237,6 +10238,13 @@ __metadata: languageName: node linkType: hard +"vscode-diff@npm:^2.1.1": + version: 2.1.1 + resolution: "vscode-diff@npm:2.1.1" + checksum: 17c2493a79d7db1c56683ea8ac14c7999d4dfb6995c3254bca402eedaabf6da7055457c423ade4ed00389983faaf7f0b027ea479e04a040c1e7153b077ae5fbf + languageName: node + linkType: hard + "vscode-jsonrpc@npm:8.2.0": version: 8.2.0 resolution: "vscode-jsonrpc@npm:8.2.0" From 655355e362eae884bdebbb975c5102b4af69afb2 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 12:10:11 -0400 Subject: [PATCH 22/42] mito-ai: fix reject ai code --- .../AiChat/ChatMessage/ChatMessage.tsx | 8 ++- .../AiChat/ChatMessage/CodeBlock.tsx | 15 ++--- .../src/Extensions/AiChat/ChatTaskpane.tsx | 60 ++++++++++++------- 3 files changed, 54 insertions(+), 29 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index 77cf1aea4..0cc081705 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -21,6 +21,8 @@ interface IChatMessageProps { isLastAiMessage: boolean operatingSystem: OperatingSystem setDisplayCodeDiff: React.Dispatch>; + acceptAICode: () => void + rejectAICode: () => void } const ChatMessage: React.FC = ({ @@ -32,7 +34,9 @@ const ChatMessage: React.FC = ({ app, isLastAiMessage, operatingSystem, - setDisplayCodeDiff + setDisplayCodeDiff, + acceptAICode, + rejectAICode }): JSX.Element | null => { if (message.role !== 'user' && message.role !== 'assistant') { // Filter out other types of messages, like system messages @@ -64,6 +68,8 @@ const ChatMessage: React.FC = ({ isLastAiMessage={isLastAiMessage} operatingSystem={operatingSystem} setDisplayCodeDiff={setDisplayCodeDiff} + acceptAICode={acceptAICode} + rejectAICode={rejectAICode} /> ) } diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index a34182832..8bfec282f 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -1,7 +1,7 @@ import React from 'react'; import PythonCode from './PythonCode'; import { INotebookTracker } from '@jupyterlab/notebook'; -import { getNotebookName, writeCodeToActiveCell } from '../../../utils/notebook'; +import { getNotebookName } from '../../../utils/notebook'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import { removeMarkdownCodeFormatting } from '../../../utils/strings'; import { JupyterFrontEnd } from '@jupyterlab/application'; @@ -19,6 +19,8 @@ interface ICodeBlockProps { isLastAiMessage: boolean, operatingSystem: OperatingSystem, setDisplayCodeDiff: React.Dispatch>; + acceptAICode: () => void, + rejectAICode: () => void } const CodeBlock: React.FC = ({ @@ -29,7 +31,9 @@ const CodeBlock: React.FC = ({ app, isLastAiMessage, operatingSystem, - setDisplayCodeDiff + setDisplayCodeDiff, + acceptAICode, + rejectAICode }): JSX.Element => { const notebookName = getNotebookName(notebookTracker) @@ -57,13 +61,10 @@ const CodeBlock: React.FC = ({
{notebookName}
- - diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index fde43a2b4..7ba22f512 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -25,8 +25,6 @@ import { StateEffect, Compartment } from '@codemirror/state'; import { codeDiffStripesExtension } from './CodeDiffDisplay'; - - // IMPORTANT: In order to improve the development experience, we allow you dispaly a // cached conversation as a starting point. Before deploying the mito-ai, we must // set USE_DEV_AI_CONVERSATION = false @@ -67,11 +65,6 @@ interface IChatTaskpaneProps { operatingSystem: OperatingSystem } -// interface IDisplayCodeDiffInfo { -// // In order to display the code diff, I need to merge the old code with the new code -// // then keep track of which lines are deleted and which lines are added/modified -// } - const ChatTaskpane: React.FC = ({ notebookTracker, rendermime, @@ -82,10 +75,13 @@ const ChatTaskpane: React.FC = ({ }) => { const textareaRef = useRef(null); const [chatHistoryManager, setChatHistoryManager] = useState(() => getDefaultChatHistoryManager()); + const chatHistoryManagerRef = useRef(chatHistoryManager); + const [input, setInput] = useState(''); const [loadingAIResponse, setLoadingAIResponse] = useState(false) - const [displayCodeDiff, setDisplayCodeDiff] = useState(undefined) - const chatHistoryManagerRef = useRef(chatHistoryManager); + + const [unifiedDiffLines, setUnifiedDiffLines] = useState(undefined) + const originalDiffedCodeRef = useRef(undefined) useEffect(() => { /* @@ -194,11 +190,14 @@ const ChatTaskpane: React.FC = ({ const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) + // Store the original code so that we can revert to it if the user rejects the AI's code + originalDiffedCodeRef.current = activeCellCode + // Temporarily write the unified code string to the active cell so we can display // the code diffs to the user. Once the user accepts or rejects the code, we'll // apply the correct version of the code. writeCodeToActiveCell(notebookTracker, unifiedCodeString) - setDisplayCodeDiff(unifiedDiffs) + setUnifiedDiffLines(unifiedDiffs) } else { updatedManager.addAIMessageFromMessageContent(apiResponse.errorMessage, true) setChatHistoryManager(updatedManager); @@ -214,16 +213,34 @@ const ChatTaskpane: React.FC = ({ const displayOptimizedChatHistory = chatHistoryManager.getDisplayOptimizedHistory() - const applyLatestCode = () => { + const acceptAICode = () => { const latestChatHistoryManager = chatHistoryManagerRef.current; const lastAIMessage = latestChatHistoryManager.getLastAIMessage() - + if (!lastAIMessage) { return } - const code = getCodeBlockFromMessage(lastAIMessage.message); + const aiGeneratedCode = getCodeBlockFromMessage(lastAIMessage.message); + if (!aiGeneratedCode) { + return + } + + _applyCode(aiGeneratedCode) + } + + const rejectAICode = () => { + const originalDiffedCode = originalDiffedCodeRef.current + if (!originalDiffedCode) { + return + } + _applyCode(originalDiffedCode) + } + + const _applyCode = (code: string) => { writeCodeToActiveCell(notebookTracker, code, true) + setUnifiedDiffLines(undefined) + originalDiffedCodeRef.current = undefined } useEffect(() => { @@ -235,14 +252,13 @@ const ChatTaskpane: React.FC = ({ */ app.commands.addCommand(COMMAND_MITO_AI_APPLY_LATEST_CODE, { execute: () => { - applyLatestCode() - setDisplayCodeDiff(undefined) + acceptAICode() } }) app.commands.addCommand(COMMAND_MITO_AI_REJECT_LATEST_CODE, { execute: () => { - setDisplayCodeDiff(undefined) + rejectAICode() } }) @@ -301,14 +317,14 @@ const ChatTaskpane: React.FC = ({ // Apply the initial configuration editorView.dispatch({ effects: StateEffect.appendConfig.of( - compartment.of(displayCodeDiff !== undefined && isActiveCodeCell? codeDiffStripesExtension({ unifiedDiffLines: displayCodeDiff }) : []) + compartment.of(unifiedDiffLines !== undefined && isActiveCodeCell? codeDiffStripesExtension({ unifiedDiffLines: unifiedDiffLines }) : []) ), }); } else { // Reconfigure the compartment editorView.dispatch({ effects: compartment.reconfigure( - displayCodeDiff !== undefined && isActiveCodeCell ? codeDiffStripesExtension({ unifiedDiffLines: displayCodeDiff }) : [] + unifiedDiffLines !== undefined && isActiveCodeCell ? codeDiffStripesExtension({ unifiedDiffLines: unifiedDiffLines }) : [] ), }); } @@ -317,12 +333,12 @@ const ChatTaskpane: React.FC = ({ } } }); - }, [displayCodeDiff, notebookTracker]); + }, [unifiedDiffLines, notebookTracker]); useEffect(() => { updateCodeCellsExtensions(); - }, [displayCodeDiff, updateCodeCellsExtensions]); + }, [unifiedDiffLines, updateCodeCellsExtensions]); const lastAIMessagesIndex = chatHistoryManager.getLastAIMessageIndex() @@ -351,7 +367,9 @@ const ChatTaskpane: React.FC = ({ app={app} isLastAiMessage={index === lastAIMessagesIndex} operatingSystem={operatingSystem} - setDisplayCodeDiff={setDisplayCodeDiff} + setDisplayCodeDiff={setUnifiedDiffLines} + acceptAICode={acceptAICode} + rejectAICode={rejectAICode} /> ) }).filter(message => message !== null)} From ad3f9992d2787c537598274fbc191f534984da26 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 12:59:41 -0400 Subject: [PATCH 23/42] tests: create mitoai testing structure --- tests/jupyter_utils/jupyterlab_utils.ts | 39 +++++++++++++++++++ .../mitosheet_utils.ts} | 39 +------------------ .../jupyterlab.dataframe-render.spec.ts | 4 +- .../jupyterlab.mitosheet.spec.ts | 6 ++- tests/jupyterlab_ui_tests/mitosheet.spec.ts | 2 +- tests/mac-setup.sh | 9 +++-- tests/mitoai_ui_tests/mitoai.spec.ts | 19 +++++++++ .../notebook.mitosheet.spec.ts | 2 - tests/package.json | 3 +- 9 files changed, 74 insertions(+), 49 deletions(-) create mode 100644 tests/jupyter_utils/jupyterlab_utils.ts rename tests/{jupyterlab_ui_tests/utils.ts => jupyter_utils/mitosheet_utils.ts} (53%) create mode 100644 tests/mitoai_ui_tests/mitoai.spec.ts diff --git a/tests/jupyter_utils/jupyterlab_utils.ts b/tests/jupyter_utils/jupyterlab_utils.ts new file mode 100644 index 000000000..aa7f183a1 --- /dev/null +++ b/tests/jupyter_utils/jupyterlab_utils.ts @@ -0,0 +1,39 @@ +import { IJupyterLabPageFixture } from "@jupyterlab/galata"; + + +export const createAndRunNotebookWithCells = async (page: IJupyterLabPageFixture, cellContents: string[]) => { + const randomFileName = `$test_file_${Math.random().toString(36).substring(2, 15)}.ipynb`; + await page.notebook.createNew(randomFileName); + + for (let i = 0; i < cellContents.length; i++) { + await page.notebook.setCell(i, 'code', cellContents[i]); + await page.notebook.runCell(i); + } + await waitForIdle(page) +} + + +export const waitForIdle = async (page: IJupyterLabPageFixture) => { + const idleLocator = page.locator('#jp-main-statusbar >> text=Idle'); + await idleLocator.waitFor(); +} + +export const waitForCodeToBeWritten = async (page: IJupyterLabPageFixture, cellIndex: number) => { + await waitForIdle(page); + const cellInput = await page.notebook.getCellInput(cellIndex); + let cellCode = (await cellInput?.innerText())?.trim(); + // We wait until there's any code in the cell + while (!/[a-zA-Z]/g.test(cellCode || '')) { + // Wait 20 ms + await page.waitForTimeout(20); + await waitForIdle(page); + const cellInput = await page.notebook.getCellInput(cellIndex); + cellCode = (await cellInput?.innerText())?.trim(); + } +} + +export const typeInNotebookCell = async (page: IJupyterLabPageFixture, cellIndex: number, cellValue: string) => { + await page.locator('.jp-Cell-inputArea').nth(cellIndex).scrollIntoViewIfNeeded(); + await page.notebook.enterCellEditingMode(cellIndex); + await page.notebook.setCell(cellIndex, 'code', cellValue); +} diff --git a/tests/jupyterlab_ui_tests/utils.ts b/tests/jupyter_utils/mitosheet_utils.ts similarity index 53% rename from tests/jupyterlab_ui_tests/utils.ts rename to tests/jupyter_utils/mitosheet_utils.ts index 44c2df57b..390d9d73b 100644 --- a/tests/jupyterlab_ui_tests/utils.ts +++ b/tests/jupyter_utils/mitosheet_utils.ts @@ -1,4 +1,5 @@ import { IJupyterLabPageFixture } from "@jupyterlab/galata"; +import { waitForIdle } from "./jupyterlab_utils"; export const dfCreationCode = `import pandas as pd df = pd.DataFrame({'a': [1], 'b': [4]})\n`; @@ -11,17 +12,6 @@ os.environ['MITO_CONFIG_DISABLE_TOURS'] = 'True' type ToolbarButton = 'Insert' | 'Delete' -export const createAndRunNotebookWithCells = async (page: IJupyterLabPageFixture, cellContents: string[]) => { - const randomFileName = `$test_file_${Math.random().toString(36).substring(2, 15)}.ipynb`; - await page.notebook.createNew(randomFileName); - - for (let i = 0; i < cellContents.length; i++) { - await page.notebook.setCell(i, 'code', cellContents[i]); - await page.notebook.runCell(i); - } - await waitForIdle(page) -} - // If the test just interacts with the mitosheet, and not JupyterLab export const createNewMitosheetOnlyTest = async (page: IJupyterLabPageFixture, firstCellCode: string) => { const randomFileName = `$test_file_${Math.random().toString(36).substring(2, 15)}.ipynb`; @@ -45,35 +35,10 @@ export const getNumberOfColumns = async (page: IJupyterLabPageFixture, cellNumbe return columns?.length || 0; } -export const waitForIdle = async (page: IJupyterLabPageFixture) => { - const idleLocator = page.locator('#jp-main-statusbar >> text=Idle'); - await idleLocator.waitFor(); -} - -export const waitForCodeToBeWritten = async (page: IJupyterLabPageFixture, cellIndex: number) => { - await waitForIdle(page); - const cellInput = await page.notebook.getCellInput(cellIndex); - let cellCode = (await cellInput?.innerText())?.trim(); - // We wait until there's any code in the cell - while (!/[a-zA-Z]/g.test(cellCode || '')) { - // Wait 20 ms - await page.waitForTimeout(20); - await waitForIdle(page); - const cellInput = await page.notebook.getCellInput(cellIndex); - cellCode = (await cellInput?.innerText())?.trim(); - } -} - export const updateCellValue = async (page: IJupyterLabPageFixture, cellValue: string, newCellValue: string) => { await page.locator('.mito-grid-cell', { hasText: cellValue }).scrollIntoViewIfNeeded(); await page.locator('.mito-grid-cell', { hasText: cellValue }).dblclick(); await page.locator('input#cell-editor-input').fill(newCellValue); await page.keyboard.press('Enter'); await waitForIdle(page); -}; - -export const typeInNotebookCell = async (page: IJupyterLabPageFixture, cellIndex: number, cellValue: string) => { - await page.locator('.jp-Cell-inputArea').nth(cellIndex).scrollIntoViewIfNeeded(); - await page.notebook.enterCellEditingMode(cellIndex); - await page.notebook.setCell(cellIndex, 'code', cellValue); -} +}; \ No newline at end of file diff --git a/tests/jupyterlab_ui_tests/jupyterlab.dataframe-render.spec.ts b/tests/jupyterlab_ui_tests/jupyterlab.dataframe-render.spec.ts index 80d9f9be9..08ae12a52 100644 --- a/tests/jupyterlab_ui_tests/jupyterlab.dataframe-render.spec.ts +++ b/tests/jupyterlab_ui_tests/jupyterlab.dataframe-render.spec.ts @@ -1,6 +1,6 @@ import { expect, test } from '@jupyterlab/galata'; -import { createAndRunNotebookWithCells, updateCellValue, waitForIdle } from './utils'; - +import { createAndRunNotebookWithCells, waitForIdle } from '../jupyter_utils/jupyterlab_utils'; +import { updateCellValue } from '../jupyter_utils/mitosheet_utils'; const placeholderCellText = '# Empty code cell'; test.describe.configure({ mode: 'parallel' }); diff --git a/tests/jupyterlab_ui_tests/jupyterlab.mitosheet.spec.ts b/tests/jupyterlab_ui_tests/jupyterlab.mitosheet.spec.ts index 8237076f2..1a9372cfa 100644 --- a/tests/jupyterlab_ui_tests/jupyterlab.mitosheet.spec.ts +++ b/tests/jupyterlab_ui_tests/jupyterlab.mitosheet.spec.ts @@ -1,5 +1,7 @@ -import { IJupyterLabPageFixture, expect, test } from '@jupyterlab/galata'; -import { TURN_OFF_TOURS, clickToolbarButton, createAndRunNotebookWithCells, dfCreationCode, getNumberOfColumns, typeInNotebookCell, updateCellValue, waitForCodeToBeWritten, waitForIdle } from './utils'; +import { expect, test } from '@jupyterlab/galata'; +import { TURN_OFF_TOURS, clickToolbarButton, updateCellValue, dfCreationCode, getNumberOfColumns } from '../jupyter_utils/mitosheet_utils'; +import { waitForIdle, createAndRunNotebookWithCells, typeInNotebookCell, waitForCodeToBeWritten } from '../jupyter_utils/jupyterlab_utils'; + import { Page } from '@playwright/test'; test.describe.configure({ mode: 'parallel' }); diff --git a/tests/jupyterlab_ui_tests/mitosheet.spec.ts b/tests/jupyterlab_ui_tests/mitosheet.spec.ts index ddb1f3e0c..c3be320af 100644 --- a/tests/jupyterlab_ui_tests/mitosheet.spec.ts +++ b/tests/jupyterlab_ui_tests/mitosheet.spec.ts @@ -1,5 +1,5 @@ import { expect, test } from '@jupyterlab/galata'; -import { createNewMitosheetOnlyTest, dfCreationCode, getNumberOfColumns } from './utils'; +import { createNewMitosheetOnlyTest, dfCreationCode, getNumberOfColumns } from '../jupyter_utils/mitosheet_utils'; test.describe.configure({ mode: 'parallel' }); diff --git a/tests/mac-setup.sh b/tests/mac-setup.sh index ab682fa5b..88798ecf3 100644 --- a/tests/mac-setup.sh +++ b/tests/mac-setup.sh @@ -13,10 +13,11 @@ jlpm install # install only the necessary browsers. if [ $# -eq 0 ] then - npx playwright install chromium webkit firefox chrome -else - npx playwright install $1 - npx playwright install + npx playwright install chromium webkit firefox || echo "Warning: Failed to install some browsers" + npx playwright install chrome || echo "Warning: Failed to install Chrome" + else + npx playwright install $1 || echo "Warning: Failed to install specified browser" + npx playwright install || echo "Warning: Failed to install additional browsers" fi # Install mitosheet diff --git a/tests/mitoai_ui_tests/mitoai.spec.ts b/tests/mitoai_ui_tests/mitoai.spec.ts new file mode 100644 index 000000000..7f424fab0 --- /dev/null +++ b/tests/mitoai_ui_tests/mitoai.spec.ts @@ -0,0 +1,19 @@ +import { expect, test } from '@jupyterlab/galata'; +import { createAndRunNotebookWithCells, waitForIdle } from '../jupyter_utils/jupyterlab_utils'; +import { updateCellValue } from '../jupyter_utils/mitosheet_utils'; +const placeholderCellText = '# Empty code cell'; + +test.describe.configure({ mode: 'parallel' }); + +test.describe('Dataframe renders as mitosheet', () => { + test('renders a mitosheet when hanging dataframe', async ({ page, tmpPath }) => { + + await createAndRunNotebookWithCells(page, ['import pandas as pd\ndf=pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})\ndf']); + await waitForIdle(page); + const cellOutput = await page.notebook.getCellOutput(0); + expect(await cellOutput?.innerHTML()).toContain('Home'); + + // The toolbar should be collapsed by default, so the Delete button should not be visible + expect(await cellOutput?.innerHTML()).not.toContain('Delete'); + }); +}); \ No newline at end of file diff --git a/tests/notebook_ui_tests/notebook.mitosheet.spec.ts b/tests/notebook_ui_tests/notebook.mitosheet.spec.ts index 720fa9836..66e313763 100644 --- a/tests/notebook_ui_tests/notebook.mitosheet.spec.ts +++ b/tests/notebook_ui_tests/notebook.mitosheet.spec.ts @@ -1,6 +1,4 @@ import { expect, test } from '@playwright/test'; -import { awaitResponse } from '../streamlit_ui_tests/utils'; -import { waitForIdle } from '../jupyterlab_ui_tests/utils'; test.describe('Mitosheet Jupyter Notebook integration', () => { test('renders a mitosheet.sheet() and generates code', async ({ page }) => { diff --git a/tests/package.json b/tests/package.json index 2f712e53e..f795e1b75 100644 --- a/tests/package.json +++ b/tests/package.json @@ -10,7 +10,8 @@ "test:streamlit:headed": "playwright test streamlit_ui_tests --headed", "test:dash": "playwright test dash_ui_tests --retries=3 --project=\"chromium\"", "test:jupyterlab": "playwright test jupyterlab_ui_tests --retries=3 --project=\"chromium\"", - "test:notebook": "playwright test notebook_ui_tests --retries=3 --project=\"chromium\"" + "test:notebook": "playwright test notebook_ui_tests --retries=3 --project=\"chromium\"", + "test:mitoai": "playwright test mitoai_ui_tests --retries=3 --project=\"chromium\"" }, "author": "", "license": "ISC", From c97b51596861e60a95c957c2de9d47067f8ae4d9 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 13:42:13 -0400 Subject: [PATCH 24/42] tests: build first working mito ai test --- tests/README.md | 48 ++++++++++++++++++------- tests/jupyter_utils/jupyterlab_utils.ts | 5 +++ tests/mitoai_ui_tests/mitoai.spec.ts | 45 ++++++++++++++++++----- 3 files changed, 77 insertions(+), 21 deletions(-) diff --git a/tests/README.md b/tests/README.md index 3b8107ea4..d18fd0c3d 100644 --- a/tests/README.md +++ b/tests/README.md @@ -27,18 +27,6 @@ And then run npm run test:streamlit -- --project=chromium ``` -### Dash Specific Tests - -Run -``` -python dash-test.py -``` - -Then, run the tests with -``` -npm run test:dash -- --project=chromium -``` - ### Jupyter Specific Tests First, you need to run (from the mitosheet/ directory) in this same virtual environment:: @@ -66,6 +54,42 @@ And then in a separate terminal run from the tests/ directory: npm run test:notebook -- --project=chromium ``` +### Mito AI Specific Tests + +First, you need to run (from the mito-ai/ directory) in this same virtual environment: + +``` +jupyter labextension develop . --overwrite +jupyter server extension enable --py mito-ai +``` + +Then, from the tests/directory, set your OPENAI_API_KEY environment variable: +``` +export OPENAI_API_KEY= +``` + +From the same terminal, run: +``` +jupyter lab --config jupyter_server_test_config.py +``` + +And then in a separate terminal run from the tests/ directory: +``` +npm run test:mitoai -- --project=chromium +``` + +### Dash Specific Tests + +Run +``` +python dash-test.py +``` + +Then, run the tests with +``` +npm run test:dash -- --project=chromium +``` + Add a `--headed` flag to see the test run. diff --git a/tests/jupyter_utils/jupyterlab_utils.ts b/tests/jupyter_utils/jupyterlab_utils.ts index aa7f183a1..6902f74e6 100644 --- a/tests/jupyter_utils/jupyterlab_utils.ts +++ b/tests/jupyter_utils/jupyterlab_utils.ts @@ -37,3 +37,8 @@ export const typeInNotebookCell = async (page: IJupyterLabPageFixture, cellIndex await page.notebook.enterCellEditingMode(cellIndex); await page.notebook.setCell(cellIndex, 'code', cellValue); } + +export const getCodeFromCell = async (page: IJupyterLabPageFixture, cellIndex: number) => { + const cellInput = await page.notebook.getCellInput(cellIndex); + return await cellInput?.innerText(); +} \ No newline at end of file diff --git a/tests/mitoai_ui_tests/mitoai.spec.ts b/tests/mitoai_ui_tests/mitoai.spec.ts index 7f424fab0..d35f748a6 100644 --- a/tests/mitoai_ui_tests/mitoai.spec.ts +++ b/tests/mitoai_ui_tests/mitoai.spec.ts @@ -1,19 +1,46 @@ import { expect, test } from '@jupyterlab/galata'; -import { createAndRunNotebookWithCells, waitForIdle } from '../jupyter_utils/jupyterlab_utils'; +import { createAndRunNotebookWithCells, getCodeFromCell, waitForIdle } from '../jupyter_utils/jupyterlab_utils'; import { updateCellValue } from '../jupyter_utils/mitosheet_utils'; const placeholderCellText = '# Empty code cell'; test.describe.configure({ mode: 'parallel' }); -test.describe('Dataframe renders as mitosheet', () => { - test('renders a mitosheet when hanging dataframe', async ({ page, tmpPath }) => { +test.describe('Mito AI Chat', () => { - await createAndRunNotebookWithCells(page, ['import pandas as pd\ndf=pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})\ndf']); + test('Apply AI Generated Code', async ({ page }) => { + await createAndRunNotebookWithCells(page, ['import pandas as pd\ndf=pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})']); await waitForIdle(page); - const cellOutput = await page.notebook.getCellOutput(0); - expect(await cellOutput?.innerHTML()).toContain('Home'); - // The toolbar should be collapsed by default, so the Delete button should not be visible - expect(await cellOutput?.innerHTML()).not.toContain('Delete'); + await page.getByRole('tab', { name: 'AI Chat for your JupyterLab' }).getByRole('img').click(); + expect (page.getByPlaceholder('Ask your personal Python')).toBeVisible(); + + await page.getByPlaceholder('Ask your personal Python').fill('Write the code df["C"] = [7, 8, 9]'); + await page.keyboard.press('Enter'); + await waitForIdle(page); + + await page.getByRole('button', { name: 'Apply' }).click(); + await waitForIdle(page); + + const code = await getCodeFromCell(page, 1); + expect(code).toContain('df["C"] = [7, 8, 9]'); }); -}); \ No newline at end of file + + test('Reject AI Generated Code', async ({ page }) => { + await createAndRunNotebookWithCells(page, ['import pandas as pd\ndf=pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})']); + await waitForIdle(page); + + await page.getByRole('tab', { name: 'AI Chat for your JupyterLab' }).getByRole('img').click(); + expect (page.getByPlaceholder('Ask your personal Python')).toBeVisible(); + + await page.getByPlaceholder('Ask your personal Python').fill('Write the code df["C"] = [7, 8, 9]'); + await page.keyboard.press('Enter'); + await waitForIdle(page); + + await page.getByRole('button', { name: 'Deny' }).click(); + await waitForIdle(page); + + const code = await getCodeFromCell(page, 1); + expect(code).not.toContain('df["C"] = [7, 8, 9]'); + }); +}); + From 16798e8cfee2904c5fa184d8f9a67695eefbbe66 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 13:54:51 -0400 Subject: [PATCH 25/42] mito-ai: fix bug in deny ai generated code + add test --- mito-ai/src/Extensions/AiChat/AiChatPlugin.ts | 6 ++---- mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx | 13 +++++-------- mito-ai/src/Extensions/AiChat/ChatWidget.tsx | 3 --- tests/mitoai_ui_tests/mitoai.spec.ts | 1 + 4 files changed, 8 insertions(+), 15 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts b/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts index c48fad276..725fc70a6 100644 --- a/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts +++ b/mito-ai/src/Extensions/AiChat/AiChatPlugin.ts @@ -8,7 +8,6 @@ import { INotebookTracker } from '@jupyterlab/notebook'; import { buildChatWidget } from './ChatWidget'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; -import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; import { COMMAND_MITO_AI_OPEN_CHAT } from '../../commands'; @@ -19,7 +18,7 @@ const AiChatPlugin: JupyterFrontEndPlugin = { id: 'mito_ai:plugin', description: 'AI chat for JupyterLab', autoStart: true, - requires: [INotebookTracker, ICommandPalette, IRenderMimeRegistry, IVariableManager, IEditorExtensionRegistry], + requires: [INotebookTracker, ICommandPalette, IRenderMimeRegistry, IVariableManager], optional: [ILayoutRestorer], activate: ( app: JupyterFrontEnd, @@ -27,7 +26,6 @@ const AiChatPlugin: JupyterFrontEndPlugin = { palette: ICommandPalette, rendermime: IRenderMimeRegistry, variableManager: IVariableManager, - editorExtensionRegistry: IEditorExtensionRegistry, restorer: ILayoutRestorer | null ) => { @@ -35,7 +33,7 @@ const AiChatPlugin: JupyterFrontEndPlugin = { // then call it to make a new widget const newWidget = () => { // Create a blank content widget inside of a MainAreaWidget - const chatWidget = buildChatWidget(app, notebookTracker, rendermime, variableManager, editorExtensionRegistry) + const chatWidget = buildChatWidget(app, notebookTracker, rendermime, variableManager) return chatWidget } diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 7ba22f512..31d80d924 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -18,7 +18,6 @@ import ResetIcon from '../../icons/ResetIcon'; import IconButton from '../../components/IconButton'; import { OperatingSystem } from '../../utils/user'; import { getCodeDiffsAndUnifiedCodeString, UnifiedDiffLine } from '../../utils/codeDiff'; -import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; import { CodeMirrorEditor } from '@jupyterlab/codemirror'; import { CodeCell } from '@jupyterlab/cells'; import { StateEffect, Compartment } from '@codemirror/state'; @@ -60,7 +59,6 @@ interface IChatTaskpaneProps { notebookTracker: INotebookTracker rendermime: IRenderMimeRegistry variableManager: IVariableManager - editorExtensionRegistry: IEditorExtensionRegistry app: JupyterFrontEnd operatingSystem: OperatingSystem } @@ -69,7 +67,6 @@ const ChatTaskpane: React.FC = ({ notebookTracker, rendermime, variableManager, - editorExtensionRegistry, app, operatingSystem }) => { @@ -81,7 +78,7 @@ const ChatTaskpane: React.FC = ({ const [loadingAIResponse, setLoadingAIResponse] = useState(false) const [unifiedDiffLines, setUnifiedDiffLines] = useState(undefined) - const originalDiffedCodeRef = useRef(undefined) + const originalCodeBeforeDiff = useRef(undefined) useEffect(() => { /* @@ -191,7 +188,7 @@ const ChatTaskpane: React.FC = ({ const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) // Store the original code so that we can revert to it if the user rejects the AI's code - originalDiffedCodeRef.current = activeCellCode + originalCodeBeforeDiff.current = activeCellCode || '' // Temporarily write the unified code string to the active cell so we can display // the code diffs to the user. Once the user accepts or rejects the code, we'll @@ -230,8 +227,8 @@ const ChatTaskpane: React.FC = ({ } const rejectAICode = () => { - const originalDiffedCode = originalDiffedCodeRef.current - if (!originalDiffedCode) { + const originalDiffedCode = originalCodeBeforeDiff.current + if (originalDiffedCode === undefined) { return } _applyCode(originalDiffedCode) @@ -240,7 +237,7 @@ const ChatTaskpane: React.FC = ({ const _applyCode = (code: string) => { writeCodeToActiveCell(notebookTracker, code, true) setUnifiedDiffLines(undefined) - originalDiffedCodeRef.current = undefined + originalCodeBeforeDiff.current = undefined } useEffect(() => { diff --git a/mito-ai/src/Extensions/AiChat/ChatWidget.tsx b/mito-ai/src/Extensions/AiChat/ChatWidget.tsx index 1674508ea..129903e70 100644 --- a/mito-ai/src/Extensions/AiChat/ChatWidget.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatWidget.tsx @@ -8,7 +8,6 @@ import chatIconSvg from '../../../src/icons/ChatIcon.svg' import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { getOperatingSystem } from '../../utils/user'; -import { IEditorExtensionRegistry } from '@jupyterlab/codemirror'; export const chatIcon = new LabIcon({ name: 'mito_ai', @@ -20,7 +19,6 @@ export function buildChatWidget( notebookTracker: INotebookTracker, rendermime: IRenderMimeRegistry, variableManager: IVariableManager, - editorExtensionRegistry: IEditorExtensionRegistry ) { // Get the operating system here so we don't have to do it each time the chat changes. @@ -34,7 +32,6 @@ export function buildChatWidget( rendermime={rendermime} variableManager={variableManager} operatingSystem={operatingSystem} - editorExtensionRegistry={editorExtensionRegistry} /> ) chatWidget.id = 'mito_ai'; diff --git a/tests/mitoai_ui_tests/mitoai.spec.ts b/tests/mitoai_ui_tests/mitoai.spec.ts index d35f748a6..ef1fb3af7 100644 --- a/tests/mitoai_ui_tests/mitoai.spec.ts +++ b/tests/mitoai_ui_tests/mitoai.spec.ts @@ -41,6 +41,7 @@ test.describe('Mito AI Chat', () => { const code = await getCodeFromCell(page, 1); expect(code).not.toContain('df["C"] = [7, 8, 9]'); + expect(code?.trim()).toBe("") }); }); From ba0cfa40637b8d5a13512273bc1d2115c5c338b3 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 14:32:34 -0400 Subject: [PATCH 26/42] tests: add more tests --- .../CellToolbarButtonsPlugin.tsx | 3 +- tests/jupyter_utils/jupyterlab_utils.ts | 5 ++ tests/mitoai_ui_tests/mitoai.spec.ts | 67 +++++++++++++++---- tests/mitoai_ui_tests/utils.ts | 17 +++++ 4 files changed, 78 insertions(+), 14 deletions(-) create mode 100644 tests/mitoai_ui_tests/utils.ts diff --git a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx index ef6bb877d..0224a1d6a 100644 --- a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx +++ b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx @@ -3,7 +3,7 @@ import { JupyterFrontEndPlugin } from '@jupyterlab/application'; import { INotebookTracker } from '@jupyterlab/notebook'; -import { COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; +import { COMMAND_MITO_AI_OPEN_CHAT, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; import { LabIcon } from '@jupyterlab/ui-components'; import LightbulbIcon from '../../../src/icons/LightbulbIcon.svg' @@ -37,6 +37,7 @@ const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { update the AI Chat History Manager so that we can generate an AI optimzied message and a display optimized message. */ + app.commands.execute(COMMAND_MITO_AI_OPEN_CHAT) app.commands.execute(COMMAND_MITO_AI_SEND_MESSAGE, { input: `Explain this code` }); }, isVisible: () => notebookTracker.activeCell?.model.type === 'code' && app.commands.hasCommand(COMMAND_MITO_AI_SEND_MESSAGE) diff --git a/tests/jupyter_utils/jupyterlab_utils.ts b/tests/jupyter_utils/jupyterlab_utils.ts index 6902f74e6..3a162db42 100644 --- a/tests/jupyter_utils/jupyterlab_utils.ts +++ b/tests/jupyter_utils/jupyterlab_utils.ts @@ -41,4 +41,9 @@ export const typeInNotebookCell = async (page: IJupyterLabPageFixture, cellIndex export const getCodeFromCell = async (page: IJupyterLabPageFixture, cellIndex: number) => { const cellInput = await page.notebook.getCellInput(cellIndex); return await cellInput?.innerText(); +} + +export const selectCell = async (page: IJupyterLabPageFixture, cellIndex: number) => { + const cell = await page.notebook.getCell(cellIndex); + await cell?.click(); } \ No newline at end of file diff --git a/tests/mitoai_ui_tests/mitoai.spec.ts b/tests/mitoai_ui_tests/mitoai.spec.ts index ef1fb3af7..989ed19c5 100644 --- a/tests/mitoai_ui_tests/mitoai.spec.ts +++ b/tests/mitoai_ui_tests/mitoai.spec.ts @@ -1,6 +1,7 @@ import { expect, test } from '@jupyterlab/galata'; -import { createAndRunNotebookWithCells, getCodeFromCell, waitForIdle } from '../jupyter_utils/jupyterlab_utils'; +import { createAndRunNotebookWithCells, getCodeFromCell, selectCell, waitForIdle } from '../jupyter_utils/jupyterlab_utils'; import { updateCellValue } from '../jupyter_utils/mitosheet_utils'; +import { sendMessageToMitoAI, waitForMitoAILoadingToDisappear } from './utils'; const placeholderCellText = '# Empty code cell'; test.describe.configure({ mode: 'parallel' }); @@ -11,12 +12,7 @@ test.describe('Mito AI Chat', () => { await createAndRunNotebookWithCells(page, ['import pandas as pd\ndf=pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})']); await waitForIdle(page); - await page.getByRole('tab', { name: 'AI Chat for your JupyterLab' }).getByRole('img').click(); - expect (page.getByPlaceholder('Ask your personal Python')).toBeVisible(); - - await page.getByPlaceholder('Ask your personal Python').fill('Write the code df["C"] = [7, 8, 9]'); - await page.keyboard.press('Enter'); - await waitForIdle(page); + await sendMessageToMitoAI(page, 'Write the code df["C"] = [7, 8, 9]'); await page.getByRole('button', { name: 'Apply' }).click(); await waitForIdle(page); @@ -29,12 +25,7 @@ test.describe('Mito AI Chat', () => { await createAndRunNotebookWithCells(page, ['import pandas as pd\ndf=pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})']); await waitForIdle(page); - await page.getByRole('tab', { name: 'AI Chat for your JupyterLab' }).getByRole('img').click(); - expect (page.getByPlaceholder('Ask your personal Python')).toBeVisible(); - - await page.getByPlaceholder('Ask your personal Python').fill('Write the code df["C"] = [7, 8, 9]'); - await page.keyboard.press('Enter'); - await waitForIdle(page); + await sendMessageToMitoAI(page, 'Write the code df["C"] = [7, 8, 9]'); await page.getByRole('button', { name: 'Deny' }).click(); await waitForIdle(page); @@ -43,5 +34,55 @@ test.describe('Mito AI Chat', () => { expect(code).not.toContain('df["C"] = [7, 8, 9]'); expect(code?.trim()).toBe("") }); + + test('Code Diffs are applied', async ({ page }) => { + await createAndRunNotebookWithCells(page, ['print(1)']); + await waitForIdle(page); + + await sendMessageToMitoAI(page, 'Remove the code print(1) and add the code print(2)', 0); + + await expect(page.locator('.cm-codeDiffRemovedStripe')).toBeVisible(); + await expect(page.locator('.cm-codeDiffInsertedStripe')).toBeVisible(); + }); + + test('No Code blocks are displayed when active cell is empty', async ({ page }) => { + await createAndRunNotebookWithCells(page, []); + await waitForIdle(page); + + await sendMessageToMitoAI(page, 'Add print (1)'); + + // Since the active cell is empty, there should only be one code message part container. + // It should be in the AI response message, which means that it is not in the user's message. + const codeMessagePartContainersCount = await page.locator('.code-message-part-container').count(); + expect(codeMessagePartContainersCount).toBe(1); + }); + + test('Errors have fix with AI button', async ({ page }) => { + await createAndRunNotebookWithCells(page, ['print(1']); + await waitForIdle(page); + + await page.getByRole('button', { name: 'Fix Error in AI Chat' }).click(); + await waitForIdle(page); + + await waitForMitoAILoadingToDisappear(page); + + await page.getByRole('button', { name: 'Apply' }).click(); + await waitForIdle(page); + + const code = await getCodeFromCell(page, 0); + expect(code).toContain('print(1)'); + }); + + test('Code cells have Explain Code button', async ({ page }) => { + await createAndRunNotebookWithCells(page, ['print(1)']); + await waitForIdle(page); + + await selectCell(page, 0); + await page.getByRole('button', { name: 'Explain code in AI Chat' }).click(); + + // Check that the message "Explain this code" exists in the AI chat + await expect(page.getByText('Explain this code')).toBeVisible(); + + }); }); diff --git a/tests/mitoai_ui_tests/utils.ts b/tests/mitoai_ui_tests/utils.ts new file mode 100644 index 000000000..9ba41a61c --- /dev/null +++ b/tests/mitoai_ui_tests/utils.ts @@ -0,0 +1,17 @@ +import { IJupyterLabPageFixture } from "@jupyterlab/galata"; +import { selectCell } from "../jupyter_utils/jupyterlab_utils"; + +export const waitForMitoAILoadingToDisappear = async (page: IJupyterLabPageFixture) => { + const mitoAILoadingLocator = page.locator('.chat-loading-message'); + await mitoAILoadingLocator.waitFor({ state: 'hidden' }); +} + +export const sendMessageToMitoAI = async (page: IJupyterLabPageFixture, message: string, activeCellIndex?: number) => { + if (activeCellIndex) { + await selectCell(page, activeCellIndex); + } + await page.getByRole('tab', { name: 'AI Chat for your JupyterLab' }).getByRole('img').click(); + await page.getByPlaceholder('Ask your personal Python').fill(message); + await page.keyboard.press('Enter'); + await waitForMitoAILoadingToDisappear(page); +} \ No newline at end of file From fdceb208b06d3307479b73b437661a4122715ab5 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 14:36:43 -0400 Subject: [PATCH 27/42] .github: update testing workflow for mito-ai --- .github/workflows/test-mitosheet-frontend.yml | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/.github/workflows/test-mitosheet-frontend.yml b/.github/workflows/test-mitosheet-frontend.yml index e23bb5695..a5da6e6b6 100644 --- a/.github/workflows/test-mitosheet-frontend.yml +++ b/.github/workflows/test-mitosheet-frontend.yml @@ -113,6 +113,58 @@ jobs: path: tests/playwright-report/ retention-days: 14 + test-mitoai-frontend-jupyterlab: + runs-on: ubuntu-20.04 + timeout-minutes: 60 + strategy: + matrix: + python-version: ['3.8', '3.10', '3.11'] + fail-fast: false + + steps: + - name: Cancel Previous Runs + uses: styfle/cancel-workflow-action@0.7.0 + with: + access_token: ${{ github.token }} + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: | + mito-ai/setup.py + tests/requirements.txt + - uses: actions/setup-node@v3 + with: + node-version: 18 + cache: 'npm' + cache-dependency-path: mito-ai/package-lock.json + - name: Install dependencies + run: | + cd tests + bash mac-setup.sh + - name: Setup JupyterLab + run: | + cd tests + source venv/bin/activate + cd ../mito-ai + jupyter labextension develop . --overwrite + jupyter server extension enable --py mito-ai + - name: Start a server and run tests + run: | + cd tests + source venv/bin/activate + jupyter lab --config jupyter_server_test_config.py & + npm run test:mitoai + - name: Upload test-results + uses: actions/upload-artifact@v3 + if: failure() + with: + name: mitoai-jupyterlab-playwright-report-${{ matrix.python-version }} + path: tests/playwright-report/ + retention-days: 14 + test-mitosheet-frontend-streamlit: timeout-minutes: 60 strategy: From f437a1b0a9d575e64567518bd6f21f3f1fabefb2 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 14:50:12 -0400 Subject: [PATCH 28/42] tests: add mito-ai to setup scripts --- tests/mac-setup.sh | 11 +++++++---- tests/windows-setup.bat | 8 ++++++-- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/mac-setup.sh b/tests/mac-setup.sh index 88798ecf3..b0054b576 100644 --- a/tests/mac-setup.sh +++ b/tests/mac-setup.sh @@ -20,12 +20,15 @@ if [ $# -eq 0 ] npx playwright install || echo "Warning: Failed to install additional browsers" fi -# Install mitosheet -cd ../mitosheet -# Install Python dependencies +# Install mitosheet and build JS +cd ../mitosheet pip install -e ".[test]" +jlpm install +jlpm run build -# Install the npm dependences for Mitosheet, and build JS +# Install mito-ai and build JS +cd ../mito-ai +pip install -e ".[test]" jlpm install jlpm run build \ No newline at end of file diff --git a/tests/windows-setup.bat b/tests/windows-setup.bat index 50968951f..959c89e92 100644 --- a/tests/windows-setup.bat +++ b/tests/windows-setup.bat @@ -16,10 +16,14 @@ REM install only the necessary browsers. if [%1]==[] npx playwright install chromium firefox chrome msedge else npx playwright install %1 -REM Navigate to the mitosheet directory and install dependencies +REM Navigate to the mitosheet directory and install dependencies + build JS cd ../mitosheet pip install -e ".[test]" +jlpm install +jlpm run build -REM Install and build npm packages +REM Navigate to the mito-ai directory and install dependencies + build JS +cd ../mito-a +pip install -e ".[test]" jlpm install jlpm run build \ No newline at end of file From 90bb7a7fea20019728141196d911a8aa199b35ce Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 15:01:32 -0400 Subject: [PATCH 29/42] .github: add openai key env var --- .github/workflows/test-mitosheet-frontend.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test-mitosheet-frontend.yml b/.github/workflows/test-mitosheet-frontend.yml index a5da6e6b6..b4ab9a7ac 100644 --- a/.github/workflows/test-mitosheet-frontend.yml +++ b/.github/workflows/test-mitosheet-frontend.yml @@ -157,6 +157,8 @@ jobs: source venv/bin/activate jupyter lab --config jupyter_server_test_config.py & npm run test:mitoai + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} - name: Upload test-results uses: actions/upload-artifact@v3 if: failure() From 599e549b3394fb85a27663b90d8cd3cbf4ae4d07 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 17:06:42 -0400 Subject: [PATCH 30/42] mito-ai: create separate prompt for each type of interaction --- .../Extensions/AiChat/ChatHistoryManager.tsx | 148 +++++++++++++--- .../AiChat/ChatMessage/ChatMessage.tsx | 7 +- .../src/Extensions/AiChat/ChatTaskpane.tsx | 165 +++++++++--------- .../CellToolbarButtonsPlugin.tsx | 12 +- .../ErrorMimeRendererPlugin.tsx | 4 +- mito-ai/src/commands.tsx | 4 +- 6 files changed, 215 insertions(+), 125 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index ca7f6f238..6034ec7ab 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -1,9 +1,11 @@ import OpenAI from "openai"; -import { Variable } from "../VariableManager/VariableInspector"; +import { IVariableManager } from "../VariableManager/VariableManagerPlugin"; +import { INotebookTracker } from '@jupyterlab/notebook'; +import { getActiveCellCode } from "../../utils/notebook"; export interface IDisplayOptimizedChatHistory { message: OpenAI.Chat.ChatCompletionMessageParam - error: boolean + type: 'openai message' | 'connection error' } export interface IChatHistory { @@ -33,12 +35,25 @@ export interface IChatHistory { */ export class ChatHistoryManager { private history: IChatHistory; + private variableManager: IVariableManager; + private notebookTracker: INotebookTracker; - constructor(initialHistory?: IChatHistory) { + constructor(variableManager: IVariableManager, notebookTracker: INotebookTracker, initialHistory?: IChatHistory) { + // Initialize the history this.history = initialHistory || { aiOptimizedChatHistory: [], displayOptimizedChatHistory: [] }; + + // Save the variable manager + this.variableManager = variableManager; + + // Save the notebook tracker + this.notebookTracker = notebookTracker; + } + + createDuplicateChatHistoryManager(): ChatHistoryManager { + return new ChatHistoryManager(this.variableManager, this.notebookTracker, this.history); } getHistory(): IChatHistory { @@ -53,17 +68,13 @@ export class ChatHistoryManager { return this.history.displayOptimizedChatHistory; } - addUserMessage(input: string, activeCellCode?: string, variables?: Variable[]): void { + addGenericUserPromptedMessage(input: string): void { - const displayMessage: OpenAI.Chat.ChatCompletionMessageParam = { - role: 'user', - content: `\`\`\`python -${activeCellCode} -\`\`\` -${input}`}; + const variables = this.variableManager.variables + const activeCellCode = getActiveCellCode(this.notebookTracker) - const aiMessage: OpenAI.Chat.ChatCompletionMessageParam = { + const aiOptimizedMessage: OpenAI.Chat.ChatCompletionMessageParam = { role: 'user', content: `You have access to the following variables: @@ -107,16 +118,89 @@ Use the pd.to_datetime function to convert the issue_date column to datetime. - +Code in the active code cell: + +\`\`\`python +${activeCellCode} +\`\`\` + +Your task: ${input}`}; + + this.history.displayOptimizedChatHistory.push({message: getDisplayedOptimizedUserMessage(input, activeCellCode), type: 'openai message'}); + this.history.aiOptimizedChatHistory.push(aiOptimizedMessage); + } + + addDebugErrorMessage(errorMessage: string): void { + + const activeCellCode = getActiveCellCode(this.notebookTracker) + + const aiOptimizedPrompt = `You just ran the active code cell and received an error. Return the full code cell with the error corrected and a short explanation of the error. + Remember that you are executing code inside a Jupyter notebook. That means you will have persistent state issues where variables from previous cells or previous code executions might still affect current code. When those errors occur, here are a few possible solutions: 1. Restarting the kernel to reset the environment if a function or variable has been unintentionally overwritten. 2. Identify which cell might need to be rerun to properly initialize the function or variable that is causing the issue. - + For example, if an error occurs because the built-in function 'print' is overwritten by an integer, you should return the code cell with the modification to the print function removed and also return an explanation that tell the user to restart their kernel. Do not add new comments to the code cell, just return the code cell with the modification removed. - + When a user hits an error because of a persistent state issue, tell them how to resolve it. - + + +Code in the active code cell: + +\`\`\`python +print(y) +\`\`\` + +Error Message: +NameError: name 'y' is not defined + +Output: + +\`\`\`python +y = 10 +print(y) +\`\`\` + +The variable y has not yet been created.Define the variable y before printing it. + + +Code in the active code cell: + +\`\`\`python +${activeCellCode} +\`\`\` + +Error Message: + +${errorMessage} + +Output: +` + + this.history.displayOptimizedChatHistory.push({message: getDisplayedOptimizedUserMessage(errorMessage, activeCellCode), type: 'openai message'}); + this.history.aiOptimizedChatHistory.push({role: 'user', content: aiOptimizedPrompt}); + } + + addExplainCodeMessage(): void { + const activeCellCode = getActiveCellCode(this.notebookTracker) + + const aiOptimizedPrompt = `Explain the code in the active code cell to me like I have a basic understanding of Python. Don't explain each line, but instead explain the overall logic of the code. + + + +Code in the active code cell: + +\`\`\`python +def multiply(x, y): + return x * y +\`\`\` + +Output: + +This code creates a function called \`multiply\` that takes two arguments \`x\` and \`y\`, and returns the product of \`x\` and \`y\`. + + Code in the active code cell: @@ -124,13 +208,16 @@ Code in the active code cell: ${activeCellCode} \`\`\` -Your task: ${input}`}; +Output: +` - this.history.displayOptimizedChatHistory.push({message: displayMessage, error: false}); - this.history.aiOptimizedChatHistory.push(aiMessage); + this.history.displayOptimizedChatHistory.push({message: getDisplayedOptimizedUserMessage('Explain this code', activeCellCode), type: 'openai message'}); + this.history.aiOptimizedChatHistory.push({role: 'user', content: aiOptimizedPrompt}); } - addAIMessageFromResponse(message: OpenAI.Chat.Completions.ChatCompletionMessage, error: boolean=false): void { + + + addAIMessageFromResponse(message: OpenAI.Chat.Completions.ChatCompletionMessage, mitoAIConnectionError: boolean=false): void { if (message.content === null) { return } @@ -139,19 +226,19 @@ Your task: ${input}`}; role: 'assistant', content: message.content } - this._addAIMessage(aiMessage, error) + this._addAIMessage(aiMessage, mitoAIConnectionError) } - addAIMessageFromMessageContent(message: string, error: boolean=false): void { + addAIMessageFromMessageContent(message: string, mitoAIConnectionError: boolean=false): void { const aiMessage: OpenAI.Chat.ChatCompletionMessageParam = { role: 'assistant', content: message } - this._addAIMessage(aiMessage, error) + this._addAIMessage(aiMessage, mitoAIConnectionError) } - _addAIMessage(aiMessage: OpenAI.Chat.ChatCompletionMessageParam, error: boolean=false): void { - this.history.displayOptimizedChatHistory.push({message: aiMessage, error: error}); + _addAIMessage(aiMessage: OpenAI.Chat.ChatCompletionMessageParam, mitoAIConnectionError: boolean=false): void { + this.history.displayOptimizedChatHistory.push({message: aiMessage, type: mitoAIConnectionError ? 'connection error' : 'openai message'}); this.history.aiOptimizedChatHistory.push(aiMessage); } @@ -160,7 +247,7 @@ Your task: ${input}`}; role: 'system', content: message } - this.history.displayOptimizedChatHistory.push({message: systemMessage, error: false}); + this.history.displayOptimizedChatHistory.push({message: systemMessage, type: 'openai message'}); this.history.aiOptimizedChatHistory.push(systemMessage); } @@ -185,4 +272,15 @@ Your task: ${input}`}; const displayOptimizedChatHistory = this.getDisplayOptimizedHistory() return displayOptimizedChatHistory[lastAIMessagesIndex] } +} + + +const getDisplayedOptimizedUserMessage = (input: string, activeCellCode?: string): OpenAI.Chat.ChatCompletionMessageParam => { + return { + role: 'user', + content: `\`\`\`python +${activeCellCode} +\`\`\` + +${input}`}; } \ No newline at end of file diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index 0cc081705..beddb453c 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -14,7 +14,7 @@ import { UnifiedDiffLine } from '../../../utils/codeDiff'; interface IChatMessageProps { message: OpenAI.Chat.ChatCompletionMessageParam messageIndex: number - error: boolean + mitoAIConnectionError: boolean notebookTracker: INotebookTracker rendermime: IRenderMimeRegistry app: JupyterFrontEnd @@ -28,7 +28,7 @@ interface IChatMessageProps { const ChatMessage: React.FC = ({ message, messageIndex, - error, + mitoAIConnectionError, notebookTracker, rendermime, app, @@ -50,7 +50,6 @@ const ChatMessage: React.FC = ({ "message", {"message-user" : message.role === 'user'}, {'message-assistant' : message.role === 'assistant'}, - {'message-error': error} )}> {messageContentParts.map(messagePart => { if (messagePart.startsWith(PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE)) { @@ -75,7 +74,7 @@ const ChatMessage: React.FC = ({ } } else { return ( -

{error && }{messagePart}

+

{mitoAIConnectionError && }{messagePart}

) } })} diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 31d80d924..896dcb84c 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -1,18 +1,17 @@ import React, { useCallback, useEffect, useRef, useState } from 'react'; -import OpenAI from 'openai'; import '../../../style/ChatTaskpane.css'; import { classNames } from '../../utils/classNames'; import { INotebookTracker } from '@jupyterlab/notebook'; import { getActiveCellCode, writeCodeToActiveCell } from '../../utils/notebook'; import ChatMessage from './ChatMessage/ChatMessage'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; -import { ChatHistoryManager, IChatHistory } from './ChatHistoryManager'; +import { ChatHistoryManager } from './ChatHistoryManager'; import { requestAPI } from '../../utils/handler'; import { IVariableManager } from '../VariableManager/VariableManagerPlugin'; import LoadingDots from '../../components/LoadingDots'; import { JupyterFrontEnd } from '@jupyterlab/application'; import { getCodeBlockFromMessage, removeMarkdownCodeFormatting } from '../../utils/strings'; -import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_REJECT_LATEST_CODE, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; +import { COMMAND_MITO_AI_APPLY_LATEST_CODE, COMMAND_MITO_AI_REJECT_LATEST_CODE, COMMAND_MITO_AI_SEND_DEBUG_ERROR_MESSAGE, COMMAND_MITO_AI_SEND_EXPLAIN_CODE_MESSAGE } from '../../commands'; import { ReadonlyPartialJSONObject } from '@lumino/coreutils'; import ResetIcon from '../../icons/ResetIcon'; import IconButton from '../../components/IconButton'; @@ -23,36 +22,11 @@ import { CodeCell } from '@jupyterlab/cells'; import { StateEffect, Compartment } from '@codemirror/state'; import { codeDiffStripesExtension } from './CodeDiffDisplay'; +const getDefaultChatHistoryManager = (notebookTracker: INotebookTracker, variableManager: IVariableManager): ChatHistoryManager => { -// IMPORTANT: In order to improve the development experience, we allow you dispaly a -// cached conversation as a starting point. Before deploying the mito-ai, we must -// set USE_DEV_AI_CONVERSATION = false -// TODO: Write a test to ensure USE_DEV_AI_CONVERSATION is false -const USE_DEV_AI_CONVERSATION = false - -const getDefaultChatHistoryManager = (): ChatHistoryManager => { - - if (USE_DEV_AI_CONVERSATION) { - const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ - { role: 'system', content: 'You are an expert Python programmer.' }, - { role: 'user', content: "```python x = 5\ny=10\nx+y``` update x to 10" }, - { role: 'assistant', content: "```python x = 10\ny=10\nx+y```" }, - { role: 'user', content: "```python x = 5\ny=10\nx+y``` Explain what this code does to me" }, - { role: 'assistant', content: "This code defines two variables, x and y. Variables are named buckets that store a value. ```python x = 5\ny=10``` It then adds them together ```python x+y``` Let me know if you want me to further explain any of those concepts" } - ] - - const chatHistory: IChatHistory = { - aiOptimizedChatHistory: [...messages], - displayOptimizedChatHistory: [...messages].map(message => ({ message: message, error: false })) - } - - return new ChatHistoryManager(chatHistory) - - } else { - const chatHistoryManager = new ChatHistoryManager() - chatHistoryManager.addSystemMessage('You are an expert Python programmer.') - return chatHistoryManager - } + const chatHistoryManager = new ChatHistoryManager(variableManager, notebookTracker) + chatHistoryManager.addSystemMessage('You are an expert Python programmer.') + return chatHistoryManager } interface IChatTaskpaneProps { @@ -71,7 +45,7 @@ const ChatTaskpane: React.FC = ({ operatingSystem }) => { const textareaRef = useRef(null); - const [chatHistoryManager, setChatHistoryManager] = useState(() => getDefaultChatHistoryManager()); + const [chatHistoryManager, setChatHistoryManager] = useState(() => getDefaultChatHistoryManager(notebookTracker, variableManager)); const chatHistoryManagerRef = useRef(chatHistoryManager); const [input, setInput] = useState(''); @@ -122,55 +96,82 @@ const ChatTaskpane: React.FC = ({ adjustHeight(); }, [input]); + + + const getDuplicateChatHistoryManager = () => { + + /* + We use getDuplicateChatHistoryManager() instead of directly accessing the state variable because + the COMMAND_MITO_AI_SEND_MESSAGE is registered in a useEffect on initial render, which + would otherwise always use the initial state values. By using a function, we ensure we always + get the most recent chat history, even when the command is executed later. + */ + return chatHistoryManagerRef.current.createDuplicateChatHistoryManager() + } + /* Send a message with a specific input, clearing what is currently in the chat input. This is useful when we want to send the error message from the MIME renderer directly to the AI chat. */ - const sendMessageWithInput = async (input: string) => { - _sendMessage(input) + const sendDebugErrorMessage = (errorMessage: string) => { + + const newChatHistoryManager = getDuplicateChatHistoryManager() + newChatHistoryManager.addGenericUserPromptedMessage(input) + _sendMessageToOpenAI(newChatHistoryManager) + } + + const sendExplainCodeMessage = () => { + const newChatHistoryManager = getDuplicateChatHistoryManager() + newChatHistoryManager.addExplainCodeMessage() + _sendMessageToOpenAI(newChatHistoryManager) } /* - Send a message with the text currently in the chat input. + Send whatever message is currently in the chat input */ - const sendMessageFromChat = async () => { - _sendMessage(input) - } + const sendGenericUserPromptedMessage = async () => { - const getChatHistoryManager = () => { - return chatHistoryManagerRef.current - } + const newChatHistoryManager = getDuplicateChatHistoryManager() + newChatHistoryManager.addGenericUserPromptedMessage(input) + + setChatHistoryManager(newChatHistoryManager) - const _sendMessage = async (input: string) => { + const aiMessage = await _sendMessageToOpenAI(newChatHistoryManager) + + if (!aiMessage) { + return + } - const variables = variableManager.variables const activeCellCode = getActiveCellCode(notebookTracker) - /* - 1. Access ChatHistoryManager via a function: - We use getChatHistoryManager() instead of directly accessing the state variable because - the COMMAND_MITO_AI_SEND_MESSAGE is registered in a useEffect on initial render, which - would otherwise always use the initial state values. By using a function, we ensure we always - get the most recent chat history, even when the command is executed later. + // Extract the code from the AI's message and then calculate the code diffs + const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); + const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); + const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) - 2. Create a new ChatHistoryManager instance: - We create a copy of the current chat history and use it to initialize a new ChatHistoryManager to - trigger a re-render in React, as simply appending to the existing ChatHistoryManager - (an immutable object) wouldn't be detected as a state change. - */ - const currentChatHistory = getChatHistoryManager().getHistory() - const updatedManager = new ChatHistoryManager(currentChatHistory); - updatedManager.addUserMessage(input, activeCellCode, variables) + // Store the original code so that we can revert to it if the user rejects the AI's code + originalCodeBeforeDiff.current = activeCellCode || '' + + // Temporarily write the unified code string to the active cell so we can display + // the code diffs to the user. Once the user accepts or rejects the code, we'll + // apply the correct version of the code. + writeCodeToActiveCell(notebookTracker, unifiedCodeString) + setUnifiedDiffLines(unifiedDiffs) + } + + const _sendMessageToOpenAI = async (newChatHistoryManager: ChatHistoryManager) => { setInput(''); setLoadingAIResponse(true) + let aiRespone = undefined + try { const apiResponse = await requestAPI('mito_ai/completion', { method: 'POST', body: JSON.stringify({ - messages: updatedManager.getAIOptimizedHistory() + messages: newChatHistoryManager.getAIOptimizedHistory() }) }); @@ -179,34 +180,21 @@ const ChatTaskpane: React.FC = ({ const response = apiResponse.response; const aiMessage = response.choices[0].message; - updatedManager.addAIMessageFromResponse(aiMessage); - setChatHistoryManager(updatedManager); - - // Extract the code from the AI's message and then calculate the code diffs - const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); - const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); - const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) - - // Store the original code so that we can revert to it if the user rejects the AI's code - originalCodeBeforeDiff.current = activeCellCode || '' - - // Temporarily write the unified code string to the active cell so we can display - // the code diffs to the user. Once the user accepts or rejects the code, we'll - // apply the correct version of the code. - writeCodeToActiveCell(notebookTracker, unifiedCodeString) - setUnifiedDiffLines(unifiedDiffs) + newChatHistoryManager.addAIMessageFromResponse(aiMessage); + setChatHistoryManager(newChatHistoryManager); + aiRespone = aiMessage } else { - updatedManager.addAIMessageFromMessageContent(apiResponse.errorMessage, true) - setChatHistoryManager(updatedManager); + newChatHistoryManager.addAIMessageFromMessageContent(apiResponse.errorMessage, true) + setChatHistoryManager(newChatHistoryManager); } - - setLoadingAIResponse(false) } catch (error) { console.error('Error calling OpenAI API:', error); } finally { setLoadingAIResponse(false) + return aiRespone } - }; + + } const displayOptimizedChatHistory = chatHistoryManager.getDisplayOptimizedHistory() @@ -275,13 +263,20 @@ const ChatTaskpane: React.FC = ({ Add a new command to the JupyterLab command registry that sends the current chat message. We use this to automatically send the message when the user adds an error to the chat. */ - app.commands.addCommand(COMMAND_MITO_AI_SEND_MESSAGE, { + app.commands.addCommand(COMMAND_MITO_AI_SEND_DEBUG_ERROR_MESSAGE, { execute: (args?: ReadonlyPartialJSONObject) => { if (args?.input) { - sendMessageWithInput(args.input.toString()) + sendDebugErrorMessage(args.input.toString()) } } }) + + app.commands.addCommand(COMMAND_MITO_AI_SEND_EXPLAIN_CODE_MESSAGE, { + execute: () => { + sendExplainCodeMessage() + } + }) + }, []) // Create a WeakMap to store compartments per code cell @@ -348,7 +343,7 @@ const ChatTaskpane: React.FC = ({ icon={} title="Clear the chat history" onClick={() => { - setChatHistoryManager(getDefaultChatHistoryManager()) + setChatHistoryManager(getDefaultChatHistoryManager(notebookTracker, variableManager)) }} />
@@ -357,7 +352,7 @@ const ChatTaskpane: React.FC = ({ return ( = ({ // shift + enter to add a new line. if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); - sendMessageFromChat(); + sendGenericUserPromptedMessage() } }} /> diff --git a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx index 0224a1d6a..9ab6fb713 100644 --- a/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx +++ b/mito-ai/src/Extensions/CellToolbarButtons/CellToolbarButtonsPlugin.tsx @@ -3,7 +3,7 @@ import { JupyterFrontEndPlugin } from '@jupyterlab/application'; import { INotebookTracker } from '@jupyterlab/notebook'; -import { COMMAND_MITO_AI_OPEN_CHAT, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; +import { COMMAND_MITO_AI_OPEN_CHAT, COMMAND_MITO_AI_SEND_EXPLAIN_CODE_MESSAGE } from '../../commands'; import { LabIcon } from '@jupyterlab/ui-components'; import LightbulbIcon from '../../../src/icons/LightbulbIcon.svg' @@ -31,16 +31,12 @@ const CellToolbarButtonsPlugin: JupyterFrontEndPlugin = { execute: () => { /* In order to click on the cell toolbar button, that cell must be the active cell, - so the Ai Chat taskpane will take care of providing the cell context. - - TODO: In the future, instead of directly calling COMMAND_MITO_AI_SEND_MESSAGE, we should - update the AI Chat History Manager so that we can generate an AI optimzied message and a - display optimized message. + so the ChatHistoryManager will take care of providing the cell context. */ app.commands.execute(COMMAND_MITO_AI_OPEN_CHAT) - app.commands.execute(COMMAND_MITO_AI_SEND_MESSAGE, { input: `Explain this code` }); + app.commands.execute(COMMAND_MITO_AI_SEND_EXPLAIN_CODE_MESSAGE); }, - isVisible: () => notebookTracker.activeCell?.model.type === 'code' && app.commands.hasCommand(COMMAND_MITO_AI_SEND_MESSAGE) + isVisible: () => notebookTracker.activeCell?.model.type === 'code' && app.commands.hasCommand(COMMAND_MITO_AI_SEND_EXPLAIN_CODE_MESSAGE) }); } }; diff --git a/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx b/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx index d61e095c0..8e34c7ace 100644 --- a/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx +++ b/mito-ai/src/Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin.tsx @@ -4,7 +4,7 @@ import { JupyterFrontEnd, JupyterFrontEndPlugin } from '@jupyterlab/application' import { IRenderMimeRegistry} from '@jupyterlab/rendermime'; import { IRenderMime } from '@jupyterlab/rendermime-interfaces'; import { Widget } from '@lumino/widgets'; -import { COMMAND_MITO_AI_OPEN_CHAT, COMMAND_MITO_AI_SEND_MESSAGE } from '../../commands'; +import { COMMAND_MITO_AI_OPEN_CHAT, COMMAND_MITO_AI_SEND_DEBUG_ERROR_MESSAGE } from '../../commands'; import MagicWandIcon from '../../icons/MagicWand'; import '../../../style/ErrorMimeRendererPlugin.css' @@ -81,7 +81,7 @@ class AugmentedStderrRenderer extends Widget implements IRenderMime.IRenderer { openChatInterfaceWithError(model: IRenderMime.IMimeModel): void { const conciseErrorMessage = this.getErrorString(model); this.app.commands.execute(COMMAND_MITO_AI_OPEN_CHAT) - this.app.commands.execute(COMMAND_MITO_AI_SEND_MESSAGE, { input: conciseErrorMessage }); + this.app.commands.execute(COMMAND_MITO_AI_SEND_DEBUG_ERROR_MESSAGE, { input: conciseErrorMessage }); } /* diff --git a/mito-ai/src/commands.tsx b/mito-ai/src/commands.tsx index 4fcc20e70..44d7ad90b 100644 --- a/mito-ai/src/commands.tsx +++ b/mito-ai/src/commands.tsx @@ -3,4 +3,6 @@ const MITO_AI = 'mito_ai' export const COMMAND_MITO_AI_OPEN_CHAT = `${MITO_AI}:open-chat` export const COMMAND_MITO_AI_APPLY_LATEST_CODE = `${MITO_AI}:apply-latest-code` export const COMMAND_MITO_AI_REJECT_LATEST_CODE = `${MITO_AI}:reject-latest-code` -export const COMMAND_MITO_AI_SEND_MESSAGE = `${MITO_AI}:send-message` \ No newline at end of file +export const COMMAND_MITO_AI_SEND_MESSAGE = `${MITO_AI}:send-message` +export const COMMAND_MITO_AI_SEND_EXPLAIN_CODE_MESSAGE = `${MITO_AI}:send-explain-code-message` +export const COMMAND_MITO_AI_SEND_DEBUG_ERROR_MESSAGE = `${MITO_AI}:send-debug-error-message` \ No newline at end of file From 835ee0c23a9d1124ba02a8b462e6dd465f74c1bd Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 25 Oct 2024 18:25:17 -0400 Subject: [PATCH 31/42] mito-ai: cleanup --- .../src/Extensions/AiChat/ChatTaskpane.tsx | 61 ++++++++++++------- .../src/Extensions/AiChat/CodeDiffDisplay.tsx | 1 + 2 files changed, 40 insertions(+), 22 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 896dcb84c..093e52a92 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -21,6 +21,7 @@ import { CodeMirrorEditor } from '@jupyterlab/codemirror'; import { CodeCell } from '@jupyterlab/cells'; import { StateEffect, Compartment } from '@codemirror/state'; import { codeDiffStripesExtension } from './CodeDiffDisplay'; +import OpenAI from "openai"; const getDefaultChatHistoryManager = (notebookTracker: INotebookTracker, variableManager: IVariableManager): ChatHistoryManager => { @@ -116,15 +117,27 @@ const ChatTaskpane: React.FC = ({ */ const sendDebugErrorMessage = (errorMessage: string) => { + // Step 1: Add the user's message to the chat history const newChatHistoryManager = getDuplicateChatHistoryManager() - newChatHistoryManager.addGenericUserPromptedMessage(input) + newChatHistoryManager.addDebugErrorMessage(errorMessage) + + // Step 2: Send the message to the AI _sendMessageToOpenAI(newChatHistoryManager) + + // Step 3: Update the code diff stripes + updateCodeDiffStripes } const sendExplainCodeMessage = () => { + + // Step 1: Add the user's message to the chat history const newChatHistoryManager = getDuplicateChatHistoryManager() newChatHistoryManager.addExplainCodeMessage() + + // Step 2: Send the message to the AI _sendMessageToOpenAI(newChatHistoryManager) + + // Step 3: No post processing step needed for explaining code. } /* @@ -132,32 +145,15 @@ const ChatTaskpane: React.FC = ({ */ const sendGenericUserPromptedMessage = async () => { + // Step 1: Add the user's message to the chat history const newChatHistoryManager = getDuplicateChatHistoryManager() newChatHistoryManager.addGenericUserPromptedMessage(input) - setChatHistoryManager(newChatHistoryManager) - + // Step 2: Send the message to the AI const aiMessage = await _sendMessageToOpenAI(newChatHistoryManager) - if (!aiMessage) { - return - } - - const activeCellCode = getActiveCellCode(notebookTracker) - - // Extract the code from the AI's message and then calculate the code diffs - const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); - const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); - const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) - - // Store the original code so that we can revert to it if the user rejects the AI's code - originalCodeBeforeDiff.current = activeCellCode || '' - - // Temporarily write the unified code string to the active cell so we can display - // the code diffs to the user. Once the user accepts or rejects the code, we'll - // apply the correct version of the code. - writeCodeToActiveCell(notebookTracker, unifiedCodeString) - setUnifiedDiffLines(unifiedDiffs) + // Step 3: Update the code diff stripes + updateCodeDiffStripes(aiMessage) } const _sendMessageToOpenAI = async (newChatHistoryManager: ChatHistoryManager) => { @@ -193,7 +189,28 @@ const ChatTaskpane: React.FC = ({ setLoadingAIResponse(false) return aiRespone } + } + + const updateCodeDiffStripes = (aiMessage: OpenAI.ChatCompletionMessage | undefined) => { + if (!aiMessage) { + return + } + + const activeCellCode = getActiveCellCode(notebookTracker) + // Extract the code from the AI's message and then calculate the code diffs + const aiGeneratedCode = getCodeBlockFromMessage(aiMessage); + const aiGeneratedCodeCleaned = removeMarkdownCodeFormatting(aiGeneratedCode || ''); + const { unifiedCodeString, unifiedDiffs } = getCodeDiffsAndUnifiedCodeString(activeCellCode, aiGeneratedCodeCleaned) + + // Store the original code so that we can revert to it if the user rejects the AI's code + originalCodeBeforeDiff.current = activeCellCode || '' + + // Temporarily write the unified code string to the active cell so we can display + // the code diffs to the user. Once the user accepts or rejects the code, we'll + // apply the correct version of the code. + writeCodeToActiveCell(notebookTracker, unifiedCodeString) + setUnifiedDiffLines(unifiedDiffs) } const displayOptimizedChatHistory = chatHistoryManager.getDisplayOptimizedHistory() diff --git a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx index 3d90e54d8..0c1f9e152 100644 --- a/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx +++ b/mito-ai/src/Extensions/AiChat/CodeDiffDisplay.tsx @@ -102,3 +102,4 @@ export function codeDiffStripesExtension(options: { unifiedDiffLines?: UnifiedDi showStripes ]; } + From 256bdcfc5b1d7e50da2084da5e497e897e0540f0 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 28 Oct 2024 12:25:33 -0400 Subject: [PATCH 32/42] mito-ai: cleanup --- .../Extensions/AiChat/ChatHistoryManager.tsx | 19 +++++++++++++++++++ .../src/Extensions/AiChat/ChatTaskpane.tsx | 13 +++++++------ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index 6034ec7ab..b5b1f3da8 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -83,6 +83,7 @@ ${variables?.map(variable => `${JSON.stringify(variable, null, 2)}\n`).join('')} Complete the task below. Decide what variables to use and what changes you need to make to the active code cell. Only return the full new active code cell and a concise explanation of the changes you made. + Do not: - Use the word "I" - Include multiple approaches in your response @@ -136,6 +137,22 @@ Your task: ${input}`}; const aiOptimizedPrompt = `You just ran the active code cell and received an error. Return the full code cell with the error corrected and a short explanation of the error. + + +Do not: +- Use the word "I" +- Include multiple approaches in your response +- Recreate variables that already exist + +Do: +- Use the variables that you have access to +- Keep as much of the original code as possible +- Ask for more context if you need it. + + + + + Remember that you are executing code inside a Jupyter notebook. That means you will have persistent state issues where variables from previous cells or previous code executions might still affect current code. When those errors occur, here are a few possible solutions: 1. Restarting the kernel to reset the environment if a function or variable has been unintentionally overwritten. 2. Identify which cell might need to be rerun to properly initialize the function or variable that is causing the issue. @@ -144,6 +161,8 @@ For example, if an error occurs because the built-in function 'print' is overwri When a user hits an error because of a persistent state issue, tell them how to resolve it. + + Code in the active code cell: diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 093e52a92..c702f13d5 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -46,10 +46,11 @@ const ChatTaskpane: React.FC = ({ operatingSystem }) => { const textareaRef = useRef(null); + const [input, setInput] = useState(''); + const [chatHistoryManager, setChatHistoryManager] = useState(() => getDefaultChatHistoryManager(notebookTracker, variableManager)); const chatHistoryManagerRef = useRef(chatHistoryManager); - const [input, setInput] = useState(''); const [loadingAIResponse, setLoadingAIResponse] = useState(false) const [unifiedDiffLines, setUnifiedDiffLines] = useState(undefined) @@ -115,17 +116,17 @@ const ChatTaskpane: React.FC = ({ This is useful when we want to send the error message from the MIME renderer directly to the AI chat. */ - const sendDebugErrorMessage = (errorMessage: string) => { + const sendDebugErrorMessage = async (errorMessage: string) => { // Step 1: Add the user's message to the chat history const newChatHistoryManager = getDuplicateChatHistoryManager() newChatHistoryManager.addDebugErrorMessage(errorMessage) // Step 2: Send the message to the AI - _sendMessageToOpenAI(newChatHistoryManager) + const aiMessage = await _sendMessageToOpenAI(newChatHistoryManager) // Step 3: Update the code diff stripes - updateCodeDiffStripes + updateCodeDiffStripes(aiMessage) } const sendExplainCodeMessage = () => { @@ -143,7 +144,7 @@ const ChatTaskpane: React.FC = ({ /* Send whatever message is currently in the chat input */ - const sendGenericUserPromptedMessage = async () => { + const sendChatInputMessage = async () => { // Step 1: Add the user's message to the chat history const newChatHistoryManager = getDuplicateChatHistoryManager() @@ -399,7 +400,7 @@ const ChatTaskpane: React.FC = ({ // shift + enter to add a new line. if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); - sendGenericUserPromptedMessage() + sendChatInputMessage() } }} /> From b5e2d375f502ce5050fb83f7f8826d62459e0052 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 28 Oct 2024 12:47:12 -0400 Subject: [PATCH 33/42] mito-ai: cleanup --- mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx | 7 +++++-- mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx index b5b1f3da8..2e1f83de7 100644 --- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx @@ -15,7 +15,10 @@ export interface IChatHistory { aiOptimizedChatHistory: OpenAI.Chat.ChatCompletionMessageParam[] // The display optimized chat history is what we display to the user. Each message - // is a subset of the corresponding message in aiOptimizedChatHistory. + // is a subset of the corresponding message in aiOptimizedChatHistory. Note that in the + // displayOptimizedChatHistory, we also include connection error messages so that we can + // display them in the chat interface. For example, if the user does not have an API key set, + // we add a message to the chat ui that tells them to set an API key. displayOptimizedChatHistory: IDisplayOptimizedChatHistory[] } @@ -68,7 +71,7 @@ export class ChatHistoryManager { return this.history.displayOptimizedChatHistory; } - addGenericUserPromptedMessage(input: string): void { + addChatInputMessage(input: string): void { const variables = this.variableManager.variables diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index c702f13d5..85d5caf80 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -148,7 +148,7 @@ const ChatTaskpane: React.FC = ({ // Step 1: Add the user's message to the chat history const newChatHistoryManager = getDuplicateChatHistoryManager() - newChatHistoryManager.addGenericUserPromptedMessage(input) + newChatHistoryManager.addChatInputMessage(input) // Step 2: Send the message to the AI const aiMessage = await _sendMessageToOpenAI(newChatHistoryManager) From 1005901120cb92b2cb610989eb3ce717d8e65caa Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 28 Oct 2024 12:49:09 -0400 Subject: [PATCH 34/42] .github: run tests on mito-ai directory changes --- .github/workflows/test-mitosheet-frontend.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-mitosheet-frontend.yml b/.github/workflows/test-mitosheet-frontend.yml index b4ab9a7ac..8b83e02a4 100644 --- a/.github/workflows/test-mitosheet-frontend.yml +++ b/.github/workflows/test-mitosheet-frontend.yml @@ -5,11 +5,13 @@ on: branches: [ dev ] paths: - 'mitosheet/**' + - 'tests/**' + - 'mito-ai/**' pull_request: paths: - 'mitosheet/**' - 'tests/**' - + - 'mito-ai/**' jobs: test-mitosheet-frontend-jupyterlab: runs-on: ubuntu-20.04 From 7b35246d5bf6d0dc13e7f53ee5445e7cd412b525 Mon Sep 17 00:00:00 2001 From: Nawaz Gafar Date: Fri, 1 Nov 2024 13:43:54 -0400 Subject: [PATCH 35/42] Added check for isLastAiMessage --- .gitignore | 1 + mito-ai/Untitled.ipynb | 2 +- .../AiChat/ChatMessage/CodeBlock.tsx | 28 +++++++++++-------- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index 272bd37c5..6130d31b7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +scratch* # Mac Specific diff --git a/mito-ai/Untitled.ipynb b/mito-ai/Untitled.ipynb index 774b7bb0c..5e9bc1cef 100644 --- a/mito-ai/Untitled.ipynb +++ b/mito-ai/Untitled.ipynb @@ -122,7 +122,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.1" + "version": "3.10.4" } }, "nbformat": 4, diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index 8bfec282f..0988ae726 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -24,18 +24,18 @@ interface ICodeBlockProps { } const CodeBlock: React.FC = ({ - code, - role, - rendermime, - notebookTracker, - app, + code, + role, + rendermime, + notebookTracker, + app, isLastAiMessage, operatingSystem, setDisplayCodeDiff, acceptAICode, rejectAICode }): JSX.Element => { - + const notebookName = getNotebookName(notebookTracker) const copyCodeToClipboard = () => { @@ -61,12 +61,16 @@ const CodeBlock: React.FC = ({
{notebookName}
- - + {isLastAiMessage && ( + <> + + + + )} Date: Fri, 1 Nov 2024 14:53:57 -0400 Subject: [PATCH 36/42] Moved mito-ai tests into seperate workflow --- .github/workflows/test-mito-ai.yml | 64 +++++++++++++++++++ .github/workflows/test-mitosheet-frontend.yml | 56 ---------------- 2 files changed, 64 insertions(+), 56 deletions(-) create mode 100644 .github/workflows/test-mito-ai.yml diff --git a/.github/workflows/test-mito-ai.yml b/.github/workflows/test-mito-ai.yml new file mode 100644 index 000000000..14004d9c6 --- /dev/null +++ b/.github/workflows/test-mito-ai.yml @@ -0,0 +1,64 @@ +name: Test - Mito AI + +on: + push: + branches: [ dev ] + paths: + - 'mito-ai/**' + pull_request: + paths: + - 'mito-ai/**' +jobs: + test-mitoai-frontend-jupyterlab: + runs-on: ubuntu-20.04 + timeout-minutes: 60 + strategy: + matrix: + python-version: ['3.8', '3.10', '3.11'] + fail-fast: false + + steps: + - name: Cancel Previous Runs + uses: styfle/cancel-workflow-action@0.7.0 + with: + access_token: ${{ github.token }} + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: | + mito-ai/setup.py + tests/requirements.txt + - uses: actions/setup-node@v3 + with: + node-version: 18 + cache: 'npm' + cache-dependency-path: mito-ai/package-lock.json + - name: Install dependencies + run: | + cd tests + bash mac-setup.sh + - name: Setup JupyterLab + run: | + cd tests + source venv/bin/activate + cd ../mito-ai + jupyter labextension develop . --overwrite + jupyter server extension enable --py mito-ai + - name: Start a server and run tests + run: | + cd tests + source venv/bin/activate + jupyter lab --config jupyter_server_test_config.py & + npm run test:mitoai + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + - name: Upload test-results + uses: actions/upload-artifact@v3 + if: failure() + with: + name: mitoai-jupyterlab-playwright-report-${{ matrix.python-version }} + path: tests/playwright-report/ + retention-days: 14 \ No newline at end of file diff --git a/.github/workflows/test-mitosheet-frontend.yml b/.github/workflows/test-mitosheet-frontend.yml index 8b83e02a4..6eedac2d7 100644 --- a/.github/workflows/test-mitosheet-frontend.yml +++ b/.github/workflows/test-mitosheet-frontend.yml @@ -6,12 +6,10 @@ on: paths: - 'mitosheet/**' - 'tests/**' - - 'mito-ai/**' pull_request: paths: - 'mitosheet/**' - 'tests/**' - - 'mito-ai/**' jobs: test-mitosheet-frontend-jupyterlab: runs-on: ubuntu-20.04 @@ -115,60 +113,6 @@ jobs: path: tests/playwright-report/ retention-days: 14 - test-mitoai-frontend-jupyterlab: - runs-on: ubuntu-20.04 - timeout-minutes: 60 - strategy: - matrix: - python-version: ['3.8', '3.10', '3.11'] - fail-fast: false - - steps: - - name: Cancel Previous Runs - uses: styfle/cancel-workflow-action@0.7.0 - with: - access_token: ${{ github.token }} - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - cache: pip - cache-dependency-path: | - mito-ai/setup.py - tests/requirements.txt - - uses: actions/setup-node@v3 - with: - node-version: 18 - cache: 'npm' - cache-dependency-path: mito-ai/package-lock.json - - name: Install dependencies - run: | - cd tests - bash mac-setup.sh - - name: Setup JupyterLab - run: | - cd tests - source venv/bin/activate - cd ../mito-ai - jupyter labextension develop . --overwrite - jupyter server extension enable --py mito-ai - - name: Start a server and run tests - run: | - cd tests - source venv/bin/activate - jupyter lab --config jupyter_server_test_config.py & - npm run test:mitoai - env: - OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} - - name: Upload test-results - uses: actions/upload-artifact@v3 - if: failure() - with: - name: mitoai-jupyterlab-playwright-report-${{ matrix.python-version }} - path: tests/playwright-report/ - retention-days: 14 - test-mitosheet-frontend-streamlit: timeout-minutes: 60 strategy: From 6d9f937e54330875acfe791f26a96615e28fd750 Mon Sep 17 00:00:00 2001 From: Nawaz Gafar Date: Fri, 1 Nov 2024 15:14:37 -0400 Subject: [PATCH 37/42] Updated margin and padding --- mito-ai/style/PythonCode.css | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mito-ai/style/PythonCode.css b/mito-ai/style/PythonCode.css index 589cf6b5b..4fd47f80d 100644 --- a/mito-ai/style/PythonCode.css +++ b/mito-ai/style/PythonCode.css @@ -2,8 +2,8 @@ flex-grow: 1; height: 100%; width: 100%; - margin: 0; - padding: 10px; + margin: 0 !important; + padding: 10px !important; font-size: 12px; font-family: Consolas,Monaco,Lucida Console,Liberation Mono,DejaVu Sans Mono,Bitstream Vera Sans Mono,Courier New; white-space: nowrap; From dd9f13a6031f5592ec148e40cfd062613046768a Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 1 Nov 2024 15:34:39 -0400 Subject: [PATCH 38/42] mito-ai: update font size and line wrap of chat code --- mito-ai/style/PythonCode.css | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mito-ai/style/PythonCode.css b/mito-ai/style/PythonCode.css index 4fd47f80d..ca74116dc 100644 --- a/mito-ai/style/PythonCode.css +++ b/mito-ai/style/PythonCode.css @@ -4,9 +4,8 @@ width: 100%; margin: 0 !important; padding: 10px !important; - font-size: 12px; - font-family: Consolas,Monaco,Lucida Console,Liberation Mono,DejaVu Sans Mono,Bitstream Vera Sans Mono,Courier New; - white-space: nowrap; + font-size: 12px !important; + white-space: pre !important; overflow-x: auto; } From 0748089dd50c2043be5a1715c80c97187aff89e8 Mon Sep 17 00:00:00 2001 From: Nawaz Gafar Date: Fri, 1 Nov 2024 16:43:53 -0400 Subject: [PATCH 39/42] Added white-space styling for code class --- mito-ai/style/PythonCode.css | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mito-ai/style/PythonCode.css b/mito-ai/style/PythonCode.css index ca74116dc..dd0157ccc 100644 --- a/mito-ai/style/PythonCode.css +++ b/mito-ai/style/PythonCode.css @@ -5,10 +5,13 @@ margin: 0 !important; padding: 10px !important; font-size: 12px !important; - white-space: pre !important; overflow-x: auto; } +code { + white-space: pre !important; +} + .code-message-part-python-code .jp-RenderedHTMLCommon > *:last-child { /* Remove the default Jupyter ending margin From 0ce41421d5ad02578216297637e48d17f864fd9a Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Fri, 1 Nov 2024 17:07:47 -0400 Subject: [PATCH 40/42] mito-ai: add more specific class name --- mito-ai/style/PythonCode.css | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mito-ai/style/PythonCode.css b/mito-ai/style/PythonCode.css index dd0157ccc..1048c4141 100644 --- a/mito-ai/style/PythonCode.css +++ b/mito-ai/style/PythonCode.css @@ -8,7 +8,7 @@ overflow-x: auto; } -code { +.code-message-part-python-code code { white-space: pre !important; } From 1f7e62b9b195c80d2a910e6bccd0562dde54177e Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 4 Nov 2024 09:29:24 -0500 Subject: [PATCH 41/42] mito-ai: update openai pinning --- mito-ai/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mito-ai/setup.py b/mito-ai/setup.py index 7d9150b6b..4ef22f397 100644 --- a/mito-ai/setup.py +++ b/mito-ai/setup.py @@ -86,7 +86,7 @@ def get_data_files_from_data_files_spec( packages=find_packages(), install_requires=[ "jupyterlab>=4.0.0,<5", - "openai", + "openai>=1.0.0" ], extras_require = { 'deploy': [ From 4990f53004d8d8cf15cb2f5f0a8902213c721662 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Mon, 4 Nov 2024 09:59:41 -0500 Subject: [PATCH 42/42] mito-ai: fix small chat input bug --- mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 85d5caf80..b97bf6a58 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -87,11 +87,11 @@ const ChatTaskpane: React.FC = ({ if (!textarea) { return } - textarea.style.height = 'auto'; + textarea.style.minHeight = 'auto'; // The height should be 20 at minimum to support the placeholder - const height = textarea.scrollHeight < 20 ? 20 : textarea.scrollHeight - textarea.style.height = `${height}px`; + const minHeight = textarea.scrollHeight < 20 ? 20 : textarea.scrollHeight + textarea.style.minHeight = `${minHeight}px`; }; useEffect(() => {