add test and code improvements

Author: Darshan808

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

@@ -28,18 +28,11 @@ 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,7 +118,6 @@ 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)
 
@@ -138,9 +137,7 @@ async def handle_exc(self, e: Exception, request: InlineCompletionRequest):
         )
         self.reply(
             InlineCompletionReply(
-                list=InlineCompletionList(
-                    items=[{"error": {"message": title}, "insertText": ""}]
-                ),
+                list=InlineCompletionList(items=[]),
                 error=error,
                 reply_to=request.number,
             )

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

@@ -6,8 +6,6 @@
 import pytest
 from jupyter_ai.completions.handlers.default import DefaultInlineCompletionHandler
 from jupyter_ai.completions.models import (
-    CompletionError,
-    InlineCompletionList,
     InlineCompletionReply,
     InlineCompletionRequest,
     InlineCompletionStreamChunk,
@@ -29,10 +27,16 @@ def __init__(self, **kwargs):
         if "responses" not in kwargs:
             kwargs["responses"] = ["Test response"]
         super().__init__(**kwargs)
+    
+    def _acall(self, *args, **kwargs):
+        if self.raise_exc:
+            raise Exception("Test exception")
+        else:
+            return super()._call(*args, **kwargs)
 
 
 class MockCompletionHandler(DefaultInlineCompletionHandler):
-    def __init__(self, lm_provider=None, lm_provider_params=None):
+    def __init__(self, lm_provider=None, lm_provider_params=None, raise_exc=False):
         self.request = HTTPServerRequest()
         self.application = Application()
         self.messages = []
@@ -53,25 +57,6 @@ def reply(
     ) -> None:
         self.messages.append(message)
 
-    async def handle_exc(self, e: Exception, _request: InlineCompletionRequest):
-        # raise all exceptions during testing rather
-        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
 def inline_handler() -> MockCompletionHandler:
     return MockCompletionHandler()
@@ -216,19 +201,13 @@ async def test_handle_request_with_error(inline_handler):
         lm_provider_params={
             "model_id": "model",
             "responses": ["test"],
-        },
+            "raise_exc": True
+        }
     )
     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
+    error = inline_handler.messages[-1].dict().get('error', None)
+    assert error is not None

packages/jupyter-ai/src/completions/provider.ts

@@ -122,15 +122,18 @@ export class JaiInlineProvider
             label: 'Show Traceback',
             callback: () => {
               showErrorMessage('Inline completion failed on the server side', {
-                message: error.traceback
+                message: `${error.title}\n${error.traceback}`
               });
             }
           }
         ]
       });
-      throw new Error(
-        `Inline completion failed: ${error.type}\n${error.traceback}`
-      );
+      const items = [{
+        error: {message: error.title},
+        insertText: ""
+      }]
+      console.log({items});
+      return { items }
     }
     return result.list;
   }

packages/jupyter-ai/src/completions/types.ts

@@ -35,6 +35,7 @@ export namespace AiCompleterService {
   export type CompletionError = {
     type: string;
     traceback: string;
+    title: string;
   };
 
   export type InlineCompletionReply = {