Merge pull request #1465 from mito-ds/upgrade-msg-redesign

Added alert block for nicer error messages
mito-ds · Jan 3, 2025 · 85b34ec · 85b34ec
2 parents 917460a + f21c10e
commit 85b34ec
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 10 deletions.
diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx
@@ -11,6 +11,7 @@ type PromptType = 'chat' | 'smartDebug' | 'codeExplain' | 'system'
 export interface IDisplayOptimizedChatHistory {
     message: OpenAI.Chat.ChatCompletionMessageParam
     type: 'openai message' | 'connection error',
+    mitoAIConnectionErrorType?: string | null,
     codeCellID: string | undefined
 }
 
@@ -187,7 +188,8 @@ export class ChatHistoryManager {
     addAIMessageFromResponse(
         messageContent: string | null, 
         promptType: PromptType, 
-        mitoAIConnectionError: boolean=false
+        mitoAIConnectionError: boolean=false,
+        mitoAIConnectionErrorType: string | null = null
     ): void {
         if (messageContent === null) {
             return
@@ -201,20 +203,22 @@ export class ChatHistoryManager {
             role: 'assistant',
             content: messageContent
         }
-        this._addAIMessage(aiMessage, promptType, mitoAIConnectionError)
+        this._addAIMessage(aiMessage, promptType, mitoAIConnectionError, mitoAIConnectionErrorType)
     }
 
     _addAIMessage(
         aiMessage: OpenAI.Chat.ChatCompletionMessageParam, 
         promptType: PromptType, 
-        mitoAIConnectionError: boolean=false
+        mitoAIConnectionError: boolean=false,
+        mitoAIConnectionErrorType?: string | null
     ): void {
         const activeCellID = getActiveCellID(this.notebookTracker)
 
         this.history.displayOptimizedChatHistory.push(
             {
                 message: aiMessage, 
                 type: mitoAIConnectionError ? 'connection error' : 'openai message',
+                mitoAIConnectionErrorType: mitoAIConnectionErrorType,
                 codeCellID: activeCellID
             }
         );

diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/AlertBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/AlertBlock.tsx
@@ -0,0 +1,35 @@
+import React, { useEffect, useState } from 'react';
+import ErrorIcon from '../../../icons/ErrorIcon';
+
+
+interface IAlertBlockProps {
+    content: string;
+    mitoAIConnectionErrorType: string | null;
+}
+
+const AlertBlock: React.FC<IAlertBlockProps> = ({ content, mitoAIConnectionErrorType }) => {
+    const [message, setMessage] = useState<string | JSX.Element>("");
+
+    useEffect(() => {
+        if (mitoAIConnectionErrorType === "mito_server_free_tier_limit_reached") {
+            const message = (
+                <p>
+                    You've reached the free tier limit for Mito AI. <a href="https://www.trymito.io/plans" target="_blank">Upgrade to Pro for unlimited uses</a> or supply your own OpenAI API key.
+                </p>
+            );
+            setMessage(message);
+        }
+        else {
+            setMessage(content);
+        }
+    }, [content]);
+
+    return (
+        <div className="chat-message-alert">
+            <span style={{ marginRight: '4px' }}><ErrorIcon /></span>
+            {message}
+        </div>
+    );
+};
+
+export default AlertBlock;
diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
@@ -3,10 +3,10 @@ import OpenAI from 'openai';
 import { classNames } from '../../../utils/classNames';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import CodeBlock from './CodeBlock';
+import AlertBlock from './AlertBlock';
 import MarkdownBlock from './MarkdownBlock';
 import { INotebookTracker } from '@jupyterlab/notebook';
 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';
 import PencilIcon from '../../../icons/Pencil';
@@ -24,6 +24,7 @@ interface IChatMessageProps {
     codeCellID: string | undefined
     messageIndex: number
     mitoAIConnectionError: boolean
+    mitoAIConnectionErrorType: string | null
     notebookTracker: INotebookTracker
     renderMimeRegistry: IRenderMimeRegistry
     app: JupyterFrontEnd
@@ -41,6 +42,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
     message,
     messageIndex,
     mitoAIConnectionError,
+    mitoAIConnectionErrorType,
     notebookTracker,
     renderMimeRegistry,
     isLastAiMessage,
@@ -163,11 +165,14 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                     }
                                 }}
                             >
-                                {mitoAIConnectionError && <span style={{ marginRight: '4px' }}><ErrorIcon /></span>}
-                                <MarkdownBlock
-                                    markdown={messagePart}
-                                    renderMimeRegistry={renderMimeRegistry}
-                                />
+                                {mitoAIConnectionError ? (
+                                    <AlertBlock content={messagePart} mitoAIConnectionErrorType={mitoAIConnectionErrorType} />
+                                ) : (
+                                    <MarkdownBlock
+                                        markdown={messagePart}
+                                        renderMimeRegistry={renderMimeRegistry}
+                                    />
+                                )}
                             </p>
                             {message.role === 'user' && (
                                 <div style={{ display: 'flex', justifyContent: 'flex-end', marginTop: '4px' }}>

diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
@@ -232,7 +232,8 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
                   ? aiResponse.error.hint
                   : `${aiResponse.error.error_type}: ${aiResponse.error.title}`,
                 promptType,
-                true
+                true,
+                aiResponse.error.title
               );
               setChatHistoryManager(newChatHistoryManager);
             } else {
@@ -544,6 +545,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
                             message={displayOptimizedChat.message}
                             codeCellID={displayOptimizedChat.codeCellID}
                             mitoAIConnectionError={displayOptimizedChat.type === 'connection error'}
+                            mitoAIConnectionErrorType={displayOptimizedChat.mitoAIConnectionErrorType || null}
                             messageIndex={index}
                             notebookTracker={notebookTracker}
                             renderMimeRegistry={renderMimeRegistry}

diff --git a/mito-ai/style/ChatTaskpane.css b/mito-ai/style/ChatTaskpane.css
@@ -149,4 +149,16 @@
 .code-block-toolbar button {
     height: 16px;
     font-size: 12px !important;
+}
+
+.chat-message-alert {
+    background-color: var(--purple-300);
+    border-radius: 5px;
+    padding: 10px;
+    border: 1px solid var(--purple-500);
+}
+
+.chat-message-alert a {
+    color: var(--purple-700);
+    text-decoration: underline;
 }