From 486b2fee678f1baaf40a40ebdb84cafe31edf110 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Tue, 7 Nov 2023 20:46:13 -0800 Subject: [PATCH] Fixes %ai error for SM endpoints. --- packages/jupyter-ai-magics/jupyter_ai_magics/magics.py | 8 +++++--- packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py | 3 ++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py index 5266b071d..8763566f4 100644 --- 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 index a3b14bbb0..cadd41f4a 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -47,11 +47,11 @@ 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" @@ -59,6 +59,7 @@ 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]