add exception handling for update and learn

ion2088 · Feb 19, 2024 · ca81808 · ca81808
1 parent 653c44b
commit ca81808
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 24 deletions.
diff --git a/src/firedust/assistant.py b/src/firedust/assistant.py
@@ -302,7 +302,7 @@ def delete(
         if not response.is_success:
             raise APIError(
                 code=response.status_code,
-                message=f"An error occured while deleting the assistant with id {id}: {response.text}",
+                message=f"An error occured while deleting the assistant with id {assistant_id}: {response.text}",
             )
         LOG.info(f"Successfully deleted assistant {assistant_id}.")
 
@@ -360,16 +360,21 @@ def name(self, name: str) -> None:
         Args:
             name (str): The new name of the assistant.
         """
+        response = self.api_client.put(
+            f"/assistant/{self.assistant.config.id}/update/name/{name}",
+        )
+        if not response.is_success:
+            raise APIError(
+                code=response.status_code,
+                message=f"Failed to update the name of the assistant with id {self.assistant.config.id}: {response.text}",
+            )
         self.assistant._config = AssistantConfig(
             id=self.assistant._config.id,
             name=name,
             instructions=self.assistant._config.instructions,
             inference=self.assistant._config.inference,
             memory=self.assistant._config.memory,
         )
-        self.api_client.put(
-            f"/assistant/{self.assistant.config.id}/update/name/{name}",
-        )
 
     def instructions(self, instructions: List[str]) -> None:
         """
@@ -378,20 +383,25 @@ def instructions(self, instructions: List[str]) -> None:
         Args:
             instructions (List[str]): The new instructions of the assistant.
         """
+        response = self.api_client.post(
+            "/assistant/update/instructions",
+            data={
+                "assistant_id": str(self.assistant.config.id),
+                "instructions": instructions,
+            },
+        )
+        if not response.is_success:
+            raise APIError(
+                code=response.status_code,
+                message=f"Failed to update the instructions of the assistant with id {self.assistant.config.id}: {response.text}",
+            )
         self.assistant._config = AssistantConfig(
             id=self.assistant._config.id,
             name=self.assistant._config.name,
             instructions=instructions,
             inference=self.assistant._config.inference,
             memory=self.assistant._config.memory,
         )
-        self.api_client.post(
-            "/assistant/update/instructions",
-            data={
-                "assistant_id": str(self.assistant.config.id),
-                "instructions": instructions,
-            },
-        )
 
     def inference(self, inference_config: InferenceConfig) -> None:
         """
@@ -400,17 +410,22 @@ def inference(self, inference_config: InferenceConfig) -> None:
         Args:
             inference_config (InferenceConfig): The new inference configuration.
         """
+        response = self.api_client.post(
+            "/assistant/update/inference",
+            data={
+                "assistant_id": str(self.assistant.config.id),
+                "inference": inference_config.model_dump(),
+            },
+        )
+        if not response.is_success:
+            raise APIError(
+                code=response.status_code,
+                message=f"Failed to update the inference configuration of the assistant with id {self.assistant.config.id}: {response.text}",
+            )
         self.assistant._config = AssistantConfig(
             id=self.assistant._config.id,
             name=self.assistant._config.name,
             instructions=self.assistant._config.instructions,
             inference=inference_config,
             memory=self.assistant._config.memory,
         )
-        self.api_client.post(
-            "/assistant/update/inference",
-            data={
-                "assistant_id": str(self.assistant.config.id),
-                "inference": inference_config.model_dump(),
-            },
-        )
diff --git a/src/firedust/learning/_base.py b/src/firedust/learning/_base.py
@@ -2,6 +2,7 @@
 from typing import Union
 
 from firedust._utils.api import APIClient
+from firedust._utils.errors import APIError
 from firedust._utils.types.assistant import AssistantConfig
 
 
@@ -33,10 +34,12 @@ def fast(self, text: str) -> None:
         Args:
             text (str): The text to learn.
         """
-        self.api_client.post(
+        response = self.api_client.post(
             "learn/fast",
             data={"assistant_id": self.assistant.id, "text": text},
         )
+        if not response.is_success:
+            raise APIError(f"Failed to teach the assistant: {response.text}")
 
     def from_pdf(self, pdf: Union[str, Path]) -> None:
         """
@@ -46,10 +49,12 @@ def from_pdf(self, pdf: Union[str, Path]) -> None:
             pdf (Union[str, Path]): The path to the PDF file.
         """
         with open(pdf, "rb") as f:
-            self.api_client.post(
+            response = self.api_client.post(
                 "learn/pdf",
                 data={"assistant_id": self.assistant.id, "pdf": f},
             )
+            if not response.is_success:
+                raise APIError(f"Failed to teach the assistant: {response.text}")
 
     def from_url(self, url: str) -> None:
         """
@@ -58,10 +63,12 @@ def from_url(self, url: str) -> None:
         Args:
             url (str): The URL to the resource.
         """
-        self.api_client.post(
+        response = self.api_client.post(
             "learn/url",
             data={"assistant_id": self.assistant.id, "url": url},
         )
+        if not response.is_success:
+            raise APIError(f"Failed to teach the assistant: {response.text}")
 
     def from_image(self, image: Union[str, Path]) -> None:
         """
@@ -71,10 +78,12 @@ def from_image(self, image: Union[str, Path]) -> None:
             image (Union[str, Path]): The path to the image.
         """
         with open(image, "rb") as f:
-            self.api_client.post(
+            response = self.api_client.post(
                 "learn/image",
                 data={"assistant_id": self.assistant.id, "image": f},
             )
+            if not response.is_success:
+                raise APIError(f"Failed to teach the assistant: {response.text}")
 
     def from_audio(self, audio: Union[str, Path]) -> None:
         """
@@ -84,10 +93,12 @@ def from_audio(self, audio: Union[str, Path]) -> None:
             audio (Union[str, Path]): The path to the audio file.
         """
         with open(audio, "rb") as f:
-            self.api_client.post(
+            response = self.api_client.post(
                 "learn/audio",
                 data={"assistant_id": self.assistant.id, "audio": f},
             )
+            if not response.is_success:
+                raise APIError(f"Failed to teach the assistant: {response.text}")
 
     def from_video(self, video: Union[str, Path]) -> None:
         """
@@ -97,7 +108,9 @@ def from_video(self, video: Union[str, Path]) -> None:
             video (Union[str, Path]): The path to the video file.
         """
         with open(video, "rb") as f:
-            self.api_client.post(
+            response = self.api_client.post(
                 f"assistant/{self.assistant.id}/learn/video",
                 data={"video": f},
             )
+            if not response.is_success:
+                raise APIError(f"Failed to teach the assistant: {response.text}")