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

enkryptai_sdk/datasets.py

@@ -128,15 +128,21 @@ class DatasetClient(BaseClient):
         response["dataset_name"] = dataset_name
         return DatasetSummary.from_dict(response)
 
-    def list_datasets(self):
+    def list_datasets(self, status: str = None):
         """
         Get a list of all available dataset tasks.
 
+        Args:
+            status (str): Filter the list of dataset tasks by status
+
         Returns:
             dict: Response from the API containing the list of dataset tasks
         """
 
-        response = self._request("GET", "/datasets/list-tasks")
+        url = "/datasets/list-tasks"
+        if status:
+            url += f"?status={status}"
+        response = self._request("GET", url)
         if response.get("error"):
             raise DatasetClientError(response["error"])
         return DatasetCollection.from_dict(response)

enkryptai_sdk/dto/__init__.py

@@ -21,6 +21,7 @@ __all__ = [
     "Location",
     "Metadata",
     "DEFAULT_REDTEAM_CONFIG",
+    "DEFAULT_REDTEAM_CONFIG_WITH_SAVED_MODEL",
     "ADVANCED_REDTEAM_TESTS",
     "DETAIL_MODEL_CONFIG",
     "DatasetConfig",

enkryptai_sdk/dto/red_team.py

@@ -73,6 +73,28 @@ class StatisticItem(BaseDTO):
         return d
 
 
+@dataclass
+class StatisticItemWithTestType(BaseDTO):
+    success_percentage: float
+    total: int
+    test_type: Optional[str] = None
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "StatisticItemWithTestType":
+        return cls(
+            success_percentage=data.get("success(%)", 0.0),
+            total=data.get("total", 0),
+            test_type=data.get("test_type", None),
+        )
+
+    def to_dict(self) -> Dict:
+        d = super().to_dict()
+        # Special handling for success percentage key
+        d["success(%)"] = d.pop("success_percentage")
+        return d
+
+
 @dataclass
 class ResultSummary(BaseDTO):
     test_date: str
@@ -86,7 +108,7 @@ class ResultSummary(BaseDTO):
     test_type: Dict[str, StatisticItem]
     nist_category: Dict[str, StatisticItem]
     scenario: Dict[str, StatisticItem]
-    category: Dict[str, StatisticItem]
+    category: Dict[str, StatisticItemWithTestType]
     attack_method: Dict[str, StatisticItem]
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
@@ -98,6 +120,13 @@ class ResultSummary(BaseDTO):
                 for key, value in item.items():
                     result[key] = StatisticItem.from_dict(value)
             return result
+        
+        def convert_stat_test_type_list(stat_list: List[Dict]) -> Dict[str, StatisticItemWithTestType]:
+            result = {}
+            for item in stat_list:
+                for key, value in item.items():
+                    result[key] = StatisticItemWithTestType.from_dict(value)
+            return result
 
         return cls(
             test_date=data.get("test_date", ""),
@@ -111,20 +140,23 @@ class ResultSummary(BaseDTO):
             test_type=convert_stat_list(data.get("test_type", [])),
             nist_category=convert_stat_list(data.get("nist_category", [])),
             scenario=convert_stat_list(data.get("scenario", [])),
-            category=convert_stat_list(data.get("category", [])),
+            category=convert_stat_test_type_list(data.get("category", [])),
             attack_method=convert_stat_list(data.get("attack_method", [])),
         )
 
     def to_dict(self) -> Dict:
         def convert_stat_dict(stat_dict: Dict[str, StatisticItem]) -> List[Dict]:
             return [{key: value.to_dict()} for key, value in stat_dict.items()]
+        
+        def convert_stat_test_type_dict(stat_dict: Dict[str, StatisticItemWithTestType]) -> List[Dict]:
+            return [{key: value.to_dict()} for key, value in stat_dict.items()]
 
         d = super().to_dict()
         # Convert stat dictionaries to lists of dictionaries
         d["test_type"] = convert_stat_dict(self.test_type)
         d["nist_category"] = convert_stat_dict(self.nist_category)
         d["scenario"] = convert_stat_dict(self.scenario)
-        d["category"] = convert_stat_dict(self.category)
+        d["category"] = convert_stat_test_type_dict(self.category)
         d["attack_method"] = convert_stat_dict(self.attack_method)
         return d
 
@@ -132,12 +164,17 @@ class ResultSummary(BaseDTO):
 @dataclass
 class RedTeamResultSummary(BaseDTO):
     summary: ResultSummary
+    task_status: Optional[str] = None
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
     @classmethod
     def from_dict(cls, data: Dict) -> "RedTeamResultSummary":
         if not data or "summary" not in data:
             return cls(summary=ResultSummary.from_dict({}))
+        
+        if "task_status" in data:
+            return cls(summary=ResultSummary.from_dict({}), task_status=data["task_status"])
+
         return cls(summary=ResultSummary.from_dict(data["summary"]))
 
     def to_dict(self) -> Dict:
@@ -174,42 +211,47 @@ class TestResult(BaseDTO):
 @dataclass
 class RedTeamResultDetails(BaseDTO):
     details: List[TestResult]
+    task_status: Optional[str] = None
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
     @classmethod
     def from_dict(cls, data: Dict) -> "RedTeamResultDetails":
         if not data or "details" not in data:
             return cls(details=[])
+        
+        if "task_status" in data:
+            return cls(details=[], task_status=data["task_status"])
 
-        details = []
-        for result in data["details"]:
+        # details = []
+        # for result in data["details"]:
             # Convert eval_tokens dict to TestEvalTokens object
-            eval_tokens = TestEvalTokens(**result["eval_tokens"])
+            # eval_tokens = TestEvalTokens(**result["eval_tokens"])
 
             # Create a copy of the result dict and replace eval_tokens
-            result_copy = dict(result)
-            result_copy["eval_tokens"] = eval_tokens
+            # result_copy = dict(result["details"])
+            # result_copy["eval_tokens"] = eval_tokens
 
             # Create TestResult object
-            test_result = TestResult(**result_copy)
-            details.append(test_result)
+            # test_result = TestResult(**result_copy)
+            # details.append(test_result)
 
-        return cls(details=details)
+        return cls(details=data["details"])
 
     def to_dict(self) -> Dict:
         return {
-            "details": [
-                {**result.to_dict(), "eval_tokens": result.eval_tokens.to_dict()}
-                for result in self.details
-            ]
+            # "details": [
+            #     {**result.to_dict(), "eval_tokens": result.eval_tokens.to_dict()}
+            #     for result in self.details
+            # ]
+            "details": self.details
         }
 
-    def to_dataframe(self) -> pd.DataFrame:
-        data = [
-            {**result.to_dict(), "eval_tokens": result.eval_tokens.to_dict()}
-            for result in self.details
-        ]
-        return pd.DataFrame(data)
+    # def to_dataframe(self) -> pd.DataFrame:
+    #     data = [
+    #         {**result.to_dict(), "eval_tokens": result.eval_tokens.to_dict()}
+    #         for result in self.details
+    #     ]
+    #     return pd.DataFrame(data)
 
 
 @dataclass
@@ -330,7 +372,6 @@ class RedteamModelHealthResponse(BaseDTO):
 class RedTeamConfig(BaseDTO):
     test_name: str = "Test Name"
     dataset_name: str = "standard"
-    model_saved_name: str = "gpt-4o-mini"
 
     redteam_test_configurations: RedTeamTestConfigurations = field(
         default_factory=RedTeamTestConfigurations
@@ -363,41 +404,33 @@ class RedTeamConfig(BaseDTO):
         )
 
 
-# @dataclass
-# class RedTeamConfigWithSavedModel(BaseDTO):
-#     test_name: str = "Test Name"
-#     dataset_name: str = "standard"
-#     model_saved_name: str = "gpt-4o-mini"
-
-#     redteam_test_configurations: RedTeamTestConfigurations = field(
-#         default_factory=RedTeamTestConfigurations
-#     )
-#     target_model_configuration: TargetModelConfiguration = field(
-#         default_factory=TargetModelConfiguration
-#     )
-
-#     _extra_fields: Dict[str, Any] = field(default_factory=dict)
-
-#     def to_dict(self) -> dict:
-#         d = asdict(self)
-#         d["redteam_test_configurations"] = self.redteam_test_configurations.to_dict()
-#         d["target_model_configuration"] = self.target_model_configuration.to_dict()
-#         return d
-
-#     @classmethod
-#     def from_dict(cls, data: dict):
-#         data = data.copy()
-#         test_configs = RedTeamTestConfigurations.from_dict(
-#             data.pop("redteam_test_configurations", {})
-#         )
-#         target_config = TargetModelConfiguration.from_dict(
-#             data.pop("target_model_configuration", {})
-#         )
-#         return cls(
-#             **data,
-#             redteam_test_configurations=test_configs,
-#             target_model_configuration=target_config,
-#         )
+@dataclass
+class RedTeamConfigWithSavedModel(BaseDTO):
+    test_name: str = "Test Name"
+    dataset_name: str = "standard"
+    model_saved_name: str = "gpt-4o-mini"
+
+    redteam_test_configurations: RedTeamTestConfigurations = field(
+        default_factory=RedTeamTestConfigurations
+    )
+
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["redteam_test_configurations"] = self.redteam_test_configurations.to_dict()
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        data = data.copy()
+        test_configs = RedTeamTestConfigurations.from_dict(
+            data.pop("redteam_test_configurations", {})
+        )
+        return cls(
+            **data,
+            redteam_test_configurations=test_configs,
+        )
 
 
 @dataclass
@@ -411,3 +444,4 @@ class RedTeamTaskList(BaseDTO):
 
 # Default configurations
 DEFAULT_REDTEAM_CONFIG = RedTeamConfig()
+DEFAULT_REDTEAM_CONFIG_WITH_SAVED_MODEL = RedTeamConfigWithSavedModel()

enkryptai_sdk/red_team.py

@@ -5,7 +5,7 @@ from .dto import (
     RedTeamModelHealthConfig,
     RedteamModelHealthResponse,
     RedTeamConfig,
-    # RedTeamConfigWithSavedModel,
+    RedTeamConfigWithSavedModel,
     RedTeamResponse,
     RedTeamResultSummary,
     RedTeamResultDetails,
@@ -103,22 +103,7 @@ class RedTeamClient(BaseClient):
             },
         }
 
-        saved_model = config.model_saved_name
-        if saved_model:
-            headers = {
-                "X-Enkrypt-Model": saved_model,
-                "Content-Type": "application/json",
-            }
-            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:
+        if config.target_model_configuration:
             payload["target_model_configuration"] = (
                 config.target_model_configuration.to_dict()
             )
@@ -133,9 +118,48 @@ class RedTeamClient(BaseClient):
             return RedTeamResponse.from_dict(response)
         else:
             raise RedTeamClientError(
-                "Please use a saved model or provide a target model configuration"
+                "Please provide a target model configuration"
             )
 
+    def add_task_with_saved_model(
+        self,
+        config: RedTeamConfigWithSavedModel,
+        model_saved_name: str,
+    ):
+        """
+        Add a new red teaming task using a saved model.
+        """
+        if not model_saved_name:
+            raise RedTeamClientError("Please provide a model_saved_name")
+
+        config = RedTeamConfigWithSavedModel.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 = {
+            # "async": config.async_enabled,
+            "dataset_name": config.dataset_name,
+            "test_name": config.test_name,
+            "redteam_test_configurations": {
+                k: v.to_dict() for k, v in test_configs.items()
+            },
+        }
+
+        headers = {
+            "X-Enkrypt-Model": model_saved_name,
+            "Content-Type": "application/json",
+        }
+        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)
+
     def status(self, task_id: str = None, test_name: str = None):
         """
         Get the status of a specific red teaming task.
@@ -253,6 +277,40 @@ class RedTeamClient(BaseClient):
             raise RedTeamClientError(response["error"])
         print(f"Response: {response}")
         return RedTeamResultSummary.from_dict(response)
+    
+    def get_result_summary_test_type(self, task_id: str = None, test_name: str = None, test_type: str = None):
+        """
+        Get the summary of results for a specific red teaming task for a specific test type.
+
+        Args:
+            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
+            test_type (str, optional): The type of test to get results for
+
+        Returns:
+            dict: The summary of the task results for the specified test type
+            
+        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")
+        
+        if not test_type:
+            raise RedTeamClientError("test_type must be provided")
+        
+        headers = {}
+        if task_id:
+            headers["X-Enkrypt-Task-ID"] = task_id
+        if test_name:
+            headers["X-Enkrypt-Test-Name"] = test_name
+
+        url = f"/redteam/v2/results/summary/{test_type}"
+        response = self._request("GET", url, headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        print(f"Response: {response}")
+        return RedTeamResultSummary.from_dict(response)
 
     def get_result_details(self, task_id: str = None, test_name: str = None):
         """
@@ -277,13 +335,60 @@ class RedTeamClient(BaseClient):
         if test_name:
             headers["X-Enkrypt-Test-Name"] = test_name
             
-        response = self._request("GET", "/redteam/results/details", headers=headers)
+        response = self._request("GET", "/redteam/v2/results/details", headers=headers)
         if response.get("error"):
             raise RedTeamClientError(response["error"])
         return RedTeamResultDetails.from_dict(response)
+    
+    def get_result_details_test_type(self, task_id: str = None, test_name: str = None, test_type: str = None):
+        """
+        Get the detailed results for a specific red teaming task for a specific test type.
+
+        Args:
+            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
+            test_type (str, optional): The type of 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
+        """
+        if not task_id and not test_name:
+            raise RedTeamClientError("Either task_id or test_name must be provided")
+        
+        if not test_type:
+            raise RedTeamClientError("test_type must be provided")
+        
+        headers = {}
+        if task_id:
+            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}"
+        response = self._request("GET", url, headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        return RedTeamResultDetails.from_dict(response)
+
+    def get_task_list(self, status: str = None):
+        """
+        Get a list of red teaming tasks.
+
+        Args:
+            status (str, optional): The status of the tasks to retrieve
+
+        Returns:
+            dict: The list of tasks
+        """
+        url = "/redteam/list-tasks"
+        if status:
+            url += f"?status={status}"
 
-    def get_task_list(self):
-        response = self._request("GET", "/redteam/list-tasks")
+        response = self._request("GET", url)
         if response.get("error"):
             raise RedTeamClientError(response["error"])
         return RedTeamTaskList.from_dict(response)