rowan-mcp 0.1.0__py3-none-any.whl

rowan_mcp/functions/multistage_opt.py

@@ -0,0 +1,171 @@
+"""
+Rowan multistage optimization function for geometry optimization.
+"""
+
+import os
+import logging
+import time
+from typing import Optional, List, Dict, Any
+
+try:
+    import rowan
+except ImportError:
+    rowan = None
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+# Setup API key
+api_key = os.getenv("ROWAN_API_KEY")
+if rowan and api_key:
+    rowan.api_key = api_key
+
+def log_rowan_api_call(workflow_type: str, **kwargs):
+    """Log Rowan API calls with detailed parameters."""
+    
+    # Special handling for long-running calculations
+    if workflow_type in ["multistage_opt", "conformer_search"]:
+        ping_interval = kwargs.get('ping_interval', 5)
+        blocking = kwargs.get('blocking', True)
+        if blocking:
+            if workflow_type == "multistage_opt":
+                logger.info(f" Multi-stage optimization may take several minutes...")
+            else:
+                logger.info(f" Conformer search may take several minutes...")
+            logger.info(f" Progress will be checked every {ping_interval} seconds")
+        else:
+            logger.info(f" {workflow_type.replace('_', ' ').title()} submitted without waiting")
+    
+    try:
+        start_time = time.time()
+        result = rowan.compute(workflow_type=workflow_type, **kwargs)
+        api_time = time.time() - start_time
+        
+        if isinstance(result, dict) and 'uuid' in result:
+            job_status = result.get('status', result.get('object_status', 'Unknown'))
+            status_names = {0: "Queued", 1: "Running", 2: "Completed", 3: "Failed", 4: "Stopped", 5: "Awaiting Queue"}
+            status_text = status_names.get(job_status, f"Unknown ({job_status})")
+        
+        return result
+        
+    except Exception as e:
+        api_time = time.time() - start_time
+        raise e
+
+def rowan_multistage_opt(
+    name: str,
+    molecule: str,
+    mode: str = "rapid",
+    solvent: Optional[str] = None,
+    xtb_preopt: bool = False,
+    constraints: Optional[List[Dict[str, Any]]] = None,
+    transition_state: bool = False,
+    frequencies: bool = False,
+    folder_uuid: Optional[str] = None,
+    blocking: bool = True,
+    ping_interval: int = 30
+) -> str:
+    """Run multi-level geometry optimization using stjames MultiStageOptWorkflow.
+    
+    Performs hierarchical optimization using multiple levels of theory based on mode. There are only four valid modes:
+    
+    **RAPID** (default): GFN2-xTB optimization → r²SCAN-3c single point
+    **RECKLESS**: GFN-FF optimization → GFN2-xTB single point  
+    **CAREFUL**: GFN2-xTB preopt → r²SCAN-3c optimization → ωB97X-3c single point
+    **METICULOUS**: GFN2-xTB preopt → r²SCAN-3c opt → ωB97X-3c opt → ωB97M-D3BJ/def2-TZVPPD single point
+    
+    This follows the exact stjames MultiStageOptWorkflow implementation for:
+    - High accuracy final structures
+    - Efficient computational cost  
+    - Reliable convergence across chemical space
+    
+    Args:
+        name: Name for the calculation
+        molecule: Molecule SMILES string
+        mode: Optimization mode - "rapid", "reckless", "careful", or "meticulous" (default: "rapid")
+        solvent: Solvent for single point calculation (e.g., "water", "dmso", "acetone")
+        xtb_preopt: Whether to include xTB pre-optimization step (default: False)
+        constraints: List of optimization constraints (default: None)
+        transition_state: Whether this is a transition state optimization (default: False)
+        frequencies: Whether to calculate vibrational frequencies (default: False)
+        folder_uuid: Optional folder UUID for organization
+        blocking: Whether to wait for completion (default: True)
+        ping_interval: Check status interval in seconds (default: 30)
+    
+    Returns:
+        Comprehensive optimization results following MultiStageOptWorkflow format
+        
+    Example:
+        # Basic optimization
+        result = rowan_multistage_opt("aspirin_opt", "CC(=O)Oc1ccccc1C(=O)O")
+        
+        # With solvent and frequency analysis
+        result = rowan_multistage_opt(
+            "aspirin_water", 
+            "CC(=O)Oc1ccccc1C(=O)O",
+            mode="careful",
+            solvent="water",
+            frequencies=True
+        )
+    """
+    
+    # Validate mode parameter - only allow Rowan's supported modes
+    valid_modes = ["rapid", "reckless", "careful", "meticulous"]
+    if mode.lower() not in valid_modes:
+        error_msg = f"Invalid mode '{mode}'. Must be one of: {', '.join(valid_modes)}"
+        logger.error(f"Mode validation failed: {error_msg}")
+        return f"Error: {error_msg}"
+    
+    # Prepare parameters following stjames MultiStageOptWorkflow structure
+    params = {
+        "name": name,
+        "molecule": molecule,
+        "initial_molecule": molecule,  # Required by MultiStageOptWorkflow
+        "mode": mode.lower(),
+        "folder_uuid": folder_uuid,
+        "blocking": blocking,
+        "ping_interval": ping_interval
+    }
+    
+    # Add optional parameters if specified
+    if solvent:
+        params["solvent"] = solvent
+        
+    if xtb_preopt:
+        params["xtb_preopt"] = xtb_preopt
+        
+    if constraints:
+        params["constraints"] = constraints
+        
+    if transition_state:
+        params["transition_state"] = transition_state
+        
+    if frequencies:
+        params["frequencies"] = frequencies
+    
+    # Submit to Rowan using multistage_opt workflow
+    result = log_rowan_api_call(
+        workflow_type="multistage_opt",
+        **params
+    )
+    return str(result)
+
+def test_rowan_multistage_opt():
+    """Test the rowan_multistage_opt function with stjames parameters."""
+    try:
+        # Test with stjames-compatible parameters
+        result = rowan_multistage_opt(
+            name="test_stjames_opt", 
+            molecule="CCO", 
+            mode="rapid",
+            blocking=False
+        )
+        print("✅ Multistage optimization test successful!")
+        print(f"Result length: {len(result)} characters")
+        return True
+    except Exception as e:
+        print(f"Multistage optimization test failed: {e}")
+        return False
+
+if __name__ == "__main__":
+    test_rowan_multistage_opt() 

rowan_mcp/functions/pdb_handler.py

@@ -0,0 +1,200 @@
+"""
+PDB handler for Rowan docking - attempts to handle PDB input intelligently
+"""
+
+import os
+import logging
+import requests
+from typing import Optional, Dict, Any, Tuple
+import rowan
+
+logger = logging.getLogger(__name__)
+
+def parse_pdb_content(pdb_content: str) -> Dict[str, Any]:
+    """
+    Parse PDB content into a structured format that might work with Rowan.
+    
+    This attempts to extract key information from PDB files that the 
+    Rowan API might accept.
+    """
+    lines = pdb_content.split('\n')
+    
+    # Extract basic information
+    header = ""
+    title = ""
+    atoms = []
+    
+    for line in lines:
+        if line.startswith("HEADER"):
+            header = line[10:50].strip()
+        elif line.startswith("TITLE"):
+            title += line[10:].strip() + " "
+        elif line.startswith("ATOM") or line.startswith("HETATM"):
+            atoms.append(line)
+    
+    # Try to create a structure that might work with Rowan's PDB format
+    pdb_data = {
+        "description": {
+            "title": title.strip() or header,
+            "classification": header
+        },
+        "experiment": {
+            "technique": "X-RAY DIFFRACTION"  # Default assumption
+        },
+        "geometry": {},
+        "contents": pdb_content,  # Keep full content
+        "name": f"pdb_upload_{len(pdb_content)}"
+    }
+    
+    return pdb_data
+
+def attempt_pdb_upload(pdb_id: str, pdb_content: str, folder_uuid: Optional[str] = None) -> Tuple[Optional[str], str]:
+    """
+    Attempt to upload PDB content to Rowan.
+    
+    Returns:
+        Tuple of (uuid or None, status message)
+    """
+    try:
+        # First, try the Protein.create method if it exists
+        if hasattr(rowan, 'Protein') and hasattr(rowan.Protein, 'create'):
+            logger.info(f"Attempting to create protein from PDB {pdb_id}")
+            result = rowan.Protein.create(
+                name=f"{pdb_id}_upload",
+                pdb_content=pdb_content,
+                folder_uuid=folder_uuid
+            )
+            if result and 'uuid' in result:
+                return result['uuid'], f"Successfully uploaded PDB {pdb_id}"
+        
+        # Try alternative upload methods
+        if hasattr(rowan, 'upload_pdb'):
+            result = rowan.upload_pdb(pdb_content, name=pdb_id)
+            if result and 'uuid' in result:
+                return result['uuid'], f"Successfully uploaded PDB {pdb_id}"
+        
+        # Try to create a file upload
+        if hasattr(rowan, 'File') and hasattr(rowan.File, 'create'):
+            result = rowan.File.create(
+                name=f"{pdb_id}.pdb",
+                content=pdb_content,
+                file_type="pdb",
+                folder_uuid=folder_uuid
+            )
+            if result and 'uuid' in result:
+                return result['uuid'], f"Uploaded PDB {pdb_id} as file"
+                
+    except Exception as e:
+        logger.error(f"Failed to upload PDB {pdb_id}: {str(e)}")
+    
+    return None, f"Cannot upload PDB through API - manual upload required"
+
+def handle_pdb_input(pdb_input: str, folder_uuid: Optional[str] = None) -> Dict[str, Any]:
+    """
+    Handle PDB input (either PDB ID or UUID) and return information for docking.
+    
+    Args:
+        pdb_input: Either a PDB ID (like "4Z18") or a UUID
+        folder_uuid: Optional folder for uploads
+        
+    Returns:
+        Dictionary with:
+        - success: bool
+        - uuid: str (if successful)
+        - message: str (status/error message)
+        - pdb_id: str (if it was a PDB ID)
+        - instructions: str (what to do next)
+    """
+    # Check if it's a UUID (36 chars with dashes)
+    if len(pdb_input) == 36 and pdb_input.count('-') == 4:
+        return {
+            "success": True,
+            "uuid": pdb_input,
+            "message": "Using provided protein UUID",
+            "instructions": None
+        }
+    
+    # Assume it's a PDB ID
+    pdb_id = pdb_input.upper()
+    
+    # Fetch PDB content
+    try:
+        url = f"https://files.rcsb.org/download/{pdb_id}.pdb"
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+        pdb_content = response.text
+        logger.info(f"Successfully fetched PDB {pdb_id} ({len(pdb_content)} characters)")
+    except Exception as e:
+        return {
+            "success": False,
+            "uuid": None,
+            "message": f"Failed to fetch PDB {pdb_id}: {str(e)}",
+            "pdb_id": pdb_id,
+            "instructions": "Please verify the PDB ID is correct"
+        }
+    
+    # Attempt to upload
+    uuid, upload_status = attempt_pdb_upload(pdb_id, pdb_content, folder_uuid)
+    
+    if uuid:
+        return {
+            "success": True,
+            "uuid": uuid,
+            "message": upload_status,
+            "pdb_id": pdb_id,
+            "instructions": None
+        }
+    
+    # Upload failed - provide instructions
+    instructions = f"""
+To use PDB {pdb_id} for docking:
+
+1. Download the PDB file:
+   https://files.rcsb.org/download/{pdb_id}.pdb
+
+2. Upload to your Rowan account:
+   https://labs.rowansci.com
+
+3. Get the protein UUID from your account
+
+4. Use the UUID in your docking call:
+   rowan_docking(
+       target_uuid="your-protein-uuid-here",
+       ...
+   )
+"""
+    
+    return {
+        "success": False,
+        "uuid": None,
+        "message": upload_status,
+        "pdb_id": pdb_id,
+        "instructions": instructions
+    }
+
+def enhance_docking_with_pdb_handler(
+    target_input: Optional[str] = None,
+    target_uuid: Optional[str] = None,
+    **kwargs
+) -> Tuple[Optional[str], str]:
+    """
+    Enhanced target handling for docking that accepts PDB IDs.
+    
+    Returns:
+        Tuple of (target_uuid or None, message)
+    """
+    # If UUID already provided, use it
+    if target_uuid:
+        return target_uuid, "Using provided target UUID"
+    
+    # If target_input provided, try to handle it
+    if target_input:
+        result = handle_pdb_input(target_input, kwargs.get('folder_uuid'))
+        
+        if result['success']:
+            return result['uuid'], result['message']
+        else:
+            # Return None but with helpful message
+            return None, result['instructions'] or result['message']
+    
+    return None, "No protein target specified"

rowan_mcp/functions/pka.py

@@ -0,0 +1,137 @@
+"""
+Calculate pKa values for molecules using Rowan API.
+"""
+
+import os
+import rowan
+from typing import Optional
+
+# Set up logging
+import logging
+logger = logging.getLogger(__name__)
+
+# Configure rowan API key
+if not hasattr(rowan, 'api_key') or not rowan.api_key:
+    api_key = os.getenv("ROWAN_API_KEY")
+    if api_key:
+        rowan.api_key = api_key
+        logger.info("🔑 Rowan API key configured")
+    else:
+        logger.error("No ROWAN_API_KEY found in environment")
+
+def log_rowan_api_call(workflow_type: str, **kwargs):
+    """Log Rowan API calls and let Rowan handle its own errors."""
+    
+    # Simple logging for calculations
+    logger.info(f" Starting {workflow_type.replace('_', ' ')}...")
+    
+    # Let Rowan handle everything - no custom error handling
+    return rowan.compute(workflow_type=workflow_type, **kwargs)
+
+def rowan_pka(
+    name: str,
+    molecule: str,
+    folder_uuid: Optional[str] = None,
+    blocking: bool = True,
+    ping_interval: int = 5
+) -> str:
+    """Calculate pKa values for molecules.
+    
+    Args:
+        name: Name for the calculation
+        molecule: Molecule SMILES string or common name
+        folder_uuid: UUID of folder to organize calculation in
+        blocking: Whether to wait for completion (default: True)
+        ping_interval: How often to check status in seconds (default: 5)
+    
+    Returns:
+        pKa calculation results
+    """
+    try:
+        result = log_rowan_api_call(
+            workflow_type="pka",
+            name=name,
+            molecule=molecule,
+            folder_uuid=folder_uuid,
+            blocking=blocking,
+            ping_interval=ping_interval
+        )
+        
+        # Format results based on whether we waited or not
+        if blocking:
+            # We waited for completion - format actual results
+            status = result.get('status', result.get('object_status', 'Unknown'))
+            
+            if status == 2:  # Completed successfully
+                formatted = f" pKa calculation for '{name}' completed successfully!\n\n"
+            elif status == 3:  # Failed
+                formatted = f" pKa calculation for '{name}' failed!\n\n"
+            else:
+                formatted = f" pKa calculation for '{name}' finished with status {status}\n\n"
+                
+            formatted += f" Molecule: {molecule}\n"
+            formatted += f" Job UUID: {result.get('uuid', 'N/A')}\n"
+            formatted += f" Status: {status}\n"
+            
+            # Try to extract pKa results
+            if isinstance(result, dict) and 'object_data' in result and result['object_data']:
+                data = result['object_data']
+                
+                # Extract pKa values
+                if 'strongest_acid' in data:
+                    if data['strongest_acid'] is not None:
+                        formatted += f" Strongest Acid pKa: {data['strongest_acid']:.2f}\n"
+                    else:
+                        formatted += f" Strongest Acid pKa: N/A (no acidic sites found)\n"
+                        
+                if 'strongest_base' in data:
+                    if data['strongest_base'] is not None:
+                        formatted += f" Strongest Base pKa: {data['strongest_base']:.2f}\n"
+                    else:
+                        formatted += f" Strongest Base pKa: N/A (no basic sites found)\n"
+                if 'pka_values' in data and isinstance(data['pka_values'], list):
+                    formatted += f" All pKa values: {', '.join([f'{val:.2f}' for val in data['pka_values']])}\n"
+                
+                # Additional properties if available
+                if 'ionizable_sites' in data:
+                    formatted += f" Ionizable sites found: {data['ionizable_sites']}\n"
+            
+            # Basic guidance
+            if status == 2:
+                formatted += f"\n Use rowan_workflow_management(action='retrieve', workflow_uuid='{result.get('uuid')}') for detailed data\n"
+        else:
+            # Non-blocking mode - just submission confirmation
+            formatted = f" pKa calculation for '{name}' submitted!\n\n"
+            formatted += f" Molecule: {molecule}\n"
+            formatted += f" Job UUID: {result.get('uuid', 'N/A')}\n"
+            formatted += f" Status: {result.get('status', 'Submitted')}\n"
+        
+        return formatted
+        
+    except Exception as e:
+        error_response = {
+            "error": f"pKa calculation failed: {str(e)}",
+            "name": name,
+            "molecule": molecule
+        }
+        return str(error_response)
+
+
+def test_rowan_pka():
+    """Test the rowan_pka function."""
+    try:
+        # Test with minimal parameters
+        result = rowan_pka(
+            name="test_pka_water",
+            molecule="O"
+        )
+        print("✅ pKa test successful!")
+        print(f"Result: {result}")
+        return True
+    except Exception as e:
+        print(f"pKa test failed: {e}")
+        return False
+
+
+if __name__ == "__main__":
+    test_rowan_pka()