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

{enkryptai_sdk-1.0.11/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.11
+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.11 → 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.11", # 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.11 → enkryptai_sdk-1.0.12}/src/enkryptai_sdk/red_team.py

@@ -106,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:
@@ -153,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 = {
@@ -196,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:
@@ -252,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.11 → 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.11
+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.11 → 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 {
@@ -162,6 +164,23 @@ def sample_model_config():
         "model_version": test_model_version,
         "testing_for": "foundationModels",
         "model_name": model_name,
+        "model_config": {
+            "model_provider": model_provider,
+            "endpoint_url": model_endpoint_url,
+            "apikey": OPENAI_API_KEY,
+            "input_modalities": ["text"],
+            "output_modalities": ["text"],
+        },
+    }
+
+
+@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",
@@ -193,6 +212,7 @@ def sample_model_config():
         },
     }
 
+
 @pytest.fixture
 def sample_deployment_config():
     return {
@@ -945,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
@@ -962,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")
@@ -972,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
@@ -984,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
@@ -996,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)
@@ -1233,7 +1267,7 @@ def test_get_health(redteam_client):
 
 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")
@@ -1297,7 +1331,7 @@ def test_saved_model_health(redteam_client):
 
 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")
@@ -1439,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)