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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.11
+Version: 1.0.13
 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.13}/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.13", # 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.13}/src/enkryptai_sdk/dto/models.py

@@ -166,6 +166,7 @@ class ModelConfigDetails(BaseDTO):
     auth_data: AuthData = field(default_factory=AuthData)
     input_modalities: List[InputModality] = field(default_factory=list)
     output_modalities: List[OutputModality] = field(default_factory=list)
+    custom_curl_command: Optional[str] = None
     custom_headers: List[CustomHeader] = field(default_factory=list)
     custom_payload: Dict[str, Any] = field(default_factory=dict)
     custom_response_content_type: Optional[str] = ""

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

@@ -359,6 +359,7 @@ class TargetModelConfiguration(BaseDTO):
     model_api_key: Optional[str] = None
     input_modalities: List[InputModality] = field(default_factory=list)
     output_modalities: List[OutputModality] = field(default_factory=list)
+    custom_curl_command: Optional[str] = None
     custom_headers: List[CustomHeader] = field(default_factory=list)
     custom_payload: Dict[str, Any] = field(default_factory=dict)
     custom_response_content_type: Optional[str] = ""

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

@@ -77,6 +77,7 @@ class ModelClient(BaseClient):
                 "tools": config.model_config.tools,
                 "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_curl_command": config.model_config.custom_curl_command,
                 "custom_headers": custom_headers,
                 "custom_payload": config.model_config.custom_payload,
                 "custom_response_content_type": config.model_config.custom_response_content_type,
@@ -225,6 +226,7 @@ class ModelClient(BaseClient):
                 "tools": config.model_config.tools,
                 "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_curl_command": config.model_config.custom_curl_command,
                 "custom_headers": custom_headers,
                 "custom_payload": config.model_config.custom_payload,
                 "custom_response_content_type": config.model_config.custom_response_content_type,

{enkryptai_sdk-1.0.11 → enkryptai_sdk-1.0.13}/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.13/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.13
 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.13}/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
@@ -972,11 +1006,16 @@ 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)
+    # -------------------------------
+    # NOTE: Using custom model here is breaking AI Proxy as model_saved_name is a global variable
+    # -------------------------------
+    # model_saved_name = test_custom_model_saved_name
     model_saved_name = test_model_saved_name
     model_version = test_model_version
 
     # Now test the get_model method
     model = model_client.get_model(model_saved_name=model_saved_name, model_version=model_version)
+    # print("\nGET Model: ", model)
     assert model is not None
     # assert hasattr(model, "model_id")
     assert model.model_saved_name == model_saved_name
@@ -984,10 +1023,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 +1035,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)
@@ -1006,7 +1045,7 @@ def test_modify_model(model_client, sample_model_config):
 # Tests for Deployment Endpoints
 # ----------------------------
 
-def test_add_deployment_success(deployment_client, sample_deployment_config):
+def test_add_deployment(deployment_client, sample_deployment_config):
     response = deployment_client.add_deployment(config=sample_deployment_config)
     assert response is not None
     assert hasattr(response, "from_dict")
@@ -1097,7 +1136,7 @@ def test_401_err_chat(ai_proxy_err_client, sample_chat_body):
 # Tests for Dataset Endpoints
 # ----------------------------
 
-def test_add_dataset_success(dataset_client, sample_dataset_config):
+def test_add_dataset(dataset_client, sample_dataset_config):
     print("\n\nTesting add_dataset")
     response = dataset_client.add_dataset(config=sample_dataset_config)
     print("\nAdd Dataset Response: ", response)
@@ -1233,7 +1272,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 +1336,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 +1478,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)