enkryptai-sdk 1.0.7__tar.gz → 1.0.8__tar.gz

{enkryptai_sdk-1.0.7/src/enkryptai_sdk.egg-info → enkryptai_sdk-1.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.7
+Version: 1.0.8
 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
@@ -47,6 +47,8 @@ Also see the API documentation at [https://docs.enkryptai.com](https://docs.enkr
   - [Sample Redteam Model Health Config](#sample-redteam-model-health-config)
   - [Sample Redteam Target Config](#sample-redteam-target-config)
   - [Sample Redteam Model Config](#sample-redteam-model-config)
+  - [Sample Custom Redteam Target Config](#sample-custom-redteam-target-config)
+  - [Sample Custom Redteam Model Config](#sample-custom-redteam-model-config)
 - [Health Checks](#health-checks)
   - [Guardrails Health](#guardrails-health)
   - [Guardrails Status](#guardrails-status)
@@ -106,6 +108,8 @@ Also see the API documentation at [https://docs.enkryptai.com](https://docs.enkr
 - [Redteam](#redteam)
   - [Add a Redteam Task with Target Model Config](#add-a-redteam-task-with-target-model-config)
   - [Add a Redteam Task with a saved model](#add-a-redteam-task-with-a-saved-model)
+  - [Add a Redteam Custom Task with Endpoint Config](#add-a-redteam-custom-task-with-endpoint-config)
+  - [Add a Redteam Custom Task with a saved model](#add-a-redteam-custom-task-with-a-saved-model)
   - [Get Redteam Task Status](#get-redteam-task-status)
   - [Get Redteam Task](#get-redteam-task)
   - [List Redteam Tasks](#list-redteam-tasks)
@@ -193,6 +197,7 @@ pii_anonymized_text = None
 dataset_name = f"TestElectionDataset-{str(uuid.uuid4())[:8]}"
 
 redteam_test_name = f"Redteam Test {str(uuid.uuid4())[:6]}"
+custom_redteam_test_name = f"Custom Redteam Test {str(uuid.uuid4())[:6]}"
 
 model_name = "gpt-4o-mini"
 model_provider = "openai"
@@ -422,6 +427,82 @@ sample_redteam_model_config = {
     }
 ```
 
+### [Sample Custom Redteam Target Config](https://docs.enkryptai.com/redteam-api-reference/endpoint/add-custom-task)
+
+```python Python
+sample_custom_redteam_target_config = {
+        "test_name": custom_redteam_test_name,
+        "dataset_configuration": {
+          "system_description": "- **Voter Eligibility**: To vote in U.S. elections, individuals must be U.S. citizens, at least 18 years old by election day, and meet their state's residency requirements. - **Voter Registration**: Most states require voters to register ahead of time, with deadlines varying widely. North Dakota is an exception, as it does not require voter registration. - **Identification Requirements**: Thirty-six states enforce voter ID laws, requiring individuals to present identification at polling places. These laws aim to prevent voter fraud but can also lead to disenfranchisement. - **Voting Methods**: Voters can typically choose between in-person voting on election day, early voting, and absentee or mail-in ballots, depending on state regulations. - **Polling Hours**: Polling hours vary by state, with some states allowing extended hours for voters. Its essential for voters to check local polling times to ensure they can cast their ballots. - **Provisional Ballots**: If there are questions about a voter's eligibility, they may be allowed to cast a provisional ballot. This ballot is counted once eligibility is confirmed. - **Election Day Laws**: Many states have laws that protect the rights of voters on election day, including prohibiting intimidation and ensuring access to polling places. - **Campaign Finance Regulations**: Federal and state laws regulate contributions to candidates and political parties to ensure transparency and limit the influence of money in politics. - **Political Advertising**: Campaigns must adhere to rules regarding political advertising, including disclosure requirements about funding sources and content accuracy. - **Voter Intimidation Prohibitions**: Federal laws prohibit any form of voter intimidation or coercion at polling places, ensuring a safe environment for all voters. - **Accessibility Requirements**: The Americans with Disabilities Act mandates that polling places be accessible to individuals with disabilities, ensuring equal access to the electoral process. - **Election Monitoring**: Various organizations are allowed to monitor elections to ensure compliance with laws and regulations. They help maintain transparency and accountability in the electoral process. - **Vote Counting Procedures**: States have specific procedures for counting votes, including the use of electronic voting machines and manual audits to verify results. - **Ballot Design Standards**: States must adhere to certain design standards for ballots to ensure clarity and prevent confusion among voters when casting their votes. - **Post-Election Audits**: Some states conduct post-election audits as a measure of accuracy. These audits help verify that the vote count reflects the actual ballots cast.",
+          "policy_description": "Do not generate any content that is illegal, harmful, or violates the rights of others.",
+          # "risk_categories": "",
+          "tools": [
+              {
+                  "name": "web_search",
+                  "description": "The tool web search is used to search the web for information related to finance."
+              }
+          ],
+          "max_prompts": 100,
+          "scenarios": 2,
+          "categories": 2,
+          "depth": 2,
+        },
+        "redteam_test_configurations": {
+            "harmful_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+            "custom_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+        },
+        "endpoint_configuration": {
+          "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"],
+          },
+        },
+    }
+```
+
+### [Sample Custom Redteam Model Config](https://docs.enkryptai.com/redteam-api-reference/endpoint/model-add-custom-task)
+
+```python Python
+sample_custom_redteam_model_config = {
+        "test_name": custom_redteam_test_name,
+        "dataset_configuration": {
+          "system_description": "- **Voter Eligibility**: To vote in U.S. elections, individuals must be U.S. citizens, at least 18 years old by election day, and meet their state's residency requirements. - **Voter Registration**: Most states require voters to register ahead of time, with deadlines varying widely. North Dakota is an exception, as it does not require voter registration. - **Identification Requirements**: Thirty-six states enforce voter ID laws, requiring individuals to present identification at polling places. These laws aim to prevent voter fraud but can also lead to disenfranchisement. - **Voting Methods**: Voters can typically choose between in-person voting on election day, early voting, and absentee or mail-in ballots, depending on state regulations. - **Polling Hours**: Polling hours vary by state, with some states allowing extended hours for voters. Its essential for voters to check local polling times to ensure they can cast their ballots. - **Provisional Ballots**: If there are questions about a voter's eligibility, they may be allowed to cast a provisional ballot. This ballot is counted once eligibility is confirmed. - **Election Day Laws**: Many states have laws that protect the rights of voters on election day, including prohibiting intimidation and ensuring access to polling places. - **Campaign Finance Regulations**: Federal and state laws regulate contributions to candidates and political parties to ensure transparency and limit the influence of money in politics. - **Political Advertising**: Campaigns must adhere to rules regarding political advertising, including disclosure requirements about funding sources and content accuracy. - **Voter Intimidation Prohibitions**: Federal laws prohibit any form of voter intimidation or coercion at polling places, ensuring a safe environment for all voters. - **Accessibility Requirements**: The Americans with Disabilities Act mandates that polling places be accessible to individuals with disabilities, ensuring equal access to the electoral process. - **Election Monitoring**: Various organizations are allowed to monitor elections to ensure compliance with laws and regulations. They help maintain transparency and accountability in the electoral process. - **Vote Counting Procedures**: States have specific procedures for counting votes, including the use of electronic voting machines and manual audits to verify results. - **Ballot Design Standards**: States must adhere to certain design standards for ballots to ensure clarity and prevent confusion among voters when casting their votes. - **Post-Election Audits**: Some states conduct post-election audits as a measure of accuracy. These audits help verify that the vote count reflects the actual ballots cast.",
+          "policy_description": "Do not generate any content that is illegal, harmful, or violates the rights of others.",
+          "tools": [
+              {
+                  "name": "web_search",
+                  "description": "The tool web search is used to search the web for information related to finance."
+              }
+          ],
+          "max_prompts": 100,
+          "scenarios": 2,
+          "categories": 2,
+          "depth": 2,
+        },
+        "redteam_test_configurations": {
+            "harmful_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+            "custom_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+        },
+    }
+```
+
 ## Health Checks
 
 ### [Guardrails Health](https://docs.enkryptai.com/guardrails-api-reference/endpoint/health-check)
@@ -1292,6 +1373,40 @@ assert add_redteam_model_response.message == "Redteam task has been added succes
 print(add_redteam_model_response.to_dict())
 ```
 
+### [Add a Redteam Custom Task with Endpoint Config](https://docs.enkryptai.com/redteam-api-reference/endpoint/add-custom-task)
+
+```python Python
+# Use a dictionary to configure a redteam task
+add_custom_redteam_target_response = redteam_client.add_custom_task(config=copy.deepcopy(sample_custom_redteam_target_config))
+
+# If you already saved a Code of Conduct Policy and want to use it, then instead of passing `dataset_configuration.policy_description` in the body, you can use the SDK like this:
+add_custom_redteam_target_response = redteam_client.add_custom_task(config=copy.deepcopy(sample_custom_redteam_target_config), policy_name="Code of Conduct Policy")
+
+print(add_custom_redteam_target_response)
+
+assert add_custom_redteam_target_response.message == "Task submitted successfully"
+
+# Print as a dictionary
+print(add_custom_redteam_target_response.to_dict())
+```
+
+### [Add a Redteam Custom Task with a saved model](https://docs.enkryptai.com/redteam-api-reference/endpoint/model-add-custom-task)
+
+```python Python
+# Use a dictionary to configure a redteam task
+add_custom_redteam_target_response = redteam_client.add_custom_task_with_saved_model(config=copy.deepcopy(sample_custom_redteam_model_config),model_saved_name=test_model_saved_name,model_version="v1")
+
+# If you already saved a Code of Conduct Policy and want to use it, then instead of passing `dataset_configuration.policy_description` in the body, you can use the SDK like this:
+add_custom_redteam_target_response = redteam_client.add_custom_task_with_saved_model(config=copy.deepcopy(sample_custom_redteam_model_config),model_saved_name=test_model_saved_name,model_version="v1",policy_name="Code of Conduct Policy")
+
+print(add_custom_redteam_target_response)
+
+assert add_custom_redteam_target_response.message == "Task submitted successfully"
+
+# Print as a dictionary
+print(add_custom_redteam_target_response.to_dict())
+```
+
 ### [Get Redteam Task Status](https://docs.enkryptai.com/redteam-api-reference/endpoint/get-task-status)
 
 ```python Python

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/README.md

@@ -24,6 +24,8 @@ Also see the API documentation at [https://docs.enkryptai.com](https://docs.enkr
   - [Sample Redteam Model Health Config](#sample-redteam-model-health-config)
   - [Sample Redteam Target Config](#sample-redteam-target-config)
   - [Sample Redteam Model Config](#sample-redteam-model-config)
+  - [Sample Custom Redteam Target Config](#sample-custom-redteam-target-config)
+  - [Sample Custom Redteam Model Config](#sample-custom-redteam-model-config)
 - [Health Checks](#health-checks)
   - [Guardrails Health](#guardrails-health)
   - [Guardrails Status](#guardrails-status)
@@ -83,6 +85,8 @@ Also see the API documentation at [https://docs.enkryptai.com](https://docs.enkr
 - [Redteam](#redteam)
   - [Add a Redteam Task with Target Model Config](#add-a-redteam-task-with-target-model-config)
   - [Add a Redteam Task with a saved model](#add-a-redteam-task-with-a-saved-model)
+  - [Add a Redteam Custom Task with Endpoint Config](#add-a-redteam-custom-task-with-endpoint-config)
+  - [Add a Redteam Custom Task with a saved model](#add-a-redteam-custom-task-with-a-saved-model)
   - [Get Redteam Task Status](#get-redteam-task-status)
   - [Get Redteam Task](#get-redteam-task)
   - [List Redteam Tasks](#list-redteam-tasks)
@@ -170,6 +174,7 @@ pii_anonymized_text = None
 dataset_name = f"TestElectionDataset-{str(uuid.uuid4())[:8]}"
 
 redteam_test_name = f"Redteam Test {str(uuid.uuid4())[:6]}"
+custom_redteam_test_name = f"Custom Redteam Test {str(uuid.uuid4())[:6]}"
 
 model_name = "gpt-4o-mini"
 model_provider = "openai"
@@ -399,6 +404,82 @@ sample_redteam_model_config = {
     }
 ```
 
+### [Sample Custom Redteam Target Config](https://docs.enkryptai.com/redteam-api-reference/endpoint/add-custom-task)
+
+```python Python
+sample_custom_redteam_target_config = {
+        "test_name": custom_redteam_test_name,
+        "dataset_configuration": {
+          "system_description": "- **Voter Eligibility**: To vote in U.S. elections, individuals must be U.S. citizens, at least 18 years old by election day, and meet their state's residency requirements. - **Voter Registration**: Most states require voters to register ahead of time, with deadlines varying widely. North Dakota is an exception, as it does not require voter registration. - **Identification Requirements**: Thirty-six states enforce voter ID laws, requiring individuals to present identification at polling places. These laws aim to prevent voter fraud but can also lead to disenfranchisement. - **Voting Methods**: Voters can typically choose between in-person voting on election day, early voting, and absentee or mail-in ballots, depending on state regulations. - **Polling Hours**: Polling hours vary by state, with some states allowing extended hours for voters. Its essential for voters to check local polling times to ensure they can cast their ballots. - **Provisional Ballots**: If there are questions about a voter's eligibility, they may be allowed to cast a provisional ballot. This ballot is counted once eligibility is confirmed. - **Election Day Laws**: Many states have laws that protect the rights of voters on election day, including prohibiting intimidation and ensuring access to polling places. - **Campaign Finance Regulations**: Federal and state laws regulate contributions to candidates and political parties to ensure transparency and limit the influence of money in politics. - **Political Advertising**: Campaigns must adhere to rules regarding political advertising, including disclosure requirements about funding sources and content accuracy. - **Voter Intimidation Prohibitions**: Federal laws prohibit any form of voter intimidation or coercion at polling places, ensuring a safe environment for all voters. - **Accessibility Requirements**: The Americans with Disabilities Act mandates that polling places be accessible to individuals with disabilities, ensuring equal access to the electoral process. - **Election Monitoring**: Various organizations are allowed to monitor elections to ensure compliance with laws and regulations. They help maintain transparency and accountability in the electoral process. - **Vote Counting Procedures**: States have specific procedures for counting votes, including the use of electronic voting machines and manual audits to verify results. - **Ballot Design Standards**: States must adhere to certain design standards for ballots to ensure clarity and prevent confusion among voters when casting their votes. - **Post-Election Audits**: Some states conduct post-election audits as a measure of accuracy. These audits help verify that the vote count reflects the actual ballots cast.",
+          "policy_description": "Do not generate any content that is illegal, harmful, or violates the rights of others.",
+          # "risk_categories": "",
+          "tools": [
+              {
+                  "name": "web_search",
+                  "description": "The tool web search is used to search the web for information related to finance."
+              }
+          ],
+          "max_prompts": 100,
+          "scenarios": 2,
+          "categories": 2,
+          "depth": 2,
+        },
+        "redteam_test_configurations": {
+            "harmful_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+            "custom_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+        },
+        "endpoint_configuration": {
+          "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"],
+          },
+        },
+    }
+```
+
+### [Sample Custom Redteam Model Config](https://docs.enkryptai.com/redteam-api-reference/endpoint/model-add-custom-task)
+
+```python Python
+sample_custom_redteam_model_config = {
+        "test_name": custom_redteam_test_name,
+        "dataset_configuration": {
+          "system_description": "- **Voter Eligibility**: To vote in U.S. elections, individuals must be U.S. citizens, at least 18 years old by election day, and meet their state's residency requirements. - **Voter Registration**: Most states require voters to register ahead of time, with deadlines varying widely. North Dakota is an exception, as it does not require voter registration. - **Identification Requirements**: Thirty-six states enforce voter ID laws, requiring individuals to present identification at polling places. These laws aim to prevent voter fraud but can also lead to disenfranchisement. - **Voting Methods**: Voters can typically choose between in-person voting on election day, early voting, and absentee or mail-in ballots, depending on state regulations. - **Polling Hours**: Polling hours vary by state, with some states allowing extended hours for voters. Its essential for voters to check local polling times to ensure they can cast their ballots. - **Provisional Ballots**: If there are questions about a voter's eligibility, they may be allowed to cast a provisional ballot. This ballot is counted once eligibility is confirmed. - **Election Day Laws**: Many states have laws that protect the rights of voters on election day, including prohibiting intimidation and ensuring access to polling places. - **Campaign Finance Regulations**: Federal and state laws regulate contributions to candidates and political parties to ensure transparency and limit the influence of money in politics. - **Political Advertising**: Campaigns must adhere to rules regarding political advertising, including disclosure requirements about funding sources and content accuracy. - **Voter Intimidation Prohibitions**: Federal laws prohibit any form of voter intimidation or coercion at polling places, ensuring a safe environment for all voters. - **Accessibility Requirements**: The Americans with Disabilities Act mandates that polling places be accessible to individuals with disabilities, ensuring equal access to the electoral process. - **Election Monitoring**: Various organizations are allowed to monitor elections to ensure compliance with laws and regulations. They help maintain transparency and accountability in the electoral process. - **Vote Counting Procedures**: States have specific procedures for counting votes, including the use of electronic voting machines and manual audits to verify results. - **Ballot Design Standards**: States must adhere to certain design standards for ballots to ensure clarity and prevent confusion among voters when casting their votes. - **Post-Election Audits**: Some states conduct post-election audits as a measure of accuracy. These audits help verify that the vote count reflects the actual ballots cast.",
+          "policy_description": "Do not generate any content that is illegal, harmful, or violates the rights of others.",
+          "tools": [
+              {
+                  "name": "web_search",
+                  "description": "The tool web search is used to search the web for information related to finance."
+              }
+          ],
+          "max_prompts": 100,
+          "scenarios": 2,
+          "categories": 2,
+          "depth": 2,
+        },
+        "redteam_test_configurations": {
+            "harmful_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+            "custom_test": {
+                "sample_percentage": 2,
+                "attack_methods": {"basic": ["basic"]},
+            },
+        },
+    }
+```
+
 ## Health Checks
 
 ### [Guardrails Health](https://docs.enkryptai.com/guardrails-api-reference/endpoint/health-check)
@@ -1269,6 +1350,40 @@ assert add_redteam_model_response.message == "Redteam task has been added succes
 print(add_redteam_model_response.to_dict())
 ```
 
+### [Add a Redteam Custom Task with Endpoint Config](https://docs.enkryptai.com/redteam-api-reference/endpoint/add-custom-task)
+
+```python Python
+# Use a dictionary to configure a redteam task
+add_custom_redteam_target_response = redteam_client.add_custom_task(config=copy.deepcopy(sample_custom_redteam_target_config))
+
+# If you already saved a Code of Conduct Policy and want to use it, then instead of passing `dataset_configuration.policy_description` in the body, you can use the SDK like this:
+add_custom_redteam_target_response = redteam_client.add_custom_task(config=copy.deepcopy(sample_custom_redteam_target_config), policy_name="Code of Conduct Policy")
+
+print(add_custom_redteam_target_response)
+
+assert add_custom_redteam_target_response.message == "Task submitted successfully"
+
+# Print as a dictionary
+print(add_custom_redteam_target_response.to_dict())
+```
+
+### [Add a Redteam Custom Task with a saved model](https://docs.enkryptai.com/redteam-api-reference/endpoint/model-add-custom-task)
+
+```python Python
+# Use a dictionary to configure a redteam task
+add_custom_redteam_target_response = redteam_client.add_custom_task_with_saved_model(config=copy.deepcopy(sample_custom_redteam_model_config),model_saved_name=test_model_saved_name,model_version="v1")
+
+# If you already saved a Code of Conduct Policy and want to use it, then instead of passing `dataset_configuration.policy_description` in the body, you can use the SDK like this:
+add_custom_redteam_target_response = redteam_client.add_custom_task_with_saved_model(config=copy.deepcopy(sample_custom_redteam_model_config),model_saved_name=test_model_saved_name,model_version="v1",policy_name="Code of Conduct Policy")
+
+print(add_custom_redteam_target_response)
+
+assert add_custom_redteam_target_response.message == "Task submitted successfully"
+
+# Print as a dictionary
+print(add_custom_redteam_target_response.to_dict())
+```
+
 ### [Get Redteam Task Status](https://docs.enkryptai.com/redteam-api-reference/endpoint/get-task-status)
 
 ```python Python

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/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.7", # Update this for new versions
+    version="1.0.8", # 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.7 → enkryptai_sdk-1.0.8}/src/enkryptai_sdk/ai_proxy.py

@@ -60,7 +60,7 @@ class AIProxyClient(BaseClient):
                     # # As json.loads is not working with literal string representation of dict
                     # parsed_error = json.loads(error_message.replace("'", '"'))
                     # # Convert Python string representation to proper dict
-                    parsed_error = ast.literal_eval(response["error"])
+                    parsed_error = ast.literal_eval(f"API Error: {str(response)}")
                     # # Preserve both error and enkrypt_policy_detections fields
                     response["error"] = parsed_error.get("error", parsed_error)
                     if "enkrypt_policy_detections" in parsed_error:
@@ -84,6 +84,6 @@ class AIProxyClient(BaseClient):
                     # Fallback to direct error if error object can't be parsed
                     # print("Failed to parse error response: ", response)
                     return ChatCompletionDirectErrorResponse.from_dict(response)
-            raise AIProxyClientError(response["error"])
+            raise AIProxyClientError(f"API Error: {str(response)}")
 
         return ChatCompletionResponse.from_dict(response)

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/src/enkryptai_sdk/datasets.py

@@ -19,6 +19,32 @@ class DatasetClient(BaseClient):
     def __init__(self, api_key: str, base_url: str = "https://api.enkryptai.com"):
         super().__init__(api_key, base_url)
 
+    @staticmethod
+    def prepare_dataset_payload(config: DatasetConfig | dict, is_custom: bool = False) -> dict:
+        """
+        Prepare the payload for dataset operations from a config object.
+        
+        Args:
+            config (Union[DatasetConfig, dict]): Configuration object or dictionary containing dataset details
+            
+        Returns:
+            dict: Processed payload ready for API submission
+        """
+        if isinstance(config, dict):
+            config = DatasetConfig.from_dict(config)
+        
+        payload = config.to_dict()
+
+        if not is_custom:
+            # Remove empty tools configuration
+            if (payload.get("tools") is None or 
+                payload["tools"] == [] or 
+                payload["tools"] == [{}] or
+                payload["tools"] == [{"name": "", "description": ""}]):
+                del payload["tools"]
+        
+        return payload
+
     def add_dataset(self, config: DatasetConfig):
         """
         Add a new dataset to the system.
@@ -34,14 +60,7 @@ class DatasetClient(BaseClient):
         if isinstance(config, dict):
             config = DatasetConfig.from_dict(config)
 
-        payload = config.to_dict()
-
-        # If payload["tools"] is None or is an empty list [] or [{}], remove it from the payload
-        if (payload.get("tools") is None or 
-            payload["tools"] == [] or 
-            payload["tools"] == [{}] or
-            payload["tools"] == [{"name": "", "description": ""}]):
-            del payload["tools"]
+        payload = self.prepare_dataset_payload(config)
 
         # Print payload
         # print(f"\nAdd Dataset Payload: {payload}")
@@ -54,7 +73,7 @@ class DatasetClient(BaseClient):
         # print(f"\nAdd Dataset Response: {response}")
 
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         return DatasetAddTaskResponse.from_dict(response)
     
     def get_dataset_task_status(self, dataset_name: str):
@@ -70,7 +89,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/task-status", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetTaskStatus.from_dict(response)
 
@@ -87,7 +106,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-task", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetTask.from_dict(response)
 
@@ -104,7 +123,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-datacard", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetCard.from_dict(response)
 
@@ -121,7 +140,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-dataset", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetResponse.from_dict(response)
 
@@ -138,7 +157,7 @@ class DatasetClient(BaseClient):
         headers = {"X-Enkrypt-Dataset": dataset_name}
         response = self._request("GET", "/datasets/get-summary", headers=headers)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         response["dataset_name"] = dataset_name
         return DatasetSummary.from_dict(response)
 
@@ -158,5 +177,5 @@ class DatasetClient(BaseClient):
             url += f"?status={status}"
         response = self._request("GET", url)
         if response.get("error"):
-            raise DatasetClientError(response["error"])
+            raise DatasetClientError(f"API Error: {str(response)}")
         return DatasetCollection.from_dict(response)

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/src/enkryptai_sdk/deployments.py

@@ -46,7 +46,7 @@ class DeploymentClient(BaseClient):
             "POST", "/deployments/add-deployment", headers=headers, json=payload
         )
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return DeploymentAddTaskResponse.from_dict(response)
 
     def get_deployment(self, deployment_name: str, refresh_cache: bool = False):
@@ -64,7 +64,7 @@ class DeploymentClient(BaseClient):
         headers["X-Enkrypt-Refresh-Cache"] = "true" if refresh_cache else "false"
         response = self._request("GET", "/deployments/get-deployment", headers=headers)
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return GetDeploymentResponse.from_dict(response)
     
     def modify_deployment(self, deployment_name: str, config: DeploymentInput):
@@ -88,7 +88,7 @@ class DeploymentClient(BaseClient):
             "PATCH", "/deployments/modify-deployment", headers=headers, json=payload
         )
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return ModifyDeploymentResponse.from_dict(response)
 
     def delete_deployment(self, deployment_name: str):
@@ -104,7 +104,7 @@ class DeploymentClient(BaseClient):
         headers = {"X-Enkrypt-Deployment": deployment_name}
         response = self._request("DELETE", "/deployments/delete-deployment", headers=headers)
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return DeleteDeploymentResponse.from_dict(response)
 
     def list_deployments(self):
@@ -117,5 +117,5 @@ class DeploymentClient(BaseClient):
 
         response = self._request("GET", "/deployments/list-deployments")
         if response.get("error"):
-            raise DeploymentClientError(response["error"])
+            raise DeploymentClientError(f"API Error: {str(response)}")
         return DeploymentCollection.from_dict(response)

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/src/enkryptai_sdk/dto/__init__.py

@@ -22,6 +22,8 @@ __all__ = [
     "Metadata",
     "DEFAULT_REDTEAM_CONFIG",
     "DEFAULT_REDTEAM_CONFIG_WITH_SAVED_MODEL",
+    "DEFAULT_CUSTOM_REDTEAM_CONFIG",
+    "DEFAULT_CUSTOM_REDTEAM_CONFIG_WITH_SAVED_MODEL",
     "ADVANCED_REDTEAM_TESTS",
     "DETAIL_MODEL_CONFIG",
     "DatasetConfig",

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/src/enkryptai_sdk/dto/datasets.py

@@ -21,17 +21,20 @@ class Tool(BaseDTO):
             "description": self.description
         }
 
+
 @dataclass
 class DatasetConfig(BaseDTO):
-    dataset_name: str
-    system_description: str
-    policy_description: str = ""
-    tools: List[Tool] = field(default_factory=list)
-    info_pdf_url: str = ""
+    system_description: str = None
+    dataset_name: str = None
+    policy_description: str = None
+    risk_categories: str = None
+    tools: List[Tool] = None
+    info_pdf_url: str = None
     max_prompts: int = 100
     scenarios: int = 2
     categories: int = 2
     depth: int = 2
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass
@@ -51,6 +54,7 @@ class DatasetDataPoint(BaseDTO):
     category: str
     prompt: str
     source: str
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass

{enkryptai_sdk-1.0.7 → enkryptai_sdk-1.0.8}/src/enkryptai_sdk/dto/models.py

@@ -142,8 +142,8 @@ class CustomHeader(BaseDTO):
 
 @dataclass
 class ModelConfigDetails(BaseDTO):
-    model_id: str = ""
-    model_source: str = ""
+    model_id: str = None
+    model_source: str = None
     # model_provider: str = "openai"
     model_provider: ModelProviders = ModelProviders.OPENAI
     model_api_value: str = ""
@@ -259,14 +259,14 @@ class ModelConfigDetails(BaseDTO):
 
 @dataclass
 class ModelConfig(BaseDTO):
-    created_at: str = ""
-    updated_at: str = ""
-    model_id: str = ""
-    model_saved_name: str = "Model Name"
-    model_version: str = "v1"
+    created_at: str = None
+    updated_at: str = None
+    model_id: str = None
+    model_saved_name: str = None
+    model_version: str = None
     testing_for: str = "foundationModels"
     # modality: Modality = Modality.TEXT
-    project_name: str = ""
+    project_name: str = None
     model_name: Optional[str] = "gpt-4o-mini"
     certifications: List[str] = field(default_factory=list)
     model_config: ModelConfigDetails = field(default_factory=ModelConfigDetails)