tooluniverse 0.2.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

tooluniverse/base_tool.py

@@ -1,11 +1,134 @@
 from .utils import extract_function_call_json, evaluate_function_call
+import json
+from pathlib import Path
+from typing import no_type_check
+
+
+class ToolExecutionError(Exception):
+    """Base exception for tool execution errors."""
+
+
+class ValidationError(Exception):
+    """Exception raised when input validation fails."""
+
+
+class AuthenticationError(Exception):
+    """Exception raised when authentication fails."""
+
+
+class RateLimitError(Exception):
+    """Exception raised when API rate limit is exceeded."""
+
 
 class BaseTool:
     def __init__(self, tool_config):
-        self.tool_config = tool_config
+        self.tool_config = self._apply_defaults(tool_config)
+
+    @classmethod
+    def get_default_config_file(cls):
+        """
+        Get the path to the default configuration file for this tool type.
+
+        This method uses a robust path resolution strategy that works across
+        different installation scenarios:
+
+        1. Installed packages: Uses importlib.resources for proper package
+           resource access
+        2. Development mode: Falls back to file-based path resolution
+        3. Legacy Python: Handles importlib.resources and importlib_resources
+
+        Override this method in subclasses to specify a custom defaults file.
+
+        Returns:
+            Path or resource object pointing to the defaults file
+        """
+        tool_type = cls.__name__
 
-    def run(self):
-        pass
+        # Use importlib.resources for robust path resolution across different
+        # installation methods
+        try:
+            import importlib.resources as pkg_resources
+        except ImportError:
+            # Fallback for Python < 3.9
+            import importlib_resources as pkg_resources
+
+        try:
+            # Try to use package resources first (works with installed
+            # packages). Use the newer files() API
+            data_files = pkg_resources.files("tooluniverse.data")
+            defaults_file = data_files / f"{tool_type.lower()}_defaults.json"
+
+            # For compatibility, convert to a regular Path if possible
+            if hasattr(defaults_file, "resolve"):
+                return defaults_file.resolve()
+            else:
+                # For older Python versions or special cases, return resource
+                # path
+                return defaults_file
+
+        except (FileNotFoundError, ModuleNotFoundError, AttributeError):
+            # Fallback to file-based path resolution for development/local use
+            current_dir = Path(__file__).parent
+            defaults_file = current_dir / "data" / f"{tool_type.lower()}_defaults.json"
+            return defaults_file
+
+    @classmethod
+    def load_defaults_from_file(cls):
+        """Load defaults from the configuration file"""
+        defaults_file = cls.get_default_config_file()
+
+        # Handle both regular Path objects and importlib resource objects
+        try:
+            # Check if it's a regular Path object
+            if hasattr(defaults_file, "exists") and not defaults_file.exists():
+                return {}
+
+            # Try to read the file (works for both Path and resource objects)
+            if hasattr(defaults_file, "read_text"):
+                # Resource object with read_text method
+                content = defaults_file.read_text(encoding="utf-8")
+                data = json.loads(content)
+            else:
+                # Regular file path
+                with open(defaults_file, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+
+            # Look for defaults under the tool type key
+            tool_type = cls.__name__
+            return data.get(f"{tool_type.lower()}_defaults", {})
+
+        except (FileNotFoundError, json.JSONDecodeError):
+            # File doesn't exist or invalid JSON, return empty defaults
+            return {}
+        except Exception as e:
+            print(f"Warning: Could not load defaults for {cls.__name__}: {e}")
+            return {}
+
+    def _apply_defaults(self, tool_config):
+        """Apply default configuration to the tool config"""
+        # Load defaults from file
+        defaults = self.load_defaults_from_file()
+
+        if not defaults:
+            # No defaults available, return original config
+            return tool_config
+
+        # Create merged configuration by starting with defaults
+        merged_config = defaults.copy()
+
+        # Override with tool-specific configuration
+        merged_config.update(tool_config)
+
+        return merged_config
+
+    @no_type_check
+    def run(self, arguments=None):
+        """Execute the tool.
+
+        The default BaseTool implementation accepts an optional arguments
+        mapping to align with most concrete tool implementations which expect
+        a dictionary of inputs.
+        """
 
     def check_function_call(self, function_call_json):
         if isinstance(function_call_json, str):
@@ -22,11 +145,11 @@ class BaseTool:
         list: List of required parameters for the given endpoint.
         """
         required_params = []
-        parameters = self.tool_config.get('parameter', {}).get('properties', {})
+        parameters = self.tool_config.get("parameter", {}).get("properties", {})
 
         # Check each parameter to see if it is required
         for param, details in parameters.items():
-            if details.get('required', False):
+            if details.get("required", False):
                 required_params.append(param.lower())
 
-        return required_params
+        return required_params

tooluniverse/boltz_tool.py

@@ -0,0 +1,207 @@
+import os
+import pprint
+import subprocess
+import tempfile
+import yaml
+import json
+import shutil
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+
+
+@register_tool("Boltz2DockingTool")
+class Boltz2DockingTool(BaseTool):
+    """
+    Tool to perform protein-ligand docking and affinity prediction using the local Boltz-2 model.
+    This tool constructs a YAML input file, runs the `boltz predict` command,
+    and parses the output to return the predicted structure and affinity.
+    """
+
+    def __init__(self, tool_config: dict):
+        """
+        Initializes the BoltzDockingTool.
+        Checks if the 'boltz' command is available in the system's PATH.
+        """
+        super().__init__(tool_config)
+        if not shutil.which("boltz"):
+            raise EnvironmentError(
+                "The 'boltz' command is not found. "
+                "Please ensure the 'boltz' package is installed and accessible in the system's PATH. "
+                "Installation guide: https://github.com/jwohlwend/boltz"
+            )
+
+    def _build_yaml_input(self, arguments: dict) -> dict:
+        """Constructs the YAML data structure for the Boltz input."""
+        protein_sequence = arguments.get("protein_sequence")
+        ligands = arguments.get("ligands", [])
+
+        # The first ligand is assumed to be the binder for affinity prediction
+        if not ligands:
+            raise ValueError(
+                "At least one ligand must be provided in the 'ligands' list."
+            )
+
+        binder_id = ligands[0].get("id")
+        if not binder_id:
+            raise ValueError("The first ligand in the list must have a valid 'id'.")
+
+        # --- Sequences Section ---
+        sequences = [{"protein": {"id": "A", "sequence": protein_sequence}}]
+
+        for i, ligand_data in enumerate(ligands):
+            chain_id = ligand_data.get("id")
+            if not chain_id:
+                raise ValueError(f"Ligand at index {i} must have an 'id' key.")
+
+            entry = {"id": chain_id}
+            if "smiles" in ligand_data:
+                entry["smiles"] = ligand_data["smiles"]
+            elif "ccd" in ligand_data:
+                entry["ccd"] = ligand_data["ccd"]
+            else:
+                raise ValueError(
+                    f"Ligand at index {i} must have a 'smiles' or 'ccd' key."
+                )
+            sequences.append({"ligand": entry})
+
+        # --- Properties Section (for Affinity) ---
+        properties = [{"affinity": {"binder": binder_id}}]
+
+        # --- Final YAML Structure ---
+        yaml_input = {"version": 1, "sequences": sequences, "properties": properties}
+
+        # Add optional fields
+        if "constraints" in arguments:
+            yaml_input["constraints"] = arguments["constraints"]
+        if "templates" in arguments:
+            yaml_input["templates"] = arguments["templates"]
+
+        return yaml_input
+
+    def run(self, arguments: dict | None = None, timeout: int = 1200) -> dict:
+        """
+        Executes the Boltz prediction.
+
+        Args:
+            arguments (dict): A dictionary containing the necessary inputs.
+                - protein_sequence (str): The amino acid sequence of the protein.
+                - ligands (list[dict]): A list of ligands, each with a 'smiles' or 'ccd' key.
+                - constraints (list[dict], optional): Covalent bonds or other constraints.
+                - templates (list[dict], optional): Structural templates.
+                - other optional boltz CLI flags (e.g., 'recycling_steps').
+            timeout (int): The maximum time in seconds to wait for the Boltz command to complete.
+
+        Returns:
+            dict: A dictionary containing the path to the predicted structure and affinity data, or an error.
+        """
+        arguments = arguments or {}
+        if not arguments.get("protein_sequence"):
+            return {"error": "The 'protein_sequence' parameter is required."}
+
+        # Create a temporary directory to store input and output files
+        with tempfile.TemporaryDirectory() as temp_dir:
+            input_filename = "boltz_input"
+            input_yaml_path = os.path.join(temp_dir, f"{input_filename}.yaml")
+            output_dir = os.path.join(temp_dir, "results")
+            os.makedirs(output_dir, exist_ok=True)
+
+            # Build and write the input YAML file
+            yaml_data = self._build_yaml_input(arguments)
+            with open(input_yaml_path, "w") as f:
+                yaml.dump(yaml_data, f, sort_keys=False)
+
+            # Construct the command-line arguments for Boltz
+            command = [
+                "boltz",
+                "predict",
+                input_yaml_path,
+                "--out_dir",
+                output_dir,
+                "--use_msa_server",
+                "--override",  # Override existing results if any
+            ]
+
+            # Add optional command-line flags from arguments
+            for key in [
+                "recycling_steps",
+                "diffusion_samples",
+                "sampling_steps",
+                "step_scale",
+            ]:
+                if key in arguments:
+                    command.extend([f"--{key}", str(arguments[key])])
+
+            if arguments.get("use_potentials", False):
+                command.append("--use_potentials")
+
+            # Execute the Boltz command
+            subprocess.run(
+                command,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+                check=True,  # Will raise CalledProcessError on non-zero exit codes
+            )
+
+            # --- Parse the output files ---
+            # 1. locate the Boltz run folder under your out_dir
+            root_dirs = [
+                d
+                for d in os.listdir(output_dir)
+                if os.path.isdir(os.path.join(output_dir, d))
+            ]
+            if not root_dirs:
+                return {"error": "No Boltz run folder found under out_dir"}
+            if len(root_dirs) > 1:
+                # you could pick the latest by timestamp instead of the first
+                run_dir_name = sorted(root_dirs)[-1]
+            else:
+                run_dir_name = root_dirs[0]
+
+            run_root = os.path.join(output_dir, run_dir_name)
+
+            # 2. now point at predictions/<input_filename>
+            prediction_folder = os.path.join(run_root, "predictions", input_filename)
+            results = {}
+
+            # 3. structure .cif
+            if arguments.get("return_structure", False):
+                structure_file = os.path.join(
+                    prediction_folder, f"{input_filename}_model_0.cif"
+                )
+                if os.path.exists(structure_file):
+                    with open(structure_file, "r") as f:
+                        results["predicted_structure"] = f.read()
+                    results["structure_format"] = "cif"
+                else:
+                    results["structure_error"] = (
+                        f"Missing {os.path.basename(structure_file)}"
+                    )
+
+            # 4. affinity .json
+            affinity_file = os.path.join(
+                prediction_folder, f"affinity_{input_filename}.json"
+            )
+            if os.path.exists(affinity_file):
+                with open(affinity_file, "r") as f:
+                    results["affinity_prediction"] = json.load(f)
+            else:
+                results["affinity_error"] = f"Missing {os.path.basename(affinity_file)}"
+
+            return results
+
+
+if __name__ == "__main__":
+    # Example usage
+    tool = Boltz2DockingTool(tool_config={})
+    query = {
+        "protein_sequence": "ACDEFGHIKLMNPQRSTVWY",
+        "ligands": [
+            {"id": "LIG1", "smiles": "C1=CC=CC=C1"},
+        ],
+        "use_potentials": False,
+        "diffusion_samples": 1,
+        "return_structure": False,
+    }
+    result = tool.run(query)
+    pprint.pprint(result)

tooluniverse/chem_tool.py

@@ -0,0 +1,192 @@
+import requests
+from urllib.parse import quote
+
+# from rdkit import Chem
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+from indigo import Indigo
+
+
+@register_tool("ChEMBLTool")
+class ChEMBLTool(BaseTool):
+    """
+    Tool to search for molecules similar to a given compound name or SMILES using the ChEMBL Web Services API.
+    """
+
+    def __init__(self, tool_config, base_url="https://www.ebi.ac.uk/chembl/api/data"):
+        super().__init__(tool_config)
+        self.base_url = base_url
+        self.indigo = Indigo()
+
+    def run(self, arguments):
+        query = arguments.get("query")
+        similarity_threshold = arguments.get("similarity_threshold", 80)
+        max_results = arguments.get("max_results", 20)
+
+        if not query:
+            return {"error": "`query` parameter is required."}
+        return self._search_similar_molecules(query, similarity_threshold, max_results)
+
+    def get_chembl_id_by_name(self, compound_name):
+        """
+        Search ChEMBL for a compound by name and return the ChEMBL ID of the first match.
+        """
+        headers = {"Accept": "application/json"}
+        search_url = f"{self.base_url}/molecule/search.json?q={quote(compound_name)}"
+        print(search_url)
+        response = requests.get(search_url, headers=headers)
+        response.raise_for_status()
+        results = response.json().get("molecules", [])
+        if not results or not isinstance(results, list):
+            return {"error": "No valid results found for the compound name."}
+        if not results:
+            return {"error": "No results found for the compound name."}
+        top_molecules = results[:3]  # Get the top 3 results
+        chembl_ids = [
+            molecule.get("molecule_chembl_id")
+            for molecule in top_molecules
+            if molecule.get("molecule_chembl_id")
+        ]
+        if not chembl_ids:
+            return {"error": "No ChEMBL IDs found for the compound name."}
+        return {"chembl_ids": chembl_ids}
+
+    def get_smiles_pref_name_by_chembl_id(self, query):
+        """
+        Given a ChEMBL ID, return a dict with canonical SMILES and preferred name.
+        """
+        headers = {"Accept": "application/json"}
+        if query.upper().startswith("CHEMBL"):
+            molecule_url = f"{self.base_url}/molecule/{quote(query)}.json"
+            response = requests.get(molecule_url, headers=headers)
+            response.raise_for_status()
+            molecule = response.json()
+            if not molecule or not isinstance(molecule, dict):
+                return {"error": "No valid molecule found for the given ChEMBL ID."}
+            molecule_structures = molecule.get("molecule_structures")
+            if not molecule_structures or not isinstance(molecule_structures, dict):
+                return {
+                    "error": "Molecule structures not found or invalid for the ChEMBL ID."
+                }
+            smiles = molecule_structures.get("canonical_smiles")
+            pref_name = molecule.get("pref_name")
+            if not smiles:
+                return {"error": "SMILES not found for the given ChEMBL ID."}
+            return {"smiles": smiles, "pref_name": pref_name}
+        else:
+            return None
+
+    def get_chembl_smiles_pref_name_id_by_name(self, compound_name):
+        """
+        Search ChEMBL for a compound by name and return a list of dicts with ChEMBL ID, canonical SMILES, and preferred name for the top 5 matches.
+        """
+        headers = {"Accept": "application/json"}
+        search_url = f"{self.base_url}/molecule/search.json?q={quote(compound_name)}"
+        response = requests.get(search_url, headers=headers)
+        response.raise_for_status()
+        results = response.json().get("molecules", [])
+        if not results or not isinstance(results, list):
+            return {"error": "No valid results found for the compound name."}
+        top_molecules = results[:5]
+        output = []
+        for molecule in top_molecules:
+            chembl_id = molecule.get("molecule_chembl_id", None)
+            molecule_structures = molecule.get("molecule_structures", {})
+            if molecule_structures is not None:
+                smiles = molecule_structures.get("canonical_smiles", None)
+            else:
+                smiles = None
+            pref_name = molecule.get("pref_name")
+            if chembl_id and smiles:
+                output.append(
+                    {"chembl_id": chembl_id, "smiles": smiles, "pref_name": pref_name}
+                )
+            elif chembl_id and not smiles:
+                smiles_pre_name_dict = self.get_smiles_pref_name_by_chembl_id(chembl_id)
+                if (
+                    isinstance(smiles_pre_name_dict, dict)
+                    and "error" not in smiles_pre_name_dict
+                ):
+                    output.append(
+                        {
+                            "chembl_id": chembl_id,
+                            "smiles": smiles_pre_name_dict["smiles"],
+                            "pref_name": smiles_pre_name_dict.get("pref_name"),
+                        }
+                    )
+        if not output:
+            return {"error": "No ChEMBL IDs or SMILES found for the compound name."}
+        return output
+
+    def _search_similar_molecules(self, query, similarity_threshold, max_results):
+        headers = {"Accept": "application/json"}
+
+        smiles_info_list = []
+
+        # If the query looks like a ChEMBL ID, fetch its SMILES and pref_name
+        if isinstance(query, str) and query.upper().startswith("CHEMBL"):
+            result = self.get_smiles_pref_name_by_chembl_id(query)
+            if isinstance(result, dict) and "error" in result:
+                return result
+            smiles_info_list.append(
+                {
+                    "chembl_id": query,
+                    "smiles": result["smiles"],
+                    "pref_name": result.get("pref_name"),
+                }
+            )
+
+        # If not a ChEMBL ID, use get_chembl_smiles_pref_name_id_by_name to get info
+        if len(smiles_info_list) == 0 and isinstance(query, str):
+            results = self.get_chembl_smiles_pref_name_id_by_name(query)
+            if isinstance(results, dict) and "error" in results:
+                return results
+            for item in results:
+                smiles_info_list.append(item)
+
+        if len(smiles_info_list) == 0:
+            return {"error": "SMILES representation not found for the compound."}
+
+        results_list = []
+        for info in smiles_info_list:
+            smiles = info["smiles"]
+            pref_name = info.get("pref_name")
+            chembl_id = info.get("chembl_id")
+            mol = self.indigo.loadMolecule(smiles)
+            if mol is None:
+                return {"error": "Failed to load molecule with Indigo."}
+
+            encoded_smiles = quote(smiles)
+            similarity_url = f"{self.base_url}/similarity/{encoded_smiles}/{similarity_threshold}.json?limit={max_results}"
+            sim_response = requests.get(similarity_url, headers=headers)
+            sim_response.raise_for_status()
+            sim_results = sim_response.json().get("molecules", [])
+            similar_molecules = []
+            for mol in sim_results:
+                sim_chembl_id = mol.get("molecule_chembl_id")
+                sim_pref_name = mol.get("pref_name", "N/A")
+                mol_structures = mol.get("molecule_structures", {})
+                if mol_structures is None:
+                    continue
+                mol_smiles = mol_structures.get("canonical_smiles", "N/A")
+                similarity = mol.get("similarity", "N/A")
+                similar_molecules.append(
+                    {
+                        "chembl_id": sim_chembl_id,
+                        "pref_name": sim_pref_name,
+                        "smiles": mol_smiles,
+                        "similarity": similarity,
+                    }
+                )
+            if len(similar_molecules) == 0:
+                continue
+            results_list.append(
+                {
+                    "chembl_id": chembl_id,
+                    "pref_name": pref_name,
+                    "smiles": smiles,
+                    "similar_molecules": similar_molecules,
+                }
+            )
+
+        return results_list

tooluniverse/compose_scripts/__init__.py

@@ -0,0 +1 @@
+# Initialize the compose_scripts package