enkryptai-sdk 0.1.4__tar.gz → 0.1.6__tar.gz

{enkryptai_sdk-0.1.4/src/enkryptai_sdk.egg-info → enkryptai_sdk-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: enkryptai-sdk
-Version: 0.1.4
+Version: 0.1.6
 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-0.1.4 → enkryptai_sdk-0.1.6}/setup.py

@@ -8,7 +8,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
-    version="0.1.4",
+    version="0.1.6",
     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-0.1.6/src/enkryptai_sdk/__init__.py

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

{enkryptai_sdk-0.1.4 → enkryptai_sdk-0.1.6}/src/enkryptai_sdk/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-0.1.6/src/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-0.1.6/src/enkryptai_sdk/dto/models.py

@@ -0,0 +1,215 @@
+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 ModelResponse:
+    message: Optional[str] = None
+    data: Optional[Dict] = None
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@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()

enkryptai_sdk-0.1.6/src/enkryptai_sdk/dto/red_team.py

@@ -0,0 +1,196 @@
+from dataclasses import dataclass, field, asdict
+from typing import Dict, List, Optional
+import json
+
+
+@dataclass
+class RedTeamResponse:
+    task_id: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class RedTeamTaskStatus:
+    status: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class RedTeamTaskDetails:
+    created_at: Optional[str] = None
+    model_name: Optional[str] = None
+    status: Optional[str] = None
+    task_id: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class RedTeamResultSummary:
+    test_date: Optional[str] = None
+    test_name: Optional[str] = None
+    dataset_name: Optional[str] = None
+    model_name: Optional[str] = None
+    model_endpoint_url: Optional[str] = None
+    model_source: Optional[str] = None
+    model_provider: Optional[str] = None
+    risk_score: Optional[float] = None
+    test_type: Optional[List] = None
+    nist_category: Optional[List] = None
+    scenario: Optional[List] = None
+    category: Optional[List] = None
+    attack_method: Optional[List] = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class RedTeamResultDetails:  # To Be Updated
+    details: Optional[Dict] = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class AttackMethods:
+    basic: List[str] = field(default_factory=lambda: ["basic"])
+    advanced: Dict[str, List[str]] = field(
+        default_factory=lambda: {"static": ["single_shot"], "dynamic": ["iterative"]}
+    )
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class TestConfig:
+    sample_percentage: int = 100
+    attack_methods: AttackMethods = field(default_factory=AttackMethods)
+
+    def to_dict(self) -> dict:
+        return {
+            "sample_percentage": self.sample_percentage,
+            "attack_methods": self.attack_methods.to_dict(),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        attack_methods = AttackMethods.from_dict(data.pop("attack_methods", {}))
+        return cls(**data, attack_methods=attack_methods)
+
+
+@dataclass
+class RedTeamTestConfigurations:
+    # Basic tests
+    bias_test: TestConfig = field(default=None)
+    cbrn_test: TestConfig = field(default=None)
+    insecure_code_test: TestConfig = field(default=None)
+    toxicity_test: TestConfig = field(default=None)
+    harmful_test: TestConfig = field(default=None)
+    # Advanced tests
+    adv_info_test: TestConfig = field(default=None)
+    adv_bias_test: TestConfig = field(default=None)
+    adv_command_test: TestConfig = field(default=None)
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**{k: TestConfig.from_dict(v) for k, v in data.items()})
+
+
+@dataclass
+class TargetModelConfiguration:
+    testing_for: str = "LLM"
+    model_name: str = "gpt-4o-mini"
+    model_version: Optional[str] = None
+    system_prompt: str = ""
+    conversation_template: str = ""
+    model_source: str = ""
+    model_provider: str = "openai"
+    model_endpoint_url: str = "https://api.openai.com/v1/chat/completions"
+    model_api_key: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
+@dataclass
+class RedTeamConfig:
+    test_name: str = "Test Name"
+    dataset_name: str = "standard"
+    model_name: str = "gpt-4o-mini"
+    redteam_test_configurations: RedTeamTestConfigurations = field(
+        default_factory=RedTeamTestConfigurations
+    )
+    target_model_configuration: TargetModelConfiguration = field(
+        default_factory=TargetModelConfiguration
+    )
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["redteam_test_configurations"] = self.redteam_test_configurations.to_dict()
+        d["target_model_configuration"] = self.target_model_configuration.to_dict()
+        return d
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        data = data.copy()
+        test_configs = RedTeamTestConfigurations.from_dict(
+            data.pop("redteam_test_configurations", {})
+        )
+        target_config = TargetModelConfiguration.from_dict(
+            data.pop("target_model_configuration", {})
+        )
+        return cls(
+            **data,
+            redteam_test_configurations=test_configs,
+            target_model_configuration=target_config,
+        )
+
+    @classmethod
+    def from_json(cls, json_str: str):
+        return cls.from_dict(json.loads(json_str))
+
+
+# Default configurations
+DEFAULT_REDTEAM_CONFIG = RedTeamConfig()

enkryptai_sdk-0.1.6/src/enkryptai_sdk/models.py

@@ -0,0 +1,160 @@
+import urllib3
+from .dto import ModelConfig, ModelResponse
+from urllib.parse import urlparse, urlsplit
+
+
+class ModelClientError(Exception):
+    """
+    A custom exception for ModelClient errors.
+    """
+
+    pass
+
+
+class ModelClient:
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com:443"):
+        self.api_key = api_key
+        self.base_url = base_url
+        self.http = urllib3.PoolManager()
+        self.headers = {"apikey": self.api_key}
+
+    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 health(self):
+        return self._request("GET", "/models/health")
+
+    def add_model(self, config: ModelConfig):
+        """
+        Add a new model configuration to the system.
+
+        Args:
+            config (ModelConfig): Configuration object containing model details
+
+        Returns:
+            dict: Response from the API containing the added model details
+        """
+        headers = {"Content-Type": "application/json"}
+        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 = ""
+
+        # Determine paths based on the endpoint
+        paths = {
+            "completions": f"/{remaining_path}" if remaining_path else "",
+            "chat": "",
+        }
+
+        payload = {
+            "model_saved_name": config.model_saved_name,
+            "testing_for": config.testing_for,
+            "model_name": config.model_name,
+            "model_type": config.model_type,
+            "certifications": config.certifications,
+            "model_config": {
+                "is_compatible_with": config.model_config.is_compatible_with,
+                "model_version": config.model_config.model_version,
+                "hosting_type": config.model_config.hosting_type,
+                "model_source": config.model_config.model_source,
+                "model_provider": config.model_config.model_provider,
+                "system_prompt": config.model_config.system_prompt,
+                "conversation_template": config.model_config.conversation_template,
+                "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}/{paths['completions']}",  # 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 []
+                ),
+                "default_request_options": config.model_config.default_request_options,
+            },
+        }
+        try:
+            response = self._request(
+                "POST", "/models/add-model", headers=headers, json=payload
+            )
+            if response.get("error"):
+                raise ModelClientError(response["error"])
+            return ModelResponse.from_dict(response)
+        except Exception as e:
+            raise ModelClientError(str(e))
+
+    def get_model(self, model_id: str) -> ModelConfig:
+        """
+        Get model configuration by model ID.
+
+        Args:
+            model_id (str): ID of the model to retrieve
+
+        Returns:
+            ModelConfig: Configuration object containing model details
+        """
+        headers = {"X-Enkrypt-Model": model_id}
+        response = self._request("GET", "/models/get-model", headers=headers)
+        if response.get("error"):
+            raise ModelClientError(response["error"])
+        return ModelConfig.from_dict(response)
+
+    def get_model_list(self):
+        """
+        Get a list of all available models.
+
+        Returns:
+            dict: Response from the API containing the list of models
+        """
+        try:
+            return self._request("GET", "/models/list-models")
+        except Exception as e:
+            return {"error": str(e)}
+
+    def delete_model(self, model_id: str):
+        """
+        Delete a specific model from the system.
+
+        Args:
+            model_id (str): The identifier or name of the model to delete
+
+        Returns:
+            dict: Response from the API containing the deletion status
+        """
+        headers = {"X-Enkrypt-Model": model_id}
+        return self._request("DELETE", "/models/delete-model", headers=headers)

enkryptai_sdk-0.1.6/src/enkryptai_sdk/red_team.py

@@ -0,0 +1,195 @@
+import urllib3
+from .dto import (
+    RedTeamConfig,
+    RedTeamResponse,
+    RedTeamResultSummary,
+    RedTeamResultDetails,
+    RedTeamTaskStatus,
+    RedTeamTaskDetails,
+)
+
+
+class RedTeamClientError(Exception):
+    """
+    A custom exception for Red Team errors.
+    """
+
+    pass
+
+
+class RedTeamClient:
+    """
+    A client for interacting with the Red Team API.
+    """
+
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
+        self.api_key = api_key
+        self.base_url = base_url
+        self.http = urllib3.PoolManager()
+        self.headers = {"apikey": self.api_key}
+
+    def _request(self, method, endpoint, 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 get_model(self, model):
+        models = self._request("GET", "/models/list-models")
+        if model in models:
+            return model
+        else:
+            return None
+
+    def add_task(
+        self,
+        config: RedTeamConfig,
+    ):
+        """
+        Add a new red teaming task.
+        """
+        config = RedTeamConfig.from_dict(config)
+        test_configs = config.redteam_test_configurations.to_dict()
+        # Remove None or empty test configurations
+        test_configs = {k: v for k, v in test_configs.items() if v is not None}
+
+        payload = {
+            # "async": config.async_enabled,
+            "dataset_name": config.dataset_name,
+            "test_name": config.test_name,
+            "redteam_test_configurations": test_configs,
+        }
+
+        model = config.model_name
+        saved_model = self.get_model(model)
+
+        if saved_model:
+            print("saved model found")
+            headers = {
+                "X-Enkrypt-Model": saved_model,
+                "Content-Type": "application/json",
+            }
+            payload["location"] = {"storage": "supabase", "container_name": "supabase"}
+            return self._request(
+                "POST",
+                "/redteam/v2/model/add-task",
+                headers=headers,
+                json=payload,
+            )
+        elif config.target_model_configuration:
+            payload["target_model_configuration"] = (
+                config.target_model_configuration.to_dict()
+            )
+            # print(payload)
+            response = self._request(
+                "POST",
+                "/redteam/v2/add-task",
+                json=payload,
+            )
+            if response.get("error"):
+                raise RedTeamClientError(response["error"])
+            return RedTeamResponse.from_dict(response)
+        else:
+            raise RedTeamClientError(
+                "Please use a saved model or provide a target model configuration"
+            )
+
+    def status(self, task_id: str):
+        """
+        Get the status of a specific red teaming task.
+
+        Args:
+            task_id (str): The ID of the task to check status
+
+        Returns:
+            dict: The task status information
+        """
+        headers = {"X-Enkrypt-Task-ID": task_id}
+
+        response = self._request("GET", "/redteam/task-status", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        return RedTeamTaskStatus.from_dict(response)
+
+    def cancel_task(self, task_id: str):
+        """
+        Cancel a specific red teaming task.
+
+        Args:
+            task_id (str): The ID of the task to cancel
+        """
+        raise RedTeamClientError(
+            "This feature is currently under development. Please check our documentation "
+            "at https://docs.enkrypt.ai for updates or contact support@enkrypt.ai for assistance."
+        )
+
+    def get_task(self, task_id: str):
+        """
+        Get the status and details of a specific red teaming task.
+
+        Args:
+            task_id (str): The ID of the task to retrieve
+
+        Returns:
+            dict: The task details and status
+        """
+        headers = {"X-Enkrypt-Task-ID": task_id}
+
+        response = self._request("GET", "/redteam/get-task", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        if response.get("data").get("job_id "):
+            response["data"]["task_id"] = response["data"].pop("job_id")
+        return RedTeamTaskDetails.from_dict(response["data"])
+
+    def get_result_summary(self, task_id: str):
+        """
+        Get the summary of results for a specific red teaming task.
+
+        Args:
+            task_id (str): The ID of the task to get results for
+
+        Returns:
+            dict: The summary of the task results
+        """
+        headers = {"X-Enkrypt-Task-ID": task_id}
+
+        response = self._request("GET", "/redteam/results/summary", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        return RedTeamResultSummary.from_dict(response["summary"])
+
+    def get_result_details(self, task_id: str):
+        """
+        Get the detailed results for a specific red teaming task.
+
+        Args:
+            task_id (str): The ID of the task to get detailed results for
+
+        Returns:
+            dict: The detailed task results
+        """
+        # TODO: Update the response to be updated
+        headers = {"X-Enkrypt-Task-ID": task_id}
+        response = self._request("GET", "/redteam/results/details", headers=headers)
+        if response.get("error"):
+            raise RedTeamClientError(response["error"])
+        return RedTeamResultDetails.from_dict(response["details"])

{enkryptai_sdk-0.1.4 → enkryptai_sdk-0.1.6/src/enkryptai_sdk.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: enkryptai-sdk
-Version: 0.1.4
+Version: 0.1.6
 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-0.1.4 → enkryptai_sdk-0.1.6}/src/enkryptai_sdk.egg-info/SOURCES.txt

@@ -12,6 +12,9 @@ src/enkryptai_sdk.egg-info/PKG-INFO
 src/enkryptai_sdk.egg-info/SOURCES.txt
 src/enkryptai_sdk.egg-info/dependency_links.txt
 src/enkryptai_sdk.egg-info/top_level.txt
+src/enkryptai_sdk/dto/__init__.py
+src/enkryptai_sdk/dto/models.py
+src/enkryptai_sdk/dto/red_team.py
 tests/test_all.py
 tests/test_basic.py
 tests/test_detect_policy.py

enkryptai_sdk-0.1.4/src/enkryptai_sdk/__init__.py

@@ -1,5 +0,0 @@
-from .guardrails import GuardrailsClient
-from .config import GuardrailsConfig
-from .evals import EvalsClient
-
-__all__ = ["GuardrailsClient", "GuardrailsConfig", "EvalsClient"]