Fixes %ai error for SM endpoints.

jupyterlab · Nov 8, 2023 · 026b7ea · 026b7ea
1 parent a951dea
commit 026b7ea
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py
@@ -392,9 +392,11 @@ def handle_error(self, args: ErrorArgs):
 
         prompt = f"Explain the following error:\n\n{last_error}"
         # Set CellArgs based on ErrorArgs
-        cell_args = CellArgs(
-            type="root", model_id=args.model_id, format=args.format, reset=False
-        )
+        values = args.dict()
+        values["type"] = "root"
+        values["reset"] = False
+        cell_args = CellArgs(**values)
+
         return self.run_ai_cell(cell_args, prompt)
 
     def _append_exchange_openai(self, prompt: str, output: str):

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py
@@ -47,18 +47,19 @@ class CellArgs(BaseModel):
     model_id: str
     format: FORMAT_CHOICES_TYPE
     reset: bool
+    model_parameters: Optional[str]
     # The following parameters are required only for SageMaker models
     region_name: Optional[str]
     request_schema: Optional[str]
     response_path: Optional[str]
-    model_parameters: Optional[str]
 
 
 # Should match CellArgs, but without "reset"
 class ErrorArgs(BaseModel):
     type: Literal["error"] = "error"
     model_id: str
     format: FORMAT_CHOICES_TYPE
+    model_parameters: Optional[str]
     # The following parameters are required only for SageMaker models
     region_name: Optional[str]
     request_schema: Optional[str]