enkryptai-sdk 1.0.7__py3-none-any.whl → 1.0.9__py3-none-any.whl

enkryptai_sdk/__init__.py

@@ -1,6 +1,7 @@
 from .evals import EvalsClient
 from .config import GuardrailsConfig
 from .guardrails import GuardrailsClient, GuardrailsClientError
+from .coc import CoCClient, CoCClientError
 from .models import ModelClient, ModelClientError
 from .red_team import RedTeamClient, RedTeamClientError
 from .datasets import DatasetClient, DatasetClientError
@@ -11,6 +12,8 @@ __all__ = [
     "GuardrailsClient",
     "GuardrailsClientError",
     "GuardrailsConfig",
+    "CoCClient",
+    "CoCClientError",
     "EvalsClient",
     "ModelClient",
     "RedTeamClient",

enkryptai_sdk/ai_proxy.py

@@ -60,7 +60,7 @@ class AIProxyClient(BaseClient):
                     # # As json.loads is not working with literal string representation of dict
                     # parsed_error = json.loads(error_message.replace("'", '"'))
                     # # Convert Python string representation to proper dict
-                    parsed_error = ast.literal_eval(response["error"])
+                    parsed_error = ast.literal_eval(f"API Error: {str(response)}")
                     # # Preserve both error and enkrypt_policy_detections fields
                     response["error"] = parsed_error.get("error", parsed_error)
                     if "enkrypt_policy_detections" in parsed_error:
@@ -84,6 +84,6 @@ class AIProxyClient(BaseClient):
                     # Fallback to direct error if error object can't be parsed
                     # print("Failed to parse error response: ", response)
                     return ChatCompletionDirectErrorResponse.from_dict(response)
-            raise AIProxyClientError(response["error"])
+            raise AIProxyClientError(f"API Error: {str(response)}")
 
         return ChatCompletionResponse.from_dict(response)

enkryptai_sdk/base.py

@@ -11,17 +11,65 @@ class BaseClient:
         self.http = urllib3.PoolManager()
         self.headers = {"apikey": self.api_key}
 
-    def _request(self, method, endpoint, payload=None, headers=None, **kwargs):
+    # def _request(self, method, endpoint, payload=None, headers=None, **kwargs):
+    #     url = self.base_url + endpoint
+    #     request_headers = {
+    #         "Accept-Encoding": "gzip",  # Add required gzip encoding
+    #         **self.headers,
+    #     }
+    #     if headers:
+    #         request_headers.update(headers)
+
+    #     try:
+    #         response = self.http.request(method, url, headers=request_headers, **kwargs)
+
+    #         if response.status >= 400:
+    #             error_data = (
+    #                 response.json()
+    #                 if response.data
+    #                 else {"message": f"HTTP {response.status}"}
+    #             )
+    #             error_message = error_data.get("message", str(error_data))
+    #             raise urllib3.exceptions.HTTPError(error_message)
+    #         return response.json()
+    #     except urllib3.exceptions.HTTPError as e:
+    #         return {"error": str(e)}
+
+    def _request(self, method, endpoint, payload=None, headers=None, form_data=None, **kwargs):
         url = self.base_url + endpoint
         request_headers = {
-            "Accept-Encoding": "gzip",  # Add required gzip encoding
+            "Accept-Encoding": "gzip",
             **self.headers,
         }
         if headers:
             request_headers.update(headers)
 
         try:
-            response = self.http.request(method, url, headers=request_headers, **kwargs)
+            if form_data:
+                # Handle multipart form data
+                fields = {}
+                for key, value in form_data.items():
+                    if isinstance(value, tuple):
+                        # Handle file upload tuple (filename, content, content_type)
+                        filename, file_data, content_type = value
+                        fields[key] = (filename, file_data, content_type)
+                    else:
+                        fields[key] = value
+
+                response = self.http.request(
+                    method,
+                    url,
+                    headers=request_headers,
+                    fields=fields,
+                    **kwargs
+                )
+            else:
+                response = self.http.request(
+                    method,
+                    url,
+                    headers=request_headers,
+                    **kwargs
+                )
 
             if response.status >= 400:
                 error_data = (
@@ -34,3 +82,4 @@ class BaseClient:
             return response.json()
         except urllib3.exceptions.HTTPError as e:
             return {"error": str(e)}
+

enkryptai_sdk/coc.py

@@ -0,0 +1,203 @@
+import os
+# import requests
+from .base import BaseClient
+from .dto import (
+    CoCPolicyData,
+    CoCPolicyResponse,
+    CoCDeletePolicyResponse,
+    CoCListPoliciesResponse,
+)
+
+
+class CoCClientError(Exception):
+    """
+    A custom exception for CoCClient errors.
+    """
+
+    pass
+
+
+class CoCClient(BaseClient):
+    """
+    A client for interacting with Enkrypt AI CoC API endpoints.
+    """
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com:443"):
+        super().__init__(api_key, base_url)
+
+    def add_policy(self, policy_name, policy_rules=None, total_rules=None, policy_file=None, policy_text=None):
+        """
+        Create a new policy with policy_rules.
+        
+        Parameters:
+        - policy_name (str): Name of the policy
+        - policy_rules (List or Str): List of rules for the policy
+        - total_rules (int): Total number of rules
+        - policy_file (str, optional): Path to the policy file (PDF)
+        - policy_text (str, optional): Policy text content
+
+        Returns:
+        - CoCPolicyResponse
+
+        Raises:
+        - CoCClientError: If validation fails or API returns an error
+        """
+        try:
+            if not policy_file and not policy_text:
+                raise CoCClientError("Must provide either policy_file or policy_text")
+            if policy_file and policy_text:
+                raise CoCClientError("Cannot provide both policy_file and policy_text")
+            
+            if isinstance(policy_rules, list):
+                policy_rules = "\n".join(policy_rules)
+            elif isinstance(policy_rules, str):
+                policy_rules = policy_rules.strip()
+            else:
+                raise CoCClientError("policy_rules must be a string or list of strings")
+
+            form_data = {
+                'name': policy_name,
+                'policy_rules': policy_rules,
+                'total_rules': total_rules
+            }
+
+            if policy_file:
+                # Normalize file path and check existence
+                file_path = os.path.abspath(policy_file)
+                file_name = os.path.basename(file_path)
+
+                if not os.path.exists(file_path):
+                    raise CoCClientError(f"File not found: {file_path}")
+
+                # Check file extension
+                if not file_path.lower().endswith('.pdf'):
+                    raise CoCClientError("Only PDF files are supported")
+                                    
+                with open(file_path, 'rb') as f:
+                    file_content = f.read()
+                    form_data['policy_file'] = (file_name, file_content, 'application/pdf')
+            else:
+                form_data['policy_text'] = policy_text
+
+            response = self._request(
+                "POST",
+                "/code-of-conduct/add-policy",
+                form_data=form_data
+            )
+
+            if isinstance(response, dict) and response.get("error"):
+                raise CoCClientError(f"API Error: {str(response)}")
+                
+            return CoCPolicyResponse.from_dict(response)
+        except Exception as e:
+            raise CoCClientError(str(e))
+        
+    def get_policy(self, policy_name):
+        """
+        Retrieve an existing policy by providing its header identifier.
+        """
+        headers = {"X-Enkrypt-Policy": policy_name}
+
+        try:
+            response = self._request("GET", "/code-of-conduct/get-policy", headers=headers)
+            if response.get("error"):
+                raise CoCClientError(f"API Error: {str(response)}")
+            return CoCPolicyData.from_dict(response)
+        except Exception as e:
+            raise CoCClientError(str(e))
+
+    def modify_policy(self, policy_name, policy_rules=None, total_rules=None, policy_file=None, policy_text=None):
+        """
+        Modify a policy with policy_rules.
+        
+        Parameters:
+        - policy_name (str): Name of the policy
+        - policy_rules (List or Str): List of rules for the policy
+        - total_rules (int): Total number of rules
+        - policy_file (str, optional): Path to the policy file (PDF)
+        - policy_text (str, optional): Policy text content
+
+        Returns:
+        - CoCPolicyResponse
+
+        Raises:
+        - CoCClientError: If validation fails or API returns an error
+        """
+        try:
+            if not policy_file and not policy_text:
+                raise CoCClientError("Must provide either policy_file or policy_text")
+            if policy_file and policy_text:
+                raise CoCClientError("Cannot provide both policy_file and policy_text")
+            
+            if isinstance(policy_rules, list):
+                policy_rules = "\n".join(policy_rules)
+            elif isinstance(policy_rules, str):
+                policy_rules = policy_rules.strip()
+            else:
+                raise CoCClientError("policy_rules must be a string or list of strings")
+
+            form_data = {
+                'name': policy_name,
+                'policy_rules': policy_rules,
+                'total_rules': total_rules
+            }
+
+            if policy_file:
+                # Normalize file path and check existence
+                file_path = os.path.abspath(policy_file)
+                file_name = os.path.basename(file_path)
+
+                if not os.path.exists(file_path):
+                    raise CoCClientError(f"File not found: {file_path}")
+
+                # Check file extension
+                if not file_path.lower().endswith('.pdf'):
+                    raise CoCClientError("Only PDF files are supported")
+                                    
+                with open(file_path, 'rb') as f:
+                    file_content = f.read()
+                    form_data['policy_file'] = (file_name, file_content, 'application/pdf')
+            else:
+                form_data['policy_text'] = policy_text
+
+            headers = {"X-Enkrypt-Policy": policy_name}
+
+            response = self._request(
+                "PATCH",
+                "/code-of-conduct/modify-policy",
+                form_data=form_data,
+                headers=headers
+            )
+
+            if isinstance(response, dict) and response.get("error"):
+                raise CoCClientError(f"API Error: {str(response)}")
+                
+            return CoCPolicyResponse.from_dict(response)
+        except Exception as e:
+            raise CoCClientError(str(e))
+    
+    def delete_policy(self, policy_name):
+        """
+        Delete a policy.
+        """
+        headers = {"X-Enkrypt-Policy": policy_name}
+
+        try:
+            response = self._request("DELETE", "/code-of-conduct/delete-policy", headers=headers)
+            if response.get("error"):
+                raise CoCClientError(f"API Error: {str(response)}")
+            return CoCDeletePolicyResponse.from_dict(response)
+        except Exception as e:
+            raise CoCClientError(str(e))
+
+    def get_policy_list(self):
+        """
+        List all policies.
+        """
+
+        try:
+            response = self._request("GET", "/code-of-conduct/list-policies")
+            if isinstance(response, dict) and response.get("error"):
+                raise CoCClientError(f"API Error: {str(response)}")
+            return CoCListPoliciesResponse.from_dict(response)
+        except Exception as e:
+            raise CoCClientError(str(e))

enkryptai_sdk/datasets.py

@@ -19,6 +19,32 @@ class DatasetClient(BaseClient):
     def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
         super().__init__(api_key, base_url)
 
+    @staticmethod
+    def prepare_dataset_payload(config: DatasetConfig | dict, is_custom: bool = False) -> dict:
+        """
+        Prepare the payload for dataset operations from a config object.
+        
+        Args:
+            config (Union[DatasetConfig, dict]): Configuration object or dictionary containing dataset details
+            
+        Returns:
+            dict: Processed payload ready for API submission
+        """
+        if isinstance(config, dict):
+            config = DatasetConfig.from_dict(config)
+        
+        payload = config.to_dict()
+
+        if not is_custom:
+            # Remove empty tools configuration
+            if (payload.get("tools") is None or 
+                payload["tools"] == [] or 
+                payload["tools"] == [{}] or
+                payload["tools"] == [{"name": "", "description": ""}]):
+                del payload["tools"]
+        
+        return payload
+
     def add_dataset(self, config: DatasetConfig):
         """
         Add a new dataset to the system.
@@ -34,14 +60,7 @@ class DatasetClient(BaseClient):
         if isinstance(config, dict):
             config = DatasetConfig.from_dict(config)
 
-        payload = config.to_dict()
-
-        # If payload["tools"] is None or is an empty list [] or [{}], remove it from the payload
-        if (payload.get("tools") is None or 
-            payload["tools"] == [] or 
-            payload["tools"] == [{}] or
-            payload["tools"] == [{"name": "", "description": ""}]):
-            del payload["tools"]
+        payload = self.prepare_dataset_payload(config)
 
         # Print payload
         # print(f"\nAdd Dataset Payload: {payload}")
@@ -54,7 +73,7 @@ class DatasetClient(BaseClient):
         # print(f"\nAdd Dataset Response: {response}")
 
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         return DatasetAddTaskResponse.from_dict(response)
     
     def get_dataset_task_status(self, dataset_name: str):
@@ -70,7 +89,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/task-status", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetTaskStatus.from_dict(response)
 
@@ -87,7 +106,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-task", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetTask.from_dict(response)
 
@@ -104,7 +123,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-datacard", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetCard.from_dict(response)
 
@@ -121,7 +140,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-dataset", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetResponse.from_dict(response)
 
@@ -138,7 +157,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-summary", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetSummary.from_dict(response)
 
@@ -158,5 +177,5 @@ class DatasetClient(BaseClient):
             url += f"?status={status}"
         response = self._request("GET", url)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         return DatasetCollection.from_dict(response)

enkryptai_sdk/deployments.py

@@ -46,7 +46,7 @@ class DeploymentClient(BaseClient):
             "POST", "/deployments/add-deployment", headers=headers, json=payload
         )
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return DeploymentAddTaskResponse.from_dict(response)
 
     def get_deployment(self, deployment_name: str, refresh_cache: bool = False):
@@ -64,7 +64,7 @@ class DeploymentClient(BaseClient):
         headers["X-Enkrypt-Refresh-Cache"] = "true" if refresh_cache else "false"
         response = self._request("GET", "/deployments/get-deployment", headers=headers)
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return GetDeploymentResponse.from_dict(response)
     
     def modify_deployment(self, deployment_name: str, config: DeploymentInput):
@@ -88,7 +88,7 @@ class DeploymentClient(BaseClient):
             "PATCH", "/deployments/modify-deployment", headers=headers, json=payload
         )
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return ModifyDeploymentResponse.from_dict(response)
 
     def delete_deployment(self, deployment_name: str):
@@ -104,7 +104,7 @@ class DeploymentClient(BaseClient):
         headers = {"X-Enkrypt-Deployment": deployment_name}
         response = self._request("DELETE", "/deployments/delete-deployment", headers=headers)
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return DeleteDeploymentResponse.from_dict(response)
 
     def list_deployments(self):
@@ -117,5 +117,5 @@ class DeploymentClient(BaseClient):
 
         response = self._request("GET", "/deployments/list-deployments")
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return DeploymentCollection.from_dict(response)

enkryptai_sdk/dto/__init__.py

@@ -4,6 +4,7 @@ from .datasets import *
 from .deployments import *
 from .ai_proxy import *
 from .guardrails import *
+from .coc import *
 
 
 __all__ = [
@@ -22,6 +23,8 @@ __all__ = [
     "Metadata",
     "DEFAULT_REDTEAM_CONFIG",
     "DEFAULT_REDTEAM_CONFIG_WITH_SAVED_MODEL",
+    "DEFAULT_CUSTOM_REDTEAM_CONFIG",
+    "DEFAULT_CUSTOM_REDTEAM_CONFIG_WITH_SAVED_MODEL",
     "ADVANCED_REDTEAM_TESTS",
     "DETAIL_MODEL_CONFIG",
     "DatasetConfig",
@@ -79,8 +82,16 @@ __all__ = [
     "GuardrailsRelevancyResponse",
     "GuardrailsPolicyRequest",
     "GuardrailsPolicyData",
-    "GuardrailsaPolicyResponse",
+    "GuardrailsPolicyResponse",
     "GuardrailsDeletePolicyData",
     "GuardrailsDeletePolicyResponse",
     "GuardrailsListPoliciesResponse",
+    "GuardrailsPolicyAtomizerRequest",
+    "GuardrailsPolicyAtomizerResponse",
+    "CoCPolicyRequest",
+    "CoCPolicyData",
+    "CoCPolicyResponse",
+    "CoCDeletePolicyData",
+    "CoCDeletePolicyResponse",
+    "CoCListPoliciesResponse",
 ]

enkryptai_sdk/dto/coc.py

@@ -0,0 +1,159 @@
+from enum import Enum
+from .base import BaseDTO
+from dataclasses import dataclass, field
+from typing import List, Dict, Any, Set, Optional, BinaryIO
+
+
+@dataclass
+class CoCPolicyData(BaseDTO):
+    created_at: str
+    name: str
+    updated_at: str
+    policy_id: int
+    project_name: str = "default"
+    policy_rules: str = ""
+    total_rules: int = 0
+    pdf_name: str = ""
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CoCPolicyData":
+        return cls(
+            created_at=data.get("created_at", ""),
+            name=data.get("name", ""),
+            updated_at=data.get("updated_at", ""),
+            policy_id=data.get("policy_id", 0),
+            project_name=data.get("project_name", "default"),
+            policy_rules=data.get("policy_rules", ""),
+            total_rules=data.get("total_rules", 0),
+            pdf_name=data.get("pdf_name", "")
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "created_at": self.created_at,
+            "name": self.name,
+            "updated_at": self.updated_at,
+            "policy_id": self.policy_id,
+            "project_name": self.project_name,
+            "policy_rules": self.policy_rules,
+            "total_rules": self.total_rules,
+            "pdf_name": self.pdf_name
+        }
+        result.update(self._extra_fields)
+        return result
+    
+    def get_rules_list(self) -> List[str]:
+        """
+        Get the policy rules as a list of strings.
+        
+        Returns:
+            List[str]: List of individual policy rules
+        """
+        if not self.policy_rules:
+            return []
+        return [rule.strip() for rule in self.policy_rules.split('\n') if rule.strip()]
+    
+    def __str__(self) -> str:
+        """
+        String representation of the response.
+        
+        Returns:
+            str: A formatted string 
+        """
+        return (
+            f"Policy Name: {self.name}\n"
+            f"Created At: {self.created_at}\n"
+            f"Updated At: {self.updated_at}\n"
+            f"Policy ID: {self.policy_id}\n"
+            f"total_rules: {self.total_rules}\n"
+        )
+
+
+@dataclass
+class CoCPolicyResponse(BaseDTO):
+    message: str
+    data: CoCPolicyData
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CoCPolicyResponse":
+        policy_data = data.get("data", {})
+        return cls(
+            message=data.get("message", ""),
+            data=CoCPolicyData.from_dict(policy_data)
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "message": self.message,
+            "data": self.data.to_dict()
+        }
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class CoCDeletePolicyData(BaseDTO):
+    policy_id: int
+    project_name: str = ""
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CoCDeletePolicyData":
+        return cls(
+            policy_id=data.get("policy_id", 0),
+            project_name=data.get("project_name", "")
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "policy_id": self.policy_id,
+            "project_name": self.project
+        }
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class CoCDeletePolicyResponse(BaseDTO):
+    message: str
+    data: CoCDeletePolicyData
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CoCDeletePolicyResponse":
+        policy_data = data.get("data", {})
+        return cls(
+            message=data.get("message", ""),
+            data=CoCDeletePolicyData.from_dict(policy_data)
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "message": self.message,
+            "data": self.data.to_dict()
+        }
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class CoCListPoliciesResponse(BaseDTO):
+    policies: List[CoCPolicyData] = field(default_factory=list)
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CoCListPoliciesResponse":
+        policies_data = data.get("policies", [])
+        return cls(
+            policies=[CoCPolicyData.from_dict(policy) for policy in policies_data]
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "policies": [policy.to_dict() for policy in self.policies]
+        }
+        result.update(self._extra_fields)
+        return result
+