enkryptai-sdk 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

enkryptai_sdk/__init__.py

@@ -1,13 +1,25 @@
-from .guardrails import GuardrailsClient
-from .config import GuardrailsConfig
 from .evals import EvalsClient
-from .models import ModelClient
-from .red_team import RedTeamClient
+from .config import GuardrailsConfig
+from .guardrails import GuardrailsClient, GuardrailsClientError
+from .models import ModelClient, ModelClientError
+from .red_team import RedTeamClient, RedTeamClientError
+from .datasets import DatasetClient, DatasetClientError
+from .deployments import DeploymentClient, DeploymentClientError
+from .ai_proxy import AIProxyClient, AIProxyClientError
 
 __all__ = [
     "GuardrailsClient",
+    "GuardrailsClientError",
     "GuardrailsConfig",
     "EvalsClient",
     "ModelClient",
     "RedTeamClient",
+    "DatasetClient",
+    "DeploymentClient",
+    "ModelClientError",
+    "RedTeamClientError",
+    "DatasetClientError",
+    "DeploymentClientError",
+    "AIProxyClient",
+    "AIProxyClientError",
 ]

enkryptai_sdk/ai_proxy.py

@@ -0,0 +1,70 @@
+from .base import BaseClient
+from .dto import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionErrorResponse,
+    ChatCompletionDirectErrorResponse,
+)
+
+
+class AIProxyClientError(Exception):
+    pass
+
+
+class AIProxyClient(BaseClient):
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
+        super().__init__(api_key, base_url)
+
+    def chat(self, chat_body: ChatCompletionRequest, deployment_name: str, tags: str = None, refresh_cache: bool = False, return_error: bool = False):
+        """
+        Get chat completion response for a given prompt.
+
+        Args:
+            chat_body (ChatCompletionRequest): Configuration object containing chat details
+
+        Returns:
+            dict: Response from the API containing the chat completion response
+        """
+        headers = {"Content-Type": "application/json"}
+
+        if deployment_name is None:
+            raise AIProxyClientError("Deployment name is required")
+        headers["X-Enkrypt-Deployment"] = deployment_name
+
+        if tags is not None:
+            headers["X-Enkrypt-Tags"] = tags
+
+        if refresh_cache is not None:
+            headers["X-Enkrypt-Refresh-Cache"] = str(refresh_cache).lower()
+
+        if isinstance(chat_body, dict):
+            chat_body = ChatCompletionRequest.from_dict(chat_body)
+
+        payload = chat_body.to_dict()
+
+        response = self._request(
+            "POST", "/ai-proxy/chat/completions", headers=headers, json=payload
+        )
+
+        if response.get("error"):
+            error_message = response["error"]
+            is_json_str = error_message.startswith("{") and error_message.endswith("}")
+            # Try to parse nested JSON error if it's a string
+            if isinstance(error_message, str) and is_json_str:
+                import json
+                try:
+                    error_message = json.loads(error_message.replace("'", '"'))
+                    response["error"] = error_message.get("error", error_message)
+                except json.JSONDecodeError:
+                    # If parsing fails, keep the original error
+                    pass
+                    
+            # print("Error in response: ", response)
+            if return_error:
+                if is_json_str:
+                    return ChatCompletionErrorResponse.from_dict(response)
+                else:
+                    return ChatCompletionDirectErrorResponse.from_dict(response)
+            raise AIProxyClientError(response["error"])
+
+        return ChatCompletionResponse.from_dict(response)

enkryptai_sdk/base.py

@@ -0,0 +1,36 @@
+import urllib3
+from typing import Dict, Any
+
+
+class BaseClient:
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
+        if api_key is None:
+            raise ValueError("API key is required")
+        self.api_key = api_key
+        self.base_url = base_url.rstrip('/')
+        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)}

enkryptai_sdk/datasets.py

@@ -0,0 +1,142 @@
+from .base import BaseClient
+from .dto import (
+    DatasetConfig,
+    DatasetCollection,
+    DatasetSummary,
+    DatasetResponse,
+    DatasetCard,
+    DatasetTaskStatus,
+    DatasetTask,
+    DatasetAddTaskResponse,
+)
+
+
+class DatasetClientError(Exception):
+    pass
+
+
+class DatasetClient(BaseClient):
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
+        super().__init__(api_key, base_url)
+
+    def add_dataset(self, config: DatasetConfig):
+        """
+        Add a new dataset to the system.
+
+        Args:
+            config (DatasetConfig): Configuration object containing dataset details
+
+        Returns:
+            dict: Response from the API containing the added dataset details
+        """
+        headers = {"Content-Type": "application/json"}
+
+        if isinstance(config, dict):
+            config = DatasetConfig.from_dict(config)
+
+        payload = config.to_dict()
+
+        response = self._request(
+            "POST", "/datasets/add-task", headers=headers, json=payload
+        )
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        return DatasetAddTaskResponse.from_dict(response)
+    
+    def get_dataset_task_status(self, dataset_name: str):
+        """
+        Get dataset task status for a specific dataset task.
+
+        Args:
+            dataset_name (str): The name of the dataset
+
+        Returns:
+            dict: Response from the API containing the dataset task status
+        """
+        headers = {"X-Enkrypt-Dataset": dataset_name}
+        response = self._request("GET", "/datasets/task-status", headers=headers)
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        response["dataset_name"] = dataset_name
+        return DatasetTaskStatus.from_dict(response)
+
+    def get_dataset_task(self, dataset_name: str):
+        """
+        Get detailed information for a specific dataset task.
+
+        Args:
+            dataset_name (str): The name of the dataset
+
+        Returns:
+            dict: Response from the API containing the dataset task information
+        """
+        headers = {"X-Enkrypt-Dataset": dataset_name}
+        response = self._request("GET", "/datasets/get-task", headers=headers)
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        response["dataset_name"] = dataset_name
+        return DatasetTask.from_dict(response)
+
+    def get_datacard(self, dataset_name: str):
+        """
+        Get datacard information for a specific dataset.
+
+        Args:
+            dataset_name (str): The name of the dataset
+
+        Returns:
+            dict: Response from the API containing the datacard information
+        """
+        headers = {"X-Enkrypt-Dataset": dataset_name}
+        response = self._request("GET", "/datasets/get-datacard", headers=headers)
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        response["dataset_name"] = dataset_name
+        return DatasetCard.from_dict(response)
+
+    def get_dataset(self, dataset_name: str):
+        """
+        Get detailed information for a specific dataset.
+
+        Args:
+            dataset_name (str): The name of the dataset
+
+        Returns:
+            dict: Response from the API containing the dataset information
+        """
+        headers = {"X-Enkrypt-Dataset": dataset_name}
+        response = self._request("GET", "/datasets/get-dataset", headers=headers)
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        response["dataset_name"] = dataset_name
+        return DatasetResponse.from_dict(response)
+
+    def get_summary(self, dataset_name: str):
+        """
+        Get summary information for a specific dataset.
+
+        Args:
+            dataset_name (str): The name of the dataset
+
+        Returns:
+            dict: Response from the API containing the dataset summary
+        """
+        headers = {"X-Enkrypt-Dataset": dataset_name}
+        response = self._request("GET", "/datasets/get-summary", headers=headers)
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        response["dataset_name"] = dataset_name
+        return DatasetSummary.from_dict(response)
+
+    def list_datasets(self):
+        """
+        Get a list of all available dataset tasks.
+
+        Returns:
+            dict: Response from the API containing the list of dataset tasks
+        """
+
+        response = self._request("GET", "/datasets/list-tasks")
+        if response.get("error"):
+            raise DatasetClientError(response["error"])
+        return DatasetCollection.from_dict(response)

enkryptai_sdk/deployments.py

@@ -0,0 +1,121 @@
+from .base import BaseClient
+from .dto import (
+    # InputGuardrailBlockType,
+    # OutputGuardrailBlockType,
+    # InputGuardrailsAdditionalConfig,
+    # OutputGuardrailsAdditionalConfig,
+    # InputGuardrailsPolicy,
+    # OutputGuardrailsPolicy,
+    DeploymentInput,
+    DeploymentAddTaskResponse,
+    GetDeploymentResponse,
+    ModifyDeploymentResponse,
+    # DeleteDeploymentData,
+    DeleteDeploymentResponse,
+    # DeploymentSummary,
+    DeploymentCollection,
+)
+
+
+class DeploymentClientError(Exception):
+    pass
+
+
+class DeploymentClient(BaseClient):
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
+        super().__init__(api_key, base_url)
+
+    def add_deployment(self, config: DeploymentInput):
+        """
+        Add a new deployment to the system.
+
+        Args:
+            config (DeploymentInput): Configuration object containing deployment details
+
+        Returns:
+            dict: Response from the API containing the added deployment details
+        """
+        headers = {"Content-Type": "application/json"}
+
+        if isinstance(config, dict):
+            config = DeploymentInput.from_dict(config)
+
+        payload = config.to_dict()
+
+        response = self._request(
+            "POST", "/deployments/add-deployment", headers=headers, json=payload
+        )
+        if response.get("error"):
+            raise DeploymentClientError(response["error"])
+        return DeploymentAddTaskResponse.from_dict(response)
+
+    def get_deployment(self, deployment_name: str, refresh_cache: bool = False):
+        """
+        Get detailed information for a specific deployment.
+
+        Args:
+            deployment_name (str): The name of the deployment
+            refresh_cache (bool): Whether to refresh the cache
+
+        Returns:
+            dict: Response from the API containing the deployment information
+        """
+        headers = {"X-Enkrypt-Deployment": deployment_name}
+        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"])
+        return GetDeploymentResponse.from_dict(response)
+    
+    def modify_deployment(self, deployment_name: str, config: DeploymentInput):
+        """
+        Modify an existing deployment in the system.
+
+        Args:
+            config (DeploymentInput): Configuration object containing deployment details
+
+        Returns:
+            dict: Response from the API containing the modified deployment details
+        """
+        headers = {"Content-Type": "application/json", "X-Enkrypt-Deployment": deployment_name}
+
+        if isinstance(config, dict):
+            config = DeploymentInput.from_dict(config)
+
+        payload = config.to_dict()
+
+        response = self._request(
+            "PATCH", "/deployments/modify-deployment", headers=headers, json=payload
+        )
+        if response.get("error"):
+            raise DeploymentClientError(response["error"])
+        return ModifyDeploymentResponse.from_dict(response)
+
+    def delete_deployment(self, deployment_name: str):
+        """
+        Delete a deployment from the system.
+
+        Args:
+            deployment_name (str): The name of the deployment
+
+        Returns:
+            dict: Response from the API containing the deployment deleted
+        """
+        headers = {"X-Enkrypt-Deployment": deployment_name}
+        response = self._request("DELETE", "/deployments/delete-deployment", headers=headers)
+        if response.get("error"):
+            raise DeploymentClientError(response["error"])
+        return DeleteDeploymentResponse.from_dict(response)
+
+    def list_deployments(self):
+        """
+        Get a list of all available deployments.
+
+        Returns:
+            dict: Response from the API containing the list of deployments
+        """
+
+        response = self._request("GET", "/deployments/list-deployments")
+        if response.get("error"):
+            raise DeploymentClientError(response["error"])
+        return DeploymentCollection.from_dict(response)

enkryptai_sdk/dto/__init__.py

@@ -1,7 +1,15 @@
 from .models import *
 from .red_team import *
+from .datasets import *
+from .deployments import *
+from .ai_proxy import *
+from .guardrails import *
+
 
 __all__ = [
+    "RedteamHealthResponse",
+    "RedTeamModelHealthConfig",
+    "RedteamModelHealthResponse",
     "DetailModelConfig",
     "ModelConfig",
     "RedTeamConfig",
@@ -15,4 +23,60 @@ __all__ = [
     "DEFAULT_REDTEAM_CONFIG",
     "ADVANCED_REDTEAM_TESTS",
     "DETAIL_MODEL_CONFIG",
+    "DatasetConfig",
+    "DatasetCollection",
+    "DatasetSummary",
+    "DatasetDataPoint",
+    "DatasetResponse",
+    "DatasetContent",
+    "DatasetTaskData",
+    "DatasetCard",
+    "DatasetTaskStatus",
+    "DatasetTask",
+    "DatasetAddTaskResponse",
+    "InputGuardrailBlockType",
+    "OutputGuardrailBlockType",
+    "InputGuardrailsAdditionalConfig",
+    "OutputGuardrailsAdditionalConfig",
+    "InputGuardrailsPolicy",
+    "OutputGuardrailsPolicy",
+    "DeploymentInput",
+    "DeploymentAddTaskResponse",
+    "GetDeploymentResponse",
+    "ModifyDeploymentResponse",
+    "DeleteDeploymentData",
+    "DeleteDeploymentResponse",
+    "DeploymentSummary",
+    "DeploymentCollection",
+    "ChatCompletionMessage",
+    "ChatCompletionRequest",
+    "ChatCompletionResponseMessage",
+    "ChatCompletionChoice",
+    "ChatCompletionUsage",
+    "ChatCompletionResponse",
+    "ChatCompletionError",
+    "ChatCompletionErrorResponse",
+    "ChatCompletionDirectErrorResponse",
+    "GuardrailsHealthResponse",
+    "GuardrailsModelsResponse",
+    "GuardrailDetectors",
+    "GuardrailsDetectRequest",
+    "GuardrailsPolicyDetectRequest",
+    "DetectResponseSummary",
+    "DetectResponseDetails",
+    "GuardrailsDetectResponse",
+    "GuardrailsPIIRequest",
+    "GuardrailsPIIResponse",
+    "GuardrailsHallucinationRequest",
+    "GuardrailsHallucinationResponse",
+    "GuardrailsAdherenceRequest",
+    "GuardrailsAdherenceResponse",
+    "GuardrailsRelevancyRequest",
+    "GuardrailsRelevancyResponse",
+    "GuardrailsPolicyRequest",
+    "GuardrailsPolicyData",
+    "GuardrailsaPolicyResponse",
+    "GuardrailsDeletePolicyData",
+    "GuardrailsDeletePolicyResponse",
+    "GuardrailsListPoliciesResponse",
 ]