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

enkryptai_sdk/guardrails.py

@@ -1,65 +1,122 @@
-import requests
+# import requests
+from .base import BaseClient
 from .config import GuardrailsConfig
 from .response import GuardrailsResponse, PIIResponse
+from .dto import (
+    GuardrailsHealthResponse,
+    GuardrailsModelsResponse,
+    # GuardrailDetectors,
+    # GuardrailsDetectRequest,
+    # GuardrailsPolicyDetectRequest,
+    # DetectResponseSummary,
+    # DetectResponseDetails,
+    GuardrailsDetectResponse,
+    # GuardrailsPIIRequest,
+    GuardrailsPIIResponse,
+    # GuardrailsHallucinationRequest,
+    GuardrailsHallucinationResponse,
+    # GuardrailsAdherenceRequest,
+    GuardrailsAdherenceResponse,
+    # GuardrailsRelevancyRequest,
+    GuardrailsRelevancyResponse,
+    # GuardrailsPolicyRequest,
+    GuardrailsPolicyData,
+    GuardrailsaPolicyResponse,
+    # GuardrailsDeletePolicyData,
+    GuardrailsDeletePolicyResponse,
+    GuardrailsListPoliciesResponse,
+)
 
-class GuardrailsClient:
+# ---------------------------------------
+# TODO: Use DTOs for request data as well
+# ---------------------------------------
+
+class GuardrailsClientError(Exception):
+    """
+    A custom exception for GuardrailsClient errors.
+    """
+
+    pass
+
+
+class GuardrailsClient(BaseClient):
     """
     A client for interacting with Enkrypt AI Guardrails API endpoints.
     """
+    def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com:443"):
+        super().__init__(api_key, base_url)
 
-    def __init__(self, api_key, base_url="https://api.enkryptai.com"):
-        """
-        Initializes the client.
+    # 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()
+    #     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.
+    # 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
+    #     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()
+    #     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)}
+    #     except Exception as e:
+    #         print(e)
+    #         return {"error": str(e)}
 
     # ----------------------------
     # Basic Guardrails Endpoints
     # ----------------------------
 
-    def health(self):
+    def get_health(self):
         """
         Get the health status of the service.
         """
-        return self._request("GET", "/guardrails/health")
+        try:
+            response = self._request("GET", "/guardrails/health")
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsHealthResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
 
-    def status(self):
+    def get_status(self):
         """
         Check if the API is up and running.
         """
-        return self._request("GET", "/guardrails/status")
+        try:
+            response = self._request("GET", "/guardrails/status")
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsHealthResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
 
-    def models(self):
+    def get_models(self):
         """
         Retrieve the list of models loaded by the service.
         """
-        return self._request("GET", "/guardrails/models")
+        try:
+            response = self._request("GET", "/guardrails/models")
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsModelsResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
 
     def detect(self, text, config=None):
         """
@@ -72,24 +129,32 @@ class GuardrailsClient:
           If not provided, defaults to injection attack detection only.
 
         Returns:
-        - JSON response from the API.
+        - Response from the API.
         """
         # Use injection attack config by default if none provided
         if config is None:
             config = GuardrailsConfig.injection_attack()
 
-        # Allow passing in either a dict or a GuardrailsConfig instance.
+        # Allow passing in either a dict or a GuardrailsConfig or GuardrailDetectors instance
         if hasattr(config, "as_dict"):
             config = config.as_dict()
+        if hasattr(config, "to_dict"):
+            config = config.to_dict()
             
         payload = {
             "text": text,
             "detectors": config
         }
-        response_body = self._request("POST", "/guardrails/detect", json=payload)
-        return GuardrailsResponse(response_body)
 
-    def pii(self, text, mode, key="null", entities=None):
+        try:
+            response = self._request("POST", "/guardrails/detect", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsDetectResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+        
+    def pii(self, text, mode="request", key="null", entities=None):
         """
         Detects Personally Identifiable Information (PII) and can de-anonymize it.
         """
@@ -99,14 +164,72 @@ class GuardrailsClient:
             "key": key,
             "entities": entities
         }
-        response_body = self._request("POST", "/guardrails/pii", json=payload)
-        return PIIResponse(response_body)
+
+        try:
+            response = self._request("POST", "/guardrails/pii", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsPIIResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+
+    def hallucination(self, request_text, response_text, context=""):
+        """
+        Detects hallucination in the response text.
+        """
+        payload = {
+            "request_text": request_text,
+            "response_text": response_text,
+            "context": context
+        }
+
+        try:
+            response = self._request("POST", "/guardrails/hallucination", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsHallucinationResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+        
+    def adherence(self, llm_answer, context):
+        """
+        Check the adherence of an LLM answer to the provided context.
+        """
+        payload = {
+            "llm_answer": llm_answer,
+            "context": context
+        }
+
+        try:
+            response = self._request("POST", "/guardrails/adherence", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsAdherenceResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+
+    def relevancy(self, question, llm_answer):
+        """
+        Check the relevancy of an LLM answer to the provided question.
+        """
+        payload = {
+            "question": question,
+            "llm_answer": llm_answer
+        }
+
+        try:
+            response = self._request("POST", "/guardrails/relevancy", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsRelevancyResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
 
     # ----------------------------
     # Guardrails Policy Endpoints
     # ----------------------------
 
-    def add_policy(self, name, config, description="guardrails policy"):
+    def add_policy(self, policy_name, config, description="guardrails policy"):
         """
         Create a new policy with custom configurations.
         
@@ -119,64 +242,78 @@ class GuardrailsClient:
         # Allow passing in either a dict or a GuardrailsConfig instance
         if hasattr(config, "as_dict"):
             config = config.as_dict()
+        if hasattr(config, "to_dict"):
+            config = config.to_dict()
             
         payload = {
-            "name": name,
+            "name": policy_name,
             "description": description,
             "detectors": config
         }
-        
+
         try:
-            return self._request("POST", "/guardrails/add-policy", json=payload)
-        
+            response = self._request("POST", "/guardrails/add-policy", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsaPolicyResponse.from_dict(response)
         except Exception as e:
-            print(e)
-            return {"error": str(e)}
+            raise GuardrailsClientError(str(e))
 
     def get_policy(self, policy_name):
         """
         Retrieve an existing policy by providing its header identifier.
         """
         headers = {"X-Enkrypt-Policy": policy_name}
+
         try:
-            return self._request("GET", "/guardrails/get-policy", headers=headers)
+            response = self._request("GET", "/guardrails/get-policy", headers=headers)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsPolicyData.from_dict(response)
         except Exception as e:
-            print(e)
-            return {"error": str(e)}
+            raise GuardrailsClientError(str(e))
 
-    def modify_policy(self, policy_name, config, name=None, description="guardrails policy"):
+    def modify_policy(self, policy_name, config, new_policy_name=None, description="guardrails policy"):
         """
         Modify an existing policy.
         """
         # Allow passing in either a dict or a GuardrailsConfig instance
         if hasattr(config, "as_dict"):
             config = config.as_dict()
-            
-        if name is None:
-            name = policy_name
-            
+        if hasattr(config, "to_dict"):
+            config = config.to_dict()
+
+        if new_policy_name is None:
+            new_policy_name = policy_name
+
         headers = {"X-Enkrypt-Policy": policy_name}
         payload = {
             "detectors": config,
-            "name": name,
+            "name": new_policy_name,
             "description": description
         }
+
         try:
-            return self._request("PATCH", "/guardrails/modify-policy", headers=headers, json=payload)
+            response = self._request("PATCH", "/guardrails/modify-policy", headers=headers, json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsaPolicyResponse.from_dict(response)
         except Exception as e:
-            print(e)
-            return {"error": str(e)}
+            raise GuardrailsClientError(str(e))
 
     def delete_policy(self, policy_name):
         """
         Delete a policy.
         """
         headers = {"X-Enkrypt-Policy": policy_name}
+
         try:
-            return self._request("DELETE", "/guardrails/delete-policy", headers=headers)
+            response = self._request("DELETE", "/guardrails/delete-policy", headers=headers)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsDeletePolicyResponse.from_dict(response)
         except Exception as e:
-            print(e)
-            return {"error": str(e)}
+            raise GuardrailsClientError(str(e))
 
     def policy_detect(self, policy_name, text):
         """
@@ -184,12 +321,24 @@ class GuardrailsClient:
         """
         headers = {"X-Enkrypt-Policy": policy_name}
         payload = {"text": text}
-        
+
         try:
-            
-            response_body = self._request("POST", "/guardrails/policy/detect", headers=headers, json=payload)
-            return GuardrailsResponse(response_body)
-        
+            response = self._request("POST", "/guardrails/policy/detect", headers=headers, json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsDetectResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+
+    def get_policy_list(self):
+        """
+        List all policies.
+        """
+
+        try:
+            response = self._request("GET", "/guardrails/list-policies")
+            if response.get("error"):
+                raise GuardrailsClientError(response["error"])
+            return GuardrailsListPoliciesResponse.from_dict(response)
         except Exception as e:
-            print(e)
-            return {"error": str(e)}
+            raise GuardrailsClientError(str(e))

enkryptai_sdk/guardrails_old.py

@@ -0,0 +1,195 @@
+import requests
+from .config import GuardrailsConfig
+from .response import GuardrailsResponse, PIIResponse
+
+class GuardrailsClient:
+    """
+    A client for interacting with Enkrypt AI Guardrails 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 Guardrails Endpoints
+    # ----------------------------
+
+    def health(self):
+        """
+        Get the health status of the service.
+        """
+        return self._request("GET", "/guardrails/health")
+
+    def status(self):
+        """
+        Check if the API is up and running.
+        """
+        return self._request("GET", "/guardrails/status")
+
+    def models(self):
+        """
+        Retrieve the list of models loaded by the service.
+        """
+        return self._request("GET", "/guardrails/models")
+
+    def detect(self, text, config=None):
+        """
+        Detects prompt injection, toxicity, NSFW content, PII, hallucination, and more.
+
+        Parameters:
+        - text (str): The text to analyze.
+        - guardrails_config (dict or GuardrailsConfig, optional): A configuration for detectors.
+          If a GuardrailsConfig instance is provided, its underlying dictionary will be used.
+          If not provided, defaults to injection attack detection only.
+
+        Returns:
+        - JSON response from the API.
+        """
+        # Use injection attack config by default if none provided
+        if config is None:
+            config = GuardrailsConfig.injection_attack()
+
+        # Allow passing in either a dict or a GuardrailsConfig instance.
+        if hasattr(config, "as_dict"):
+            config = config.as_dict()
+            
+        payload = {
+            "text": text,
+            "detectors": config
+        }
+        response_body = self._request("POST", "/guardrails/detect", json=payload)
+        return GuardrailsResponse(response_body)
+
+    def pii(self, text, mode, key="null", entities=None):
+        """
+        Detects Personally Identifiable Information (PII) and can de-anonymize it.
+        """
+        payload = {
+            "text": text,
+            "mode": mode,
+            "key": key,
+            "entities": entities
+        }
+        response_body = self._request("POST", "/guardrails/pii", json=payload)
+        return PIIResponse(response_body)
+
+    # ----------------------------
+    # Guardrails Policy Endpoints
+    # ----------------------------
+
+    def add_policy(self, name, config, description="guardrails policy"):
+        """
+        Create a new policy with custom configurations.
+        
+        Args:
+            name (str): Name of the policy
+            config (dict or GuardrailsConfig): Configuration for the policy detectors.
+                If a GuardrailsConfig instance is provided, its underlying dictionary will be used.
+            description (str, optional): Description of the policy. Defaults to "guardrails policy"
+        """
+        # Allow passing in either a dict or a GuardrailsConfig instance
+        if hasattr(config, "as_dict"):
+            config = config.as_dict()
+            
+        payload = {
+            "name": name,
+            "description": description,
+            "detectors": config
+        }
+        
+        try:
+            return self._request("POST", "/guardrails/add-policy", json=payload)
+        
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}
+
+    def get_policy(self, policy_name):
+        """
+        Retrieve an existing policy by providing its header identifier.
+        """
+        headers = {"X-Enkrypt-Policy": policy_name}
+        try:
+            return self._request("GET", "/guardrails/get-policy", headers=headers)
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}
+
+    def modify_policy(self, policy_name, config, name=None, description="guardrails policy"):
+        """
+        Modify an existing policy.
+        """
+        # Allow passing in either a dict or a GuardrailsConfig instance
+        if hasattr(config, "as_dict"):
+            config = config.as_dict()
+            
+        if name is None:
+            name = policy_name
+            
+        headers = {"X-Enkrypt-Policy": policy_name}
+        payload = {
+            "detectors": config,
+            "name": name,
+            "description": description
+        }
+        try:
+            return self._request("PATCH", "/guardrails/modify-policy", headers=headers, json=payload)
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}
+
+    def delete_policy(self, policy_name):
+        """
+        Delete a policy.
+        """
+        headers = {"X-Enkrypt-Policy": policy_name}
+        try:
+            return self._request("DELETE", "/guardrails/delete-policy", headers=headers)
+        except Exception as e:
+            print(e)
+            return {"error": str(e)}
+
+    def policy_detect(self, policy_name, text):
+        """
+        Apply a specific policy to detect and filter content.
+        """
+        headers = {"X-Enkrypt-Policy": policy_name}
+        payload = {"text": text}
+        
+        try:
+            
+            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)}