beekeeper-monitors-watsonx 1.1.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

beekeeper/monitors/watsonx/base.py

@@ -202,6 +202,40 @@ class WatsonxExternalPromptMonitor(PromptMonitor):
 
         return created_detached_pta.to_dict()["asset_id"]
 
+    def _delete_detached_prompt(self, detached_pta_id: str) -> None:
+        from ibm_aigov_facts_client import (  # type: ignore
+            AIGovFactsClient,
+            CloudPakforDataConfig,
+        )
+
+        try:
+            if hasattr(self, "_fact_cpd_creds") and self._fact_cpd_creds:
+                cpd_creds = CloudPakforDataConfig(**self._fact_cpd_creds)
+
+                aigov_client = AIGovFactsClient(
+                    container_id=self._container_id,
+                    container_type=self._container_type,
+                    cloud_pak_for_data_configs=cpd_creds,
+                    disable_tracing=True,
+                )
+
+            else:
+                aigov_client = AIGovFactsClient(
+                    api_key=self._api_key,
+                    container_id=self._container_id,
+                    container_type=self._container_type,
+                    disable_tracing=True,
+                    region=self.region.factsheet,
+                )
+
+        except Exception as e:
+            logging.error(
+                f"Error connecting to IBM watsonx.governance (factsheets): {e}",
+            )
+            raise
+
+        suppress_output(aigov_client.assets.delete_prompt_asset, detached_pta_id)
+
     def _create_deployment_pta(self, asset_id: str, name: str, model_id: str) -> str:
         from ibm_watsonx_ai import APIClient, Credentials  # type: ignore
 
@@ -239,6 +273,30 @@ class WatsonxExternalPromptMonitor(PromptMonitor):
 
         return wml_client.deployments.get_uid(created_deployment)
 
+    def _delete_deployment_pta(self, deployment_id: str):
+        from ibm_watsonx_ai import APIClient, Credentials  # type: ignore
+
+        try:
+            if hasattr(self, "_wml_cpd_creds") and self._wml_cpd_creds:
+                creds = Credentials(**self._wml_cpd_creds)
+
+                wml_client = APIClient(creds)
+                wml_client.set.default_space(self.space_id)
+
+            else:
+                creds = Credentials(
+                    url=self.region.watsonxai,
+                    api_key=self._api_key,
+                )
+                wml_client = APIClient(creds)
+                wml_client.set.default_space(self.space_id)
+
+        except Exception as e:
+            logging.error(f"Error connecting to IBM watsonx.ai Runtime: {e}")
+            raise
+
+        suppress_output(wml_client.deployments.delete, deployment_id)
+
     @deprecated(
         reason="'add_prompt_observer()' is deprecated and will be removed in a future version. Use 'create_prompt_monitor()' instead.",
         version="1.0.5",
@@ -381,6 +439,8 @@ class WatsonxExternalPromptMonitor(PromptMonitor):
             ```
         """
         task_id = TaskType.from_value(task_id).value
+        rollback_stack = []
+
         # DEPRECATION NOTICE
         if input_text is not None:
             warnings.warn(
@@ -493,11 +553,14 @@ class WatsonxExternalPromptMonitor(PromptMonitor):
             prompt_details,
             detached_asset_details,
         )
+        rollback_stack.append(lambda: self._delete_detached_prompt(detached_pta_id))
+
         deployment_id = None
         if self._container_type == "space":
             deployment_id = suppress_output(
                 self._create_deployment_pta, detached_pta_id, name, model_id
             )
+            rollback_stack.append(lambda: self._delete_deployment_pta(deployment_id))
 
         monitors = {
             "generative_ai_quality": {
@@ -558,10 +621,17 @@ class WatsonxExternalPromptMonitor(PromptMonitor):
 
         generative_ai_monitor_details = generative_ai_monitor_details.result._to_dict()
 
+        wos_status = generative_ai_monitor_details.get("status", {})
+        if wos_status.get("state") == "ERROR":
+            for rollback_step in reversed(rollback_stack):
+                rollback_step()
+            raise Exception(wos_status.get("failure"))
+
+
         return {
             "detached_prompt_template_asset_id": detached_pta_id,
             "deployment_id": deployment_id,
-            "subscription_id": generative_ai_monitor_details["subscription_id"],
+            "subscription_id": generative_ai_monitor_details.get("subscription_id", None),
         }
 
     def store_payload_records(
@@ -938,6 +1008,40 @@ class WatsonxPromptMonitor(PromptMonitor):
 
         return created_pta.to_dict()["asset_id"]
 
+    def _delete_prompt(self, pta_id: str) -> None:
+        from ibm_aigov_facts_client import (  # type: ignore
+            AIGovFactsClient,
+            CloudPakforDataConfig,
+        )
+
+        try:
+            if hasattr(self, "_fact_cpd_creds") and self._fact_cpd_creds:
+                cpd_creds = CloudPakforDataConfig(**self._fact_cpd_creds)
+
+                aigov_client = AIGovFactsClient(
+                    container_id=self._container_id,
+                    container_type=self._container_type,
+                    cloud_pak_for_data_configs=cpd_creds,
+                    disable_tracing=True,
+                )
+
+            else:
+                aigov_client = AIGovFactsClient(
+                    api_key=self._api_key,
+                    container_id=self._container_id,
+                    container_type=self._container_type,
+                    disable_tracing=True,
+                    region=self.region.factsheet,
+                )
+
+        except Exception as e:
+            logging.error(
+                f"Error connecting to IBM watsonx.governance (factsheets): {e}",
+            )
+            raise
+
+        suppress_output(aigov_client.assets.delete_prompt_asset, pta_id)
+
     def _create_deployment_pta(self, asset_id: str, name: str, model_id: str) -> str:
         from ibm_watsonx_ai import APIClient, Credentials  # type: ignore
 
@@ -976,6 +1080,30 @@ class WatsonxPromptMonitor(PromptMonitor):
 
         return wml_client.deployments.get_uid(created_deployment)
 
+    def _delete_deployment_pta(self, deployment_id: str):
+        from ibm_watsonx_ai import APIClient, Credentials  # type: ignore
+
+        try:
+            if hasattr(self, "_wml_cpd_creds") and self._wml_cpd_creds:
+                creds = Credentials(**self._wml_cpd_creds)
+
+                wml_client = APIClient(creds)
+                wml_client.set.default_space(self.space_id)
+
+            else:
+                creds = Credentials(
+                    url=self.region.watsonxai,
+                    api_key=self._api_key,
+                )
+                wml_client = APIClient(creds)
+                wml_client.set.default_space(self.space_id)
+
+        except Exception as e:
+            logging.error(f"Error connecting to IBM watsonx.ai Runtime: {e}")
+            raise
+
+        suppress_output(wml_client.deployments.delete, deployment_id)
+
     @deprecated(
         reason="'add_prompt_observer()' is deprecated and will be removed in a future version. Use 'create_prompt_monitor()' instead.",
         version="1.0.5",
@@ -1085,6 +1213,8 @@ class WatsonxPromptMonitor(PromptMonitor):
             ```
         """
         task_id = TaskType.from_value(task_id).value
+        rollback_stack = []
+
         # DEPRECATION NOTICE
         if input_text is not None:
             warnings.warn(
@@ -1178,11 +1308,14 @@ class WatsonxPromptMonitor(PromptMonitor):
         pta_id = suppress_output(
             self._create_prompt_template, prompt_details, asset_details
         )
+        rollback_stack.append(lambda: self._delete_detached_prompt(pta_id))
+
         deployment_id = None
         if self._container_type == "space":
             deployment_id = suppress_output(
                 self._create_deployment_pta, pta_id, name, model_id
             )
+            rollback_stack.append(lambda: self._delete_deployment_pta(deployment_id))
 
         monitors = {
             "generative_ai_quality": {
@@ -1243,10 +1376,16 @@ class WatsonxPromptMonitor(PromptMonitor):
 
         generative_ai_monitor_details = generative_ai_monitor_details._to_dict()
 
+        wos_status = generative_ai_monitor_details.get("status", {})
+        if wos_status.get("state") == "ERROR":
+            for rollback_step in reversed(rollback_stack):
+                rollback_step()
+            raise Exception(wos_status.get("failure"))
+
         return {
             "prompt_template_asset_id": pta_id,
             "deployment_id": deployment_id,
-            "subscription_id": generative_ai_monitor_details["subscription_id"],
+            "subscription_id": generative_ai_monitor_details.get("subscription_id", None),
         }
 
     def store_payload_records(

beekeeper/monitors/watsonx/supporting_classes/enums.py

@@ -72,6 +72,19 @@ class Region(str, Enum):
 
 
 class TaskType(Enum):
+    """
+    Supported IBM watsonx.governance tasks.
+
+    Attributes:
+        QUESTION_ANSWERING (str): "question_answering"
+        SUMMARIZATION (str): "summarization"
+        RETRIEVAL_AUGMENTED_GENERATION (str): "retrieval_augmented_generation"
+        CLASSIFICATION (str): "classification"
+        GENERATION (str): "generation"
+        CODE (str): "code"
+        EXTRACTION (str): "extraction"
+    """
+
     QUESTION_ANSWERING = "question_answering"
     SUMMARIZATION = "summarization"
     RETRIEVAL_AUGMENTED_GENERATION = "retrieval_augmented_generation"

{beekeeper_monitors_watsonx-1.1.1.dist-info → beekeeper_monitors_watsonx-1.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: beekeeper-monitors-watsonx
-Version: 1.1.1
+Version: 1.1.2
 Summary: beekeeper monitors watsonx extension
 Author-email: Leonardo Furnielis <leonardofurnielis@outlook.com>
 License: Apache-2.0

{beekeeper_monitors_watsonx-1.1.1.dist-info → beekeeper_monitors_watsonx-1.1.2.dist-info}/RECORD

@@ -1,12 +1,12 @@
 beekeeper/monitors/watsonx/__init__.py,sha256=iJv6D68IT00ZC40TNSVYtqyFTen9sWoDqUtxvVVJjOE,789
-beekeeper/monitors/watsonx/base.py,sha256=AVgYwk-LgjjZ1RTAkk-8f5_T9pJ-jQoeDm_NX1uaX5I,57089
+beekeeper/monitors/watsonx/base.py,sha256=gx_iduLDy1TmMSM-jw9OyUNKltxV8xUeMz-POPhjT6k,62174
 beekeeper/monitors/watsonx/custom_metric.py,sha256=XdJ26lb0y_P1nirJE_krpKONuqJjP3sWFApoG-DR6hk,23959
 beekeeper/monitors/watsonx/supporting_classes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 beekeeper/monitors/watsonx/supporting_classes/credentials.py,sha256=x4rYoOFvx0pWDhFZfuy6fM0rj7JCivaSYn_jYFXlV8U,5190
-beekeeper/monitors/watsonx/supporting_classes/enums.py,sha256=R86yqE_KNHjrrFwqy1tCA2a6UAXwmIzCNxw1vnsH32w,2865
+beekeeper/monitors/watsonx/supporting_classes/enums.py,sha256=Kp8_-YYQ2Z2ha9_pYyF7F7DaZOVFJoV4YSoCqoEgbOk,3274
 beekeeper/monitors/watsonx/supporting_classes/metric.py,sha256=iERXRi0iBFHITFaLtonoV97nNUCPXbieD7Z1wX6bvX8,3370
 beekeeper/monitors/watsonx/utils/data_utils.py,sha256=qBLYtHGY0MJ0JJ8BpFDT2YIjA3QOYJQNemLvpA3DMz0,1252
 beekeeper/monitors/watsonx/utils/instrumentation.py,sha256=ztgR1kZ9h-JvASzRA47AYHdc-Isv33EQum9XBLg-dQk,525
-beekeeper_monitors_watsonx-1.1.1.dist-info/METADATA,sha256=5F1R9-ocW3kbbJQfjzFNEuVMiUYiH2B47EeMCEW7yes,716
-beekeeper_monitors_watsonx-1.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-beekeeper_monitors_watsonx-1.1.1.dist-info/RECORD,,
+beekeeper_monitors_watsonx-1.1.2.dist-info/METADATA,sha256=9NZNF7aRduO-j9Up_BfNbgIRWrBb67zbcGnTvylHk_E,716
+beekeeper_monitors_watsonx-1.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+beekeeper_monitors_watsonx-1.1.2.dist-info/RECORD,,