Error handling and messaging when the chat service doesn't work (jupy…

…terlab#88) * WIP: fix missing env * Added error handling when chat provider key is missing
dbelgrod · Apr 18, 2023 · aaafdec · aaafdec
1 parent b9178cc
commit aaafdec
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 32 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -1,4 +1,5 @@
 import asyncio
+import os
 import queue
 from langchain.memory import ConversationBufferWindowMemory
 from jupyter_ai.actors.default import DefaultActor 
@@ -13,7 +14,7 @@
 from importlib_metadata import entry_points
 import inspect
 from .engine import BaseModelEngine
-from jupyter_ai_magics.providers import ChatOpenAIProvider
+from jupyter_ai_magics.providers import ChatOpenAINewProvider, ChatOpenAIProvider
 
 import ray
 from ray.util.queue import Queue
@@ -89,6 +90,9 @@ def initialize_settings(self):
         self.settings["ai_default_tasks"] = default_tasks
         self.log.info("Registered all default tasks.")
 
+        if ChatOpenAINewProvider.auth_strategy.name not in os.environ:
+            raise EnvironmentError(f"`{ChatOpenAINewProvider.auth_strategy.name}` value not set in environment. For chat to work, this value should be provided.")
+
         ## load OpenAI provider
         self.settings["openai_chat"] = ChatOpenAIProvider(model_id="gpt-3.5-turbo")
 
@@ -104,6 +108,7 @@ def initialize_settings(self):
         # store chat messages in memory for now
         self.settings["chat_history"] = []
 
+
         reply_queue = Queue()
         self.settings["reply_queue"] = reply_queue
 
@@ -136,4 +141,4 @@ def initialize_settings(self):
 
         reply_processor = ReplyProcessor(self.settings['chat_handlers'], reply_queue, log=self.log)        
         loop = asyncio.get_event_loop()
-        loop.create_task(reply_processor.start())
+        loop.create_task(reply_processor.start())            
diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py
@@ -81,16 +81,8 @@ async def get(self, id=None):
 class ChatHistoryHandler(BaseAPIHandler):
     """Handler to return message history"""
 
-    _chat_provider = None  
     _messages = []
 
-    @property
-    def chat_provider(self):
-        actor = ray.get_actor("default")
-        self._chat_provider = actor.chat_provider
-
-        return self._chat_provider
-
     @property
     def chat_history(self):
         return self.settings["chat_history"]
@@ -104,13 +96,6 @@ async def get(self):
         history = ChatHistory(messages=self.chat_history)
         self.finish(history.json())
 
-    @tornado.web.authenticated
-    async def delete(self):
-        self.chat_provider.memory.chat_memory.clear()
-        self.chat_history = []
-        self.set_status(204)
-        self.finish()
-
 
 class ChatHandler(
     JupyterHandler,
@@ -122,7 +107,6 @@ class ChatHandler(
 
     _chat_provider = None  
     _chat_message_queue = None 
-    _reply_queue = None
 
     @property
     def chat_message_queue(self):
@@ -150,13 +134,6 @@ def chat_client(self) -> ChatClient:
     @property
     def chat_history(self) -> List[ChatMessage]:
         return self.settings["chat_history"]
-
-    @property
-    def reply_queue(self):
-        if not self._reply_queue:
-            self._reply_queue = self.settings["reply_queue"]
-
-        return self._reply_queue
 
     def initialize(self):
         self.log.debug("Initializing websocket connection %s", self.request.path)

diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts
@@ -40,6 +40,7 @@ export class ChatHandler implements IDisposable {
         (token ? `?token=${encodeURIComponent(token)}` : '');
 
       const socket = (this._socket = new WebSocket(url));
+      socket.onerror = (e) => reject(e);
       socket.onmessage = msg =>
         msg.data && this._onMessage(JSON.parse(msg.data));
 

diff --git a/packages/jupyter-ai/src/components/chat.tsx b/packages/jupyter-ai/src/components/chat.tsx
@@ -31,8 +31,13 @@ function ChatBody({ chatHandler }: ChatBodyProps): JSX.Element {
    */
   useEffect(() => {
     async function fetchHistory() {
-      const history = await chatHandler.getHistory();
-      setMessages(history.messages);
+      try {
+        const history = await chatHandler.getHistory();
+        setMessages(history.messages);
+      } catch (e) {
+
+      }
+
     }
 
     fetchHistory();

diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts
@@ -4,13 +4,14 @@ import {
   ILayoutRestorer
 } from '@jupyterlab/application';
 
-import { IWidgetTracker } from '@jupyterlab/apputils';
+import { IWidgetTracker, ReactWidget } from '@jupyterlab/apputils';
 import { IDocumentWidget } from '@jupyterlab/docregistry';
 import { IGlobalAwareness } from '@jupyterlab/collaboration';
 import type { Awareness } from 'y-protocols/awareness';
 import { buildChatSidebar } from './widgets/chat-sidebar';
 import { SelectionWatcher } from './selection-watcher';
 import { ChatHandler } from './chat_handler';
+import { buildErrorWidget } from './widgets/chat-error';
 
 export type DocumentTracker = IWidgetTracker<IDocumentWidget>;
 
@@ -35,10 +36,15 @@ const plugin: JupyterFrontEndPlugin<void> = {
      * Initialize chat handler, open WS connection
      */
     const chatHandler = new ChatHandler();
-    await chatHandler.initialize();
-
-    const chatWidget = buildChatSidebar(selectionWatcher, chatHandler, globalAwareness);
-
+
+    let chatWidget: ReactWidget | null = null;
+    try {
+      await chatHandler.initialize();
+      chatWidget = buildChatSidebar(selectionWatcher, chatHandler, globalAwareness);
+    } catch (e) {
+      chatWidget = buildErrorWidget()
+    }
+
     /**
      * Add Chat widget to right sidebar
      */

diff --git a/packages/jupyter-ai/src/widgets/chat-error.tsx b/packages/jupyter-ai/src/widgets/chat-error.tsx
@@ -0,0 +1,37 @@
+import React from 'react';
+import { ReactWidget } from '@jupyterlab/apputils';
+
+import { chatIcon } from '../icons';
+import { Alert, Box } from '@mui/material';
+import { JlThemeProvider } from '../components/jl-theme-provider';
+
+export function buildErrorWidget() {
+    const ErrorWidget = ReactWidget.create(
+        <JlThemeProvider>
+            <Box
+                sx={{
+                    width: '100%',
+                    height: '100%',
+                    boxSizing: 'border-box',
+                    background: 'var(--jp-layout-color0)',
+                    display: 'flex',
+                    flexDirection: 'column'
+                }}
+            >
+                <Box sx={{ padding: 4 }}>
+                    <Alert severity="error">
+                        There seems to be a problem with the Chat backend, please look at the
+                        JupyterLab server logs or contact your administrator to correct this
+                        problem.
+                    </Alert>
+                </Box>
+            </Box>
+        </JlThemeProvider>
+    );
+    ErrorWidget.id = 'jupyter-ai::chat';
+    ErrorWidget.title.icon = chatIcon;
+    ErrorWidget.title.caption = 'Jupyter AI Chat'; // TODO: i18n
+
+    return ErrorWidget
+}
+