enkryptai-sdk 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

enkryptai_sdk/dto/models.py

@@ -14,6 +14,19 @@ class Modality(Enum):
         return self.value
 
 
+@dataclass
+class ModelResponse:
+    message: Optional[str] = None
+    data: Optional[Dict] = None
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
 @dataclass
 class EndpointConfig:
     scheme: str = "https"

enkryptai_sdk/models.py

@@ -1,8 +1,16 @@
 import urllib3
-from .dto import ModelConfig, ModelDetailConfig
+from .dto import ModelConfig, ModelResponse
 from urllib.parse import urlparse, urlsplit
 
 
+class ModelClientError(Exception):
+    """
+    A custom exception for ModelClient errors.
+    """
+
+    pass
+
+
 class ModelClient:
     def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com:443"):
         self.api_key = api_key
@@ -23,14 +31,13 @@ class ModelClient:
             response = self.http.request(method, url, headers=request_headers, **kwargs)
 
             if response.status >= 400:
-                error_response = (
+                error_data = (
                     response.json()
                     if response.data
                     else {"message": f"HTTP {response.status}"}
                 )
-                raise urllib3.exceptions.HTTPError(
-                    f"HTTP {response.status}: {error_response}"
-                )
+                error_message = error_data.get("message", str(error_data))
+                raise urllib3.exceptions.HTTPError(error_message)
             return response.json()
         except urllib3.exceptions.HTTPError as e:
             return {"error": str(e)}
@@ -101,8 +108,15 @@ class ModelClient:
                 "default_request_options": config.model_config.default_request_options,
             },
         }
-        print(payload)
-        return self._request("POST", "/models/add-model", headers=headers, json=payload)
+        try:
+            response = self._request(
+                "POST", "/models/add-model", headers=headers, json=payload
+            )
+            if response.get("error"):
+                raise ModelClientError(response["error"])
+            return ModelResponse.from_dict(response)
+        except Exception as e:
+            raise ModelClientError(str(e))
 
     def get_model(self, model_id: str) -> ModelConfig:
         """
@@ -116,6 +130,8 @@ class ModelClient:
         """
         headers = {"X-Enkrypt-Model": model_id}
         response = self._request("GET", "/models/get-model", headers=headers)
+        if response.get("error"):
+            raise ModelClientError(response["error"])
         return ModelConfig.from_dict(response)
 
     def get_model_list(self):

enkryptai_sdk/red_team.py

@@ -9,7 +9,7 @@ from .dto import (
 )
 
 
-class RedTeamError(Exception):
+class RedTeamClientError(Exception):
     """
     A custom exception for Red Team errors.
     """
@@ -39,18 +39,17 @@ class RedTeamClient:
 
         try:
             response = self.http.request(method, url, headers=request_headers, **kwargs)
+
             if response.status >= 400:
-                error_response = (
+                error_data = (
                     response.json()
                     if response.data
                     else {"message": f"HTTP {response.status}"}
                 )
-                raise urllib3.exceptions.HTTPError(
-                    f"HTTP {response.status}: {error_response}"
-                )
+                error_message = error_data.get("message", str(error_data))
+                raise urllib3.exceptions.HTTPError(error_message)
             return response.json()
         except urllib3.exceptions.HTTPError as e:
-            print(f"Request failed: {e}")
             return {"error": str(e)}
 
     def get_model(self, model):
@@ -105,9 +104,11 @@ class RedTeamClient:
                 "/redteam/v2/add-task",
                 json=payload,
             )
+            if response.get("error"):
+                raise RedTeamClientError(response["error"])
             return RedTeamResponse.from_dict(response)
         else:
-            raise RedTeamError(
+            raise RedTeamClientError(
                 "Please use a saved model or provide a target model configuration"
             )
 
@@ -124,6 +125,8 @@ class RedTeamClient:
         headers = {"X-Enkrypt-Task-ID": task_id}
 
         response = self._request("GET", "/redteam/task-status", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
         return RedTeamTaskStatus.from_dict(response)
 
     def cancel_task(self, task_id: str):
@@ -133,7 +136,7 @@ class RedTeamClient:
         Args:
             task_id (str): The ID of the task to cancel
         """
-        raise RedTeamError(
+        raise RedTeamClientError(
             "This feature is currently under development. Please check our documentation "
             "at https://docs.enkrypt.ai for updates or contact support@enkrypt.ai for assistance."
         )
@@ -151,7 +154,10 @@ class RedTeamClient:
         headers = {"X-Enkrypt-Task-ID": task_id}
 
         response = self._request("GET", "/redteam/get-task", headers=headers)
-        response["data"]["task_id"] = response["data"].pop("job_id")
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        if response.get("data").get("job_id "):
+            response["data"]["task_id"] = response["data"].pop("job_id")
         return RedTeamTaskDetails.from_dict(response["data"])
 
     def get_result_summary(self, task_id: str):
@@ -167,6 +173,8 @@ class RedTeamClient:
         headers = {"X-Enkrypt-Task-ID": task_id}
 
         response = self._request("GET", "/redteam/results/summary", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
         return RedTeamResultSummary.from_dict(response["summary"])
 
     def get_result_details(self, task_id: str):
@@ -182,4 +190,6 @@ class RedTeamClient:
         # TODO: Update the response to be updated
         headers = {"X-Enkrypt-Task-ID": task_id}
         response = self._request("GET", "/redteam/results/details", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
         return RedTeamResultDetails.from_dict(response["details"])

{enkryptai_sdk-0.1.5.dist-info → enkryptai_sdk-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: enkryptai-sdk
-Version: 0.1.5
+Version: 0.1.6
 Summary: A Python SDK with guardrails and red teaming functionality for API interactions
 Home-page: https://github.com/enkryptai/enkryptai-sdk
 Author: Enkrypt AI Team

{enkryptai_sdk-0.1.5.dist-info → enkryptai_sdk-0.1.6.dist-info}/RECORD

@@ -2,14 +2,14 @@ enkryptai_sdk/__init__.py,sha256=PJq9O9HIhE4OLQsMfYVbdrymg6fBQ0owKGzCTnt1yJk,299
 enkryptai_sdk/config.py,sha256=IpB8_aO4zXdvv061v24oh83oyJ5Tp1QBQTzeuW4h9QY,8828
 enkryptai_sdk/evals.py,sha256=SAQ3GM2ppV-uPLWR1NdOa0fjiEWY5UqzdA8bXFZVl5I,2589
 enkryptai_sdk/guardrails.py,sha256=SgzPZkTzbAPD9XfmYNG6M1-TrzbhDHpAkI3FjnVWS_s,6434
-enkryptai_sdk/models.py,sha256=Yv9HfwxT1D4j1yTf1sfysIKAypUKxeN-NCEv_oit020,5446
-enkryptai_sdk/red_team.py,sha256=Pf3z1jRrE4jhZrZYVIJVyUEx70latucg-vCpQdLoKMA,5838
+enkryptai_sdk/models.py,sha256=CyUv5qOE2F9G8DcSLyr3J4U3VhnRQOZw_8qFB9k25C4,5894
+enkryptai_sdk/red_team.py,sha256=NmXNfbZMYqqS2w7Dfj-8VAMvs1Xs986X1bnC8FK_lJQ,6328
 enkryptai_sdk/response.py,sha256=43JRubzgGCpoVxYNzBZY0AlUgLbfcXD_AwD7wU3qY9o,4086
 enkryptai_sdk/dto/__init__.py,sha256=Ngo7pcuOeWNG0tiSWrXf14NCPBfpD6FBXw0R4jDwLPU,377
-enkryptai_sdk/dto/models.py,sha256=uJevi8Fj67upBVNOcqLJJEPu7DWb8OLKOGJ3hZdS2C4,6144
+enkryptai_sdk/dto/models.py,sha256=W3L-yloW4QsanHJg6Gg3CqNAibKKCFb50kdTKLAgRp0,6377
 enkryptai_sdk/dto/red_team.py,sha256=LupAXG_2izEbPQct1MS_vWOcov06bvnDnnbs_4zFBfY,5234
-enkryptai_sdk-0.1.5.dist-info/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-enkryptai_sdk-0.1.5.dist-info/METADATA,sha256=LYZo7NOhMHIvBrEyv5a627S-0B33On9sG5TJ6452qZs,7482
-enkryptai_sdk-0.1.5.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-enkryptai_sdk-0.1.5.dist-info/top_level.txt,sha256=s2X9UJJwvJamNmr6ZXWyyQe60sXtQGWFuaBYfhgHI_4,14
-enkryptai_sdk-0.1.5.dist-info/RECORD,,
+enkryptai_sdk-0.1.6.dist-info/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+enkryptai_sdk-0.1.6.dist-info/METADATA,sha256=vRbv-pl5hYfjZZkAm7KhqvnrOduVYeWfMRZp99hL-8M,7482
+enkryptai_sdk-0.1.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+enkryptai_sdk-0.1.6.dist-info/top_level.txt,sha256=s2X9UJJwvJamNmr6ZXWyyQe60sXtQGWFuaBYfhgHI_4,14
+enkryptai_sdk-0.1.6.dist-info/RECORD,,