workbench 0.8.172__py3-none-any.whl → 0.8.174__py3-none-any.whl

workbench/utils/chem_utils/toxicity.py

@@ -0,0 +1,250 @@
+"""Toxicity detection utilities for molecular compounds"""
+
+from typing import List, Optional, Tuple
+from rdkit import Chem
+from rdkit.Chem import Mol
+from rdkit.Chem import FunctionalGroups as FG
+
+# Precompiled SMARTS patterns for custom toxic functional groups
+toxic_smarts_patterns = [
+    ("C(=S)N"),  # Dithiocarbamate
+    ("P(=O)(O)(O)O"),  # Phosphate Ester
+    ("[As](=O)(=O)-[OH]"),  # Arsenic Oxide
+    ("[C](Cl)(Cl)(Cl)"),  # Trichloromethyl
+    ("[Cr](=O)(=O)=O"),  # Chromium(VI)
+    ("[N+](C)(C)(C)(C)"),  # Quaternary Ammonium
+    ("[Se][Se]"),  # Diselenide
+    ("c1c(Cl)c(Cl)c(Cl)c1"),  # Trichlorinated Aromatic Ring
+    ("[CX3](=O)[CX4][Cl,Br,F,I]"),  # Halogenated Carbonyl
+    ("[P+](C*)(C*)(C*)(C*)"),  # Phosphonium Group
+    ("NC(=S)c1c(Cl)cccc1Cl"),  # Chlorobenzene Thiocarbamate
+    ("NC(=S)Nc1ccccc1"),  # Phenyl Thiocarbamate
+    ("S=C1NCCN1"),  # Thiourea Derivative
+]
+compiled_toxic_smarts = [Chem.MolFromSmarts(smarts) for smarts in toxic_smarts_patterns]
+
+# Precompiled SMARTS patterns for exemptions
+exempt_smarts_patterns = [
+    "c1ccc(O)c(O)c1",  # Phenols
+]
+compiled_exempt_smarts = [Chem.MolFromSmarts(smarts) for smarts in exempt_smarts_patterns]
+
+# Load functional group hierarchy once during initialization
+fgroup_hierarchy = FG.BuildFuncGroupHierarchy()
+
+
+def contains_heavy_metals(mol: Mol) -> bool:
+    """
+    Check if a molecule contains any heavy metals (broad filter).
+
+    Args:
+        mol: RDKit molecule object.
+
+    Returns:
+        bool: True if any heavy metals are detected, False otherwise.
+    """
+    heavy_metals = {"Zn", "Cu", "Fe", "Mn", "Co", "Pb", "Hg", "Cd", "As"}
+    return any(atom.GetSymbol() in heavy_metals for atom in mol.GetAtoms())
+
+
+def halogen_toxicity_score(mol: Mol) -> Tuple[int, int]:
+    """
+    Calculate the halogen count and toxicity threshold for a molecule.
+
+    Args:
+        mol: RDKit molecule object.
+
+    Returns:
+        Tuple[int, int]: (halogen_count, halogen_threshold), where the threshold
+        scales with molecule size (minimum of 2 or 20% of atom count).
+    """
+    # Define halogens and count their occurrences
+    halogens = {"Cl", "Br", "I", "F"}
+    halogen_count = sum(1 for atom in mol.GetAtoms() if atom.GetSymbol() in halogens)
+
+    # Define threshold: small molecules tolerate fewer halogens
+    # Threshold scales with molecule size to account for reasonable substitution
+    molecule_size = mol.GetNumAtoms()
+    halogen_threshold = max(2, int(molecule_size * 0.2))  # Minimum 2, scaled by 20% of molecule size
+
+    return halogen_count, halogen_threshold
+
+
+def toxic_elements(mol: Mol) -> Optional[List[str]]:
+    """
+    Identifies toxic elements or specific forms of elements in a molecule.
+
+    Args:
+        mol: RDKit molecule object.
+
+    Returns:
+        Optional[List[str]]: List of toxic elements or specific forms if found, otherwise None.
+
+    Notes:
+        Halogen toxicity logic integrates with `halogen_toxicity_score` and scales thresholds
+        based on molecule size.
+    """
+    # Always toxic elements (heavy metals and known toxic single elements)
+    always_toxic = {"Pb", "Hg", "Cd", "As", "Be", "Tl", "Sb"}
+    toxic_found = set()
+
+    for atom in mol.GetAtoms():
+        symbol = atom.GetSymbol()
+        formal_charge = atom.GetFormalCharge()
+
+        # Check for always toxic elements
+        if symbol in always_toxic:
+            toxic_found.add(symbol)
+
+        # Conditionally toxic nitrogen (positively charged)
+        if symbol == "N" and formal_charge > 0:
+            # Exclude benign quaternary ammonium (e.g., choline-like structures)
+            if mol.HasSubstructMatch(Chem.MolFromSmarts("[N+](C)(C)(C)C")):  # Example benign structure
+                continue
+            toxic_found.add("N+")
+
+        # Halogen toxicity: Uses halogen_toxicity_score to flag excessive halogenation
+        if symbol in {"Cl", "Br", "I", "F"}:
+            halogen_count, halogen_threshold = halogen_toxicity_score(mol)
+            if halogen_count > halogen_threshold:
+                toxic_found.add(symbol)
+
+    return list(toxic_found) if toxic_found else None
+
+
+def toxic_groups(mol: Chem.Mol) -> Optional[List[str]]:
+    """
+    Check if a molecule contains known toxic functional groups using RDKit's functional groups and SMARTS patterns.
+
+    Args:
+        mol (rdkit.Chem.Mol): The molecule to evaluate.
+
+    Returns:
+        Optional[List[str]]: List of SMARTS patterns for toxic groups if found, otherwise None.
+    """
+    toxic_smarts_matches = []
+
+    # Use RDKit's functional group definitions
+    toxic_group_names = ["Nitro", "Azide", "Alcohol", "Aldehyde", "Halogen", "TerminalAlkyne"]
+    for group_name in toxic_group_names:
+        group_node = next(node for node in fgroup_hierarchy if node.label == group_name)
+        if mol.HasSubstructMatch(Chem.MolFromSmarts(group_node.smarts)):
+            toxic_smarts_matches.append(group_node.smarts)  # Use group_node's SMARTS directly
+
+    # Check for custom precompiled toxic SMARTS patterns
+    for smarts, compiled in zip(toxic_smarts_patterns, compiled_toxic_smarts):
+        if mol.HasSubstructMatch(compiled):  # Use precompiled SMARTS
+            toxic_smarts_matches.append(smarts)
+
+    # Special handling for N+
+    if mol.HasSubstructMatch(Chem.MolFromSmarts("[N+]")):
+        if not mol.HasSubstructMatch(Chem.MolFromSmarts("C[N+](C)(C)C")):  # Exclude benign
+            toxic_smarts_matches.append("[N+]")  # Append as SMARTS
+
+    # Exempt stabilizing functional groups using precompiled patterns
+    for compiled in compiled_exempt_smarts:
+        if mol.HasSubstructMatch(compiled):
+            return None
+
+    return toxic_smarts_matches if toxic_smarts_matches else None
+
+
+if __name__ == "__main__":
+    print("Running toxicity detection tests...")
+
+    # Test molecules with descriptions
+    test_molecules = {
+        # Safe molecules
+        "water": ("O", "Water - should be safe"),
+        "benzene": ("c1ccccc1", "Benzene - simple aromatic"),
+        "glucose": ("C(C1C(C(C(C(O1)O)O)O)O)O", "Glucose - sugar"),
+        "ethanol": ("CCO", "Ethanol - simple alcohol"),
+        # Heavy metal containing
+        "lead_acetate": ("CC(=O)[O-].CC(=O)[O-].[Pb+2]", "Lead acetate - contains Pb"),
+        "mercury_chloride": ("Cl[Hg]Cl", "Mercury chloride - contains Hg"),
+        "arsenic_trioxide": ("O=[As]O[As]=O", "Arsenic trioxide - contains As"),
+        # Halogenated compounds
+        "chloroform": ("C(Cl)(Cl)Cl", "Chloroform - trichloromethyl"),
+        "ddt": ("c1ccc(cc1)C(c2ccc(cc2)Cl)C(Cl)(Cl)Cl", "DDT - heavily chlorinated"),
+        "fluorobenzene": ("Fc1ccccc1", "Fluorobenzene - single halogen"),
+        # Nitrogen compounds
+        "nitrobenzene": ("c1ccc(cc1)[N+](=O)[O-]", "Nitrobenzene - nitro group"),
+        "choline": ("C[N+](C)(C)CCO", "Choline - benign quaternary ammonium"),
+        "toxic_quat": ("[N+](C)(C)(C)(C)", "Toxic quaternary ammonium"),
+        # Phenol (exempt)
+        "catechol": ("c1ccc(O)c(O)c1", "Catechol - phenol, should be exempt"),
+        # Phosphate
+        "phosphate": ("P(=O)(O)(O)O", "Phosphate ester - toxic pattern"),
+    }
+
+    # Test 1: Heavy Metals Detection
+    print("\n1. Testing heavy metals detection...")
+    for name, (smiles, desc) in test_molecules.items():
+        mol = Chem.MolFromSmiles(smiles)
+        if mol:
+            has_metals = contains_heavy_metals(mol)
+            expected = name in ["lead_acetate", "mercury_chloride", "arsenic_trioxide"]
+            status = "✓" if has_metals == expected else "✗"
+            print(f"   {status} {name}: {has_metals} (expected: {expected})")
+
+    # Test 2: Halogen Toxicity Score
+    print("\n2. Testing halogen toxicity scoring...")
+    halogen_tests = ["chloroform", "ddt", "fluorobenzene", "benzene"]
+    for name in halogen_tests:
+        if name in test_molecules:
+            smiles, desc = test_molecules[name]
+            mol = Chem.MolFromSmiles(smiles)
+            if mol:
+                count, threshold = halogen_toxicity_score(mol)
+                print(f"   {name}: {count} halogens, threshold: {threshold}, toxic: {count > threshold}")
+
+    # Test 3: Toxic Elements
+    print("\n3. Testing toxic elements detection...")
+    for name, (smiles, desc) in test_molecules.items():
+        mol = Chem.MolFromSmiles(smiles)
+        if mol:
+            toxics = toxic_elements(mol)
+            if toxics:
+                print(f"   ⚠ {name}: {toxics}")
+            elif name in ["lead_acetate", "mercury_chloride", "arsenic_trioxide", "chloroform", "ddt"]:
+                print(f"   ✗ {name}: Should have detected toxic elements")
+            else:
+                print(f"   ✓ {name}: No toxic elements (as expected)")
+
+    # Test 4: Toxic Groups
+    print("\n4. Testing toxic functional groups...")
+    for name, (smiles, desc) in test_molecules.items():
+        mol = Chem.MolFromSmiles(smiles)
+        if mol:
+            groups = toxic_groups(mol)
+            if groups:
+                print(f"   ⚠ {name}: Found {len(groups)} toxic group(s)")
+                for g in groups[:3]:  # Show first 3 patterns
+                    print(f"      - {g[:50]}...")
+            elif name == "catechol":
+                print(f"   ✓ {name}: Exempt (phenol)")
+            elif name in ["nitrobenzene", "phosphate", "chloroform", "ethanol"]:
+                print(f"   ✗ {name}: Should have detected toxic groups")
+            else:
+                print(f"   ✓ {name}: No toxic groups")
+
+    # Test 5: Edge Cases
+    print("\n5. Testing edge cases...")
+    edge_cases = [
+        ("", "Empty SMILES"),
+        ("INVALID", "Invalid SMILES"),
+        ("C" * 100, "Very long carbon chain"),
+        ("[N+](C)(C)(C)C", "Benign quaternary ammonium"),
+    ]
+
+    for smiles, desc in edge_cases:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol:
+            metals = contains_heavy_metals(mol)
+            elements = toxic_elements(mol)
+            groups = toxic_groups(mol)
+            print(f"   {desc}: metals={metals}, elements={elements is not None}, groups={groups is not None}")
+        else:
+            print(f"   {desc}: Invalid molecule (as expected)")
+
+    print("\n✅ All toxicity detection tests completed!")

workbench/utils/chem_utils/vis.py

@@ -0,0 +1,253 @@
+"""Molecular visualization utilities for Workbench"""
+
+import logging
+import base64
+import re
+from typing import Optional, Tuple
+from rdkit import Chem
+from rdkit.Chem import AllChem, Draw
+from rdkit.Chem.Draw import rdMolDraw2D
+
+# Set up the logger
+log = logging.getLogger("workbench")
+
+
+def _is_dark(color: str) -> bool:
+    """Determine if an rgba color is dark based on RGB average.
+
+    Args:
+        color: Color in rgba(...) format
+
+    Returns:
+        True if the color is dark, False otherwise
+    """
+    match = re.match(r"rgba?\((\d+),\s*(\d+),\s*(\d+)", color)
+    if not match:
+        log.warning(f"Invalid color format: {color}, defaulting to dark")
+        return True  # Default to dark mode on error
+
+    r, g, b = map(int, match.groups())
+    return (r + g + b) / 3 < 128
+
+
+def _rgba_to_tuple(rgba: str) -> Tuple[float, float, float, float]:
+    """Convert rgba string to normalized tuple (R, G, B, A).
+
+    Args:
+        rgba: RGBA color string (e.g., "rgba(255, 0, 0, 0.5)")
+
+    Returns:
+        Normalized tuple of (R, G, B, A) with RGB in [0, 1]
+    """
+    try:
+        components = rgba.strip("rgba() ").split(",")
+        r, g, b = (int(components[i]) / 255 for i in range(3))
+        a = float(components[3]) if len(components) > 3 else 1.0
+        return r, g, b, a
+    except (IndexError, ValueError) as e:
+        log.warning(f"Error parsing color '{rgba}': {e}, using default")
+        return 0.25, 0.25, 0.25, 1.0  # Default dark grey
+
+
+def _validate_molecule(smiles: str) -> Optional[Chem.Mol]:
+    """Validate and return RDKit molecule from SMILES.
+
+    Args:
+        smiles: SMILES string
+
+    Returns:
+        RDKit molecule or None if invalid
+    """
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            log.warning(f"Invalid SMILES: {smiles}")
+        return mol
+    except Exception as e:
+        log.error(f"Error parsing SMILES '{smiles}': {e}")
+        return None
+
+
+def _configure_draw_options(options: Draw.MolDrawOptions, background: str) -> None:
+    """Configure drawing options for molecule visualization.
+
+    Args:
+        options: RDKit drawing options object
+        background: Background color string
+    """
+    if _is_dark(background):
+        rdMolDraw2D.SetDarkMode(options)
+    options.setBackgroundColour(_rgba_to_tuple(background))
+
+
+def img_from_smiles(
+    smiles: str, width: int = 500, height: int = 500, background: str = "rgba(64, 64, 64, 1)"
+) -> Optional:
+    """Generate an image of the molecule from SMILES.
+
+    Args:
+        smiles: SMILES string representing the molecule
+        width: Width of the image in pixels (default: 500)
+        height: Height of the image in pixels (default: 500)
+        background: Background color (default: dark grey)
+
+    Returns:
+        PIL Image object or None if SMILES is invalid
+    """
+    mol = _validate_molecule(smiles)
+    if not mol:
+        return None
+
+    # Set up drawing options
+    dos = Draw.MolDrawOptions()
+    _configure_draw_options(dos, background)
+
+    # Generate and return image
+    return Draw.MolToImage(mol, options=dos, size=(width, height))
+
+
+def svg_from_smiles(
+    smiles: str, width: int = 500, height: int = 500, background: str = "rgba(64, 64, 64, 1)"
+) -> Optional[str]:
+    """Generate an SVG image of the molecule from SMILES.
+
+    Args:
+        smiles: SMILES string representing the molecule
+        width: Width of the image in pixels (default: 500)
+        height: Height of the image in pixels (default: 500)
+        background: Background color (default: dark grey)
+
+    Returns:
+        Base64-encoded SVG data URI or None if SMILES is invalid
+    """
+    mol = _validate_molecule(smiles)
+    if not mol:
+        return None
+
+    # Compute 2D coordinates
+    AllChem.Compute2DCoords(mol)
+
+    # Initialize SVG drawer
+    drawer = rdMolDraw2D.MolDraw2DSVG(width, height)
+
+    # Configure drawing options
+    _configure_draw_options(drawer.drawOptions(), background)
+
+    # Draw molecule
+    drawer.DrawMolecule(mol)
+    drawer.FinishDrawing()
+
+    # Encode SVG
+    svg = drawer.GetDrawingText()
+    encoded_svg = base64.b64encode(svg.encode("utf-8")).decode("utf-8")
+    return f"data:image/svg+xml;base64,{encoded_svg}"
+
+
+def show(smiles: str, width: int = 500, height: int = 500, background: str = "rgba(64, 64, 64, 1)") -> None:
+    """Display an image of the molecule.
+
+    Args:
+        smiles: SMILES string representing the molecule
+        width: Width of the image in pixels (default: 500)
+        height: Height of the image in pixels (default: 500)
+        background: Background color (default: dark grey)
+    """
+    img = img_from_smiles(smiles, width, height, background)
+    if img:
+        img.show()
+    else:
+        log.error(f"Cannot display molecule for SMILES: {smiles}")
+
+
+if __name__ == "__main__":
+    # Test suite
+    print("Running molecular visualization tests...")
+
+    # Test molecules
+    test_molecules = {
+        "benzene": "c1ccccc1",
+        "caffeine": "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
+        "aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",
+        "invalid": "not_a_smiles",
+        "empty": "",
+    }
+
+    # Test 1: Valid SMILES image generation
+    print("\n1. Testing image generation from valid SMILES...")
+    for name, smiles in test_molecules.items():
+        if name not in ["invalid", "empty"]:
+            img = img_from_smiles(smiles, width=200, height=200)
+            status = "✓" if img else "✗"
+            print(f"   {status} {name}: {'Success' if img else 'Failed'}")
+
+    # Test 2: Invalid SMILES handling
+    print("\n2. Testing invalid SMILES handling...")
+    img = img_from_smiles(test_molecules["invalid"])
+    print(f"   {'✓' if img is None else '✗'} Invalid SMILES returns None: {img is None}")
+
+    img = img_from_smiles(test_molecules["empty"])
+    print(f"   {'✓' if img is None else '✗'} Empty SMILES returns None: {img is None}")
+
+    # Test 3: SVG generation
+    print("\n3. Testing SVG generation...")
+    for name, smiles in test_molecules.items():
+        if name not in ["invalid", "empty"]:
+            svg = svg_from_smiles(smiles, width=200, height=200)
+            is_valid = svg and svg.startswith("data:image/svg+xml;base64,")
+            status = "✓" if is_valid else "✗"
+            print(f"   {status} {name}: {'Valid SVG data URI' if is_valid else 'Failed'}")
+
+    # Test 4: Different backgrounds
+    print("\n4. Testing different background colors...")
+    backgrounds = [
+        ("Light", "rgba(255, 255, 255, 1)"),
+        ("Dark", "rgba(0, 0, 0, 1)"),
+        ("Custom", "rgba(100, 150, 200, 0.8)"),
+    ]
+
+    for bg_name, bg_color in backgrounds:
+        img = img_from_smiles(test_molecules["benzene"], background=bg_color)
+        status = "✓" if img else "✗"
+        print(f"   {status} {bg_name} background: {'Success' if img else 'Failed'}")
+
+    # Test 5: Size variations
+    print("\n5. Testing different image sizes...")
+    sizes = [(100, 100), (500, 500), (1000, 800)]
+
+    for w, h in sizes:
+        img = img_from_smiles(test_molecules["caffeine"], width=w, height=h)
+        status = "✓" if img else "✗"
+        print(f"   {status} Size {w}x{h}: {'Success' if img else 'Failed'}")
+
+    # Test 6: Color parsing functions
+    print("\n6. Testing color utility functions...")
+    test_colors = [
+        ("invalid_color", True, (0.25, 0.25, 0.25, 1.0)),  # Should use defaults
+        ("rgba(255, 255, 255, 1)", False, (1.0, 1.0, 1.0, 1.0)),
+        ("rgba(0, 0, 0, 1)", True, (0.0, 0.0, 0.0, 1.0)),
+        ("rgba(64, 64, 64, 0.5)", True, (0.251, 0.251, 0.251, 0.5)),
+        ("rgb(128, 128, 128)", False, (0.502, 0.502, 0.502, 1.0)),
+    ]
+
+    for color, expected_dark, expected_tuple in test_colors:
+        is_dark_result = _is_dark(color)
+        tuple_result = _rgba_to_tuple(color)
+
+        dark_status = "✓" if is_dark_result == expected_dark else "✗"
+        print(f"   {dark_status} is_dark('{color[:20]}...'): {is_dark_result} == {expected_dark}")
+
+        # Check tuple values with tolerance for floating point
+        tuple_match = all(abs(a - b) < 0.01 for a, b in zip(tuple_result, expected_tuple))
+        tuple_status = "✓" if tuple_match else "✗"
+        print(f"   {tuple_status} rgba_to_tuple('{color[:20]}...'): matches expected")
+
+    # Test the show function (will open image windows)
+    print("\n7. Testing show function (will open image windows)...")
+    try:
+        show(test_molecules["aspirin"])
+        show(test_molecules["aspirin"], background="rgba(220, 220, 220, 1)")
+        print("   ✓ show() function executed (check for image window)")
+    except Exception as e:
+        print(f"   ✗ show() function failed: {e}")
+
+    print("\n✅ All tests completed!")

workbench/utils/monitor_utils.py

@@ -14,7 +14,7 @@ from workbench.utils.s3_utils import read_content_from_s3
 log = logging.getLogger("workbench")
 
 
-def pull_data_capture(data_capture_path, max_files=1) -> Union[pd.DataFrame, None]:
+def pull_data_capture_for_testing(data_capture_path, max_files=1) -> Union[pd.DataFrame, None]:
     """
     Read and process captured data from S3.
 
@@ -26,7 +26,12 @@ def pull_data_capture(data_capture_path, max_files=1) -> Union[pd.DataFrame, Non
 
     Returns:
         Union[pd.DataFrame, None]: A dataframe of the captured data (or None if no data is found).
+
+    Notes:
+        This method is really only for testing and debugging.
     """
+    log.important("This method is for testing and debugging only.")
+
     # List files in the specified S3 path
     files = wr.s3.list_objects(data_capture_path)
     if not files:
@@ -64,59 +69,53 @@ def pull_data_capture(data_capture_path, max_files=1) -> Union[pd.DataFrame, Non
 def process_data_capture(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
     """
     Process the captured data DataFrame to extract input and output data.
-    Continues processing even if individual files are malformed.
+    Handles cases where input or output might not be captured.
+
     Args:
         df (DataFrame): DataFrame with captured data.
     Returns:
         tuple[DataFrame, DataFrame]: Input and output DataFrames.
     """
+
+    def parse_endpoint_data(data: dict) -> pd.DataFrame:
+        """Parse endpoint data based on encoding type."""
+        encoding = data["encoding"].upper()
+
+        if encoding == "CSV":
+            return pd.read_csv(StringIO(data["data"]))
+        elif encoding == "JSON":
+            json_data = json.loads(data["data"])
+            if isinstance(json_data, dict):
+                return pd.DataFrame({k: [v] if not isinstance(v, list) else v for k, v in json_data.items()})
+            else:
+                return pd.DataFrame(json_data)
+        else:
+            return None  # Unknown encoding
+
     input_dfs = []
     output_dfs = []
 
-    for idx, row in df.iterrows():
+    # Use itertuples() instead of iterrows() for better performance
+    for row in df.itertuples(index=True):
         try:
-            capture_data = row["captureData"]
-
-            # Check if this capture has the required fields (all or nothing)
-            if "endpointInput" not in capture_data:
-                log.warning(f"Row {idx}: No endpointInput found in capture data.")
-                continue
-
-            if "endpointOutput" not in capture_data:
-                log.critical(
-                    f"Row {idx}: No endpointOutput found in capture data. DataCapture needs to include Output capture!"
-                )
-                continue
-
-            # Process input data
-            input_data = capture_data["endpointInput"]
-            if input_data["encoding"].upper() == "CSV":
-                input_df = pd.read_csv(StringIO(input_data["data"]))
-            elif input_data["encoding"].upper() == "JSON":
-                json_data = json.loads(input_data["data"])
-                if isinstance(json_data, dict):
-                    input_df = pd.DataFrame({k: [v] if not isinstance(v, list) else v for k, v in json_data.items()})
-                else:
-                    input_df = pd.DataFrame(json_data)
-
-            # Process output data
-            output_data = capture_data["endpointOutput"]
-            if output_data["encoding"].upper() == "CSV":
-                output_df = pd.read_csv(StringIO(output_data["data"]))
-            elif output_data["encoding"].upper() == "JSON":
-                json_data = json.loads(output_data["data"])
-                if isinstance(json_data, dict):
-                    output_df = pd.DataFrame({k: [v] if not isinstance(v, list) else v for k, v in json_data.items()})
-                else:
-                    output_df = pd.DataFrame(json_data)
-
-            # If we get here, both processed successfully
-            input_dfs.append(input_df)
-            output_dfs.append(output_df)
+            capture_data = row.captureData
+
+            # Process input data if present
+            if "endpointInput" in capture_data:
+                input_df = parse_endpoint_data(capture_data["endpointInput"])
+                if input_df is not None:
+                    input_dfs.append(input_df)
+
+            # Process output data if present
+            if "endpointOutput" in capture_data:
+                output_df = parse_endpoint_data(capture_data["endpointOutput"])
+                if output_df is not None:
+                    output_dfs.append(output_df)
 
         except Exception as e:
-            log.error(f"Row {idx}: Failed to process row: {e}")
+            log.debug(f"Row {row.Index}: Failed to process row: {e}")
             continue
+
     # Combine and return results
     return (
         pd.concat(input_dfs, ignore_index=True) if input_dfs else pd.DataFrame(),
@@ -178,23 +177,6 @@ def parse_monitoring_results(results_json: str) -> Dict[str, Any]:
         return {"error": str(e)}
 
 
-"""TEMP
-                # If the status is "CompletedWithViolations", we grab the lastest
-                # violation file and add it to the result
-                if status == "CompletedWithViolations":
-                    violation_file = f"{self.monitoring_path}/
-                    {last_run['CreationTime'].strftime('%Y/%m/%d')}/constraint_violations.json"
-                    if wr.s3.does_object_exist(violation_file):
-                        violations_json = read_content_from_s3(violation_file)
-                        violations = parse_monitoring_results(violations_json)
-                        result["violations"] = violations.get("constraint_violations", [])
-                        result["violation_count"] = len(result["violations"])
-                    else:
-                        result["violations"] = []
-                        result["violation_count"] = 0
-"""
-
-
 def preprocessing_script(feature_list: list[str]) -> str:
     """
     A preprocessing script for monitoring jobs.
@@ -245,8 +227,8 @@ if __name__ == "__main__":
     from workbench.api.monitor import Monitor
 
     # Test pulling data capture
-    mon = Monitor("caco2-pappab-class-0")
-    df = pull_data_capture(mon.data_capture_path)
+    mon = Monitor("abalone-regression-rt")
+    df = pull_data_capture_for_testing(mon.data_capture_path)
     print("Data Capture:")
     print(df.head())
 
@@ -262,4 +244,4 @@ if __name__ == "__main__":
     # Test preprocessing script
     script = preprocessing_script(["feature1", "feature2", "feature3"])
     print("\nPreprocessing Script:")
-    print(script)
+    # print(script)

workbench/utils/pandas_utils.py

@@ -152,7 +152,7 @@ def compare_dataframes(df1: pd.DataFrame, df2: pd.DataFrame, display_columns: li
 
     # Check for differences in common columns
     for column in common_columns:
-        if pd.api.types.is_string_dtype(df1[column]) or pd.api.types.is_string_dtype(df2[column]):
+        if pd.api.types.is_string_dtype(df1[column]) and pd.api.types.is_string_dtype(df2[column]):
             # String comparison with NaNs treated as equal
             differences = ~(df1[column].fillna("") == df2[column].fillna(""))
         elif pd.api.types.is_float_dtype(df1[column]) or pd.api.types.is_float_dtype(df2[column]):
@@ -161,8 +161,8 @@ def compare_dataframes(df1: pd.DataFrame, df2: pd.DataFrame, display_columns: li
                 pd.isna(df1[column]) & pd.isna(df2[column])
             )
         else:
-            # Other types (e.g., int) with NaNs treated as equal
-            differences = ~(df1[column].fillna(0) == df2[column].fillna(0))
+            # Other types (int, Int64, etc.) - compare with NaNs treated as equal
+            differences = (df1[column] != df2[column]) & ~(pd.isna(df1[column]) & pd.isna(df2[column]))
 
         # If differences exist, display them
         if differences.any():