enkryptai-sdk 1.0.16__py3-none-any.whl → 1.0.17__py3-none-any.whl

enkryptai_sdk/dto/models.py

@@ -5,6 +5,7 @@ from .base import BaseDTO
 from tabulate import tabulate
 from dataclasses import dataclass, field, asdict
 from typing import Optional, List, Set, Dict, Any
+from .red_team import ModelAuthTypeEnum, CustomHeader, ModelJwtConfig
 
 
 # class Modality(Enum):
@@ -126,25 +127,6 @@ class OutputModality(str, Enum):
     # code = "code"
 
 
-@dataclass
-class CustomHeader(BaseDTO):
-    key: str
-    value: str
-
-    @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "CustomHeader":
-        return cls(
-            key=data.get("key", ""),
-            value=data.get("value", "")
-        )
-
-    def to_dict(self) -> Dict[str, Any]:
-        return {
-            "key": self.key,
-            "value": self.value
-        }
-
-
 @dataclass
 class ModelConfigDetails(BaseDTO):
     model_id: str = None
@@ -166,6 +148,8 @@ class ModelConfigDetails(BaseDTO):
     apikey: Optional[str] = None
     paths: Optional[PathsConfig] = None
     tools: List[Dict[str, str]] = field(default_factory=list)
+    model_auth_type: Optional[ModelAuthTypeEnum] = ModelAuthTypeEnum.APIKEY
+    model_jwt_config: Optional[ModelJwtConfig] = None
     auth_data: AuthData = field(default_factory=AuthData)
     input_modalities: List[InputModality] = field(default_factory=list)
     output_modalities: List[OutputModality] = field(default_factory=list)
@@ -186,6 +170,12 @@ class ModelConfigDetails(BaseDTO):
         if "custom_headers" in data:
             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"])
+
         # Convert input_modalities strings to enum values
         if "input_modalities" in data:
             data["input_modalities"] = [InputModality(m) for m in data["input_modalities"]]
@@ -251,6 +241,9 @@ class ModelConfigDetails(BaseDTO):
 
     def to_dict(self):
         d = super().to_dict()
+        d["model_auth_type"] = self.model_auth_type.value
+        if self.model_jwt_config:
+            d["model_jwt_config"] = self.model_jwt_config.to_dict()
         # Handle AuthData specifically
         d["auth_data"] = self.auth_data.to_dict()
         # Handle CustomHeader list

enkryptai_sdk/dto/red_team.py

@@ -20,6 +20,64 @@ class RiskGuardrailDetectorsEnum(str, Enum):
     # SYSTEM_PROMPT = "system_prompt"
 
 
+class ModelAuthTypeEnum(str, Enum):
+    APIKEY = "apikey"
+    JWT = "jwt"
+
+
+class ModelJwtMethodEnum(str, Enum):
+    POST = "POST"
+    GET = "GET"
+
+
+@dataclass
+class CustomHeader(BaseDTO):
+    key: str
+    value: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CustomHeader":
+        return cls(
+            key=data.get("key", ""),
+            value=data.get("value", "")
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "key": self.key,
+            "value": self.value
+        }
+
+
+@dataclass
+class ModelJwtConfig(BaseDTO):
+    jwt_method: ModelJwtMethodEnum = ModelJwtMethodEnum.POST
+    jwt_url: str = ""
+    jwt_headers: List[CustomHeader] = field(default_factory=list)
+    jwt_body: str = ""
+    jwt_response_key: str = ""
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ModelJwtConfig":
+        return cls(
+            jwt_method=ModelJwtMethodEnum(data.get("jwt_method", ModelJwtMethodEnum.POST)),
+            jwt_url=data.get("jwt_url", ""),
+            jwt_headers=[CustomHeader.from_dict(header) for header in data.get("jwt_headers", [])],
+            jwt_body=data.get("jwt_body", ""),
+            jwt_response_key=data.get("jwt_response_key", ""),
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "jwt_method": self.jwt_method.value,
+            "jwt_url": self.jwt_url,
+            "jwt_headers": [header.to_dict() for header in self.jwt_headers],
+            "jwt_body": self.jwt_body,
+            "jwt_response_key": self.jwt_response_key,
+        }
+
+
 @dataclass
 class RedteamHealthResponse(BaseDTO):
     status: str
@@ -95,6 +153,7 @@ class RedTeamTaskDetails(BaseDTO):
 
     @classmethod
     def from_dict(cls, data: Dict) -> "RedTeamTaskDetails":
+        # print(f"RedTeamTaskDetails data: {data}")
         return cls(
             created_at=data.get("created_at"),
             model_saved_name=data.get("model_saved_name"),
@@ -368,8 +427,21 @@ class RedTeamTestConfigurations(BaseDTO):
     # Advanced tests
     adv_info_test: TestConfig = field(default=None)
     adv_bias_test: TestConfig = field(default=None)
+    adv_tool_test: TestConfig = field(default=None)
     adv_command_test: TestConfig = field(default=None)
+    adv_pii_test: TestConfig = field(default=None)
+    adv_competitor_test: TestConfig = field(default=None)
+    # Custom tests
     custom_test: TestConfig = field(default=None)
+    # Agents tests
+    alignment_and_governance_test: TestConfig = field(default=None)
+    input_and_content_integrity_test: TestConfig = field(default=None)
+    infrastructure_and_integration_test: TestConfig = field(default=None)
+    security_and_privacy_test: TestConfig = field(default=None)
+    human_factors_and_societal_impact_test: TestConfig = field(default=None)
+    access_control_test: TestConfig = field(default=None)
+    physical_and_actuation_safety_test: TestConfig = field(default=None)
+    reliability_and_monitoring_test: TestConfig = field(default=None)
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
     @classmethod
@@ -391,25 +463,6 @@ class OutputModality(str, Enum):
     # audio = "audio"
     # video = "video"
     # code = "code"
-
-
-@dataclass
-class CustomHeader(BaseDTO):
-    key: str
-    value: str
-
-    @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "CustomHeader":
-        return cls(
-            key=data.get("key", ""),
-            value=data.get("value", "")
-        )
-
-    def to_dict(self) -> Dict[str, Any]:
-        return {
-            "key": self.key,
-            "value": self.value
-        }
     
 
 @dataclass
@@ -422,6 +475,8 @@ class TargetModelConfiguration(BaseDTO):
     rate_per_min: int = 20
     model_name: Optional[str] = "gpt-4o-mini"
     model_version: Optional[str] = None
+    model_auth_type: Optional[ModelAuthTypeEnum] = ModelAuthTypeEnum.APIKEY
+    model_jwt_config: Optional[ModelJwtConfig] = None
     model_api_key: Optional[str] = None
     input_modalities: List[InputModality] = field(default_factory=list)
     output_modalities: List[OutputModality] = field(default_factory=list)
@@ -437,10 +492,17 @@ class TargetModelConfiguration(BaseDTO):
         data = data.copy()
         if "custom_headers" in data:
             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"])
         return cls(**data)
     
     def to_dict(self) -> dict:
         d = asdict(self)
+        d["model_auth_type"] = self.model_auth_type.value
+        if self.model_jwt_config:
+            d["model_jwt_config"] = self.model_jwt_config.to_dict()
         d["custom_headers"] = [header.to_dict() for header in self.custom_headers]
         return d
 

enkryptai_sdk/models.py

@@ -80,6 +80,8 @@ class ModelClient(BaseClient):
                     "header_prefix": config.model_config.auth_data.header_prefix,
                     "space_after_prefix": config.model_config.auth_data.space_after_prefix,
                 },
+                "model_auth_type": config.model_config.model_auth_type,
+                "model_jwt_config": config.model_config.model_jwt_config,
                 "apikeys": [config.model_config.apikey] if config.model_config.apikey else [],
                 "tools": config.model_config.tools,
                 "input_modalities": [m.value if hasattr(m, 'value') else m for m in config.model_config.input_modalities],

enkryptai_sdk/red_team.py

@@ -367,6 +367,7 @@ class RedTeamClient(BaseClient):
         response = self._request("GET", "/redteam/get-task", headers=headers)
         if response.get("error"):
             raise RedTeamClientError(f"API Error: {str(response)}")
+        # print(f"RedTeamTaskDetails response: {response}")
         return RedTeamTaskDetails.from_dict(response["data"])
 
     def get_result_summary(self, task_id: str = None, test_name: str = None):

{enkryptai_sdk-1.0.16.dist-info → enkryptai_sdk-1.0.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.16
+Version: 1.0.17
 Summary: A Python SDK with guardrails and red teaming functionality for API interactions
 Home-page: https://github.com/enkryptai/enkryptai-sdk
 Author: Enkrypt AI Team

{enkryptai_sdk-1.0.16.dist-info → enkryptai_sdk-1.0.17.dist-info}/RECORD

@@ -8,8 +8,8 @@ enkryptai_sdk/deployments.py,sha256=A7XZ2JwrMod9V4_aV8bFY_Soh9E3jHdwaTuJ9BwXuyk,
 enkryptai_sdk/evals.py,sha256=BywyEgIT7xdJ58svO_sDNOMVowdB0RTGoAZPEbCnDVo,2595
 enkryptai_sdk/guardrails.py,sha256=NluimOA0gM9N3S_q47LTUeG97t9PlYqPHlZahDPkJvI,16365
 enkryptai_sdk/guardrails_old.py,sha256=SgzPZkTzbAPD9XfmYNG6M1-TrzbhDHpAkI3FjnVWS_s,6434
-enkryptai_sdk/models.py,sha256=OfADGvP2gy52_fXT5iAj81Q07iDGAEYWiKYLD_ycozU,8786
-enkryptai_sdk/red_team.py,sha256=cjN4LODbpYiECcoL0JROMcCPCzm3Ib6kXi7kQspP4hQ,19869
+enkryptai_sdk/models.py,sha256=0R0I4KOq0aDNi5utabANot-E8dT9GqiSsgrcI9RULHM,8932
+enkryptai_sdk/red_team.py,sha256=7bWNognd6TCzNQDocml-C_f6u5z1MeursSe3Rtz16Sw,19929
 enkryptai_sdk/response.py,sha256=43JRubzgGCpoVxYNzBZY0AlUgLbfcXD_AwD7wU3qY9o,4086
 enkryptai_sdk/dto/__init__.py,sha256=wHgIv_OCnVMJOys-vqImF59ifogDrMcgxVRmfNayVvc,2761
 enkryptai_sdk/dto/ai_proxy.py,sha256=clwMN4xdH8Zr55dnhilHbs-qaHRlCOrLPrij0Zd1Av0,11283
@@ -18,10 +18,10 @@ enkryptai_sdk/dto/coc.py,sha256=Lp2aat_24J4KuUg4BeJl9S39tEak8Bw15eJ4cQDrRQk,4749
 enkryptai_sdk/dto/datasets.py,sha256=RFA9CmbhD-QDDyweBq_k9iBd00b6I6SWmdP9DPNd9fc,5002
 enkryptai_sdk/dto/deployments.py,sha256=Aw4b8tDA3FYIomqDvCjblCXTagL4bT8Fx91X0SFXs40,11216
 enkryptai_sdk/dto/guardrails.py,sha256=oJQqFhsdQd_yPU187AhKse-Y4xktgmVNwwKKkzFazbg,50167
-enkryptai_sdk/dto/models.py,sha256=SDzNMwxDVIp2D1qO29ere8WgqC5n_bJuitesaoxitZo,14312
-enkryptai_sdk/dto/red_team.py,sha256=EHwDbDM0mMLDyGB_kN4KEOO65ng9Bn51q9O3IdwO5Ng,24766
-enkryptai_sdk-1.0.16.dist-info/licenses/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-enkryptai_sdk-1.0.16.dist-info/METADATA,sha256=T7ws5lOF-ynN_Z3cegMX0Ysotjp2_n7pet8Ci-ULFmI,72860
-enkryptai_sdk-1.0.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-enkryptai_sdk-1.0.16.dist-info/top_level.txt,sha256=s2X9UJJwvJamNmr6ZXWyyQe60sXtQGWFuaBYfhgHI_4,14
-enkryptai_sdk-1.0.16.dist-info/RECORD,,
+enkryptai_sdk/dto/models.py,sha256=Nk6ZQyMfDMbl5ITQWQJUZMPxQOTATA65QbKHIy7i3qM,14533
+enkryptai_sdk/dto/red_team.py,sha256=NrYj9NZX3EmrBpESiZYFUCzgXH3eConh4NhVRKT28LM,27362
+enkryptai_sdk-1.0.17.dist-info/licenses/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+enkryptai_sdk-1.0.17.dist-info/METADATA,sha256=10vTRJNUYuEu7KgixIMOvgaweistehy8_s3kf3lkXiA,72860
+enkryptai_sdk-1.0.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+enkryptai_sdk-1.0.17.dist-info/top_level.txt,sha256=s2X9UJJwvJamNmr6ZXWyyQe60sXtQGWFuaBYfhgHI_4,14
+enkryptai_sdk-1.0.17.dist-info/RECORD,,