pre-commit

jupyterlab · Jun 11, 2024 · 09f3102 · 09f3102
1 parent f9e4009
commit 09f3102
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 8 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py
@@ -1,6 +1,6 @@
 from typing import Dict, Type
 
-from jupyter_ai.models import HumanChatMessage, CellWithErrorSelection
+from jupyter_ai.models import CellWithErrorSelection, HumanChatMessage
 from jupyter_ai_magics.providers import BaseProvider
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
@@ -31,9 +31,16 @@
 """.strip()
 
 FIX_PROMPT_TEMPLATE = PromptTemplate(
-                input_variables=["extra_instructions", "cell_content", "traceback", "error_name", "error_value"],
-                template=FIX_STRING_TEMPLATE,
-            )
+    input_variables=[
+        "extra_instructions",
+        "cell_content",
+        "traceback",
+        "error_name",
+        "error_value",
+    ],
+    template=FIX_STRING_TEMPLATE,
+)
+
 
 class FixChatHandler(BaseChatHandler):
     id = "fix"
@@ -59,12 +66,15 @@ def create_llm_chain(
 
     async def process_message(self, message: HumanChatMessage):
         if not (message.selection and message.selection.type == "cell-with-error"):
-            self.reply("`/fix` be given a cell with error output. Please click on a cell with error output and retry.", message)
+            self.reply(
+                "`/fix` be given a cell with error output. Please click on a cell with error output and retry.",
+                message,
+            )
             return
-        
+
         # hint type of selection
         selection: CellWithErrorSelection = message.selection
-        
+
         # parse additional instructions specified after `/fix`
         extra_instructions = message.body[4:].strip() or "None."
 
@@ -75,6 +85,6 @@ async def process_message(self, message: HumanChatMessage):
             cell_content=selection.source,
             error_name=selection.error.name,
             error_value=selection.error.value,
-            traceback="\n".join(selection.error.traceback)
+            traceback="\n".join(selection.error.traceback),
         )
         self.reply(response, message)
diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py
@@ -13,13 +13,16 @@ class CellError(BaseModel):
     value: str
     traceback: List[str]
 
+
 class CellWithErrorSelection(BaseModel):
     type: Literal["cell-with-error"] = "cell-with-error"
     source: str
     error: CellError
 
+
 Selection = Union[CellWithErrorSelection]
 
+
 # the type of message used to chat with the agent
 class ChatRequest(BaseModel):
     prompt: str