enkryptai-sdk 1.0.26__tar.gz → 1.0.27__tar.gz

{enkryptai_sdk-1.0.26 → enkryptai_sdk-1.0.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: enkryptai-sdk
-Version: 1.0.26
+Version: 1.0.27
 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.26 → enkryptai_sdk-1.0.27}/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.26", # Update this for new versions
+    version="1.0.27", # 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.26 → enkryptai_sdk-1.0.27}/src/enkryptai_sdk/dto/__init__.py

@@ -97,6 +97,7 @@ __all__ = [
     "GuardrailsListPoliciesResponse",
     "GuardrailsPolicyAtomizerRequest",
     "GuardrailsPolicyAtomizerResponse",
+    "GuardrailsScanPdfResponse",
     "CoCPolicyRequest",
     "CoCPolicyData",
     "CoCPolicyResponse",

{enkryptai_sdk-1.0.26 → enkryptai_sdk-1.0.27}/src/enkryptai_sdk/dto/guardrails.py

@@ -1751,3 +1751,79 @@ class GuardrailsScanUrlResponse(BaseDTO):
             
         return f"URL Scan Result for {self.url}\nStatus: {status}\n{violation_str}"
 
+
+# -------------------------------------
+# Guardrails Scan PDF
+# -------------------------------------
+
+
+@dataclass
+class ScanPdfDetails(BaseDTO):
+    total_pages: int = 0
+    pages_with_detections: int = 0
+    detections: List[Dict[str, Any]] = field(default_factory=list)
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ScanPdfDetails":
+        return cls(
+            total_pages=data.get("total_pages", 0),
+            pages_with_detections=data.get("pages_with_detections", 0),
+            detections=data.get("detections", [])
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "total_pages": self.total_pages,
+            "pages_with_detections": self.pages_with_detections,
+            "detections": self.detections
+        }
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class GuardrailsScanPdfResponse(BaseDTO):
+    summary: Dict[str, int] = field(default_factory=dict)
+    details: ScanPdfDetails = field(default_factory=ScanPdfDetails)
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GuardrailsScanPdfResponse":
+        return cls(
+            summary=data.get("summary", {}),
+            details=ScanPdfDetails.from_dict(data.get("details", {}))
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "summary": self.summary,
+            "details": self.details.to_dict()
+        }
+        result.update(self._extra_fields)
+        return result
+
+    def has_detections(self) -> bool:
+        return self.details.pages_with_detections > 0 or len(self.details.detections) > 0
+
+    def is_safe(self) -> bool:
+        return not self.has_detections()
+
+    def is_attack(self) -> bool:
+        return self.has_detections()
+
+    def __str__(self) -> str:
+        status = "UNSAFE" if self.has_detections() else "SAFE"
+        flagged = [k for k, v in self.summary.items() if v > 0]
+        if flagged:
+            detection_str = f"Detections: {', '.join(flagged)}"
+        else:
+            detection_str = "No detections"
+        return (
+            f"PDF Scan Result\n"
+            f"Status: {status}\n"
+            f"Total Pages: {self.details.total_pages}\n"
+            f"Pages With Detections: {self.details.pages_with_detections}\n"
+            f"{detection_str}"
+        )
+

{enkryptai_sdk-1.0.26 → enkryptai_sdk-1.0.27}/src/enkryptai_sdk/guardrails.py

@@ -31,7 +31,8 @@ from .dto import (
     GuardrailsListPoliciesResponse,
     GuardrailsPolicyAtomizerRequest,
     GuardrailsPolicyAtomizerResponse,
-    GuardrailsScanUrlResponse
+    GuardrailsScanUrlResponse,
+    GuardrailsScanPdfResponse
 )
 
 # ---------------------------------------
@@ -292,16 +293,20 @@ class GuardrailsClient(BaseClient):
 
     def scan_url(self, url, config=None):
         """
-        Scan a URL for security threats including injection attacks and policy violations.
+        Scan a URL for security threats using any combination of available detectors.
+
+        Supports all detectors available in the detect endpoint, including:
+        injection_attack, policy_violation, toxicity, nsfw, pii, bias,
+        keyword_detector, copyright_ip, system_prompt, sponge_attack, topic_detector.
 
         Parameters:
         - url (str): The URL to scan and analyze.
         - config (dict or GuardrailsConfig, optional): A configuration for detectors.
-        If a GuardrailsConfig instance is provided, its underlying dictionary will be used.
-        If not provided, defaults to injection attack and policy violation detection.
+          If a GuardrailsConfig instance is provided, its underlying dictionary will be used.
+          If not provided, defaults to injection attack and policy violation detection.
 
         Returns:
-        - Response from the API.
+        - GuardrailsScanUrlResponse: Response from the API containing scan results.
         """
         # Use default config if none provided
         if config is None:
@@ -358,6 +363,75 @@ class GuardrailsClient(BaseClient):
         except Exception as e:
             raise GuardrailsClientError(str(e))
 
+    def scan_pdf(self, file, config=None):
+        """
+        Scan a PDF for security threats using any combination of available detectors.
+
+        Supports all detectors available in the detect endpoint, including:
+        injection_attack, policy_violation, toxicity, nsfw, pii, bias,
+        keyword_detector, copyright_ip, system_prompt, sponge_attack, topic_detector.
+
+        Parameters:
+        - file (str): Base64-encoded PDF string.
+        - config (dict or GuardrailsConfig, optional): A configuration for detectors.
+          If a GuardrailsConfig instance is provided, its underlying dictionary will be used.
+          If not provided, defaults to injection attack and policy violation detection.
+
+        Returns:
+        - GuardrailsScanPdfResponse: Response from the API containing scan results.
+        """
+        if config is None:
+            config = {
+                "injection_attack": {
+                    "enabled": True
+                },
+                "policy_violation": {
+                    "enabled": True,
+                    "policy_text": "Detect any malicious text or injection attacks",
+                    "need_explanation": True
+                }
+            }
+
+        if hasattr(config, "as_dict"):
+            config = config.as_dict()
+        if hasattr(config, "to_dict"):
+            config = config.to_dict()
+
+        payload = {
+            "file": file,
+            "detectors": config
+        }
+
+        try:
+            response = self._request("POST", "/guardrails/scan-pdf", json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(f"API Error: {str(response)}")
+            return GuardrailsScanPdfResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+
+    def policy_scan_pdf(self, policy_name, file):
+        """
+        Apply a specific policy to scan a PDF for security threats.
+
+        Parameters:
+        - policy_name (str): Name of the policy to apply.
+        - file (str): Base64-encoded PDF string.
+
+        Returns:
+        - GuardrailsScanPdfResponse: Response from the API containing scan results.
+        """
+        headers = {"X-Enkrypt-Policy": policy_name}
+        payload = {"file": file}
+
+        try:
+            response = self._request("POST", "/guardrails/policy/scan-pdf", headers=headers, json=payload)
+            if response.get("error"):
+                raise GuardrailsClientError(f"API Error: {str(response)}")
+            return GuardrailsScanPdfResponse.from_dict(response)
+        except Exception as e:
+            raise GuardrailsClientError(str(e))
+
     # ----------------------------
     # Guardrails Policy Endpoints
     # ----------------------------

{enkryptai_sdk-1.0.26 → enkryptai_sdk-1.0.27}/src/enkryptai_sdk.egg-info/PKG-INFO

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

@@ -2,7 +2,7 @@
 
 import os
 import time
-import uuid
+# import uuid  # Only used for dataset tests which are now commented out
 import pytest
 from fixtures import *
 from enkryptai_sdk import *
@@ -28,8 +28,8 @@ pii_de_anonymized_text = "My email is example@example.com. My phone number is 12
 pii_key = None
 pii_anonymized_text = None
 
-dataset_name = None
-picked_dataset_name = None
+# dataset_name = None
+# picked_dataset_name = None
 
 redteam_test_name = None
 redteam_model_test_name = None
@@ -85,10 +85,10 @@ def deployment_client():
     # You'll want to use a test API key here
     return DeploymentClient(api_key=ENKRYPT_API_KEY, base_url=ENKRYPT_BASE_URL)
 
-@pytest.fixture
-def dataset_client():
-    # You'll want to use a test API key here
-    return DatasetClient(api_key=ENKRYPT_API_KEY, base_url=ENKRYPT_BASE_URL)
+# @pytest.fixture
+# def dataset_client():
+#     # You'll want to use a test API key here
+#     return DatasetClient(api_key=ENKRYPT_API_KEY, base_url=ENKRYPT_BASE_URL)
 
 # Getting from fixtures.py
 # @pytest.fixture
@@ -250,29 +250,29 @@ def sample_deployment_config():
         },
     }
 
-@pytest.fixture
-def sample_dataset_config():
-    print("\nCreating sample dataset config")
-    global dataset_name
-    dataset_name = f"TestElectionDataset-{str(uuid.uuid4())[:8]}"
-    print("\nDataset Name: ", dataset_name)
-    return {
-        "dataset_name": dataset_name,
-        # "dataset_name": "TestElectionDataset",
-        "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": "",
-        "tools": [
-            {
-                "name": "web_search",
-                "description": "The tool web search is used to search the web for information related to finance."
-            }
-        ],
-        "info_pdf_url": "",
-        "max_prompts": 100,
-        "scenarios": 2,
-        "categories": 2,
-        "depth": 2,
-    }
+# @pytest.fixture
+# def sample_dataset_config():
+#     print("\nCreating sample dataset config")
+#     global dataset_name
+#     dataset_name = f"TestElectionDataset-{str(uuid.uuid4())[:8]}"
+#     print("\nDataset Name: ", dataset_name)
+#     return {
+#         "dataset_name": dataset_name,
+#         # "dataset_name": "TestElectionDataset",
+#         "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": "",
+#         "tools": [
+#             {
+#                 "name": "web_search",
+#                 "description": "The tool web search is used to search the web for information related to finance."
+#             }
+#         ],
+#         "info_pdf_url": "",
+#         "max_prompts": 100,
+#         "scenarios": 2,
+#         "categories": 2,
+#         "depth": 2,
+#     }
 
 # Getting Red team fixtures from fixtures.py
 
@@ -325,34 +325,34 @@ def get_task_name_from_list(redteam_client, status=None):
     return test_name
 
 
-def get_dataset_name_from_list(dataset_client, status=None):
-    """Helper function to get a dataset name from the datasets list.
-    
-    Args:
-        dataset_client: The DatasetClient instance
-        status: Optional status filter (e.g., "Finished", "Running")
-    
-    Returns:
-        str: A dataset name from the datasets list
-    """
-    print(f"\nFetching dataset with status: {status if status else 'Finished'}")
-    datasets = dataset_client.list_datasets(status=status)
-    datasets_dict = datasets.to_dict()
-    print(f"\nDatasets list retrieved with {len(datasets_dict.get('datasets', []))} datasets")
-    
-    if not datasets_dict.get("datasets"):
-        # If no datasets with specified status, try without status filter
-        if status:
-            print(f"\nNo datasets with status '{status}', fetching any dataset")
-            datasets = dataset_client.list_datasets()
-            datasets_dict = datasets.to_dict()
-    
-    if not datasets_dict.get("datasets"):
-        return None
-    
-    dataset_name = datasets_dict["datasets"][0]
-    print(f"\nSelected dataset: {dataset_name}")
-    return dataset_name
+# def get_dataset_name_from_list(dataset_client, status=None):
+#     """Helper function to get a dataset name from the datasets list.
+#
+#     Args:
+#         dataset_client: The DatasetClient instance
+#         status: Optional status filter (e.g., "Finished", "Running")
+#
+#     Returns:
+#         str: A dataset name from the datasets list
+#     """
+#     print(f"\nFetching dataset with status: {status if status else 'Finished'}")
+#     datasets = dataset_client.list_datasets(status=status)
+#     datasets_dict = datasets.to_dict()
+#     print(f"\nDatasets list retrieved with {len(datasets_dict.get('datasets', []))} datasets")
+#
+#     if not datasets_dict.get("datasets"):
+#         # If no datasets with specified status, try without status filter
+#         if status:
+#             print(f"\nNo datasets with status '{status}', fetching any dataset")
+#             datasets = dataset_client.list_datasets()
+#             datasets_dict = datasets.to_dict()
+#
+#     if not datasets_dict.get("datasets"):
+#         return None
+#
+#     dataset_name = datasets_dict["datasets"][0]
+#     print(f"\nSelected dataset: {dataset_name}")
+#     return dataset_name
 
 
 # ----------------------------
@@ -475,13 +475,13 @@ def test_batch_detect(guardrails_client, sample_detectors):
 def test_scan_url(guardrails_client):
     print("\n\nTesting scan url method")
     # Test the scan url method
-    response = guardrails_client.scan_url(url="https://www.digitalocean.com/community/tutorials/how-to-install-poetry-to-manage-python-dependencies-on-ubuntu-22-04")
+    response = guardrails_client.scan_url(url="https://example.com/?q=<script>alert('XSS')</script>")
     print("\nResponse from scan url: ", response)
     print("\nResponse data type: ", type(response))
     assert response is not None
-    assert response.url == "https://www.digitalocean.com/community/tutorials/how-to-install-poetry-to-manage-python-dependencies-on-ubuntu-22-04"
-    assert response.violations is not None
-    assert len(response.violations) > 0
+    assert response.url == "https://example.com/?q=<script>alert('XSS')</script>"
+    # assert response.violations is not None
+    # assert len(response.violations) > 0
 
 
 def test_pii_request(guardrails_client):
@@ -642,13 +642,47 @@ def test_batch_detect_with_policy(guardrails_client):
 def test_scan_url_with_policy(guardrails_client):
     print("\n\nTesting scan url method with policy")
     # Test the scan url method
-    response = guardrails_client.policy_scan_url(policy_name=test_guardrails_policy_name, url="https://www.digitalocean.com/community/tutorials/how-to-install-poetry-to-manage-python-dependencies-on-ubuntu-22-04")
+    response = guardrails_client.policy_scan_url(policy_name=test_guardrails_policy_name, url="https://example.com/?q=<script>alert('XSS')</script>")
     print("\nResponse from scan url: ", response)
     print("\nResponse data type: ", type(response))
     assert response is not None
-    assert response.url == "https://www.digitalocean.com/community/tutorials/how-to-install-poetry-to-manage-python-dependencies-on-ubuntu-22-04"
-    assert response.violations is not None
-    assert len(response.violations) > 0
+    assert response.url == "https://example.com/?q=<script>alert('XSS')</script>"
+    # assert response.violations is not None
+    # assert len(response.violations) > 0
+
+
+def test_scan_pdf(guardrails_client):
+    print("\n\nTesting scan pdf method")
+    import base64
+    # Create a minimal valid PDF for testing
+    pdf_path = "../docs/healthcare_guidelines.pdf"
+    with open(pdf_path, "rb") as f:
+        pdf_base64 = base64.b64encode(f.read()).decode("utf-8")
+    response = guardrails_client.scan_pdf(file=pdf_base64)
+    print("\nResponse from scan pdf: ", response)
+    print("\nResponse data type: ", type(response))
+    assert response is not None
+    assert hasattr(response, "summary")
+    assert hasattr(response, "details")
+    assert response.details.total_pages >= 1
+
+
+def test_scan_pdf_with_policy(guardrails_client):
+    print("\n\nTesting scan pdf method with policy")
+    import base64
+    pdf_path = "../docs/healthcare_guidelines.pdf"
+    with open(pdf_path, "rb") as f:
+        pdf_base64 = base64.b64encode(f.read()).decode("utf-8")
+    response = guardrails_client.policy_scan_pdf(
+        policy_name=test_guardrails_policy_name,
+        file=pdf_base64
+    )
+    print("\nResponse from scan pdf with policy: ", response)
+    print("\nResponse data type: ", type(response))
+    assert response is not None
+    assert hasattr(response, "summary")
+    assert hasattr(response, "details")
+    assert response.details.total_pages >= 1
 
 
 def test_atomize_policy(guardrails_client):
@@ -982,113 +1016,113 @@ def test_401_err_chat(ai_proxy_err_client, sample_chat_body):
 # Tests for Dataset Endpoints
 # ----------------------------
 
-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)
-    assert response is not None
-    assert hasattr(response, "task_id")
-    assert hasattr(response, "message")
-    response.message == "Dataset 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)
+# 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)
+#     assert response is not None
+#     assert hasattr(response, "task_id")
+#     assert hasattr(response, "message")
+#     response.message == "Dataset 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)
 
 
-def test_list_datasets(dataset_client):
-    print("\n\nTesting list_datasets")
-    # Test the list_datasets method
-    datasets = dataset_client.list_datasets(status="Finished")
-    datasets_dict = datasets.to_dict()
-    print("\nList Datasets Response: ", datasets_dict)
-    assert datasets_dict is not None
-    assert isinstance(datasets_dict, dict)
-    
-    # Get a dataset name using our helper function
-    global picked_dataset_name
-    if picked_dataset_name is None:
-        picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
-    assert picked_dataset_name is not None
-    print(f"\nSelected dataset for further tests: {picked_dataset_name}")
-
-
-def test_get_dataset_task(dataset_client):
-    print("\n\nTesting get_dataset_task")
-    global picked_dataset_name
-    if picked_dataset_name is None:
-        picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
-    assert picked_dataset_name is not None
-    print(f"\nSelected dataset for further tests: {picked_dataset_name}")
-
-    # Now test the get_dataset_task method
-    dataset_task = dataset_client.get_dataset_task(dataset_name=picked_dataset_name)
-    print("\nDataset Task: ", dataset_task)
-    assert dataset_task is not None
-    assert hasattr(dataset_task, "dataset_name")
-    assert dataset_task.dataset_name == picked_dataset_name
-    assert hasattr(dataset_task, "data")
-    data = dataset_task.data
-    assert data is not None
-    assert hasattr(data, "status")
-    status = data.status
-    assert status is not None
-    print("\nDataset Task Status: ", status)
-
-
-def test_get_dataset_task_status(dataset_client):
-    print("\n\nTesting get_dataset_task_status")
-    global picked_dataset_name
-    if picked_dataset_name is None:
-        picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
-    assert picked_dataset_name is not None
-    print(f"\nSelected dataset for further tests: {picked_dataset_name}")
-
-    # Now test the get_dataset_task_status method
-    dataset_task_status = dataset_client.get_dataset_task_status(dataset_name=picked_dataset_name)
-    print("\nDataset Task Status: ", dataset_task_status)
-    assert dataset_task_status is not None
-    assert hasattr(dataset_task_status, "dataset_name")
-    assert dataset_task_status.dataset_name == picked_dataset_name
-    assert hasattr(dataset_task_status, "status")
-
-
-def test_get_datacard(dataset_client):
-    print("\n\nTesting get_datacard")
-    global picked_dataset_name
-    if picked_dataset_name is None:
-        picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
-    assert picked_dataset_name is not None
-    print(f"\nSelected dataset for further tests: {picked_dataset_name}")
-
-    # Now test the get_datacard method
-    datacard = dataset_client.get_datacard(dataset_name=picked_dataset_name)
-    print("\nDatacard: ", datacard)
-    assert datacard is not None
-    # # Dataset might not be generated yet
-    # # TODO: How to handle this?
-    # assert hasattr(datacard, "dataset_name")
-    # assert datacard.dataset_name == picked_dataset_name
-    # assert hasattr(datacard, "description")
-
-
-def test_get_summary(dataset_client):
-    print("\n\nTesting get_summary")
-    global picked_dataset_name
-    if picked_dataset_name is None:
-        picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
-    assert picked_dataset_name is not None
-    print(f"\nSelected dataset for further tests: {picked_dataset_name}")
-
-    # Now test the get_summary method
-    summary = dataset_client.get_summary(dataset_name=picked_dataset_name)
-    print("\nsummary: ", summary)
-    assert summary is not None
-    # # Dataset might not be generated yet
-    # # TODO: How to handle this?
-    # assert hasattr(summary, "dataset_name")
-    # assert summary.dataset_name == picked_dataset_name
-    # assert hasattr(summary, "test_types")
+# def test_list_datasets(dataset_client):
+#     print("\n\nTesting list_datasets")
+#     # Test the list_datasets method
+#     datasets = dataset_client.list_datasets(status="Finished")
+#     datasets_dict = datasets.to_dict()
+#     print("\nList Datasets Response: ", datasets_dict)
+#     assert datasets_dict is not None
+#     assert isinstance(datasets_dict, dict)
+#
+#     # Get a dataset name using our helper function
+#     global picked_dataset_name
+#     if picked_dataset_name is None:
+#         picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
+#     assert picked_dataset_name is not None
+#     print(f"\nSelected dataset for further tests: {picked_dataset_name}")
+
+
+# def test_get_dataset_task(dataset_client):
+#     print("\n\nTesting get_dataset_task")
+#     global picked_dataset_name
+#     if picked_dataset_name is None:
+#         picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
+#     assert picked_dataset_name is not None
+#     print(f"\nSelected dataset for further tests: {picked_dataset_name}")
+#
+#     # Now test the get_dataset_task method
+#     dataset_task = dataset_client.get_dataset_task(dataset_name=picked_dataset_name)
+#     print("\nDataset Task: ", dataset_task)
+#     assert dataset_task is not None
+#     assert hasattr(dataset_task, "dataset_name")
+#     assert dataset_task.dataset_name == picked_dataset_name
+#     assert hasattr(dataset_task, "data")
+#     data = dataset_task.data
+#     assert data is not None
+#     assert hasattr(data, "status")
+#     status = data.status
+#     assert status is not None
+#     print("\nDataset Task Status: ", status)
+
+
+# def test_get_dataset_task_status(dataset_client):
+#     print("\n\nTesting get_dataset_task_status")
+#     global picked_dataset_name
+#     if picked_dataset_name is None:
+#         picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
+#     assert picked_dataset_name is not None
+#     print(f"\nSelected dataset for further tests: {picked_dataset_name}")
+#
+#     # Now test the get_dataset_task_status method
+#     dataset_task_status = dataset_client.get_dataset_task_status(dataset_name=picked_dataset_name)
+#     print("\nDataset Task Status: ", dataset_task_status)
+#     assert dataset_task_status is not None
+#     assert hasattr(dataset_task_status, "dataset_name")
+#     assert dataset_task_status.dataset_name == picked_dataset_name
+#     assert hasattr(dataset_task_status, "status")
+
+
+# def test_get_datacard(dataset_client):
+#     print("\n\nTesting get_datacard")
+#     global picked_dataset_name
+#     if picked_dataset_name is None:
+#         picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
+#     assert picked_dataset_name is not None
+#     print(f"\nSelected dataset for further tests: {picked_dataset_name}")
+#
+#     # Now test the get_datacard method
+#     datacard = dataset_client.get_datacard(dataset_name=picked_dataset_name)
+#     print("\nDatacard: ", datacard)
+#     assert datacard is not None
+#     # # Dataset might not be generated yet
+#     # # TODO: How to handle this?
+#     # assert hasattr(datacard, "dataset_name")
+#     # assert datacard.dataset_name == picked_dataset_name
+#     # assert hasattr(datacard, "description")
+
+
+# def test_get_summary(dataset_client):
+#     print("\n\nTesting get_summary")
+#     global picked_dataset_name
+#     if picked_dataset_name is None:
+#         picked_dataset_name = get_dataset_name_from_list(dataset_client, status="Finished")
+#     assert picked_dataset_name is not None
+#     print(f"\nSelected dataset for further tests: {picked_dataset_name}")
+#
+#     # Now test the get_summary method
+#     summary = dataset_client.get_summary(dataset_name=picked_dataset_name)
+#     print("\nsummary: ", summary)
+#     assert summary is not None
+#     # # Dataset might not be generated yet
+#     # # TODO: How to handle this?
+#     # assert hasattr(summary, "dataset_name")
+#     # assert summary.dataset_name == picked_dataset_name
+#     # assert hasattr(summary, "test_types")
 
 
 # ----------------------------