enkryptai-sdk 1.0.23__py3-none-any.whl → 1.0.25__py3-none-any.whl

enkryptai_sdk/red_team.py

@@ -6,11 +6,14 @@ from .datasets import DatasetClient
 from .dto import (
     RedteamHealthResponse,
     RedTeamModelHealthConfig,
+    RedTeamModelHealthConfigV3,
     RedteamModelHealthResponse,
     RedTeamConfig,
     RedTeamConfigWithSavedModel,
     RedTeamCustomConfig,
     RedTeamCustomConfigWithSavedModel,
+    RedTeamCustomConfigV3,
+    RedTeamCustomConfigWithSavedModelV3,
     RedTeamResponse,
     RedTeamResultSummary,
     RedTeamResultDetails,
@@ -21,7 +24,8 @@ from .dto import (
     RedTeamRiskMitigationGuardrailsPolicyResponse,
     RedTeamRiskMitigationSystemPromptConfig,
     RedTeamRiskMitigationSystemPromptResponse,
-    RedTeamFindingsResponse
+    RedTeamFindingsResponse,
+    RedTeamDownloadLinkResponse
 )
 
 
@@ -94,6 +98,40 @@ class RedTeamClient(BaseClient):
             return RedteamModelHealthResponse.from_dict(response)
         except Exception as e:
             raise RedTeamClientError(str(e))
+    
+    def check_model_health_v3(self, config: RedTeamModelHealthConfigV3):
+        """
+        Get the health status of a model using V3 format with endpoint_configuration.
+        
+        This method accepts endpoint_configuration (similar to add_custom_task) and
+        converts it internally to target_model_configuration format for backend compatibility.
+        
+        Args:
+            config (RedTeamModelHealthConfigV3): Configuration object containing endpoint_configuration
+            
+        Returns:
+            RedteamModelHealthResponse: Response from the API containing health status
+            
+        Raises:
+            RedTeamClientError: If there's an error from the API
+        """
+        try:
+            config = RedTeamModelHealthConfigV3.from_dict(config)
+            
+            # Convert endpoint_configuration to target_model_configuration
+            target_config = config.to_target_model_configuration()
+            
+            # Create the payload in the format expected by the backend
+            payload = {
+                "target_model_configuration": target_config.to_dict()
+            }
+            
+            response = self._request("POST", "/redteam/model-health", json=payload)
+            if response.get("error") not in [None, ""]:
+                raise RedTeamClientError(f"API Error: {str(response)}")
+            return RedteamModelHealthResponse.from_dict(response)
+        except Exception as e:
+            raise RedTeamClientError(str(e))
 
     def add_task(
         self,
@@ -200,6 +238,10 @@ class RedTeamClient(BaseClient):
             "redteam_test_configurations": test_configs,
         }
 
+        # Only add frameworks if provided and not empty
+        if config.frameworks:
+            payload["frameworks"] = config.frameworks
+
         if config.dataset_configuration:
             payload["dataset_configuration"] = DatasetClient.prepare_dataset_payload(
                 config.dataset_configuration, True)
@@ -254,6 +296,10 @@ class RedTeamClient(BaseClient):
             "redteam_test_configurations": test_configs,
         }
 
+        # Only add frameworks if provided and not empty
+        if config.frameworks:
+            payload["frameworks"] = config.frameworks
+
         if config.dataset_configuration:
             payload["dataset_configuration"] = DatasetClient.prepare_dataset_payload(
                 config.dataset_configuration, True)
@@ -281,6 +327,146 @@ class RedTeamClient(BaseClient):
             raise RedTeamClientError(f"API Error: {str(response)}")
         return RedTeamResponse.from_dict(response)
 
+    def add_custom_task_v3(
+        self,
+        config: RedTeamCustomConfigV3,
+        policy_name: str = None,
+    ):
+        """
+        Add a new custom red teaming task with v3 attack methods format.
+        
+        V3 format supports nested attack methods:
+        {
+            "test_name": {
+                "sample_percentage": 50,
+                "attack_methods": {
+                    "method_category": {
+                        "method_name": {
+                            "params": {}
+                        }
+                    }
+                }
+            }
+        }
+        """
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        if policy_name is not None:
+            headers["X-Enkrypt-Policy"] = policy_name
+
+        config = RedTeamCustomConfigV3.from_dict(config)
+        test_configs = config.redteam_test_configurations.to_dict()
+        # Remove None or empty test configurations
+        test_configs = {k: v for k, v in test_configs.items() if v is not None}
+
+        payload = {
+            "test_name": config.test_name,
+            "redteam_test_configurations": test_configs,
+        }
+
+        # Only add frameworks if provided and not empty
+        if config.frameworks:
+            payload["frameworks"] = config.frameworks
+
+        if config.dataset_configuration:
+            payload["dataset_configuration"] = DatasetClient.prepare_dataset_payload(
+                config.dataset_configuration, True)
+        else:
+            raise RedTeamClientError(
+                "Please provide a dataset configuration"
+            )
+
+        if config.endpoint_configuration:
+            payload["endpoint_configuration"] = ModelClient.prepare_model_payload(
+                config.endpoint_configuration, True)
+
+            response = self._request(
+                "POST",
+                "/redteam/v3/add-custom-task",
+                headers=headers,
+                json=payload,
+            )
+            if response.get("error"):
+                raise RedTeamClientError(f"API Error: {str(response)}")
+            return RedTeamResponse.from_dict(response)
+        else:
+            raise RedTeamClientError(
+                "Please provide a endpoint configuration"
+            )
+
+    def add_custom_task_with_saved_model_v3(
+        self,
+        config: RedTeamCustomConfigWithSavedModelV3,
+        model_saved_name: str,
+        model_version: str,
+        policy_name: str = None,
+    ):
+        """
+        Add a new red teaming custom task using a saved model with v3 attack methods format.
+        
+        V3 format supports nested attack methods:
+        {
+            "test_name": {
+                "sample_percentage": 50,
+                "attack_methods": {
+                    "method_category": {
+                        "method_name": {
+                            "params": {}
+                        }
+                    }
+                }
+            }
+        }
+        """
+        if not model_saved_name:
+            raise RedTeamClientError("Please provide a model_saved_name")
+        
+        if not model_version:
+            raise RedTeamClientError("Please provide a model_version. Default is 'v1'")
+
+        config = RedTeamCustomConfigWithSavedModelV3.from_dict(config)
+        test_configs = config.redteam_test_configurations.to_dict()
+        # Remove None or empty test configurations
+        test_configs = {k: v for k, v in test_configs.items() if v is not None}
+
+        payload = {
+            "test_name": config.test_name,
+            "redteam_test_configurations": test_configs,
+        }
+
+        # Only add frameworks if provided and not empty
+        if config.frameworks:
+            payload["frameworks"] = config.frameworks
+
+        if config.dataset_configuration:
+            payload["dataset_configuration"] = DatasetClient.prepare_dataset_payload(
+                config.dataset_configuration, True)
+        else:
+            raise RedTeamClientError(
+                "Please provide a dataset configuration"
+            )
+
+        headers = {
+            "X-Enkrypt-Model": model_saved_name,
+            "X-Enkrypt-Model-Version": model_version,
+            "Content-Type": "application/json",
+        }
+
+        if policy_name is not None:
+            headers["X-Enkrypt-Policy"] = policy_name
+
+        response = self._request(
+            "POST",
+            "/redteam/v3/model/add-custom-task",
+            headers=headers,
+            json=payload,
+        )
+        if response.get("error"):
+            raise RedTeamClientError(f"API Error: {str(response)}")
+        return RedTeamResponse.from_dict(response)
+
     def status(self, task_id: str = None, test_name: str = None):
         """
         Get the status of a specific red teaming task.
@@ -394,7 +580,7 @@ class RedTeamClient(BaseClient):
         if test_name:
             headers["X-Enkrypt-Test-Name"] = test_name
 
-        response = self._request("GET", "/redteam/results/summary", headers=headers)
+        response = self._request("GET", "/redteam/v3/results/summary", headers=headers)
         if response.get("error"):
             raise RedTeamClientError(f"API Error: {str(response)}")
         # print(f"Response: {response}")
@@ -427,7 +613,7 @@ class RedTeamClient(BaseClient):
         if test_name:
             headers["X-Enkrypt-Test-Name"] = test_name
 
-        url = f"/redteam/v2/results/summary/{test_type}"
+        url = f"/redteam/v3/results/summary/{test_type}"
         response = self._request("GET", url, headers=headers)
         if response.get("error"):
             raise RedTeamClientError(f"API Error: {str(response)}")
@@ -457,7 +643,7 @@ class RedTeamClient(BaseClient):
         if test_name:
             headers["X-Enkrypt-Test-Name"] = test_name
             
-        response = self._request("GET", "/redteam/v2/results/details", headers=headers)
+        response = self._request("GET", "/redteam/v3/results/details", headers=headers)
         if response.get("error"):
             raise RedTeamClientError(f"API Error: {str(response)}")
         return RedTeamResultDetails.from_dict(response)
@@ -488,9 +674,8 @@ class RedTeamClient(BaseClient):
             headers["X-Enkrypt-Task-ID"] = task_id
         if test_name:
             headers["X-Enkrypt-Test-Name"] = test_name
-            
 
-        url = f"/redteam/v2/results/details/{test_type}"
+        url = f"/redteam/v3/results/details/{test_type}"
         response = self._request("GET", url, headers=headers)
         if response.get("error"):
             raise RedTeamClientError(f"API Error: {str(response)}")
@@ -564,4 +749,35 @@ class RedTeamClient(BaseClient):
             return RedTeamFindingsResponse.from_dict(response)
         except Exception as e:
             raise RedTeamClientError(str(e))
+
+    def get_download_link(self, task_id: str = None, test_name: str = None):
+        """
+        Get a download link for red team test results.
+        
+        Args:
+            task_id (str, optional): The ID of the task to get download link for
+            test_name (str, optional): The name of the test to get download link for
+            
+        Returns:
+            RedTeamDownloadLinkResponse: Response containing download link and expiry information
+            
+        Raises:
+            RedTeamClientError: If neither task_id nor test_name is provided, or if there's an error from the API
+        """
+        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
+
+        try:
+            response = self._request("GET", "/redteam/download-link", headers=headers)
+            if response.get("error"):
+                raise RedTeamClientError(f"API Error: {str(response)}")
+            return RedTeamDownloadLinkResponse.from_dict(response)
+        except Exception as e:
+            raise RedTeamClientError(str(e))
         

enkryptai_sdk/response.py

@@ -1,3 +1,238 @@
+import math
+from typing import Dict, Any, List, Optional, Union
+
+
+class PaginationInfo:
+    """
+    A class to handle pagination information and calculations.
+    """
+    
+    def __init__(self, page: int = 1, per_page: int = 10, total_count: int = 0):
+        """
+        Initialize pagination information.
+        
+        Args:
+            page (int): Current page number (1-based)
+            per_page (int): Number of items per page
+            total_count (int): Total number of items
+        """
+        self.page = max(1, page)
+        self.per_page = max(1, min(100, per_page))  # Ensure per_page is between 1 and 100
+        self.total_count = max(0, total_count)
+        
+    @property
+    def total_pages(self) -> int:
+        """Calculate total number of pages."""
+        if self.total_count == 0:
+            return 0
+        return math.ceil(self.total_count / self.per_page)
+    
+    @property
+    def has_next(self) -> bool:
+        """Check if there's a next page."""
+        return self.page < self.total_pages
+    
+    @property
+    def has_previous(self) -> bool:
+        """Check if there's a previous page."""
+        return self.page > 1
+    
+    @property
+    def offset(self) -> int:
+        """Calculate the offset for database queries."""
+        return (self.page - 1) * self.per_page
+    
+    @property
+    def limit(self) -> int:
+        """Get the limit for database queries."""
+        return self.per_page
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert pagination info to dictionary."""
+        return {
+            "page": self.page,
+            "per_page": self.per_page,
+            "total_count": self.total_count,
+            "total_pages": self.total_pages,
+            "has_next": self.has_next,
+            "has_previous": self.has_previous
+        }
+    
+    @classmethod
+    def from_query_params(cls, query_params: Dict[str, Any], default_per_page: int = 10) -> "PaginationInfo":
+        """
+        Create PaginationInfo from query parameters.
+        
+        Args:
+            query_params (Dict[str, Any]): Dictionary containing query parameters
+            default_per_page (int): Default items per page
+            
+        Returns:
+            PaginationInfo: Pagination information object
+        """
+        try:
+            page = int(query_params.get("page", 1))
+            per_page = int(query_params.get("per_page", default_per_page))
+        except (ValueError, TypeError):
+            page = 1
+            per_page = default_per_page
+            
+        # Validate pagination parameters
+        if page < 1:
+            page = 1
+        if per_page < 1 or per_page > 100:
+            per_page = min(max(1, per_page), 100)
+            
+        return cls(page=page, per_page=per_page)
+    
+    @classmethod
+    def validate_params(cls, page: Union[str, int], per_page: Union[str, int]) -> tuple[int, int]:
+        """
+        Validate pagination parameters and return validated values.
+        
+        Args:
+            page: Page number (can be string or int)
+            per_page: Items per page (can be string or int)
+            
+        Returns:
+            tuple[int, int]: Validated (page, per_page) values
+            
+        Raises:
+            ValueError: If parameters are invalid
+        """
+        try:
+            page_num = int(page) if page is not None else 1
+            per_page_num = int(per_page) if per_page is not None else 10
+        except (ValueError, TypeError):
+            raise ValueError("Page and per_page must be valid integers")
+        
+        if page_num < 1:
+            raise ValueError("Page must be >= 1")
+        if per_page_num < 1 or per_page_num > 100:
+            raise ValueError("Per_page must be between 1 and 100")
+            
+        return page_num, per_page_num
+
+
+class PaginatedResponse(dict):
+    """
+    A wrapper class for paginated API responses that provides pagination information
+    and maintains backward compatibility with dictionary access.
+    """
+    
+    def __init__(self, data: List[Any], pagination: PaginationInfo, **kwargs):
+        """
+        Initialize the PaginatedResponse object.
+        
+        Args:
+            data (List[Any]): List of items for the current page
+            pagination (PaginationInfo): Pagination information
+            **kwargs: Additional response data
+        """
+        response_data = {
+            "data": data,
+            "pagination": pagination.to_dict(),
+            **kwargs
+        }
+        super().__init__(response_data)
+        self._data = response_data
+        self._pagination = pagination
+    
+    def get_data(self) -> List[Any]:
+        """Get the data items for the current page."""
+        return self._data.get("data", [])
+    
+    def get_pagination(self) -> Dict[str, Any]:
+        """Get pagination information."""
+        return self._data.get("pagination", {})
+    
+    def get_page(self) -> int:
+        """Get current page number."""
+        return self._pagination.page
+    
+    def get_per_page(self) -> int:
+        """Get items per page."""
+        return self._pagination.per_page
+    
+    def get_total_count(self) -> int:
+        """Get total number of items."""
+        return self._pagination.total_count
+    
+    def get_total_pages(self) -> int:
+        """Get total number of pages."""
+        return self._pagination.total_pages
+    
+    def has_next_page(self) -> bool:
+        """Check if there's a next page."""
+        return self._pagination.has_next
+    
+    def has_previous_page(self) -> bool:
+        """Check if there's a previous page."""
+        return self._pagination.has_previous
+    
+    def get_next_page_url(self, base_url: str, **query_params) -> Optional[str]:
+        """
+        Generate URL for the next page.
+        
+        Args:
+            base_url (str): Base URL for the endpoint
+            **query_params: Additional query parameters to include
+            
+        Returns:
+            Optional[str]: Next page URL or None if no next page
+        """
+        if not self.has_next_page():
+            return None
+            
+        params = {**query_params, "page": self.get_page() + 1, "per_page": self.get_per_page()}
+        query_string = "&".join([f"{k}={v}" for k, v in params.items()])
+        return f"{base_url}?{query_string}"
+    
+    def get_previous_page_url(self, base_url: str, **query_params) -> Optional[str]:
+        """
+        Generate URL for the previous page.
+        
+        Args:
+            base_url (str): Base URL for the endpoint
+            **query_params: Additional query parameters to include
+            
+        Returns:
+            Optional[str]: Previous page URL or None if no previous page
+        """
+        if not self.has_previous_page():
+            return None
+            
+        params = {**query_params, "page": self.get_page() - 1, "per_page": self.get_per_page()}
+        query_string = "&".join([f"{k}={v}" for k, v in params.items()])
+        return f"{base_url}?{query_string}"
+    
+    def get_page_urls(self, base_url: str, **query_params) -> Dict[str, Optional[str]]:
+        """
+        Generate URLs for all pagination actions.
+        
+        Args:
+            base_url (str): Base URL for the endpoint
+            **query_params: Additional query parameters to include
+            
+        Returns:
+            Dict[str, Optional[str]]: Dictionary with pagination URLs
+        """
+        return {
+            "first": f"{base_url}?{self._build_query_string(1, **query_params)}",
+            "previous": self.get_previous_page_url(base_url, **query_params),
+            "current": f"{base_url}?{self._build_query_string(self.get_page(), **query_params)}",
+            "next": self.get_next_page_url(base_url, **query_params),
+            "last": f"{base_url}?{self._build_query_string(self.get_total_pages(), **query_params)}" if self.get_total_pages() > 0 else None
+        }
+    
+    def _build_query_string(self, page: int, **query_params) -> str:
+        """Build query string for a specific page."""
+        params = {**query_params, "page": page, "per_page": self.get_per_page()}
+        return "&".join([f"{k}={v}" for k, v in params.items()])
+    
+    def __str__(self) -> str:
+        """String representation of the paginated response."""
+        return f"PaginatedResponse(page={self.get_page()}, per_page={self.get_per_page()}, total={self.get_total_count()}, items={len(self.get_data())})"
 
 
 class GuardrailsResponse(dict):
@@ -100,7 +335,25 @@ class GuardrailsResponse(dict):
             
         return f"Response Status: {status}\n{violation_str}"
     
+    def get_pagination(self) -> Optional[Dict[str, Any]]:
+        """
+        Get pagination information if available.
+        
+        Returns:
+            Optional[Dict[str, Any]]: Pagination data or None if not available
+        """
+        return self._data.get("pagination")
     
+    def is_paginated(self) -> bool:
+        """
+        Check if the response contains pagination information.
+        
+        Returns:
+            bool: True if response is paginated, False otherwise
+        """
+        return "pagination" in self._data
+
+
 class PIIResponse(dict):
     """
     A wrapper class for Enkrypt AI PII API responses that provides additional functionality
@@ -132,4 +385,22 @@ class PIIResponse(dict):
         """ 
         return self._data.get("key", "")
     
+    def get_pagination(self) -> Optional[Dict[str, Any]]:
+        """
+        Get pagination information if available.
+        
+        Returns:
+            Optional[Dict[str, Any]]: Pagination data or None if not available
+        """
+        return self._data.get("pagination")
+    
+    def is_paginated(self) -> bool:
+        """
+        Check if the response contains pagination information.
+        
+        Returns:
+            bool: True if response is paginated, False otherwise
+        """
+        return "pagination" in self._data 
+    
     

enkryptai_sdk/utils/__init__.py

@@ -0,0 +1,29 @@
+# Utils module for Enkrypt AI SDK
+
+from .pagination import (
+    PaginationInfo,
+    PaginatedResponse,
+    parse_pagination_params,
+    build_pagination_url,
+    create_paginated_response,
+    validate_pagination_params,
+    get_pagination_metadata,
+    calculate_page_info,
+    create_pagination_links,
+    apply_pagination_to_list,
+    format_pagination_response
+)
+
+__all__ = [
+    "PaginationInfo",
+    "PaginatedResponse",
+    "parse_pagination_params",
+    "build_pagination_url",
+    "create_paginated_response",
+    "validate_pagination_params",
+    "get_pagination_metadata",
+    "calculate_page_info",
+    "create_pagination_links",
+    "apply_pagination_to_list",
+    "format_pagination_response"
+]