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

enkryptai_sdk/dto/datasets.py

@@ -10,32 +10,78 @@ class Tool(BaseDTO):
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "Tool":
-        return cls(
-            name=data.get("name", ""),
-            description=data.get("description", "")
-        )
+        return cls(name=data.get("name", ""), description=data.get("description", ""))
 
     def to_dict(self) -> Dict[str, Any]:
-        return {
-            "name": self.name,
-            "description": self.description
-        }
+        return {"name": self.name, "description": self.description}
 
 
 @dataclass
 class DatasetConfig(BaseDTO):
-    system_description: str = None
-    dataset_name: str = None
-    policy_description: str = None
-    risk_categories: str = None
-    tools: List[Tool] = None
-    info_pdf_url: str = None
+    system_description: Optional[str] = None
+    dataset_name: Optional[str] = None
+    policy_description: Optional[str] = None
+    risk_categories: Optional[str] = None
+    tools: Optional[List[Tool]] = None
+    info_pdf_url: Optional[str] = None
     max_prompts: int = 100
     scenarios: int = 2
     categories: int = 2
     depth: int = 2
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetConfig":
+        # Handle tools conversion if present
+        tools = None
+        if "tools" in data and data["tools"] is not None:
+            tools = [
+                Tool.from_dict(tool) if isinstance(tool, dict) else tool
+                for tool in data["tools"]
+            ]
+
+        return cls(
+            system_description=data.get("system_description"),
+            dataset_name=data.get("dataset_name"),
+            policy_description=data.get("policy_description"),
+            risk_categories=data.get("risk_categories"),
+            tools=tools,
+            info_pdf_url=data.get("info_pdf_url"),
+            max_prompts=data.get("max_prompts", 100),
+            scenarios=data.get("scenarios", 2),
+            categories=data.get("categories", 2),
+            depth=data.get("depth", 2),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        result = {}
+
+        # Only include optional fields if they are not None
+        if self.system_description is not None:
+            result["system_description"] = self.system_description
+        if self.dataset_name is not None:
+            result["dataset_name"] = self.dataset_name
+        if self.policy_description is not None:
+            result["policy_description"] = self.policy_description
+        if self.risk_categories is not None:
+            result["risk_categories"] = self.risk_categories
+        if self.tools is not None:
+            result["tools"] = [
+                tool.to_dict() if hasattr(tool, "to_dict") else tool
+                for tool in self.tools
+            ]
+        if self.info_pdf_url is not None:
+            result["info_pdf_url"] = self.info_pdf_url
+
+        # Always include fields with default values
+        result["max_prompts"] = self.max_prompts
+        result["scenarios"] = self.scenarios
+        result["categories"] = self.categories
+        result["depth"] = self.depth
+
+        result.update(self._extra_fields)
+        return result
+
 
 @dataclass
 class DatasetCollection(BaseDTO):
@@ -63,7 +109,9 @@ class DatasetResponse(BaseDTO):
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "DatasetResponse":
-        dataset_items = [DatasetDataPoint.from_dict(item) for item in data.get("dataset", [])]
+        dataset_items = [
+            DatasetDataPoint.from_dict(item) for item in data.get("dataset", [])
+        ]
         return cls(dataset=dataset_items)
 
     def to_dict(self) -> Dict[str, Any]:
@@ -149,7 +197,9 @@ class DatasetTaskStatus(BaseDTO):
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "DatasetTaskStatus":
-        return cls(status=data.get("status", ""), dataset_name=data.get("dataset_name", ""))
+        return cls(
+            status=data.get("status", ""), dataset_name=data.get("dataset_name", "")
+        )
 
     def to_dict(self) -> Dict[str, Any]:
         return {"status": self.status, "dataset_name": self.dataset_name}

enkryptai_sdk/dto/models.py

@@ -43,6 +43,8 @@ class ModelProviders(str, Enum):
     URL = "url"
     ENKRYPTAI = "enkryptai"
     BOXAI = "boxai"
+    NUTANIX = "nutanix"
+    XACTLY = "xactly"
 
 
 @dataclass
@@ -92,11 +94,12 @@ class AuthData(BaseDTO):
 @dataclass
 class BoxAIAuthData(AuthData):
     """BoxAI-specific authentication data."""
+
     boxai_client_id: Optional[str] = None
     boxai_client_secret: Optional[str] = None
     boxai_user_id: Optional[str] = None
     boxai_default_file_id: Optional[str] = None
-    
+
     def __post_init__(self):
         # Store BoxAI fields in extra_fields for backward compatibility
         if self.boxai_client_id:
@@ -107,22 +110,32 @@ class BoxAIAuthData(AuthData):
             self._extra_fields["boxai_user_id"] = self.boxai_user_id
         if self.boxai_default_file_id:
             self._extra_fields["boxai_default_file_id"] = self.boxai_default_file_id
-    
+
     @classmethod
     def from_dict(cls, data: dict):
         # Extract BoxAI fields from extra_fields if they exist
         boxai_data = {}
         if "_extra_fields" in data:
             extra_fields = data["_extra_fields"]
-            for field in ["boxai_client_id", "boxai_client_secret", "boxai_user_id", "boxai_default_file_id"]:
+            for field in [
+                "boxai_client_id",
+                "boxai_client_secret",
+                "boxai_user_id",
+                "boxai_default_file_id",
+            ]:
                 if field in extra_fields:
                     boxai_data[field] = extra_fields[field]
-        
+
         # Merge with direct field values
-        for field in ["boxai_client_id", "boxai_client_secret", "boxai_user_id", "boxai_default_file_id"]:
+        for field in [
+            "boxai_client_id",
+            "boxai_client_secret",
+            "boxai_user_id",
+            "boxai_default_file_id",
+        ]:
             if field in data:
                 boxai_data[field] = data[field]
-        
+
         # Create the instance
         return cls(**boxai_data)
 
@@ -168,8 +181,8 @@ class OutputModality(str, Enum):
 
 @dataclass
 class ModelConfigDetails(BaseDTO):
-    model_id: str = None
-    model_source: str = None
+    model_id: Optional[str] = None
+    model_source: Optional[str] = None
     # model_provider: str = "openai"
     model_provider: ModelProviders = ModelProviders.OPENAI
     model_api_value: str = ""
@@ -207,21 +220,29 @@ class ModelConfigDetails(BaseDTO):
         data = data.copy()
 
         if "custom_headers" in data:
-            data["custom_headers"] = [CustomHeader.from_dict(h) for h in data["custom_headers"]]
+            data["custom_headers"] = [
+                CustomHeader.from_dict(h) for h 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"]
+            )
 
         # Convert input_modalities strings to enum values
         if "input_modalities" in data:
-            data["input_modalities"] = [InputModality(m) for m in data["input_modalities"]]
-        
+            data["input_modalities"] = [
+                InputModality(m) for m in data["input_modalities"]
+            ]
+
         # Convert output_modalities strings to enum values
         if "output_modalities" in data:
-            data["output_modalities"] = [OutputModality(m) for m in data["output_modalities"]]
+            data["output_modalities"] = [
+                OutputModality(m) for m in data["output_modalities"]
+            ]
 
         # Validate model_provider if present
         if "model_provider" in data:
@@ -232,11 +253,15 @@ class ModelConfigDetails(BaseDTO):
                     data["model_provider"] = ModelProviders(provider)
                 except ValueError:
                     valid_providers = [p.value for p in ModelProviders]
-                    raise ValueError(f"Invalid model_provider: '{provider}'. Must be one of: {valid_providers}")
+                    raise ValueError(
+                        f"Invalid model_provider: '{provider}'. Must be one of: {valid_providers}"
+                    )
             # If it's already an enum instance, keep it as is
             elif not isinstance(provider, ModelProviders):
                 valid_providers = [p.value for p in ModelProviders]
-                raise ValueError(f"Invalid model_provider type. Valid values: {valid_providers}")
+                raise ValueError(
+                    f"Invalid model_provider type. Valid values: {valid_providers}"
+                )
 
         # # Remove known fields that we don't want in our model
         # unwanted_fields = ["queryParams"]
@@ -305,16 +330,16 @@ class ModelConfigDetails(BaseDTO):
 
 @dataclass
 class ModelConfig(BaseDTO):
-    created_at: str = None
-    updated_at: str = None
-    created_by: str = None
-    updated_by: str = None
-    model_id: str = None
-    model_saved_name: str = None
-    model_version: str = None
+    created_at: Optional[str] = None
+    updated_at: Optional[str] = None
+    created_by: Optional[str] = None
+    updated_by: Optional[str] = None
+    model_id: Optional[str] = None
+    model_saved_name: Optional[str] = None
+    model_version: Optional[str] = None
     testing_for: str = "foundationModels"
     # modality: Modality = Modality.TEXT
-    project_name: str = None
+    project_name: Optional[str] = None
     model_name: Optional[str] = "gpt-4o-mini"
     certifications: List[str] = field(default_factory=list)
     model_config: ModelConfigDetails = field(default_factory=ModelConfigDetails)