Merge pull request #1499 from mito-ds/upgrade-notification

Update Upgrade to Pro CTAs

mito-ai/src/Extensions/AiChat/ChatMessage/AlertBlock.tsx

@@ -1,33 +1,47 @@
-import React, { useEffect, useState } from 'react';
-import ErrorIcon from '../../../icons/ErrorIcon';
+import React from 'react';
+import TextButton from '../../../components/TextButton';
+import { FREE_TIER_LIMIT_REACHED_ERROR_TITLE } from '../../../utils/errors';
+import { STRIPE_PAYMENT_LINK } from '../../../utils/stripe';
 
 
 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 = (
+    if (mitoAIConnectionErrorType === FREE_TIER_LIMIT_REACHED_ERROR_TITLE) {
+        return (
+            <div className="chat-message-alert">
+                <p>
+                    Your Mito AI free trial has ended. To continue using Mito AI, upgrade to <a href="https://www.trymito.io/plans" target="_blank">Mito Pro</a> and get access to:
+                </p>
+                <ul>
+                    <li>Unlimited AI Completions</li>
+                    <li>All Mito Spreadsheet Pro features</li>
+                </ul>
                 <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.
+                    Or supply your own Open AI Key to continue using the basic version of Mito AI. 
                 </p>
-            );
-            setMessage(message);
-        }
-        else {
-            setMessage(content);
-        }
-    }, [content]);
+                <TextButton
+                    title="Upgrade to Pro"
+                    text="Upgrade to Pro"
+                    onClick={() => {
+                        window.open(STRIPE_PAYMENT_LINK, '_blank');
+                    }}
+                    variant="purple"
+                    width="block"
+                />
+            </div>
+        );
+    }
 
     return (
         <div className="chat-message-alert">
-            <span style={{ marginRight: '4px' }}><ErrorIcon /></span>
-            {message}
+            {content}
         </div>
     );
 };

mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx

@@ -123,13 +123,15 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                             icon={PlayButtonIcon}
                                             title={'Write the Ai generated code to the active cell in the jupyter notebook, replacing the current code.'}
                                             variant='gray'
+                                            width='fit-contents'
                                         />
                                         <TextAndIconButton 
                                             onClick={() => {copyToClipboard(messagePart)}}
                                             text={'Copy'}
                                             icon={CopyIcon}
                                             title={'Copy the Ai generated code to your clipboard'}
                                             variant='gray'
+                                            width='fit-contents'
                                         />
                                     </div>
                                 }
@@ -140,12 +142,14 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                             text={`Accept code ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`}
                                             title={'Accept the Ai generated code'}
                                             variant='green'
+                                            width='fit-contents'
                                         />
                                         <TextButton 
                                             onClick={() => {rejectAICode()}}
                                             text={`Reject code ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`}
                                             title={'Reject the Ai generated code and revert to the previous version of the code cell'}
                                             variant='red'
+                                            width='fit-contents'
                                         />
                                     </div>
 

mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx

@@ -1,5 +1,7 @@
 import React, { useEffect, useRef, useState } from 'react';
 import '../../../style/ChatTaskpane.css';
+import '../../../style/button.css';
+import '../../../style/TextButton.css';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { writeCodeToCellByID, getCellCodeByID, getActiveCellID, highlightCodeCell } from '../../utils/notebook';
 import ChatMessage from './ChatMessage/ChatMessage';
@@ -405,7 +407,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         */
         app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_ACCEPT_CODE, {
             label: `Accept ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`,
-            className: 'text-and-icon-button button-green small',
+            className: 'text-button-mito-ai button-base button-green',
             caption: 'Accept Code',
             execute: () => {acceptAICode()},
             // We use the cellStateBeforeDiff because it contains the code cell ID that we want to write to
@@ -421,7 +423,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
 
         app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_REJECT_CODE, {
             label: `Reject ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`,
-            className: 'text-and-icon-button button-red small',
+            className: 'text-button-mito-ai button-base button-red',
             caption: 'Reject Code',
             execute: () => {rejectAICode()},
             isVisible: () => {

mito-ai/src/Extensions/InlineCompleter/provider.ts

@@ -22,6 +22,8 @@ import type {
   InlineCompletionStreamChunk
 } from '../../utils/websocket/models';
 import { IChatMessageMetadata } from '../AiChat/ChatHistoryManager';
+import { STRIPE_PAYMENT_LINK } from '../../utils/stripe';
+import { FREE_TIER_LIMIT_REACHED_ERROR_TITLE } from '../../utils/errors';
 
 /**
  * Mito AI inline completer
@@ -33,26 +35,26 @@ export class MitoAIInlineCompleter
   private _client: CompletionWebsocketClient;
   private _counter = 0;
   private _isDisposed = false;
-  private _settings: MitoAIInlineCompleter.ISettings =
-    MitoAIInlineCompleter.DEFAULT_SETTINGS;
+  private _settings: MitoAIInlineCompleter.ISettings = MitoAIInlineCompleter.DEFAULT_SETTINGS;
+
   // Store only one inline completion stream
-  //   Each new request should invalidate any other suggestions.
+  // Each new request should invalidate any other suggestions.
   private _currentToken = '';
-  private _currentStream: Stream<
-    MitoAIInlineCompleter,
-    InlineCompletionStreamChunk
-  > | null = null;
+  private _currentStream: Stream<MitoAIInlineCompleter, InlineCompletionStreamChunk> | null = null;
+
   /**
    * Block processing chunks while waiting for the acknowledge request
    * that will provide the unique completion token.
    */
   private _completionLock = new PromiseDelegate<void>();
-  private _fullCompletionMap = new WeakMap<
-    Stream<MitoAIInlineCompleter, InlineCompletionStreamChunk>,
-    string
-  >();
+  private _fullCompletionMap = new WeakMap<Stream<MitoAIInlineCompleter, InlineCompletionStreamChunk>, string>();
   private _variableManager: IVariableManager;
 
+  // We only want to display the free tier limit reached notification once 
+  // per session to avoid spamming the user. 
+  private _displayed_free_tier_limit_reached_notification = false;
+
+
   constructor({
     serverSettings,
     variableManager,
@@ -144,6 +146,7 @@ export class MitoAIInlineCompleter
     request: CompletionHandler.IRequest,
     context: IInlineCompletionContext
   ): Promise<IInlineCompletionList<IInlineCompletionItem>> {
+    console.log("HERE")
     if (!this.isEnabled()) {
       return Promise.reject('Mito AI completion is disabled.');
     }
@@ -201,7 +204,12 @@ export class MitoAIInlineCompleter
       }
 
       const error = result.error;
-      if (error) {
+      if (error?.title === FREE_TIER_LIMIT_REACHED_ERROR_TITLE) {
+        if (!this._displayed_free_tier_limit_reached_notification) {
+          this._notifyFreeTierLimitReached();
+          this._displayed_free_tier_limit_reached_notification = true;
+        }
+      } else if (error) {
         this._notifyCompletionFailure(error);
         throw new Error(
           `Inline completion failed: ${error.error_type}\n${error.traceback}`
@@ -279,6 +287,26 @@ export class MitoAIInlineCompleter
     return request.text.slice(request.offset);
   }
 
+  private _notifyFreeTierLimitReached() {
+    Notification.emit(`Your Mito AI free trial ended. Upgrade to Mito Pro to access advanced models and continue using Mito AI or supply your own key.`, 'error', {
+      autoClose: false,
+      actions: [
+        {
+          label: 'Upgrade to Mito Pro',
+          callback: () => {
+            window.open(STRIPE_PAYMENT_LINK, '_blank');
+          }
+        },
+        {
+          label: "Learn more",
+          callback: () => {
+            window.open("https://www.trymito.io/plans", '_blank');
+          }
+        }
+      ]
+    });
+  }
+
   private _notifyCompletionFailure(error: CompletionError) {
     Notification.emit(`Inline completion failed: ${error.error_type}`, 'error', {
       autoClose: false,
@@ -302,7 +330,11 @@ export class MitoAIInlineCompleter
     _emitter: CompletionWebsocketClient,
     chunk: ICompletionStreamChunk
   ) {
-    if (chunk.error) {
+
+    if (chunk.error?.title === FREE_TIER_LIMIT_REACHED_ERROR_TITLE) {
+      this._notifyFreeTierLimitReached();
+      this._displayed_free_tier_limit_reached_notification = true;
+    } else if (chunk.error) {
       this._notifyCompletionFailure(chunk.error);
     }
 

mito-ai/src/Extensions/status/index.tsx

@@ -9,11 +9,11 @@ import {
 } from '@jupyterlab/ui-components';
 import React from 'react';
 import { NucleusLabIcon } from '../../icons';
-import type {
-  ErrorMessage,
-  IAICapabilities
-} from '../../utils/websocket/models';
+import type { ErrorMessage, IAICapabilities } from '../../utils/websocket/models';
 import { IChatTracker, type IChatWidget } from '../AiChat/token';
+import { FREE_TIER_LIMIT_REACHED_ERROR_TITLE } from '../../utils/errors';
+import TextButton from '../../components/TextButton';
+import { STRIPE_PAYMENT_LINK } from '../../utils/stripe';
 
 /**
  * Mito AI status model
@@ -61,29 +61,40 @@ class StatusModel extends VDomModel {
  */
 class StatusPopUp extends VDomRenderer<StatusModel> {
   protected render(): JSX.Element {
+
+    let status_paragraph = <p className='mito-ai-status-ready'>Ready</p>
+    if (this.model.lastError?.title == FREE_TIER_LIMIT_REACHED_ERROR_TITLE) {
+      status_paragraph = <p className='mito-ai-status-error'>Free Trial Expired</p>
+    }
+
     return (
       <div className="mito-ai-status-popup">
-        <h4>Mito AI Status</h4>
-        <p>Provider: {this.model.capabilities?.provider ?? 'None'}</p>
-        {this.model.capabilities && (
+        <h3>Mito AI Status</h3>
+        <div className='mito-ai-status-popup-table-row'>
+          <p>Status:</p>
+          {status_paragraph}
+        </div>
+        <div className='mito-ai-status-popup-table-row'>
+          <p>Provider:</p>
+          <p>{this.model.capabilities?.provider ?? 'None'}</p>
+        </div>
+        {this.model.lastError?.title == FREE_TIER_LIMIT_REACHED_ERROR_TITLE && (
           <>
-            <p>Configuration:</p>
-            <ul>
-              <li>Model: {this.model.capabilities.configuration['model']}</li>
-              <li>
-                Max tokens:{' '}
-                {this.model.capabilities.configuration[
-                  'max_completion_tokens'
-                ] ?? 'undefined'}
-              </li>
-              <li>
-                Temperature:{' '}
-                {this.model.capabilities.configuration['temperature']}
-              </li>
-            </ul>
+            <p>
+              ⚠️ Your Mito AI free trial has ended. Upgrade to <a href="https://www.trymito.io/plans" target="_blank">Mito Pro</a> or supply your own Open AI Key to get access to more advanced models and continue using Mito AI.
+            </p>
+            <TextButton
+              title="Upgrade to Pro"
+              text="Upgrade to Pro"
+              onClick={() => {
+                window.open(STRIPE_PAYMENT_LINK, '_blank');
+              }}
+              variant='gray'
+              width='block'
+            />
           </>
         )}
-        {this.model.lastError && (
+        {this.model.lastError && this.model.lastError.title !== FREE_TIER_LIMIT_REACHED_ERROR_TITLE && (
           <>
             <p>Last error:</p>
             <ul>
@@ -125,7 +136,7 @@ class StatusItem extends ReactWidget {
         small
         title="Mito AI Status"
       >
-        <NucleusLabIcon.react tag={'span'} stylesheet={'statusBar'} />
+        Mito AI &nbsp; <NucleusLabIcon.react tag={'span'} stylesheet={'statusBar'} />
       </Button>
     );
   }

mito-ai/src/components/TextAndIconButton.tsx

@@ -12,7 +12,12 @@ interface TextAndIconButtonProps extends ButtonProps {
 const TextAndIconButton: React.FC<TextAndIconButtonProps> = ({ text, icon: Icon, onClick, title, variant }) => {
 
     return (
-        <button className={classNames("text-and-icon-button", `button-${variant}`)} onClick={onClick} title={title}>
+        <button className={classNames(
+                'text-and-icon-button',
+                'button-base',
+                `button-${variant}`,
+                `button-width-fit-contents`
+            )} onClick={onClick} title={title}>
             {text}
             <span className="text-and-icon-button__icon">
                 <Icon />

mito-ai/src/components/TextButton.tsx

@@ -1,5 +1,5 @@
 import React from 'react';
-import '../../style/TextAndIconButton.css';
+import '../../style/TextButton.css';
 import '../../style/button.css';
 import { classNames } from '../utils/classNames';
 
@@ -8,16 +8,22 @@ export interface ButtonProps {
     text: string;
     onClick: () => void;
     title: string;
-    variant: 'green' | 'red' | 'gray';
+    variant: 'green' | 'red' | 'gray' | 'purple';
+    width: 'block' | 'fit-contents';
 }
 
 // Text Button is just the basic Button Props, nothing else. 
 interface TextButtonProps extends ButtonProps {}
 
-const TextButton: React.FC<TextButtonProps> = ({ text, onClick, title, variant }) => {
+const TextButton: React.FC<TextButtonProps> = ({ text, onClick, title, variant, width }) => {
 
     return (
-        <button className={classNames("text-and-icon-button", `button-${variant}`)} onClick={onClick} title={title}>
+        <button className={classNames(
+            "text-button-mito-ai", 
+            "button-base",
+            `button-${variant}`,
+            `button-width-${width}`
+        )} onClick={onClick} title={title}>
             {text}
         </button>
     )

mito-ai/src/utils/errors.tsx

@@ -0,0 +1 @@
+export const FREE_TIER_LIMIT_REACHED_ERROR_TITLE = 'mito_server_free_tier_limit_reached'

mito-ai/src/utils/stripe.tsx

@@ -0,0 +1 @@
+export const STRIPE_PAYMENT_LINK = "https://checkout.stripe.com/c/pay/cs_live_a14mzpiTbiuc62hknBSe5JifejsWpDVdcARfa8SoYBvx2WEhINZpsXuWO7#fidkdWxOYHwnPyd1blppbHNgWjA0VFVVTFZOM108YjVLX1ByfzxNR2pqdW5yVUY3Rks1cXNDc0BWamh%2FRktKMWBfUXRtYkdNQ2tLMl1mc0RcUkJVPUZAV2Z%2FVnBEZDc0UXNAZFY1MmQ9YGQwNTVDcVNnMjFOUicpJ2N3amhWYHdzYHcnP3F3cGApJ2lkfGpwcVF8dWAnPyd2bGtiaWBabHFgaCcpJ2BrZGdpYFVpZGZgbWppYWB3dic%2FcXdwYHgl"