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

enkryptai_sdk/models.py

@@ -1,6 +1,6 @@
-import urllib3
-from .dto import ModelConfig, ModelResponse
+from .base import BaseClient
 from urllib.parse import urlparse, urlsplit
+from .dto import ModelConfig, ModelResponse, ModelCollection
 
 
 class ModelClientError(Exception):
@@ -11,41 +11,11 @@ class ModelClientError(Exception):
     pass
 
 
-class ModelClient:
+class ModelClient(BaseClient):
     def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com:443"):
-        self.api_key = api_key
-        self.base_url = base_url
-        self.http = urllib3.PoolManager()
-        self.headers = {"apikey": self.api_key}
-
-    def _request(self, method, endpoint, payload=None, headers=None, **kwargs):
-        url = self.base_url + endpoint
-        request_headers = {
-            "Accept-Encoding": "gzip",  # Add required gzip encoding
-            **self.headers,
-        }
-        if headers:
-            request_headers.update(headers)
-
-        try:
-            response = self.http.request(method, url, headers=request_headers, **kwargs)
-
-            if response.status >= 400:
-                error_data = (
-                    response.json()
-                    if response.data
-                    else {"message": f"HTTP {response.status}"}
-                )
-                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)}
-
-    def health(self):
-        return self._request("GET", "/models/health")
+        super().__init__(api_key, base_url)
 
-    def add_model(self, config: ModelConfig):
+    def add_model(self, config: ModelConfig) -> ModelResponse:
         """
         Add a new model configuration to the system.
 
@@ -69,11 +39,15 @@ class ModelClient:
             base_path = ""
             remaining_path = ""
 
-        # Determine paths based on the endpoint
-        paths = {
-            "completions": f"/{remaining_path}" if remaining_path else "",
-            "chat": "",
-        }
+        if config.model_config.paths:
+            paths = config.model_config.paths.to_dict()
+        else:
+            paths = {
+                "completions": (
+                    f"/{remaining_path.split('/')[-1]}" if remaining_path else ""
+                ),
+                "chat": f"/{remaining_path}" if remaining_path else "",
+            }
 
         payload = {
             "model_saved_name": config.model_saved_name,
@@ -82,11 +56,10 @@ class ModelClient:
             "model_type": config.model_type,
             "certifications": config.certifications,
             "model_config": {
-                "is_compatible_with": config.model_config.is_compatible_with,
+                "model_provider": config.model_config.model_provider,
                 "model_version": config.model_config.model_version,
                 "hosting_type": config.model_config.hosting_type,
                 "model_source": config.model_config.model_source,
-                "model_provider": config.model_config.model_provider,
                 "system_prompt": config.model_config.system_prompt,
                 "conversation_template": config.model_config.conversation_template,
                 "endpoint": {
@@ -94,7 +67,7 @@ class ModelClient:
                     "host": parsed_url.hostname,
                     "port": parsed_url.port
                     or (443 if parsed_url.scheme == "https" else 80),
-                    "base_path": f"/{base_path}/{paths['completions']}",  # Just v1
+                    "base_path": f"/{base_path}",  # Just v1
                 },
                 "paths": paths,
                 "auth_data": {
@@ -118,18 +91,19 @@ class ModelClient:
         except Exception as e:
             raise ModelClientError(str(e))
 
-    def get_model(self, model_id: str) -> ModelConfig:
+    def get_model(self, model_saved_name: str) -> ModelConfig:
         """
-        Get model configuration by model ID.
+        Get model configuration by model saved name.
 
         Args:
-            model_id (str): ID of the model to retrieve
+            model_saved_name (str): Saved name of the model to retrieve
 
         Returns:
             ModelConfig: Configuration object containing model details
         """
-        headers = {"X-Enkrypt-Model": model_id}
+        headers = {"X-Enkrypt-Model": model_saved_name}
         response = self._request("GET", "/models/get-model", headers=headers)
+        print(response)
         if response.get("error"):
             raise ModelClientError(response["error"])
         return ModelConfig.from_dict(response)
@@ -142,19 +116,111 @@ class ModelClient:
             dict: Response from the API containing the list of models
         """
         try:
-            return self._request("GET", "/models/list-models")
+            response = self._request("GET", "/models/list-models")
+            if response.get("error"):
+                raise ModelClientError(response["error"])
+            return ModelCollection.from_dict(response)
         except Exception as e:
             return {"error": str(e)}
 
-    def delete_model(self, model_id: str):
+    def modify_model(self, config: ModelConfig, old_model_saved_name=None) -> ModelResponse:
+        """
+        Modify an existing model in the system.
+
+        Args:
+            model_saved_name (str): The saved name of the model to modify
+            config (ModelConfig): Configuration object containing model details
+
+        Returns:
+            dict: Response from the API containing the modified model details
+        """
+        if old_model_saved_name is None:
+            old_model_saved_name = config["model_saved_name"]
+
+        headers = {"Content-Type": "application/json", "X-Enkrypt-Model": old_model_saved_name}
+        # print(config)
+        config = ModelConfig.from_dict(config)
+        # Parse endpoint_url into components
+        parsed_url = urlparse(config.model_config.endpoint_url)
+        path_parts = parsed_url.path.strip("/").split("/")
+
+        # Extract base_path and endpoint path
+        if len(path_parts) >= 1:
+            base_path = path_parts[0]  # Usually 'v1'
+            remaining_path = "/".join(path_parts[1:])  # The rest of the path
+        else:
+            base_path = ""
+            remaining_path = ""
+
+        if config.model_config.paths:
+            paths = config.model_config.paths.to_dict()
+        else:
+            # Determine paths based on the endpoint
+            paths = {
+                "completions": (
+                    f"/{remaining_path.split('/')[-1]}" if remaining_path else ""
+                ),
+                "chat": f"/{remaining_path}" if remaining_path else "",
+            }
+
+        payload = {
+            "model_saved_name": config.model_saved_name,
+            "testing_for": config.testing_for,
+            "model_name": config.model_name,
+            "model_type": config.model_type,
+            "certifications": config.certifications,
+            "model_config": {
+                "model_provider": config.model_config.model_provider,
+                "model_version": config.model_config.model_version,
+                "hosting_type": config.model_config.hosting_type,
+                "model_source": config.model_config.model_source,
+                "system_prompt": config.model_config.system_prompt,
+                "conversation_template": config.model_config.conversation_template,
+                "endpoint": {
+                    "scheme": parsed_url.scheme,
+                    "host": parsed_url.hostname,
+                    "port": parsed_url.port
+                    or (443 if parsed_url.scheme == "https" else 80),
+                    "base_path": f"/{base_path}",  # Just v1
+                },
+                "paths": paths,
+                "auth_data": {
+                    "header_name": config.model_config.auth_data.header_name,
+                    "header_prefix": config.model_config.auth_data.header_prefix,
+                    "space_after_prefix": config.model_config.auth_data.space_after_prefix,
+                },
+                "apikeys": (
+                    [config.model_config.apikey] if config.model_config.apikey else []
+                ),
+                "default_request_options": config.model_config.default_request_options,
+            },
+        }
+        try:
+            response = self._request(
+                "PATCH", "/models/modify-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 delete_model(self, model_saved_name: str) -> ModelResponse:
         """
         Delete a specific model from the system.
 
         Args:
-            model_id (str): The identifier or name of the model to delete
+            model_saved_name (str): The saved name of the model to delete
 
         Returns:
             dict: Response from the API containing the deletion status
         """
-        headers = {"X-Enkrypt-Model": model_id}
-        return self._request("DELETE", "/models/delete-model", headers=headers)
+        headers = {"X-Enkrypt-Model": model_saved_name}
+
+        try:
+            response = self._request("DELETE", "/models/delete-model", headers=headers)
+            if response.get("error"):
+                raise ModelClientError(response["error"])
+            return ModelResponse.from_dict(response)
+        except Exception as e:
+            raise ModelClientError(str(e))

enkryptai_sdk/red_team.py

@@ -1,11 +1,17 @@
 import urllib3
+from .base import BaseClient
 from .dto import (
+    RedteamHealthResponse,
+    RedTeamModelHealthConfig,
+    RedteamModelHealthResponse,
     RedTeamConfig,
+    # RedTeamConfigWithSavedModel,
     RedTeamResponse,
     RedTeamResultSummary,
     RedTeamResultDetails,
     RedTeamTaskStatus,
     RedTeamTaskDetails,
+    RedTeamTaskList,
 )
 
 
@@ -17,47 +23,64 @@ class RedTeamClientError(Exception):
     pass
 
 
-class RedTeamClient:
+class RedTeamClient(BaseClient):
     """
     A client for interacting with the Red Team API.
     """
 
     def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
-        self.api_key = api_key
-        self.base_url = base_url
-        self.http = urllib3.PoolManager()
-        self.headers = {"apikey": self.api_key}
-
-    def _request(self, method, endpoint, headers=None, **kwargs):
-        url = self.base_url + endpoint
-        request_headers = {
-            "Accept-Encoding": "gzip",  # Add required gzip encoding
-            **self.headers,
-        }
-        if headers:
-            request_headers.update(headers)
+        super().__init__(api_key, base_url)
+
+    # def get_model(self, model):
+    #     models = self._request("GET", "/models/list-models")
+    #     models = models["models"]
+    #     for _model_data in models:
+    #         if _model_data["model_saved_name"] == model:
+    #             return _model_data["model_saved_name"]
+    #     else:
+    #         return None
 
+    def get_health(self):
+        """
+        Get the health status of the service.
+        """
         try:
-            response = self.http.request(method, url, headers=request_headers, **kwargs)
-
-            if response.status >= 400:
-                error_data = (
-                    response.json()
-                    if response.data
-                    else {"message": f"HTTP {response.status}"}
-                )
-                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)}
-
-    def get_model(self, model):
-        models = self._request("GET", "/models/list-models")
-        if model in models:
-            return model
-        else:
-            return None
+            response = self._request("GET", "/redteam/health")
+            if response.get("error"):
+                raise RedTeamClientError(response["error"])
+            return RedteamHealthResponse.from_dict(response)
+        except Exception as e:
+            raise RedTeamClientError(str(e))
+
+    def check_model_health(self, config: RedTeamModelHealthConfig):
+        """
+        Get the health status of a model.
+        """
+        try:
+            config = RedTeamModelHealthConfig.from_dict(config)
+            response = self._request("POST", "/redteam/model-health", json=config.to_dict())
+            # if response.get("error"):
+            if response.get("error") not in [None, ""]:
+                raise RedTeamClientError(response["error"])
+            return RedteamModelHealthResponse.from_dict(response)
+        except Exception as e:
+            raise RedTeamClientError(str(e))
+        
+    def check_saved_model_health(self, model_saved_name: str):
+        """
+        Get the health status of a saved model.
+        """
+        try:
+            headers = {
+                "X-Enkrypt-Model": model_saved_name,
+            }
+            response = self._request("POST", "/redteam/model/model-health", headers=headers)
+            # if response.get("error"):
+            if response.get("error") not in [None, ""]:
+                raise RedTeamClientError(response["error"])
+            return RedteamModelHealthResponse.from_dict(response)
+        except Exception as e:
+            raise RedTeamClientError(str(e))
 
     def add_task(
         self,
@@ -75,25 +98,26 @@ class RedTeamClient:
             # "async": config.async_enabled,
             "dataset_name": config.dataset_name,
             "test_name": config.test_name,
-            "redteam_test_configurations": test_configs,
+            "redteam_test_configurations": {
+                k: v.to_dict() for k, v in test_configs.items()
+            },
         }
 
-        model = config.model_name
-        saved_model = self.get_model(model)
-
+        saved_model = config.model_saved_name
         if saved_model:
-            print("saved model found")
             headers = {
                 "X-Enkrypt-Model": saved_model,
                 "Content-Type": "application/json",
             }
-            payload["location"] = {"storage": "supabase", "container_name": "supabase"}
-            return self._request(
+            response = self._request(
                 "POST",
                 "/redteam/v2/model/add-task",
                 headers=headers,
                 json=payload,
             )
+            if response.get("error"):
+                raise RedTeamClientError(response["error"])
+            return RedTeamResponse.from_dict(response)
         elif config.target_model_configuration:
             payload["target_model_configuration"] = (
                 config.target_model_configuration.to_dict()
@@ -112,84 +136,154 @@ class RedTeamClient:
                 "Please use a saved model or provide a target model configuration"
             )
 
-    def status(self, task_id: str):
+    def status(self, task_id: str = None, test_name: str = None):
         """
         Get the status of a specific red teaming task.
 
         Args:
-            task_id (str): The ID of the task to check status
+            task_id (str, optional): The ID of the task to check status
+            test_name (str, optional): The name of the test to check status
 
         Returns:
             dict: The task status information
+            
+        Raises:
+            RedTeamClientError: If neither task_id nor test_name is provided, or if there's an error from the API
         """
-        headers = {"X-Enkrypt-Task-ID": task_id}
+        if not task_id and not test_name:
+            raise RedTeamClientError("Either task_id or test_name must be provided")
+        
+        headers = {}
+        if task_id:
+            headers["X-Enkrypt-Task-ID"] = task_id
+        if test_name:
+            headers["X-Enkrypt-Test-Name"] = test_name
 
         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):
+    def cancel_task(self, task_id: str = None, test_name: str = None):
         """
         Cancel a specific red teaming task.
 
         Args:
-            task_id (str): The ID of the task to cancel
+            task_id (str, optional): The ID of the task to cancel
+            test_name (str, optional): The name of the test to cancel
+
+        Returns:
+            dict: The cancellation response
+
+        Raises:
+            RedTeamClientError: If neither task_id nor test_name is provided, or if there's an error from the API
         """
         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."
         )
 
-    def get_task(self, task_id: str):
+        if not task_id and not test_name:
+            raise RedTeamClientError("Either task_id or test_name must be provided")
+
+        headers = {}
+        if task_id:
+            headers["X-Enkrypt-Task-ID"] = task_id
+        if test_name:
+            headers["X-Enkrypt-Test-Name"] = test_name
+
+        response = self._request("POST", "/redteam/cancel-task", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        return response
+
+    def get_task(self, task_id: str = None, test_name: str = None):
         """
         Get the status and details of a specific red teaming task.
 
         Args:
-            task_id (str): The ID of the task to retrieve
+            task_id (str, optional): The ID of the task to retrieve
+            test_name (str, optional): The name of the test to retrieve
 
         Returns:
             dict: The task details and status
+            
+        Raises:
+            RedTeamClientError: If neither task_id nor test_name is provided, or if there's an error from the API
         """
-        headers = {"X-Enkrypt-Task-ID": task_id}
+        if not task_id and not test_name:
+            raise RedTeamClientError("Either task_id or test_name must be provided")
+        
+        headers = {}
+        if task_id:
+            headers["X-Enkrypt-Task-ID"] = task_id
+        if test_name:
+            headers["X-Enkrypt-Test-Name"] = test_name
 
         response = self._request("GET", "/redteam/get-task", headers=headers)
         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):
+    def get_result_summary(self, task_id: str = None, test_name: str = None):
         """
         Get the summary of results for a specific red teaming task.
 
         Args:
-            task_id (str): The ID of the task to get results for
+            task_id (str, optional): The ID of the task to get results for
+            test_name (str, optional): The name of the test to get results for
 
         Returns:
             dict: The summary of the task results
+            
+        Raises:
+            RedTeamClientError: If neither task_id nor test_name is provided, or if there's an error from the API
         """
-        headers = {"X-Enkrypt-Task-ID": task_id}
+        if not task_id and not test_name:
+            raise RedTeamClientError("Either task_id or test_name must be provided")
+        
+        headers = {}
+        if task_id:
+            headers["X-Enkrypt-Task-ID"] = task_id
+        if test_name:
+            headers["X-Enkrypt-Test-Name"] = test_name
 
         response = self._request("GET", "/redteam/results/summary", headers=headers)
         if response.get("error"):
             raise RedTeamClientError(response["error"])
-        return RedTeamResultSummary.from_dict(response["summary"])
+        print(f"Response: {response}")
+        return RedTeamResultSummary.from_dict(response)
 
-    def get_result_details(self, task_id: str):
+    def get_result_details(self, task_id: str = None, test_name: str = None):
         """
         Get the detailed results for a specific red teaming task.
 
         Args:
-            task_id (str): The ID of the task to get detailed results for
+            task_id (str, optional): The ID of the task to get detailed results for
+            test_name (str, optional): The name of the test to get detailed results for
 
         Returns:
             dict: The detailed task results
+            
+        Raises:
+            RedTeamClientError: If neither task_id nor test_name is provided, or if there's an error from the API
         """
-        # TODO: Update the response to be updated
-        headers = {"X-Enkrypt-Task-ID": task_id}
+        if not task_id and not test_name:
+            raise RedTeamClientError("Either task_id or test_name must be provided")
+        
+        headers = {}
+        if task_id:
+            headers["X-Enkrypt-Task-ID"] = task_id
+        if test_name:
+            headers["X-Enkrypt-Test-Name"] = test_name
+            
         response = self._request("GET", "/redteam/results/details", headers=headers)
         if response.get("error"):
             raise RedTeamClientError(response["error"])
-        return RedTeamResultDetails.from_dict(response["details"])
+        return RedTeamResultDetails.from_dict(response)
+
+    def get_task_list(self):
+        response = self._request("GET", "/redteam/list-tasks")
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        return RedTeamTaskList.from_dict(response)