Merge pull request #1449 from mito-ds/chat-message-apply-buttons

Chat message apply buttons
mito-ds · Dec 30, 2024 · 3bbf57f · 3bbf57f
2 parents c277316 + dab94d1
commit 3bbf57f
Show file tree

Hide file tree

Showing 13 changed files with 147 additions and 40 deletions.
diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
@@ -106,6 +106,11 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                     code={messagePart}
                                     role={message.role}
                                     renderMimeRegistry={renderMimeRegistry}
+                                    previewAICode={previewAICode}
+                                    acceptAICode={acceptAICode}
+                                    rejectAICode={rejectAICode}
+                                    isLastAiMessage={isLastAiMessage}
+                                    codeReviewStatus={codeReviewStatus}
                                 />
 
                                 {isLastAiMessage && codeReviewStatus === 'chatPreview' && 

diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx
@@ -3,18 +3,34 @@ import PythonCode from './PythonCode';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import '../../../../style/CodeBlock.css'
 import copyToClipboard from '../../../utils/copyToClipboard';
+import IconButton from '../../../components/IconButton';
+import CopyIcon from '../../../icons/CopyIcon';
+import PlayButtonIcon from '../../../icons/PlayButtonIcon';
+import { CodeReviewStatus } from '../ChatTaskpane';
+import AcceptIcon from '../../../icons/AcceptIcon';
+import RejectIcon from '../../../icons/RejectIcon';
 
 
 interface ICodeBlockProps {
     code: string,
     role: 'user' | 'assistant'
     renderMimeRegistry: IRenderMimeRegistry
+    previewAICode: () => void
+    acceptAICode: () => void
+    rejectAICode: () => void
+    isLastAiMessage: boolean
+    codeReviewStatus: CodeReviewStatus
 }
 
 const CodeBlock: React.FC<ICodeBlockProps> = ({
     code,
     role,
     renderMimeRegistry,
+    previewAICode,
+    acceptAICode,
+    rejectAICode,
+    isLastAiMessage,
+    codeReviewStatus,
 }): JSX.Element => {
 
     if (role === 'user') {
@@ -31,9 +47,53 @@ const CodeBlock: React.FC<ICodeBlockProps> = ({
     if (role === 'assistant') {
         return (
             <div className='code-block-container'>
-                <div className='code-block-toolbar'>
-                    <button onClick={() => {copyToClipboard(code)}}>Copy</button>
-                </div>
+                <>
+                    {/* The code block toolbar for the last AI message */}
+                    {isLastAiMessage && 
+                        <div className='code-block-toolbar'>
+                            {codeReviewStatus === 'chatPreview' && 
+                                <IconButton 
+                                    icon={<PlayButtonIcon />}
+                                    title="Overwrite Active Cell"
+                                    onClick={() => {previewAICode()}}
+                                />
+                            }
+                            {codeReviewStatus === 'codeCellPreview' && 
+                                <IconButton 
+                                    icon={<AcceptIcon />}
+                                    title="Accept AI Generated Code"
+                                    onClick={() => {acceptAICode()}}
+                                    style={{color: 'var(--green-700)'}}
+                                />
+                            }
+                            {codeReviewStatus === 'codeCellPreview' && 
+                                <IconButton 
+                                    icon={<RejectIcon />}
+                                    title="Reject AI Generated Code"
+                                    onClick={() => {rejectAICode()}}
+                                    style={{color: 'var(--red-700)'}}
+                                />
+                            }
+                            {codeReviewStatus !== 'codeCellPreview' && 
+                                <IconButton
+                                    icon={<CopyIcon />}
+                                    title="Copy"
+                                    onClick={() => {copyToClipboard(code)}}
+                                />
+                            }
+                        </div>
+                    }
+                    {/* The code block toolbar for every other AI message */}
+                    {!isLastAiMessage && 
+                        <div className='code-block-toolbar'>
+                            <IconButton
+                                icon={<CopyIcon />}
+                                title="Copy"
+                                onClick={() => {copyToClipboard(code)}}
+                            />
+                        </div>
+                    }
+                </>
                 <PythonCode
                     code={code}
                     renderMimeRegistry={renderMimeRegistry}

diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
@@ -388,8 +388,8 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
             Register the code cell toolbar buttons for accepting and rejecting code.
         */
         app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_ACCEPT_CODE, {
-            label: `Accept AI edits ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`,
-            className: 'text-and-icon-button green',
+            label: `Accept ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`,
+            className: 'text-and-icon-button button-green small',
             caption: 'Accept Code',
             execute: () => {acceptAICode()},
             // We use the cellStateBeforeDiff because it contains the code cell ID that we want to write to
@@ -404,8 +404,8 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         });
 
         app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_REJECT_CODE, {
-            label: `Reject AI edits ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`,
-            className: 'text-and-icon-button red',
+            label: `Reject ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`,
+            className: 'text-and-icon-button button-red small',
             caption: 'Reject Code',
             execute: () => {rejectAICode()},
             isVisible: () => {

diff --git a/mito-ai/src/components/IconButton.tsx b/mito-ai/src/components/IconButton.tsx
@@ -1,16 +1,19 @@
 import React from 'react';
 import '../../style/IconButton.css';
+import '../../style/button.css';
+import { classNames } from '../utils/classNames';
 
 interface IconButtonProps {
     icon: React.ReactNode;
     onClick: () => void;
     title: string;
+    style?: React.CSSProperties;
 }
 
-const IconButton: React.FC<IconButtonProps> = ({ icon, onClick, title }) => {
+const IconButton: React.FC<IconButtonProps> = ({ icon, onClick, title, style }) => {
 
     return (
-        <button className="icon-button" onClick={onClick} title={title}>
+        <button className={classNames("icon-button")} onClick={onClick} title={title} style={style}>
             {icon}
         </button>
     )

diff --git a/mito-ai/src/components/TextAndIconButton.tsx b/mito-ai/src/components/TextAndIconButton.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import '../../style/TextAndIconButton.css';
+import '../../style/button.css';
 import { ButtonProps } from './TextButton';
 import { classNames } from '../utils/classNames';
 
@@ -11,7 +12,7 @@ interface TextAndIconButtonProps extends ButtonProps {
 const TextAndIconButton: React.FC<TextAndIconButtonProps> = ({ text, icon: Icon, onClick, title, variant }) => {
 
     return (
-        <button className={classNames("text-and-icon-button", variant)} onClick={onClick} title={title}>
+        <button className={classNames("text-and-icon-button", `button-${variant}`)} onClick={onClick} title={title}>
             {text}
             <span className="text-and-icon-button__icon">
                 <Icon />

diff --git a/mito-ai/src/components/TextButton.tsx b/mito-ai/src/components/TextButton.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import '../../style/TextAndIconButton.css';
+import '../../style/button.css';
 import { classNames } from '../utils/classNames';
 
 
@@ -16,7 +17,7 @@ interface TextButtonProps extends ButtonProps {}
 const TextButton: React.FC<TextButtonProps> = ({ text, onClick, title, variant }) => {
 
     return (
-        <button className={classNames("text-and-icon-button", variant)} onClick={onClick} title={title}>
+        <button className={classNames("text-and-icon-button", `button-${variant}`)} onClick={onClick} title={title}>
             {text}
         </button>
     )

diff --git a/mito-ai/src/icons/AcceptIcon.tsx b/mito-ai/src/icons/AcceptIcon.tsx
@@ -0,0 +1,12 @@
+import React from 'react';
+
+const AcceptIcon: React.FC = () => (
+    <svg width="1em" height="1em" viewBox="0 0 21 20" fill="none" xmlns="http://www.w3.org/2000/svg">
+        <g fill="currentColor">
+            <path d="M10.2236 0C4.70642 0 0.224487 4.48351 0.224487 10.0009C0.224487 15.5182 4.70642 20 10.2236 20C15.7408 20 20.2245 15.5181 20.2245 10.0009C20.2245 4.48368 15.741 0 10.2236 0ZM14.6192 5.46298C14.986 5.46298 15.3511 5.60114 15.6283 5.87746C16.1817 6.43091 16.1817 7.34456 15.6283 7.89886L9.56828 13.9589C9.01482 14.5123 8.10035 14.5123 7.54687 13.9589L4.85354 11.2648C4.30008 10.7114 4.30008 9.79937 4.85354 9.24518C5.407 8.69172 6.319 8.69172 6.87319 9.24518L8.55657 10.9294L13.6083 5.87851C13.8855 5.60054 14.2525 5.46298 14.6192 5.46298Z" />
+        </g>
+    </svg>
+
+);
+
+export default AcceptIcon;
diff --git a/mito-ai/src/icons/CopyIcon.tsx b/mito-ai/src/icons/CopyIcon.tsx
@@ -1,8 +1,3 @@
-
-
-
-
-
 import React from 'react';
 
 const CopyIcon: React.FC = () => (

diff --git a/mito-ai/src/icons/RejectIcon.tsx b/mito-ai/src/icons/RejectIcon.tsx
@@ -0,0 +1,12 @@
+import React from 'react';
+
+const RejectIcon: React.FC = () => (
+
+    <svg width="1em" height="1em" viewBox="0 0 21 20" fill="none" xmlns="http://www.w3.org/2000/svg">
+        <g fill="currentColor">
+            <path d="M10.2245 0C7.57229 0 5.02869 1.0536 3.15349 2.929C1.27817 4.8044 0.224487 7.348 0.224487 10C0.224487 12.652 1.27809 15.1958 3.15349 17.071C5.02889 18.9463 7.57249 20 10.2245 20C12.8765 20 15.4203 18.9464 17.2955 17.071C19.1708 15.1956 20.2245 12.652 20.2245 10C20.2245 8.2446 19.7624 6.5202 18.8847 5C18.0071 3.47984 16.7447 2.2174 15.2245 1.3398C13.7044 0.46206 11.9799 3.9978e-05 10.2245 3.9978e-05L10.2245 0ZM13.9315 12.293C14.1244 12.4794 14.2344 12.7354 14.2367 13.0036C14.2391 13.2718 14.1335 13.5298 13.9438 13.7194C13.7542 13.9091 13.4962 14.0146 13.228 14.0123C12.9598 14.01 12.7038 13.9 12.5174 13.7071L10.2244 11.4141L7.93143 13.7071C7.67737 13.9524 7.31283 14.0455 6.97229 13.9521C6.63175 13.8588 6.36565 13.5927 6.27229 13.2521C6.17893 12.9116 6.27205 12.5471 6.51737 12.293L8.81037 10L6.51737 7.707C6.27205 7.45294 6.17893 7.0884 6.27229 6.74786C6.36565 6.40732 6.63175 6.14122 6.97229 6.04786C7.31283 5.9545 7.67737 6.04763 7.93143 6.29294L10.2244 8.58594L12.5174 6.29294V6.29286C12.7715 6.04754 13.136 5.95442 13.4766 6.04778C13.8171 6.14114 14.0832 6.40724 14.1766 6.74778C14.2699 7.08832 14.1768 7.45286 13.9315 7.70692L11.6385 9.99992L13.9315 12.293Z" />
+        </g>
+    </svg>
+);
+
+export default RejectIcon;
diff --git a/mito-ai/style/ChatTaskpane.css b/mito-ai/style/ChatTaskpane.css
@@ -147,23 +147,6 @@
 }
 
 .code-block-toolbar button {
+    height: 16px;
     font-size: 12px !important;
-}
-
-.code-block-accept-button {
-    background-color: var(--green-600) !important;
-    color: white !important;
-}
-
-.code-block-accept-button:hover {
-    background-color: var(--green-700) !important;
-}
-
-.code-block-deny-button {
-    background-color: var(--red-600) !important;
-    color: white !important;
-}
-
-.code-block-deny-button:hover {
-    background-color: var(--red-700) !important;
-}
+}
diff --git a/mito-ai/style/IconButton.css b/mito-ai/style/IconButton.css
@@ -1,4 +1,7 @@
 .icon-button {
+    display: flex;
+    align-items: center;
+    justify-content: center;
     background-color: transparent;
     border: none;
     cursor: pointer;

diff --git a/mito-ai/style/button.css b/mito-ai/style/button.css
@@ -0,0 +1,32 @@
+/* 
+    Classes that can be used for any button, making it easier to keep 
+    the theme consistent. For example, use them to make a textButton or
+    textAndIconButton green.
+*/
+
+.button-green {
+    background-color: var(--green-400);
+    color: var(--green-900) !important;
+}
+
+.button-green:hover {
+    background-color: var(--green-500);
+}
+
+.button-red {
+    background-color: var(--red-400);
+    color: var(--red-900) !important;
+}
+
+.button-red:hover {
+    background-color: var(--red-500);
+}
+
+.button-gray {
+    background-color: var(--jp-layout-color2);
+    color: var(--jp-content-font-color1) !important;
+}
+
+.button-gray:hover {
+    background-color: var(--jp-layout-color3);
+}
diff --git a/tests/mitoai_ui_tests/utils.ts b/tests/mitoai_ui_tests/utils.ts
@@ -67,7 +67,7 @@ export const editMitoAIMessage = async (
 }
 
 export const clickPreviewButton = async (page: IJupyterLabPageFixture) => {
-    await page.getByRole('button', { name: 'Overwrite Active Cell' }).click();
+    await page.locator('.chat-message-buttons').getByRole('button', { name: 'Overwrite Active Cell' }).click();
     await waitForIdle(page);
 }
 
@@ -78,9 +78,9 @@ export const clickAcceptButton = async (
     { useCellToolbar = false }: { useCellToolbar?: boolean } = {useCellToolbar: false}
 ) => {
     if (useCellToolbar) {
-        await page.getByText('Accept AI edits').click();
+        await page.locator('.jp-ToolbarButtonComponent-label').filter({ hasText: 'Accept' }).click();
     } else {
-        await page.locator('.chat-taskpane').getByRole('button', { name: 'Accept code' }).click();
+        await page.locator('.chat-message-buttons').getByRole('button', { name: 'Accept code' }).click();
     }
     await waitForIdle(page);
 }
@@ -92,9 +92,9 @@ export const clickDenyButton = async (
     { useCellToolbar = false }: { useCellToolbar?: boolean } = {useCellToolbar: false}
 ) => {
     if (useCellToolbar) {
-        await page.getByText('Reject AI edits').click();
+        await page.locator('.jp-ToolbarButtonComponent-label').filter({ hasText: 'Reject' }).click();
     } else {
-        await page.locator('.chat-taskpane').getByRole('button', { name: 'Reject code' }).click();
+        await page.locator('.chat-message-buttons').getByRole('button', { name: 'Reject code' }).click();
     }
     await waitForIdle(page);
 }
Original file line number	Diff line number	Diff line change
		@@ -1,8 +1,3 @@





		import React from 'react';

		const CopyIcon: React.FC = () => (
Expand Down