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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.10
+Version: 1.0.12
 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.12}/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.12", # 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.12}/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.12}/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.12}/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.12}/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.12}/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.12}/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, ""]:
@@ -103,9 +106,7 @@ class RedTeamClient(BaseClient):
             # "async": config.async_enabled,
             "dataset_name": config.dataset_name,
             "test_name": config.test_name,
-            "redteam_test_configurations": {
-                k: v.to_dict() for k, v in test_configs.items()
-            },
+            "redteam_test_configurations": test_configs,
         }
 
         if config.target_model_configuration:
@@ -150,9 +151,7 @@ class RedTeamClient(BaseClient):
             # "async": config.async_enabled,
             "dataset_name": config.dataset_name,
             "test_name": config.test_name,
-            "redteam_test_configurations": {
-                k: v.to_dict() for k, v in test_configs.items()
-            },
+            "redteam_test_configurations": test_configs,
         }
 
         headers = {
@@ -193,9 +192,7 @@ class RedTeamClient(BaseClient):
         payload = {
             # "async": config.async_enabled,
             "test_name": config.test_name,
-            "redteam_test_configurations": {
-                k: v.to_dict() for k, v in test_configs.items()
-            },
+            "redteam_test_configurations": test_configs,
         }
 
         if config.dataset_configuration:
@@ -249,9 +246,7 @@ class RedTeamClient(BaseClient):
         payload = {
             # "async": config.async_enabled,
             "test_name": config.test_name,
-            "redteam_test_configurations": {
-                k: v.to_dict() for k, v in test_configs.items()
-            },
+            "redteam_test_configurations": test_configs,
         }
 
         if config.dataset_configuration:

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.12/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.12
 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.12}/tests/test_all_v2.py

@@ -19,6 +19,7 @@ test_guardrails_policy_name = "Test Guardrails Policy"
 model_saved_name = None
 model_version = None
 test_model_saved_name = "Test Model"
+test_custom_model_saved_name = "Test Custom Model"
 test_model_version = "v1"
 test_deployment_name = "test-deployment"
 
@@ -155,6 +156,7 @@ def sample_detectors():
         },
     }
 
+
 @pytest.fixture
 def sample_model_config():
     return {
@@ -171,6 +173,46 @@ def sample_model_config():
         },
     }
 
+
+@pytest.fixture
+def sample_custom_model_config():
+    return {
+        "model_saved_name": test_custom_model_saved_name,
+        "model_version": test_model_version,
+        "testing_for": "foundationModels",
+        "model_name": model_name,
+        "model_config": {
+            # "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",
+        },
+    }
+
+
 @pytest.fixture
 def sample_deployment_config():
     return {
@@ -235,9 +277,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 +723,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):
@@ -900,6 +965,20 @@ def test_add_model(model_client, sample_model_config):
     assert model_info["model_saved_name"] == test_model_saved_name
     assert model_info["model_version"] == test_model_version
 
+def test_add_custom_model(model_client, sample_custom_model_config):
+    print("\n\nTesting adding a new custom model")
+    response = model_client.add_model(config=sample_custom_model_config)
+    print("\nResponse from adding a new custom model: ", response)
+    print("\nResponse data type: ", type(response))
+    assert response is not None
+    # assert hasattr(response, "message")
+    assert response.message == "Model details added successfully"
+    # assert hasattr(response, "data")
+    model_info = response.data
+    assert model_info is not None
+    assert model_info["model_saved_name"] == test_custom_model_saved_name
+    assert model_info["model_version"] == test_model_version
+
 def test_list_models(model_client):
     print("\n\nTesting list_models")
     # Test the list_models method
@@ -917,7 +996,7 @@ def test_list_models(model_client):
     # assert model_info["model_version"] == test_model_version
 
 def test_get_model(model_client):
-    print("\n\nTesting get_model")
+    print("\n\nTesting get_model with custom model")
     # global model_saved_name
     # if model_saved_name is None:
     #     print("\nModel saved name is None, fetching it from list_models")
@@ -927,7 +1006,7 @@ def test_get_model(model_client):
     #     model_saved_name = model_info.model_saved_name
     #     assert model_saved_name is not None
     #     print("\nPicked model in get_model: ", model_saved_name)
-    model_saved_name = test_model_saved_name
+    model_saved_name = test_custom_model_saved_name
     model_version = test_model_version
 
     # Now test the get_model method
@@ -939,10 +1018,10 @@ def test_get_model(model_client):
     # assert hasattr(model, "status")
 
 
-def test_modify_model(model_client, sample_model_config):
-    print("\n\nTesting modifying a new model")
+def test_modify_model(model_client, sample_custom_model_config):
+    print("\n\nTesting modifying custom model")
     # Test creating a new model
-    response = model_client.modify_model(config=sample_model_config, old_model_saved_name=None, old_model_version=None)
+    response = model_client.modify_model(config=sample_custom_model_config, old_model_saved_name=None, old_model_version=None)
     print("\nResponse from modifying a new model: ", response)
     print("\nResponse data type: ", type(response))
     assert response is not None
@@ -951,7 +1030,7 @@ def test_modify_model(model_client, sample_model_config):
     # assert hasattr(response, "data")
     model_info = response.data
     assert model_info is not None
-    assert model_info["model_saved_name"] == test_model_saved_name
+    assert model_info["model_saved_name"] == test_custom_model_saved_name
     assert model_info["model_version"] == test_model_version
     print("\nSleeping for 60 seconds after model tests...")
     time.sleep(60)
@@ -1174,25 +1253,29 @@ 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):
     print("\n\nTesting check_saved_model_health")
-    response = redteam_client.check_saved_model_health(model_saved_name=test_model_saved_name, model_version=test_model_version)
+    response = redteam_client.check_saved_model_health(model_saved_name=test_custom_model_saved_name, model_version=test_model_version)
     print("\nResponse from check_saved_model_health: ", response)
     assert response is not None
     assert hasattr(response, "status")
     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 +1293,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):
@@ -1244,7 +1331,7 @@ def test_add_task_with_saved_model(redteam_client, sample_redteam_model_config):
 
 def test_add_custom_task_with_saved_model(redteam_client, sample_custom_redteam_model_config):
     print("\n\nTesting adding a new custom redteam task with saved model")
-    response = redteam_client.add_custom_task_with_saved_model(config=sample_custom_redteam_model_config, model_saved_name=test_model_saved_name, model_version=test_model_version)
+    response = redteam_client.add_custom_task_with_saved_model(config=sample_custom_redteam_model_config, model_saved_name=test_custom_model_saved_name, model_version=test_model_version)
     print("\nResponse from adding a new custom redteam task with saved model: ", response)
     assert response is not None
     assert hasattr(response, "task_id")
@@ -1386,6 +1473,14 @@ def test_delete_model(model_client):
     # assert hasattr(response, "message")
     assert response.message == "Model details deleted successfully"
 
+def test_delete_custom_model(model_client):
+    print("\n\nTesting delete_custom_model")
+    response = model_client.delete_model(model_saved_name=test_custom_model_saved_name, model_version=test_model_version)
+    print("\nResponse from delete_custom_model: ", response)
+    assert response is not None
+    # assert hasattr(response, "message")
+    assert response.message == "Model details deleted successfully"
+
 def test_delete_deployment(deployment_client):
     print("\n\nTesting delete_deployment")
     response = deployment_client.delete_deployment(deployment_name=test_deployment_name)

{enkryptai_sdk-1.0.10 → enkryptai_sdk-1.0.12}/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",
         },
     }