enkryptai-sdk 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

enkryptai_sdk/__init__.py

@@ -1,4 +1,13 @@
 from .guardrails import GuardrailsClient
-from .guardrails_config import GuardrailsConfig
+from .config import GuardrailsConfig
+from .evals import EvalsClient
+from .models import ModelClient
+from .red_team import RedTeamClient
 
-__all__ = ["GuardrailsClient", "GuardrailsConfig"]
+__all__ = [
+    "GuardrailsClient",
+    "GuardrailsConfig",
+    "EvalsClient",
+    "ModelClient",
+    "RedTeamClient",
+]

enkryptai_sdk/{guardrails_config.py → config.py}

@@ -1,17 +1,21 @@
 import copy
 
 # Base default configuration for all detectors.
-DEFAULT_CONFIG = {
+DEFAULT_GUARDRAILS_CONFIG = {
     "topic_detector": {"enabled": False, "topic": []},
     "nsfw": {"enabled": False},
     "toxicity": {"enabled": False},
     "pii": {"enabled": False, "entities": []},
     "injection_attack": {"enabled": False},
     "keyword_detector": {"enabled": False, "banned_keywords": []},
-    "policy_violation": {"enabled": False, "policy_text": "", "need_explanation": False},
+    "policy_violation": {
+        "enabled": False,
+        "policy_text": "",
+        "need_explanation": False,
+    },
     "bias": {"enabled": False},
     "copyright_ip": {"enabled": False},
-    "system_prompt": {"enabled": False, "index": "system"}
+    "system_prompt": {"enabled": False, "index": "system"},
 }
 
 
@@ -24,29 +28,30 @@ class GuardrailsConfig:
 
     def __init__(self, config=None):
         # Use a deep copy of the default to avoid accidental mutation.
-        self.config = copy.deepcopy(DEFAULT_CONFIG) if config is None else config
+        self.config = (
+            copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG) if config is None else config
+        )
 
     @classmethod
     def injection_attack(cls):
         """
         Returns a configuration instance pre-configured for injection attack detection.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["injection_attack"] = {"enabled": True}
         return cls(config)
 
     @classmethod
-    def policy_violation(cls, 
-                         policy_text: str, 
-                         need_explanation: bool = False):
+    def policy_violation(cls, policy_text: str, need_explanation: bool = False):
         """
         Returns a configuration instance pre-configured for policy violation detection.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
-        config["policy_violation"] = {"enabled": True, 
-                                      "policy_text": policy_text, 
-                                      "need_explanation": need_explanation
-                                      }
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
+        config["policy_violation"] = {
+            "enabled": True,
+            "policy_text": policy_text,
+            "need_explanation": need_explanation,
+        }
         return cls(config)
 
     @classmethod
@@ -54,7 +59,7 @@ class GuardrailsConfig:
         """
         Returns a configuration instance pre-configured for toxicity detection.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["toxicity"] = {"enabled": True}
         return cls(config)
 
@@ -63,7 +68,7 @@ class GuardrailsConfig:
         """
         Returns a configuration instance pre-configured for NSFW content detection.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["nsfw"] = {"enabled": True}
         return cls(config)
 
@@ -72,7 +77,7 @@ class GuardrailsConfig:
         """
         Returns a configuration instance pre-configured for bias detection.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["bias"] = {"enabled": True}
         return cls(config)
 
@@ -80,14 +85,14 @@ class GuardrailsConfig:
     def pii(cls, entities=None):
         """
         Returns a configuration instance pre-configured for PII detection.
-        
+
         Args:
             entities (list, optional): List of PII entity types to detect.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["pii"] = {
             "enabled": True,
-            "entities": entities if entities is not None else []
+            "entities": entities if entities is not None else [],
         }
         return cls(config)
 
@@ -95,14 +100,14 @@ class GuardrailsConfig:
     def topic(cls, topics=None):
         """
         Returns a configuration instance pre-configured for topic detection.
-        
+
         Args:
             topics (list, optional): List of topics to detect.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["topic_detector"] = {
             "enabled": True,
-            "topic": topics if topics is not None else []
+            "topic": topics if topics is not None else [],
         }
         return cls(config)
 
@@ -110,14 +115,14 @@ class GuardrailsConfig:
     def keyword(cls, keywords=None):
         """
         Returns a configuration instance pre-configured for keyword detection.
-        
+
         Args:
             keywords (list, optional): List of banned keywords to detect.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["keyword_detector"] = {
             "enabled": True,
-            "banned_keywords": keywords if keywords is not None else []
+            "banned_keywords": keywords if keywords is not None else [],
         }
         return cls(config)
 
@@ -126,7 +131,7 @@ class GuardrailsConfig:
         """
         Returns a configuration instance pre-configured for copyright/IP detection.
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
         config["copyright_ip"] = {"enabled": True}
         return cls(config)
 
@@ -134,21 +139,18 @@ class GuardrailsConfig:
     def system_prompt(cls, index="system"):
         """
         Returns a configuration instance pre-configured for system prompt detection.
-        
+
         Args:
             index (str, optional): Index name for system prompt detection. Defaults to "system".
         """
-        config = copy.deepcopy(DEFAULT_CONFIG)
-        config["system_prompt"] = {
-            "enabled": True,
-            "index": index
-        }
+        config = copy.deepcopy(DEFAULT_GUARDRAILS_CONFIG)
+        config["system_prompt"] = {"enabled": True, "index": index}
         return cls(config)
 
     def update(self, **kwargs):
         """
         Update the configuration with custom values.
-        
+
         Only keys that exist in the default configuration can be updated.
         For example:
             config.update(nsfw={"enabled": True}, toxicity={"enabled": True})
@@ -170,16 +172,16 @@ class GuardrailsConfig:
     def from_custom_config(cls, config_dict: dict):
         """
         Configure guardrails from a dictionary input.
-        
+
         Validates that the input dictionary matches the expected schema structure.
         Each key must exist in the default configuration, and its value must be a dictionary.
-        
+
         Args:
             config_dict (dict): Dictionary containing guardrails configuration
-            
+
         Returns:
             GuardrailsConfig: Returns a new GuardrailsConfig instance
-            
+
         Raises:
             ValueError: If the input dictionary contains invalid keys or malformed values
         """
@@ -189,33 +191,92 @@ class GuardrailsConfig:
                 raise ValueError(f"Unknown detector config: {key}")
             if not isinstance(value, dict):
                 raise ValueError(f"Config value for {key} must be a dictionary")
-            
+
             # Validate that all required fields exist in the default config
-            default_fields = set(DEFAULT_CONFIG[key].keys())
+            default_fields = set(DEFAULT_GUARDRAILS_CONFIG[key].keys())
             provided_fields = set(value.keys())
-            
+
             if not provided_fields.issubset(default_fields):
                 invalid_fields = provided_fields - default_fields
                 raise ValueError(f"Invalid fields for {key}: {invalid_fields}")
-                
+
             instance.config[key] = value
-            
+
         return instance
 
     def get_config(self, detector_name: str) -> dict:
         """
         Get the configuration for a specific detector.
-        
+
         Args:
             detector_name (str): Name of the detector to get configuration for
-            
+
         Returns:
             dict: Configuration dictionary for the specified detector
-            
+
         Raises:
             ValueError: If the detector name doesn't exist in the configuration
         """
         if detector_name not in self.config:
             raise ValueError(f"Unknown detector: {detector_name}")
-            
+
         return copy.deepcopy(self.config[detector_name])
+
+
+class RedTeamConfig:
+    """
+    A helper class to manage RedTeam configuration.
+    """
+
+    def __init__(self, config=None):
+        if config is None:
+            config = copy.deepcopy(DEFAULT_REDTEAM_CONFIG)
+            # Only include advanced tests if dataset is not standard
+            if config.get("dataset_name") != "standard":
+                config["redteam_test_configurations"].update(
+                    copy.deepcopy(ADVANCED_REDTEAM_TESTS)
+                )
+        self.config = config
+
+    def as_dict(self):
+        """
+        Return the underlying configuration dictionary.
+        """
+        return self.config
+
+
+class ModelConfig:
+    def __init__(self, config=None):
+        if config is None:
+            config = copy.deepcopy(DETAIL_MODEL_CONFIG)
+        self.config = config
+
+    @classmethod
+    def model_name(self, model_name: str):
+        """
+        Set the model name.
+        """
+        self.config["model_name"] = model_name
+        return self
+
+    @classmethod
+    def testing_for(self, testing_for: str):
+        """
+        Set the testing for.
+        """
+        self.config["testing_for"] = testing_for
+        return self
+
+    @classmethod
+    def model_config(self, model_config: dict):
+        """
+        Set the model config.
+        """
+        self.config["model_config"] = model_config
+        return self
+
+    def as_dict(self):
+        """
+        Return the underlying configuration dictionary.
+        """
+        return self.config

enkryptai_sdk/dto/__init__.py

@@ -0,0 +1,18 @@
+from .models import *
+from .red_team import *
+
+__all__ = [
+    "DetailModelConfig",
+    "ModelConfig",
+    "RedTeamConfig",
+    "AdvancedRedTeamTests",
+    "TestConfig",
+    "AttackMethods",
+    "RedTeamTestConfigurations",
+    "TargetModelConfiguration",
+    "Location",
+    "Metadata",
+    "DEFAULT_REDTEAM_CONFIG",
+    "ADVANCED_REDTEAM_TESTS",
+    "DETAIL_MODEL_CONFIG",
+]

enkryptai_sdk/dto/models.py

@@ -0,0 +1,202 @@
+from dataclasses import dataclass, field, asdict
+from typing import Optional, List, Set, Dict, Any
+from enum import Enum
+import json
+
+
+class Modality(Enum):
+    TEXT = "text"
+    IMAGE = "image"
+    AUDIO = "audio"
+    VIDEO = "video"
+
+    def to_dict(self):
+        return self.value
+
+
+@dataclass
+class EndpointConfig:
+    scheme: str = "https"
+    host: str = "api.openai.com"
+    port: int = 443
+    base_path: str = "v1"
+
+
+@dataclass
+class PathsConfig:
+    completions: str = "/chat/completions"
+    chat: str = "chat/completions"
+
+
+@dataclass
+class AuthData:
+    header_name: str = "Authorization"
+    header_prefix: str = "Bearer"
+    space_after_prefix: bool = True
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class ModelDetailConfig:
+    model_version: Optional[str] = None
+    model_source: str = ""
+    model_provider: str = "openai"
+    system_prompt: str = ""
+
+    endpoint_url: str = "https://api.openai.com/v1/chat/completions"
+    auth_data: AuthData = field(default_factory=AuthData)
+    api_keys: Set[Optional[str]] = field(default_factory=lambda: {None})
+
+
+@dataclass
+class DetailModelConfig:
+    model_saved_name: str = "Model Name"
+    testing_for: str = "LLM"
+    model_name: str = "gpt-4o-mini"
+    modality: Modality = Modality.TEXT
+    model_config: ModelDetailConfig = field(default_factory=ModelDetailConfig)
+
+
+@dataclass
+class ModelConfigDetails:
+    model_version: Optional[str] = None
+    model_source: str = ""
+    model_provider: str = "openai"
+    system_prompt: str = ""
+    conversation_template: str = ""
+    is_compatible_with: str = "openai"
+    hosting_type: str = "External"
+    endpoint_url: str = "https://api.openai.com/v1/chat/completions"
+    auth_data: AuthData = field(default_factory=AuthData)
+    apikey: Optional[str] = None
+    default_request_options: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        # Create a copy of the data to avoid modifying the original
+        data = data.copy()
+
+        # Remove known fields that we don't want in our model
+        unwanted_fields = ["queryParams", "paths"]
+        for field in unwanted_fields:
+            data.pop(field, None)
+
+        # Handle apikeys to apikey conversion
+        if "apikeys" in data:
+            apikeys = data.pop("apikeys")
+            if apikeys and not data.get("apikey"):
+                data["apikey"] = apikeys[0]
+
+        # Convert endpoint dict to endpoint_url if present
+        if "endpoint" in data:
+            endpoint = data.pop("endpoint")
+            scheme = endpoint.get("scheme", "https")
+            host = endpoint.get("host", "")
+            port = endpoint.get("port", "")
+            base_path = endpoint.get("base_path", "")
+
+            endpoint_url = f"{scheme}://{host}"
+            if port and port not in [80, 443]:
+                endpoint_url += f":{port}"
+            if base_path:
+                base_path = "/" + base_path.strip("/")
+                endpoint_url += base_path
+
+            data["endpoint_url"] = endpoint_url
+
+        # Handle nested AuthData
+        auth_data = data.pop("auth_data", {})
+        auth_data_obj = AuthData.from_dict(auth_data)
+
+        # Only keep fields that are defined in the dataclass
+        valid_fields = cls.__dataclass_fields__.keys()
+        filtered_data = {k: v for k, v in data.items() if k in valid_fields}
+
+        return cls(**filtered_data, auth_data=auth_data_obj)
+
+    def to_dict(self):
+        d = asdict(self)
+        # Handle AuthData specifically
+        d["auth_data"] = self.auth_data.to_dict()
+        return d
+
+    def to_json(self):
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str):
+        return cls.from_dict(json.loads(json_str))
+
+
+@dataclass
+class ModelConfig:
+    created_at: str = ""
+    updated_at: str = ""
+    model_id: str = ""
+    model_saved_name: str = "Model Name"
+    testing_for: str = "LLM"
+    model_name: str = "gpt-4o-mini"
+    model_type: str = "text_2_text"
+    modality: Modality = Modality.TEXT
+    certifications: List[str] = field(default_factory=list)
+    model_config: ModelConfigDetails = field(default_factory=ModelConfigDetails)
+
+    def to_dict(self) -> dict:
+        """Convert the ModelConfig instance to a dictionary."""
+        # First create a shallow copy of self as dict
+        d = {}
+        for field in self.__dataclass_fields__:
+            value = getattr(self, field)
+            if field == "modality":
+                d[field] = value.value
+            elif field == "model_config":
+                if isinstance(value, ModelConfigDetails):
+                    d[field] = value.to_dict()
+                else:
+                    d[field] = value
+            else:
+                d[field] = value
+        return d
+
+    def to_json(self) -> str:
+        """Convert the ModelConfig instance to a JSON string."""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        """Create a ModelConfig instance from a dictionary."""
+        # Handle nested ModelConfigDetails
+        model_config_data = data.pop("model_config", {})
+        model_config = ModelConfigDetails.from_dict(model_config_data)
+
+        # Handle Modality enum
+        modality_value = data.pop("modality", "text")
+        modality = Modality(modality_value)
+
+        return cls(**data, modality=modality, model_config=model_config)
+
+    @classmethod
+    def from_json(cls, json_str: str):
+        """Create a ModelConfig instance from a JSON string."""
+        data = json.loads(json_str)
+        return cls.from_dict(data)
+
+    def __str__(self):
+        """String representation of the ModelConfig."""
+        return f"ModelConfig(name={self.model_saved_name}, model={self.model_name})"
+
+    def __repr__(self):
+        """Detailed string representation of the ModelConfig."""
+        return (
+            f"ModelConfig({', '.join(f'{k}={v!r}' for k, v in self.to_dict().items())})"
+        )
+
+
+# Default configuration
+DETAIL_MODEL_CONFIG = ModelConfig()