enkryptai-sdk 1.0.10__tar.gz → 1.0.11__tar.gz

{enkryptai_sdk-1.0.10/src/enkryptai_sdk.egg-info → enkryptai_sdk-1.0.11}/PKG-INFO

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

@@ -9,7 +9,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
     # NOTE: Also change this in .github/workflows/test.yaml
-    version="1.0.10", # Update this for new versions
+    version="1.0.11", # Update this for new versions
     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-1.0.10 → enkryptai_sdk-1.0.11}/src/enkryptai_sdk/base.py

@@ -72,14 +72,17 @@ class BaseClient:
                 )
 
             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))
+                try:
+                    error_data = response.json()
+                    error_message = error_data.get("message", str(error_data))
+                except:
+                    error_message = response.data.decode('utf-8') if response.data else f"HTTP {response.status}"
                 raise urllib3.exceptions.HTTPError(error_message)
-            return response.json()
+            
+            try:
+                return response.json()
+            except:
+                return {"error": response.data.decode('utf-8') if response.data else "Invalid JSON response"}
         except urllib3.exceptions.HTTPError as e:
             return {"error": str(e)}
 

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11}/src/enkryptai_sdk/dto/base.py

@@ -52,7 +52,18 @@ class BaseDTO:
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the instance to a dictionary."""
-        d = {k: v for k, v in self.__dict__.items() if k != "_extra_fields"}
+        d = {}
+        for k, v in self.__dict__.items():
+            if k == "_extra_fields":
+                continue
+            if hasattr(v, "to_dict"):
+                d[k] = v.to_dict()
+            elif isinstance(v, list):
+                d[k] = [item.to_dict() if hasattr(item, "to_dict") else item for item in v]
+            elif isinstance(v, dict):
+                d[k] = {key: val.to_dict() if hasattr(val, "to_dict") else val for key, val in v.items()}
+            else:
+                d[k] = v
         d.update(self._extra_fields)
         return d
 

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11}/src/enkryptai_sdk/dto/models.py

@@ -37,6 +37,8 @@ class ModelProviders(str, Enum):
     OPENAI_COMPATIBLE = "openai_compatible"
     COHERE_COMPATIBLE = "cohere_compatible"
     ANTHROPIC_COMPATIBLE = "anthropic_compatible"
+    CUSTOM = "custom"
+    HR = "hr"
 
 
 @dataclass
@@ -247,6 +249,14 @@ class ModelConfigDetails(BaseDTO):
         d = super().to_dict()
         # Handle AuthData specifically
         d["auth_data"] = self.auth_data.to_dict()
+        # Handle CustomHeader list
+        d["custom_headers"] = [header.to_dict() for header in self.custom_headers]
+        # Handle ModelProviders enum
+        if isinstance(d["model_provider"], ModelProviders):
+            d["model_provider"] = d["model_provider"].value
+        # Handle input/output modalities
+        d["input_modalities"] = [m.value for m in self.input_modalities]
+        d["output_modalities"] = [m.value for m in self.output_modalities]
         return d
 
     def to_json(self):
@@ -289,6 +299,13 @@ class ModelConfig(BaseDTO):
 
         return cls(**data, model_config=model_config)
 
+    def to_dict(self) -> dict:
+        """Convert the ModelConfig instance to a dictionary."""
+        d = super().to_dict()
+        # Handle nested ModelConfigDetails
+        d["model_config"] = self.model_config.to_dict()
+        return d
+
     @classmethod
     def __str__(self):
         """String representation of the ModelConfig."""

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11}/src/enkryptai_sdk/dto/red_team.py

@@ -367,7 +367,15 @@ class TargetModelConfiguration(BaseDTO):
 
     @classmethod
     def from_dict(cls, data: dict):
+        data = data.copy()
+        if "custom_headers" in data:
+            data["custom_headers"] = [CustomHeader.from_dict(header) for header in data["custom_headers"]]
         return cls(**data)
+    
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["custom_headers"] = [header.to_dict() for header in self.custom_headers]
+        return d
 
 
 @dataclass

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11}/src/enkryptai_sdk/models.py

@@ -49,6 +49,9 @@ class ModelClient(BaseClient):
                 "chat": f"/{remaining_path}" if remaining_path else "",
             }
 
+        # Convert custom_headers to list of dictionaries
+        custom_headers = [header.to_dict() for header in config.model_config.custom_headers]
+
         payload = {
             "testing_for": config.testing_for,
             "model_name": config.model_name,
@@ -72,9 +75,9 @@ class ModelClient(BaseClient):
                 },
                 "apikeys": [config.model_config.apikey] if config.model_config.apikey else [],
                 "tools": config.model_config.tools,
-                "input_modalities": config.model_config.input_modalities,
-                "output_modalities": config.model_config.output_modalities,
-                "custom_headers": config.model_config.custom_headers,
+                "input_modalities": [m.value if hasattr(m, 'value') else m for m in config.model_config.input_modalities],
+                "output_modalities": [m.value if hasattr(m, 'value') else m for m in config.model_config.output_modalities],
+                "custom_headers": custom_headers,
                 "custom_payload": config.model_config.custom_payload,
                 "custom_response_content_type": config.model_config.custom_response_content_type,
                 "custom_response_format": config.model_config.custom_response_format,
@@ -189,6 +192,9 @@ class ModelClient(BaseClient):
                 "chat": f"/{remaining_path}" if remaining_path else "",
             }
 
+        # Convert custom_headers to list of dictionaries
+        custom_headers = [header.to_dict() for header in config.model_config.custom_headers]
+
         payload = {
             "model_saved_name": config.model_saved_name,
             "model_version": config.model_version,
@@ -217,9 +223,9 @@ class ModelClient(BaseClient):
                     [config.model_config.apikey] if config.model_config.apikey else []
                 ),
                 "tools": config.model_config.tools,
-                "input_modalities": config.model_config.input_modalities,
-                "output_modalities": config.model_config.output_modalities,
-                "custom_headers": config.model_config.custom_headers,
+                "input_modalities": [m.value if hasattr(m, 'value') else m for m in config.model_config.input_modalities],
+                "output_modalities": [m.value if hasattr(m, 'value') else m for m in config.model_config.output_modalities],
+                "custom_headers": custom_headers,
                 "custom_payload": config.model_config.custom_payload,
                 "custom_response_content_type": config.model_config.custom_response_content_type,
                 "custom_response_format": config.model_config.custom_response_format,

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11}/src/enkryptai_sdk/red_team.py

@@ -1,4 +1,5 @@
-import urllib3
+# import json
+# import urllib3
 from .base import BaseClient
 from .models import ModelClient
 from .datasets import DatasetClient
@@ -62,6 +63,8 @@ class RedTeamClient(BaseClient):
         """
         try:
             config = RedTeamModelHealthConfig.from_dict(config)
+            # Print the config as json string
+            # print(f"Config: {json.dumps(config.to_dict(), indent=4)}")
             response = self._request("POST", "/redteam/model-health", json=config.to_dict())
             # if response.get("error"):
             if response.get("error") not in [None, ""]:

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11/src/enkryptai_sdk.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.10
+Version: 1.0.11
 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-1.0.10 → enkryptai_sdk-1.0.11}/tests/test_all_v2.py

@@ -163,11 +163,33 @@ def sample_model_config():
         "testing_for": "foundationModels",
         "model_name": model_name,
         "model_config": {
-            "model_provider": model_provider,
+            # "model_provider": model_provider,
+            "model_provider": "custom",
             "endpoint_url": model_endpoint_url,
-            "apikey": OPENAI_API_KEY,
+            # "apikey": OPENAI_API_KEY,
             "input_modalities": ["text"],
             "output_modalities": ["text"],
+            "custom_headers": [
+                {
+                    "key": "Content-Type",
+                    "value": "application/json"
+                },
+                {
+                    "key": "Authorization",
+                    "value": "Bearer " + OPENAI_API_KEY
+                }
+            ],
+            "custom_payload": {
+                "model": model_name,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "{prompt}"
+                    }
+                ]
+            },
+            "custom_response_content_type": "json",
+            "custom_response_format": ".choices[0].message.content",
         },
     }
 
@@ -235,9 +257,31 @@ def sample_redteam_model_health_config():
             "testing_for": "foundationModels",
             "model_version": "v1",
             "model_source": "https://openai.com",
-            "model_provider": model_provider,
+            # "model_provider": model_provider,
+            "model_provider": "custom",
             "model_endpoint_url": model_endpoint_url,
-            "model_api_key": OPENAI_API_KEY,
+            # "model_api_key": OPENAI_API_KEY,
+            "custom_headers": [
+                {
+                    "key": "Content-Type",
+                    "value": "application/json"
+                },
+                {
+                    "key": "Authorization",
+                    "value": "Bearer " + OPENAI_API_KEY
+                }
+            ],
+            "custom_payload": {
+                "model": model_name,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "{prompt}"
+                    }
+                ]
+            },
+            "custom_response_content_type": "json",
+            "custom_response_format": ".choices[0].message.content",
             "system_prompt": "",
             "rate_per_min": 20,
             "input_modalities": ["text"],
@@ -659,18 +703,19 @@ def test_adherence(guardrails_client):
     assert summary.adherence_score == 0
 
 
-def test_relevancy(guardrails_client):
-    print("\n\nTesting relevancy")
-    # Test the relevancy method
-    response = guardrails_client.relevancy(llm_answer="Hello! How can I help you today? If you have any questions or need assistance with something, feel free to ask. I'm here to provide information and support. Is there something specific you'd like to know or discuss?", question="Hi")
-    print("\nResponse from relevancy: ", response)
-    print("\nResponse data type: ", type(response))
-    assert response is not None
-    assert hasattr(response, "summary")
-    summary = response.summary
-    assert summary.relevancy_score == 0
-    print("\nSleeping for 60 seconds after guardrails tests...")
-    time.sleep(60)
+# Not being used in Prod at this time
+# def test_relevancy(guardrails_client):
+#     print("\n\nTesting relevancy")
+#     # Test the relevancy method
+#     response = guardrails_client.relevancy(llm_answer="Hello! How can I help you today? If you have any questions or need assistance with something, feel free to ask. I'm here to provide information and support. Is there something specific you'd like to know or discuss?", question="Hi")
+#     print("\nResponse from relevancy: ", response)
+#     print("\nResponse data type: ", type(response))
+#     assert response is not None
+#     assert hasattr(response, "summary")
+#     summary = response.summary
+#     assert summary.relevancy_score == 0
+#     print("\nSleeping for 60 seconds after guardrails tests...")
+#     time.sleep(60)
 
 
 def test_list_policies(guardrails_client):
@@ -1174,13 +1219,16 @@ def test_get_health(redteam_client):
     assert response.status == "healthy"
 
 
-def test_model_health(redteam_client, sample_redteam_model_health_config):
-    print("\n\nTesting check_model_health")
-    response = redteam_client.check_model_health(config=sample_redteam_model_health_config)
-    print("\nResponse from check_model_health: ", response)
-    assert response is not None
-    assert hasattr(response, "status")
-    assert response.status == "healthy"
+# ---------------------------------------------------------
+# Commenting as we already test for saved model health
+# ---------------------------------------------------------
+# def test_model_health(redteam_client, sample_redteam_model_health_config):
+#     print("\n\nTesting check_model_health")
+#     response = redteam_client.check_model_health(config=sample_redteam_model_health_config)
+#     print("\nResponse from check_model_health: ", response)
+#     assert response is not None
+#     assert hasattr(response, "status")
+#     assert response.status == "healthy"
 
 
 def test_saved_model_health(redteam_client):
@@ -1192,7 +1240,8 @@ def test_saved_model_health(redteam_client):
     assert response.status == "healthy"
 
 
-# # Testing only via saved model as it should be sufficient
+# ---------------------------------------------------------
+# Commenting as we are testing with add custom task with saved model
 # ---------------------------------------------------------
 # def test_add_task_with_target_model(redteam_client, sample_redteam_target_config):
 #     print("\n\nTesting adding a new redteam task with target model")
@@ -1210,20 +1259,24 @@ def test_saved_model_health(redteam_client):
 #     time.sleep(60)
 
 
-def test_add_task_with_saved_model(redteam_client, sample_redteam_model_config):
-    print("\n\nTesting adding a new redteam task with saved model")
-    response = redteam_client.add_task_with_saved_model(config=sample_redteam_model_config,model_saved_name=test_model_saved_name, model_version=test_model_version)
-    print("\nResponse from adding a new redteam task with saved model: ", response)
-    assert response is not None
-    assert hasattr(response, "task_id")
-    assert hasattr(response, "message")
-    response.message == "Redteam task has been added successfully"
-    # Sleep for a while to let the task complete
-    # This is also useful to avoid rate limiting issues
-    print("\nSleeping for 60 seconds to let the task complete if possible ...")
-    time.sleep(60)
+# ---------------------------------------------------------
+# Commenting as we are testing with add custom task with saved model
+# ---------------------------------------------------------
+# def test_add_task_with_saved_model(redteam_client, sample_redteam_model_config):
+#     print("\n\nTesting adding a new redteam task with saved model")
+#     response = redteam_client.add_task_with_saved_model(config=sample_redteam_model_config,model_saved_name=test_model_saved_name, model_version=test_model_version)
+#     print("\nResponse from adding a new redteam task with saved model: ", response)
+#     assert response is not None
+#     assert hasattr(response, "task_id")
+#     assert hasattr(response, "message")
+#     response.message == "Redteam task has been added successfully"
+#     # Sleep for a while to let the task complete
+#     # This is also useful to avoid rate limiting issues
+#     print("\nSleeping for 60 seconds to let the task complete if possible ...")
+#     time.sleep(60)
 
 
+# ---------------------------------------------------------
 # # Testing only via saved model as it should be sufficient
 # ---------------------------------------------------------
 # def test_add_custom_task_with_target_model(redteam_client, sample_custom_redteam_target_config):

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.11}/tests/test_model.py

@@ -14,6 +14,9 @@ model_saved_name = None
 test_model_saved_name = "Test Model"
 model_version = None
 test_model_version = "v1"
+model_provider = "openai"
+model_name = "gpt-4o-mini"
+model_endpoint_url = "https://api.openai.com/v1/chat/completions"
 
 @pytest.fixture
 def model_client():
@@ -27,13 +30,35 @@ def sample_model_config():
         "model_saved_name": test_model_saved_name,
         "model_version": test_model_version,
         "testing_for": "foundationModels",
-        "model_name": "gpt-4o-mini",
+        "model_name": model_name,
         "model_config": {
-            "model_provider": "openai",
-            "endpoint_url": "https://api.openai.com/v1/chat/completions",
-            "apikey": OPENAI_API_KEY,
+            # "model_provider": model_provider,
+            "model_provider": "custom",
+            "endpoint_url": model_endpoint_url,
+            # "apikey": OPENAI_API_KEY,
             "input_modalities": ["text"],
             "output_modalities": ["text"],
+            "custom_headers": [
+                {
+                    "key": "Content-Type",
+                    "value": "application/json"
+                },
+                {
+                    "key": "Authorization",
+                    "value": "Bearer " + OPENAI_API_KEY
+                }
+            ],
+            "custom_payload": {
+                "model": model_name,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "{prompt}"
+                    }
+                ]
+            },
+            "custom_response_content_type": "json",
+            "custom_response_format": ".choices[0].message.content",
         },
     }