error animation in cells for jupyter-ai

Co-authored-by: Andrew Fulton [email protected]

Author: Darshan808

packages/jupyter-ai-magics/jupyter_ai_magics/models/completion.py

@@ -28,13 +28,15 @@ class InlineCompletionRequest(BaseModel):
     # previous cells and following cells can be used to learn the wider context
     cell_id: Optional[str]
 
+class InlineCompletionError(BaseModel):
+    message: str
 
 class InlineCompletionItem(BaseModel):
     """The inline completion suggestion to be displayed on the frontend.
 
     See JupyterLab `InlineCompletionItem` documentation for the details.
     """
-
+    error: Optional[InlineCompletionError]
     insertText: str
     filterText: Optional[str]
     isIncomplete: Optional[bool]

packages/jupyter-ai/jupyter_ai/completions/handlers/base.py

@@ -118,6 +118,7 @@ async def on_message(self, message):
         async def handle_request_and_catch():
             try:
                 await handle_request
+                # raise Exception("An error occured!")
             except Exception as e:
                 await self.handle_exc(e, request)
 
@@ -129,14 +130,15 @@ async def handle_exc(self, e: Exception, request: InlineCompletionRequest):
         `handle_stream_request()`. This base class provides a default
         implementation, which may be overridden by subclasses.
         """
+        title = e.args[0] if e.args else "Exception"
         error = CompletionError(
             type=e.__class__.__name__,
-            title=e.args[0] if e.args else "Exception",
+            title=title,
             traceback=traceback.format_exc(),
         )
         self.reply(
             InlineCompletionReply(
-                list=InlineCompletionList(items=[]),
+                list=InlineCompletionList(items=[{"error":{"message":title},"insertText":""}]),
                 error=error,
                 reply_to=request.number,
             )

packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py

@@ -1,13 +1,16 @@
 import json
 from types import SimpleNamespace
 from typing import Union
+import traceback
 
 import pytest
 from jupyter_ai.completions.handlers.default import DefaultInlineCompletionHandler
 from jupyter_ai.completions.models import (
     InlineCompletionReply,
     InlineCompletionRequest,
     InlineCompletionStreamChunk,
+    CompletionError,
+    InlineCompletionList,
 )
 from jupyter_ai_magics import BaseProvider
 from langchain_community.llms import FakeListLLM
@@ -52,7 +55,19 @@ def reply(
 
     async def handle_exc(self, e: Exception, _request: InlineCompletionRequest):
         # raise all exceptions during testing rather
-        raise e
+        title = e.args[0] if e.args else "Exception"
+        error = CompletionError(
+            type=e.__class__.__name__,
+            title=title,
+            traceback=traceback.format_exc(),
+        )
+        self.reply(
+            InlineCompletionReply(
+                list=InlineCompletionList(items=[{"error":{"message":title},"insertText":""}]),
+                error=error,
+                reply_to=_request.number,
+            )
+        )
 
 
 @fixture
@@ -191,3 +206,19 @@ async def test_handle_stream_request():
     assert third.type == "stream"
     assert third.response.insertText == "test"
     assert third.done is True
+
+async def test_handle_request_with_error(inline_handler):
+    inline_handler = MockCompletionHandler(
+        lm_provider=MockProvider,
+        lm_provider_params={
+            "model_id": "model",
+            "responses": ["test"],
+        },
+    )
+    dummy_request = InlineCompletionRequest(
+        number=1, prefix="", suffix="", mime="", stream=True
+    )
+    await inline_handler.on_message(json.dumps(dict(dummy_request)))
+    await inline_handler.tasks[0]
+    error_message = inline_handler.messages[-1].dict().get('list', {}).get('items', [{}])[0].get('error', {}).get('message', None)
+    assert error_message is not None