enkryptai-sdk 0.1.5__tar.gz → 0.1.6__tar.gz

{enkryptai_sdk-0.1.5/src/enkryptai_sdk.egg-info → enkryptai_sdk-0.1.6}/PKG-INFO

@@ -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 → enkryptai_sdk-0.1.6}/setup.py

@@ -8,7 +8,7 @@ with open(os.path.join(here, "README.md"), encoding="utf-8") as fh:
 
 setup(
     name="enkryptai-sdk",  # This is the name of your package on PyPI
-    version="0.1.5",
+    version="0.1.6",
     description="A Python SDK with guardrails and red teaming functionality for API interactions",
     long_description=long_description,
     long_description_content_type="text/markdown",

{enkryptai_sdk-0.1.5 → enkryptai_sdk-0.1.6}/src/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-0.1.5 → enkryptai_sdk-0.1.6}/src/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-0.1.5 → enkryptai_sdk-0.1.6}/src/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 → enkryptai_sdk-0.1.6/src/enkryptai_sdk.egg-info}/PKG-INFO

@@ -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