enkryptai-sdk 1.0.15__tar.gz → 1.0.17__tar.gz

{enkryptai_sdk-1.0.15/src/enkryptai_sdk.egg-info → enkryptai_sdk-1.0.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.15
+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
@@ -526,9 +526,6 @@ sample_custom_redteam_model_config = {
 
 ```python Python
 sample_redteam_risk_mitigation_guardrails_policy_config = {
-        "required_detectors": [
-            "policy_violation"
-        ],
         "redteam_summary": {
             "category": [
                 {

{enkryptai_sdk-1.0.15 → enkryptai_sdk-1.0.17}/README.md

@@ -503,9 +503,6 @@ sample_custom_redteam_model_config = {
 
 ```python Python
 sample_redteam_risk_mitigation_guardrails_policy_config = {
-        "required_detectors": [
-            "policy_violation"
-        ],
         "redteam_summary": {
             "category": [
                 {

{enkryptai_sdk-1.0.15 → enkryptai_sdk-1.0.17}/setup.py

@@ -9,7 +9,7 @@ with open(os.path.join(here, "README.md"), encoding="utf-8") as fh:
 setup(
     name="enkryptai-sdk",  # This is the name of your package on PyPI
     # NOTE: Also change this in .github/workflows/test.yaml
-    version="1.0.15", # Update this for new versions
+    version="1.0.17", # Update this for new versions
     description="A Python SDK with guardrails and red teaming functionality for API interactions",
     long_description=long_description,
     long_description_content_type="text/markdown",

{enkryptai_sdk-1.0.15 → enkryptai_sdk-1.0.17}/src/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):
@@ -40,6 +41,7 @@ class ModelProviders(str, Enum):
     CUSTOM = "custom"
     HR = "hr"
     URL = "url"
+    ENKRYPTAI = "enkryptai"
 
 
 @dataclass
@@ -92,10 +94,11 @@ class ModelDetailConfig:
     # model_provider: str = "openai"
     model_provider: ModelProviders = ModelProviders.OPENAI
     system_prompt: str = ""
-
-    endpoint_url: str = "https://api.openai.com/v1/chat/completions"
+    endpoint_url: str = ""
     auth_data: AuthData = field(default_factory=AuthData)
+    metadata: Dict[str, Any] = field(default_factory=dict)
     api_keys: Set[Optional[str]] = field(default_factory=lambda: {None})
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass
@@ -124,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
@@ -159,11 +143,13 @@ class ModelConfigDetails(BaseDTO):
     headers: str = ""
     system_prompt: str = ""
     hosting_type: str = "External"
-    endpoint_url: str = "https://api.openai.com/v1/chat/completions"
+    endpoint_url: str = ""
     model_name: Optional[str] = ""
     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)
@@ -184,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"]]
@@ -249,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-1.0.15 → enkryptai_sdk-1.0.17}/src/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
@@ -62,13 +120,60 @@ class RedTeamTaskStatus(BaseDTO):
     status: Optional[str] = None
 
 
+@dataclass
+class RedTeamTaskDetailsModelConfig(BaseDTO):
+    system_prompt: Optional[str] = None
+    model_version: Optional[str] = None
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "RedTeamTaskDetailsModelConfig":
+        return cls(
+            system_prompt=data.get("system_prompt"),
+            model_version=data.get("model_version"),
+        )
+    
+    def to_dict(self) -> Dict:
+        return {
+            "system_prompt": self.system_prompt,
+            "model_version": self.model_version,
+        }
+
+
 @dataclass
 class RedTeamTaskDetails(BaseDTO):
     created_at: Optional[str] = None
+    model_saved_name: Optional[str] = None
     model_name: Optional[str] = None
     status: Optional[str] = None
     test_name: Optional[str] = None
     task_id: Optional[str] = None
+    model_config: Optional[RedTeamTaskDetailsModelConfig] = None
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @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"),
+            model_name=data.get("model_name"),
+            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", {})),
+        )
+    
+    def to_dict(self) -> Dict:
+        return {
+            "created_at": self.created_at,
+            "model_saved_name": self.model_saved_name,
+            "model_name": self.model_name,
+            "status": self.status,
+            "test_name": self.test_name,
+            "task_id": self.task_id,
+            "model_config": self.model_config.to_dict(),
+        }
 
 
 @dataclass
@@ -322,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
@@ -345,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
@@ -376,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)
@@ -391,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
 
@@ -602,7 +710,7 @@ class RedTeamTaskList(BaseDTO):
 
 @dataclass
 class RedTeamRiskMitigationGuardrailsPolicyConfig(BaseDTO):
-    required_detectors: List[RiskGuardrailDetectorsEnum] = field(default_factory=list)
+    # required_detectors: List[RiskGuardrailDetectorsEnum] = field(default_factory=list)
     redteam_summary: ResultSummary = field(default_factory=ResultSummary)
     _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
@@ -611,14 +719,14 @@ class RedTeamRiskMitigationGuardrailsPolicyConfig(BaseDTO):
         data = data.copy()
         summary = ResultSummary.from_dict(data.pop("redteam_summary", {}))
         return cls(
-            required_detectors=[RiskGuardrailDetectorsEnum(detector) for detector in data.get("required_detectors", [])],
+            # required_detectors=[RiskGuardrailDetectorsEnum(detector) for detector in data.get("required_detectors", [])],
             redteam_summary=summary,
             _extra_fields=data,
         )
 
     def to_dict(self) -> dict:
         return {
-            "required_detectors": [detector.value for detector in self.required_detectors],
+            # "required_detectors": [detector.value for detector in self.required_detectors],
             "redteam_summary": self.redteam_summary.to_dict(),
         }
 

{enkryptai_sdk-1.0.15 → enkryptai_sdk-1.0.17}/src/enkryptai_sdk/models.py

@@ -29,17 +29,29 @@ class ModelClient(BaseClient):
         if isinstance(config, dict):
             config = ModelConfig.from_dict(config)
         
-        # Parse endpoint_url into components
-        parsed_url = urlparse(config.model_config.endpoint_url)
-        path_parts = parsed_url.path.strip("/").split("/")
-
-        # Extract base_path and endpoint path
-        if len(path_parts) >= 1:
-            base_path = path_parts[0]  # Usually 'v1'
-            remaining_path = "/".join(path_parts[1:])  # The rest of the path
-        else:
-            base_path = ""
-            remaining_path = ""
+        endpoint_data = {}
+        base_path = ""
+        remaining_path = ""
+
+        if config.model_config.endpoint_url:
+            # Parse endpoint_url into components
+            parsed_url = urlparse(config.model_config.endpoint_url)
+            path_parts = parsed_url.path.strip("/").split("/")
+
+            # Extract base_path and endpoint path
+            if len(path_parts) >= 1:
+                base_path = path_parts[0]  # Usually 'v1'
+                remaining_path = "/".join(path_parts[1:])  # The rest of the path
+            else:
+                base_path = ""
+                remaining_path = ""
+            
+            endpoint_data = {
+                    "scheme": parsed_url.scheme,
+                    "host": parsed_url.hostname,
+                    "port": parsed_url.port or (443 if parsed_url.scheme == "https" else 80),
+                    "base_path": f"/{base_path}",
+                }
 
         if config.model_config.paths:
             paths = config.model_config.paths.to_dict()
@@ -61,18 +73,15 @@ class ModelClient(BaseClient):
                 "hosting_type": config.model_config.hosting_type,
                 "model_source": config.model_config.model_source,
                 "system_prompt": config.model_config.system_prompt,
-                "endpoint": {
-                    "scheme": parsed_url.scheme,
-                    "host": parsed_url.hostname,
-                    "port": parsed_url.port or (443 if parsed_url.scheme == "https" else 80),
-                    "base_path": f"/{base_path}",
-                },
+                "endpoint": endpoint_data,
                 "paths": paths,
                 "auth_data": {
                     "header_name": config.model_config.auth_data.header_name,
                     "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],
@@ -149,92 +158,33 @@ class ModelClient(BaseClient):
         except Exception as e:
             return {"error": str(e)}
 
-    def modify_model(self, config: ModelConfig, old_model_saved_name=None, old_model_version=None) -> ModelResponse:
+    def modify_model(self, config: ModelConfig | dict, old_model_saved_name=None, old_model_version=None) -> ModelResponse:
         """
         Modify an existing model in the system.
 
         Args:
-            old_model_saved_name (str): The old saved name of the model to modify
-            old_model_version (str): The old version of the model to modify
-            config (ModelConfig): Configuration object containing model details
+            config (Union[ModelConfig, dict]): Configuration object or dictionary containing model details
+            old_model_saved_name (str, optional): The old saved name of the model to modify. Defaults to None.
+            old_model_version (str, optional): The old version of the model to modify. Defaults to None.
 
         Returns:
             dict: Response from the API containing the modified model details
         """
+        
+        temp_config = config
+        if isinstance(temp_config, dict):
+            temp_config = ModelConfig.from_dict(temp_config)
+
         if old_model_saved_name is None:
-            old_model_saved_name = config["model_saved_name"]
+            old_model_saved_name = temp_config.model_saved_name
 
         if old_model_version is None:
-            old_model_version = config["model_version"]
+            old_model_version = temp_config.model_version
 
         headers = {"Content-Type": "application/json", "X-Enkrypt-Model": old_model_saved_name, "X-Enkrypt-Model-Version": old_model_version}
-        # print(config)
-        config = ModelConfig.from_dict(config)
-        # Parse endpoint_url into components
-        parsed_url = urlparse(config.model_config.endpoint_url)
-        path_parts = parsed_url.path.strip("/").split("/")
-
-        # Extract base_path and endpoint path
-        if len(path_parts) >= 1:
-            base_path = path_parts[0]  # Usually 'v1'
-            remaining_path = "/".join(path_parts[1:])  # The rest of the path
-        else:
-            base_path = ""
-            remaining_path = ""
-
-        if config.model_config.paths:
-            paths = config.model_config.paths.to_dict()
-        else:
-            # Determine paths based on the endpoint
-            paths = {
-                "completions": (
-                    f"/{remaining_path.split('/')[-1]}" if remaining_path else ""
-                ),
-                "chat": f"/{remaining_path}" if remaining_path else "",
-            }
-
-        # Convert custom_headers to list of dictionaries
-        custom_headers = [header.to_dict() for header in config.model_config.custom_headers]
+        
+        payload = self.prepare_model_payload(temp_config)
 
-        payload = {
-            "model_saved_name": config.model_saved_name,
-            "model_version": config.model_version,
-            "testing_for": config.testing_for,
-            "model_name": config.model_name,
-            "certifications": config.certifications,
-            "model_config": {
-                "model_provider": config.model_config.model_provider,
-                "hosting_type": config.model_config.hosting_type,
-                "model_source": config.model_config.model_source,
-                "system_prompt": config.model_config.system_prompt,
-                "endpoint": {
-                    "scheme": parsed_url.scheme,
-                    "host": parsed_url.hostname,
-                    "port": parsed_url.port
-                    or (443 if parsed_url.scheme == "https" else 80),
-                    "base_path": f"/{base_path}",  # Just v1
-                },
-                "paths": paths,
-                "auth_data": {
-                    "header_name": config.model_config.auth_data.header_name,
-                    "header_prefix": config.model_config.auth_data.header_prefix,
-                    "space_after_prefix": config.model_config.auth_data.space_after_prefix,
-                },
-                "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],
-                "output_modalities": [m.value if hasattr(m, 'value') else m for m in config.model_config.output_modalities],
-                "custom_curl_command": config.model_config.custom_curl_command,
-                "custom_headers": custom_headers,
-                "custom_payload": config.model_config.custom_payload,
-                "custom_response_content_type": config.model_config.custom_response_content_type,
-                "custom_response_format": config.model_config.custom_response_format,
-                "metadata": config.model_config.metadata,
-                "default_request_options": config.model_config.default_request_options,
-            },
-        }
         try:
             response = self._request(
                 "PATCH", "/models/modify-model", headers=headers, json=payload

{enkryptai_sdk-1.0.15 → enkryptai_sdk-1.0.17}/src/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.15 → enkryptai_sdk-1.0.17/src/enkryptai_sdk.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.15
+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
@@ -526,9 +526,6 @@ sample_custom_redteam_model_config = {
 
 ```python Python
 sample_redteam_risk_mitigation_guardrails_policy_config = {
-        "required_detectors": [
-            "policy_violation"
-        ],
         "redteam_summary": {
             "category": [
                 {