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

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/evals.py

@@ -0,0 +1,84 @@
+import requests
+
+class EvalsClient:
+    """
+    A client for interacting with Enkrypt AI Evals API endpoints.
+    """
+
+    def __init__(self, api_key, base_url="https://api.enkryptai.com"):
+        """
+        Initializes the client.
+
+        Parameters:
+        - api_key (str): Your API key for authenticating with the service.
+        - base_url (str): Base URL of the API (default: "https://api.enkryptai.com").
+        """
+        self.api_key = api_key
+        self.base_url = base_url.rstrip('/')
+        self.session = requests.Session()
+
+    def _request(self, method, endpoint, headers=None, **kwargs):
+        """
+        Internal helper to send an HTTP request.
+
+        Automatically adds the API key to headers.
+        """
+        url = self.base_url + endpoint
+        headers = headers or {}
+        if 'apikey' not in headers:
+            headers['apikey'] = self.api_key
+            
+        try:
+            response = self.session.request(method, url, headers=headers, **kwargs)
+            response.raise_for_status()
+            return response.json()
+        
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}
+
+    # ----------------------------
+    # Basic Evals Endpoints
+    # ----------------------------
+
+    def check_adherence(self, llm_answer, context):
+        """
+        Checks if the LLM's answer adheres to the provided context.
+
+        Parameters:
+        - llm_answer (str): The response generated by the LLM
+        - context (str): The context against which to check the answer
+
+        Returns:
+        - JSON response from the API containing adherence analysis
+        """
+        payload = {
+            "llm_answer": llm_answer,
+            "context": context
+        }
+        try:
+            return self._request("POST", "/guardrails/adherence", json=payload)
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}
+
+    def check_relevancy(self, question, llm_answer):
+        """
+        Checks if the LLM's answer is relevant to the asked question.
+
+        Parameters:
+        - question (str): The original question asked
+        - llm_answer (str): The response generated by the LLM
+
+        Returns:
+        - JSON response from the API containing relevancy analysis
+        """
+        payload = {
+            "question": question,
+            "llm_answer": llm_answer
+        }
+        try:
+            return self._request("POST", "/guardrails/relevancy", json=payload)
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}

enkryptai_sdk/guardrails.py

@@ -1,4 +1,6 @@
 import requests
+from .config import GuardrailsConfig
+from .response import GuardrailsResponse, PIIResponse
 
 class GuardrailsClient:
     """
@@ -74,7 +76,6 @@ class GuardrailsClient:
         """
         # Use injection attack config by default if none provided
         if config is None:
-            from .guardrails_config import GuardrailsConfig
             config = GuardrailsConfig.injection_attack()
 
         # Allow passing in either a dict or a GuardrailsConfig instance.
@@ -85,7 +86,8 @@ class GuardrailsClient:
             "text": text,
             "detectors": config
         }
-        return self._request("POST", "/guardrails/detect", json=payload)
+        response_body = self._request("POST", "/guardrails/detect", json=payload)
+        return GuardrailsResponse(response_body)
 
     def pii(self, text, mode, key="null", entities=None):
         """
@@ -97,7 +99,8 @@ class GuardrailsClient:
             "key": key,
             "entities": entities
         }
-        return self._request("POST", "/guardrails/pii", json=payload)
+        response_body = self._request("POST", "/guardrails/pii", json=payload)
+        return PIIResponse(response_body)
 
     # ----------------------------
     # Guardrails Policy Endpoints
@@ -181,8 +184,12 @@ class GuardrailsClient:
         """
         headers = {"X-Enkrypt-Policy": policy_name}
         payload = {"text": text}
+        
         try:
-            return self._request("POST", "/guardrails/policy/detect", headers=headers, json=payload)
+            
+            response_body = self._request("POST", "/guardrails/policy/detect", headers=headers, json=payload)
+            return GuardrailsResponse(response_body)
+        
         except Exception as e:
             print(e)
             return {"error": str(e)}

enkryptai_sdk/models.py

@@ -0,0 +1,144 @@
+import urllib3
+from .dto import ModelConfig, ModelDetailConfig
+from urllib.parse import urlparse, urlsplit
+
+
+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_response = (
+                    response.json()
+                    if response.data
+                    else {"message": f"HTTP {response.status}"}
+                )
+                raise urllib3.exceptions.HTTPError(
+                    f"HTTP {response.status}: {error_response}"
+                )
+            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,
+            },
+        }
+        print(payload)
+        return self._request("POST", "/models/add-model", headers=headers, json=payload)
+
+    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)
+        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/red_team.py

@@ -0,0 +1,185 @@
+import urllib3
+from .dto import (
+    RedTeamConfig,
+    RedTeamResponse,
+    RedTeamResultSummary,
+    RedTeamResultDetails,
+    RedTeamTaskStatus,
+    RedTeamTaskDetails,
+)
+
+
+class RedTeamError(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_response = (
+                    response.json()
+                    if response.data
+                    else {"message": f"HTTP {response.status}"}
+                )
+                raise urllib3.exceptions.HTTPError(
+                    f"HTTP {response.status}: {error_response}"
+                )
+            return response.json()
+        except urllib3.exceptions.HTTPError as e:
+            print(f"Request failed: {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,
+            )
+            return RedTeamResponse.from_dict(response)
+        else:
+            raise RedTeamError(
+                "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)
+        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 RedTeamError(
+            "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)
+        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)
+        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)
+        return RedTeamResultDetails.from_dict(response["details"])