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

enkryptai_sdk/dto/red_team.py

@@ -9,6 +9,7 @@ from .models import ModelConfig
 from .guardrails import GuardrailDetectors
 from .common import ModelAuthTypeEnum, CustomHeader, ModelJwtConfig
 
+
 # The risk mitigation do not support all detectors, so we need to create a separate enum for them.
 class RiskGuardrailDetectorsEnum(str, Enum):
     NSFW = "nsfw"
@@ -27,17 +28,13 @@ class RiskGuardrailDetectorsEnum(str, Enum):
 class RedteamHealthResponse(BaseDTO):
     status: str
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
-    
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "RedteamHealthResponse":
-        return cls(
-            status=data.get("status", "")
-        )
-    
+        return cls(status=data.get("status", ""))
+
     def to_dict(self) -> Dict[str, Any]:
-        return {
-            "status": self.status
-        }
+        return {"status": self.status}
 
 
 @dataclass
@@ -55,7 +52,7 @@ class RedTeamResponse(BaseDTO):
             message=data.get("message"),
             data=data.get("data"),
         )
-    
+
     def to_dict(self) -> Dict:
         return super().to_dict()
 
@@ -77,7 +74,7 @@ class RedTeamTaskDetailsModelConfig(BaseDTO):
             system_prompt=data.get("system_prompt"),
             model_version=data.get("model_version"),
         )
-    
+
     def to_dict(self) -> Dict:
         return {
             "system_prompt": self.system_prompt,
@@ -108,9 +105,11 @@ class RedTeamTaskDetails(BaseDTO):
             status=data.get("status"),
             test_name=data.get("test_name"),
             task_id=data.get("task_id"),
-            model_config=RedTeamTaskDetailsModelConfig.from_dict(data.get("model_config", {})),
+            model_config=RedTeamTaskDetailsModelConfig.from_dict(
+                data.get("model_config", {})
+            ),
         )
-    
+
     def to_dict(self) -> Dict:
         return {
             "created_at": self.created_at,
@@ -191,8 +190,10 @@ 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]:
+
+        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():
@@ -221,8 +222,10 @@ class ResultSummary(BaseDTO):
     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]:
+
+        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()
@@ -248,9 +251,11 @@ class RedTeamResultSummary(BaseDTO):
     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({}), task_status=data["task_status"]
+            )
 
         return cls(summary=ResultSummary.from_dict(data["summary"]))
 
@@ -295,22 +300,22 @@ class RedTeamResultDetails(BaseDTO):
     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"]:
-            # Convert eval_tokens dict to TestEvalTokens object
-            # eval_tokens = TestEvalTokens(**result["eval_tokens"])
+        # Convert eval_tokens dict to TestEvalTokens object
+        # eval_tokens = TestEvalTokens(**result["eval_tokens"])
 
-            # Create a copy of the result dict and replace eval_tokens
-            # result_copy = dict(result["details"])
-            # result_copy["eval_tokens"] = eval_tokens
+        # Create a copy of the result dict and replace 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)
+        # Create TestResult object
+        # test_result = TestResult(**result_copy)
+        # details.append(test_result)
 
         return cls(details=data["details"])
 
@@ -336,8 +341,19 @@ class AttackMethods(BaseDTO):
     basic: List[str] = field(default_factory=lambda: ["basic"])
     advanced: Dict[str, List[str]] = field(
         default_factory=lambda: {
-            "static": ["masking", "figstep", "hades","encoding", "single_shot", "echo", "speed", "pitch", "reverb", "noise" ],
-            "dynamic": ["iterative","jood"]
+            "static": [
+                "masking",
+                "figstep",
+                "hades",
+                "encoding",
+                "single_shot",
+                "echo",
+                "speed",
+                "pitch",
+                "reverb",
+                "noise",
+            ],
+            "dynamic": ["iterative", "jood"],
         }
     )
 
@@ -450,7 +466,7 @@ class OutputModality(str, Enum):
     # audio = "audio"
     # video = "video"
     # code = "code"
-    
+
 
 @dataclass
 class TargetModelConfiguration(BaseDTO):
@@ -478,13 +494,17 @@ class TargetModelConfiguration(BaseDTO):
     def from_dict(cls, data: dict):
         data = data.copy()
         if "custom_headers" in data:
-            data["custom_headers"] = [CustomHeader.from_dict(header) for header in data["custom_headers"]]
+            data["custom_headers"] = [
+                CustomHeader.from_dict(header) for header in data["custom_headers"]
+            ]
         if "model_auth_type" in data:
             data["model_auth_type"] = ModelAuthTypeEnum(data["model_auth_type"])
         if "model_jwt_config" in data:
-            data["model_jwt_config"] = ModelJwtConfig.from_dict(data["model_jwt_config"])
+            data["model_jwt_config"] = ModelJwtConfig.from_dict(
+                data["model_jwt_config"]
+            )
         return cls(**data)
-    
+
     def to_dict(self) -> dict:
         d = asdict(self)
         d["model_auth_type"] = self.model_auth_type.value
@@ -523,9 +543,8 @@ class RedTeamModelHealthConfigV3(BaseDTO):
     V3 format for model health check that accepts endpoint_configuration
     similar to add_custom_task.
     """
-    endpoint_configuration: ModelConfig = field(
-        default_factory=ModelConfig
-    )
+
+    endpoint_configuration: ModelConfig = field(default_factory=ModelConfig)
 
     def to_dict(self) -> dict:
         d = asdict(self)
@@ -535,25 +554,27 @@ class RedTeamModelHealthConfigV3(BaseDTO):
     @classmethod
     def from_dict(cls, data: dict):
         data = data.copy()
-        endpoint_config = ModelConfig.from_dict(
-            data.pop("endpoint_configuration", {})
-        )
+        endpoint_config = ModelConfig.from_dict(data.pop("endpoint_configuration", {}))
         return cls(
             endpoint_configuration=endpoint_config,
         )
-    
+
     def to_target_model_configuration(self) -> TargetModelConfiguration:
         """
         Convert endpoint_configuration to target_model_configuration format.
         This enables the V3 format to be compatible with the existing backend API.
         """
         model_config = self.endpoint_configuration.model_config
-        
+
         return TargetModelConfiguration(
             testing_for=self.endpoint_configuration.testing_for,
             system_prompt=model_config.system_prompt,
             model_source=model_config.model_source,
-            model_provider=model_config.model_provider.value if hasattr(model_config.model_provider, 'value') else model_config.model_provider,
+            model_provider=(
+                model_config.model_provider.value
+                if hasattr(model_config.model_provider, "value")
+                else model_config.model_provider
+            ),
             model_endpoint_url=model_config.endpoint_url,
             rate_per_min=model_config.rate_per_min,
             model_name=self.endpoint_configuration.model_name,
@@ -561,8 +582,14 @@ class RedTeamModelHealthConfigV3(BaseDTO):
             model_auth_type=model_config.model_auth_type,
             model_jwt_config=model_config.model_jwt_config,
             model_api_key=model_config.apikey,
-            input_modalities=[InputModality(m) if isinstance(m, str) else m for m in model_config.input_modalities],
-            output_modalities=[OutputModality(m) if isinstance(m, str) else m for m in model_config.output_modalities],
+            input_modalities=[
+                InputModality(m) if isinstance(m, str) else m
+                for m in model_config.input_modalities
+            ],
+            output_modalities=[
+                OutputModality(m) if isinstance(m, str) else m
+                for m in model_config.output_modalities
+            ],
             custom_curl_command=model_config.custom_curl_command,
             custom_headers=model_config.custom_headers,
             custom_payload=model_config.custom_payload,
@@ -578,22 +605,22 @@ class RedteamModelHealthResponse(BaseDTO):
     error: str
     data: Optional[Dict[str, Any]] = field(default_factory=dict)
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
-    
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "RedteamModelHealthResponse":
         return cls(
             status=data.get("status", ""),
             message=data.get("message", ""),
             data=data.get("data", {}),
-            error=data.get("error", "")
+            error=data.get("error", ""),
         )
-    
+
     def to_dict(self) -> Dict[str, Any]:
         return {
             "status": self.status,
             "message": self.message,
             "data": self.data,
-            "error": self.error
+            "error": self.error,
         }
 
 
@@ -671,20 +698,22 @@ class RedTeamCustomConfig(BaseDTO):
     redteam_test_configurations: RedTeamTestConfigurations = field(
         default_factory=RedTeamTestConfigurations
     )
-    dataset_configuration: DatasetConfig = field(
-        default_factory=DatasetConfig
-    )
-    endpoint_configuration: ModelConfig = field(
-        default_factory=ModelConfig
-    )
+    dataset_configuration: Optional[DatasetConfig] = None
+    endpoint_configuration: Optional[ModelConfig] = None
 
     _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["dataset_configuration"] = self.dataset_configuration.to_dict()
-        d["endpoint_configuration"] = self.endpoint_configuration.to_dict()
+        if self.dataset_configuration is not None:
+            d["dataset_configuration"] = self.dataset_configuration.to_dict()
+        else:
+            d.pop("dataset_configuration", None)
+        if self.endpoint_configuration is not None:
+            d["endpoint_configuration"] = self.endpoint_configuration.to_dict()
+        else:
+            d.pop("endpoint_configuration", None)
         return d
 
     @classmethod
@@ -693,12 +722,18 @@ class RedTeamCustomConfig(BaseDTO):
         test_configs = RedTeamTestConfigurations.from_dict(
             data.pop("redteam_test_configurations", {})
         )
-        dataset_config = DatasetConfig.from_dict(
-            data.pop("dataset_configuration", {})
-        )
-        endpoint_config = ModelConfig.from_dict(
-            data.pop("endpoint_configuration", {})
-        )
+        dataset_config = None
+        if "dataset_configuration" in data and data["dataset_configuration"]:
+            dataset_config = DatasetConfig.from_dict(data.pop("dataset_configuration"))
+        else:
+            data.pop("dataset_configuration", None)
+
+        endpoint_config = None
+        if "endpoint_configuration" in data and data["endpoint_configuration"]:
+            endpoint_config = ModelConfig.from_dict(data.pop("endpoint_configuration"))
+        else:
+            data.pop("endpoint_configuration", None)
+
         return cls(
             **data,
             redteam_test_configurations=test_configs,
@@ -706,6 +741,7 @@ class RedTeamCustomConfig(BaseDTO):
             endpoint_configuration=endpoint_config,
         )
 
+
 @dataclass
 class RedTeamCustomConfigWithSavedModel(BaseDTO):
     test_name: str = "Test Name"
@@ -714,16 +750,17 @@ class RedTeamCustomConfigWithSavedModel(BaseDTO):
     redteam_test_configurations: RedTeamTestConfigurations = field(
         default_factory=RedTeamTestConfigurations
     )
-    dataset_configuration: DatasetConfig = field(
-        default_factory=DatasetConfig
-    )
+    dataset_configuration: Optional[DatasetConfig] = None
 
     _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["dataset_configuration"] = self.dataset_configuration.to_dict()
+        if self.dataset_configuration is not None:
+            d["dataset_configuration"] = self.dataset_configuration.to_dict()
+        else:
+            d.pop("dataset_configuration", None)
         return d
 
     @classmethod
@@ -732,9 +769,12 @@ class RedTeamCustomConfigWithSavedModel(BaseDTO):
         test_configs = RedTeamTestConfigurations.from_dict(
             data.pop("redteam_test_configurations", {})
         )
-        dataset_config = DatasetConfig.from_dict(
-            data.pop("dataset_configuration", {})
-        )
+        dataset_config = None
+        if "dataset_configuration" in data and data["dataset_configuration"]:
+            dataset_config = DatasetConfig.from_dict(data.pop("dataset_configuration"))
+        else:
+            data.pop("dataset_configuration", None)
+
         return cls(
             **data,
             redteam_test_configurations=test_configs,
@@ -749,7 +789,7 @@ class RedTeamTaskList(BaseDTO):
     def to_dataframe(self) -> pd.DataFrame:
         data = [task for task in self.tasks]
         return pd.DataFrame(data)
-    
+
 
 @dataclass
 class RedTeamRiskMitigationGuardrailsPolicyConfig(BaseDTO):
@@ -793,21 +833,21 @@ class RedTeamRiskMitigationGuardrailsPolicyResponse(BaseDTO):
 
     def to_dict(self) -> dict:
         policy_dict = self.guardrails_policy.to_dict()
-        
+
         # Remove detector entries that are disabled and have no other config
         final_policy_dict = {}
         for key, value in policy_dict.items():
             if isinstance(value, dict):
                 # Check if 'enabled' is the only key and its value is False
-                if list(value.keys()) == ['enabled'] and not value['enabled']:
+                if list(value.keys()) == ["enabled"] and not value["enabled"]:
                     continue
                 # Check for empty detectors that only have 'enabled': False
                 if not value.get("enabled", True) and len(value) == 1:
                     continue
                 # check for other empty values
-                if not any(v for k, v in value.items() if k != 'enabled'):
-                     if not value.get('enabled'):
-                         continue
+                if not any(v for k, v in value.items() if k != "enabled"):
+                    if not value.get("enabled"):
+                        continue
                 final_policy_dict[key] = value
 
         return {
@@ -862,21 +902,18 @@ class RedTeamRiskMitigationSystemPromptResponse(BaseDTO):
             "message": self.message,
         }
 
+
 @dataclass
 class RedTeamKeyFinding(BaseDTO):
     text: str
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
-    
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "RedTeamKeyFinding":
-        return cls(
-            text=data.get("text", "")
-        )
-    
+        return cls(text=data.get("text", ""))
+
     def to_dict(self) -> Dict[str, Any]:
-        result = {
-            "text": self.text
-        }
+        result = {"text": self.text}
         result.update(self._extra_fields)
         return result
 
@@ -886,21 +923,20 @@ class RedTeamFindingsResponse(BaseDTO):
     key_findings: List[RedTeamKeyFinding] = field(default_factory=list)
     message: str = ""
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
-    
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "RedTeamFindingsResponse":
         key_findings_data = data.get("key_findings", [])
-        key_findings = [RedTeamKeyFinding.from_dict(finding) for finding in key_findings_data]
-        
-        return cls(
-            key_findings=key_findings,
-            message=data.get("message", "")
-        )
-    
+        key_findings = [
+            RedTeamKeyFinding.from_dict(finding) for finding in key_findings_data
+        ]
+
+        return cls(key_findings=key_findings, message=data.get("message", ""))
+
     def to_dict(self) -> Dict[str, Any]:
         result = {
             "key_findings": [finding.to_dict() for finding in self.key_findings],
-            "message": self.message
+            "message": self.message,
         }
         result.update(self._extra_fields)
         return result
@@ -912,20 +948,20 @@ class RedTeamDownloadLinkResponse(BaseDTO):
     expiry: str = ""
     expires_at: str = ""
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
-    
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "RedTeamDownloadLinkResponse":
         return cls(
             link=data.get("link", ""),
             expiry=data.get("expiry", ""),
-            expires_at=data.get("expires_at", "")
+            expires_at=data.get("expires_at", ""),
         )
-    
+
     def to_dict(self) -> Dict[str, Any]:
         result = {
             "link": self.link,
             "expiry": self.expiry,
-            "expires_at": self.expires_at
+            "expires_at": self.expires_at,
         }
         result.update(self._extra_fields)
         return result
@@ -944,6 +980,7 @@ class AttackMethodsV3(BaseDTO):
         }
     }
     """
+
     _data: Dict[str, Dict[str, Dict[str, Any]]] = field(default_factory=dict)
 
     def to_dict(self) -> dict:
@@ -970,13 +1007,15 @@ class TestConfigV3(BaseDTO):
         attack_methods = AttackMethodsV3.from_dict(data.get("attack_methods", {}))
         return cls(
             sample_percentage=data.get("sample_percentage", 5),
-            attack_methods=attack_methods
+            attack_methods=attack_methods,
         )
 
 
 @dataclass
 class RedTeamTestConfigurationsV3(BaseDTO):
     """V3 format for red team test configurations with nested attack methods"""
+
+    version: str = "3.0"
     # Basic tests
     bias_test: TestConfigV3 = field(default=None)
     cbrn_test: TestConfigV3 = field(default=None)
@@ -1019,7 +1058,12 @@ class RedTeamTestConfigurationsV3(BaseDTO):
 
     @classmethod
     def from_dict(cls, data: dict):
-        return cls(**{k: TestConfigV3.from_dict(v) if isinstance(v, dict) else v for k, v in data.items()})
+        return cls(
+            **{
+                k: TestConfigV3.from_dict(v) if isinstance(v, dict) else v
+                for k, v in data.items()
+            }
+        )
 
 
 @dataclass
@@ -1030,20 +1074,22 @@ class RedTeamCustomConfigV3(BaseDTO):
     redteam_test_configurations: RedTeamTestConfigurationsV3 = field(
         default_factory=RedTeamTestConfigurationsV3
     )
-    dataset_configuration: DatasetConfig = field(
-        default_factory=DatasetConfig
-    )
-    endpoint_configuration: ModelConfig = field(
-        default_factory=ModelConfig
-    )
+    dataset_configuration: Optional[DatasetConfig] = None
+    endpoint_configuration: Optional[ModelConfig] = None
 
     _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["dataset_configuration"] = self.dataset_configuration.to_dict()
-        d["endpoint_configuration"] = self.endpoint_configuration.to_dict()
+        if self.dataset_configuration is not None:
+            d["dataset_configuration"] = self.dataset_configuration.to_dict()
+        else:
+            d.pop("dataset_configuration", None)
+        if self.endpoint_configuration is not None:
+            d["endpoint_configuration"] = self.endpoint_configuration.to_dict()
+        else:
+            d.pop("endpoint_configuration", None)
         return d
 
     @classmethod
@@ -1052,12 +1098,18 @@ class RedTeamCustomConfigV3(BaseDTO):
         test_configs = RedTeamTestConfigurationsV3.from_dict(
             data.pop("redteam_test_configurations", {})
         )
-        dataset_config = DatasetConfig.from_dict(
-            data.pop("dataset_configuration", {})
-        )
-        endpoint_config = ModelConfig.from_dict(
-            data.pop("endpoint_configuration", {})
-        )
+        dataset_config = None
+        if "dataset_configuration" in data and data["dataset_configuration"]:
+            dataset_config = DatasetConfig.from_dict(data.pop("dataset_configuration"))
+        else:
+            data.pop("dataset_configuration", None)
+
+        endpoint_config = None
+        if "endpoint_configuration" in data and data["endpoint_configuration"]:
+            endpoint_config = ModelConfig.from_dict(data.pop("endpoint_configuration"))
+        else:
+            data.pop("endpoint_configuration", None)
+
         return cls(
             **data,
             redteam_test_configurations=test_configs,
@@ -1074,16 +1126,17 @@ class RedTeamCustomConfigWithSavedModelV3(BaseDTO):
     redteam_test_configurations: RedTeamTestConfigurationsV3 = field(
         default_factory=RedTeamTestConfigurationsV3
     )
-    dataset_configuration: DatasetConfig = field(
-        default_factory=DatasetConfig
-    )
+    dataset_configuration: Optional[DatasetConfig] = None
 
     _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["dataset_configuration"] = self.dataset_configuration.to_dict()
+        if self.dataset_configuration is not None:
+            d["dataset_configuration"] = self.dataset_configuration.to_dict()
+        else:
+            d.pop("dataset_configuration", None)
         return d
 
     @classmethod
@@ -1092,9 +1145,12 @@ class RedTeamCustomConfigWithSavedModelV3(BaseDTO):
         test_configs = RedTeamTestConfigurationsV3.from_dict(
             data.pop("redteam_test_configurations", {})
         )
-        dataset_config = DatasetConfig.from_dict(
-            data.pop("dataset_configuration", {})
-        )
+        dataset_config = None
+        if "dataset_configuration" in data and data["dataset_configuration"]:
+            dataset_config = DatasetConfig.from_dict(data.pop("dataset_configuration"))
+        else:
+            data.pop("dataset_configuration", None)
+
         return cls(
             **data,
             redteam_test_configurations=test_configs,
@@ -1110,4 +1166,6 @@ DEFAULT_CUSTOM_REDTEAM_CONFIG = RedTeamCustomConfig()
 DEFAULT_CUSTOM_REDTEAM_CONFIG_WITH_SAVED_MODEL = RedTeamCustomConfigWithSavedModel()
 
 DEFAULT_CUSTOM_REDTEAM_CONFIG_V3 = RedTeamCustomConfigV3()
-DEFAULT_CUSTOM_REDTEAM_CONFIG_WITH_SAVED_MODEL_V3 = RedTeamCustomConfigWithSavedModelV3()
+DEFAULT_CUSTOM_REDTEAM_CONFIG_WITH_SAVED_MODEL_V3 = (
+    RedTeamCustomConfigWithSavedModelV3()
+)