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

tooluniverse/graphql_tool.py

@@ -2,9 +2,9 @@ from graphql import build_schema
 from graphql.language import parse
 from graphql.validation import validate
 from .base_tool import BaseTool
+from .tool_registry import register_tool
 import requests
 import copy
-import json
 
 
 def validate_query(query_str, schema_str):
@@ -22,39 +22,49 @@ def validate_query(query_str, schema_str):
             return True
         else:
             # Collect and return the validation errors
-            error_messages = '\n'.join(str(error)
-                                       for error in validation_errors)
+            error_messages = "\n".join(str(error) for error in validation_errors)
             return f"Query validation errors:\n{error_messages}"
     except Exception as e:
         return f"An error occurred during validation: {str(e)}"
 
+
 def remove_none_and_empty_values(json_obj):
     """Remove all key-value pairs where the value is None or an empty list"""
     if isinstance(json_obj, dict):
-        return {k: remove_none_and_empty_values(v) for k, v in json_obj.items() if v is not None and v != []}
+        return {
+            k: remove_none_and_empty_values(v)
+            for k, v in json_obj.items()
+            if v is not None and v != []
+        }
     elif isinstance(json_obj, list):
-        return [remove_none_and_empty_values(item) for item in json_obj if item is not None and item != []]
+        return [
+            remove_none_and_empty_values(item)
+            for item in json_obj
+            if item is not None and item != []
+        ]
     else:
         return json_obj
 
+
 def execute_query(endpoint_url, query, variables=None):
     response = requests.post(
-        endpoint_url, json={'query': query, 'variables': variables})
+        endpoint_url, json={"query": query, "variables": variables}
+    )
     try:
         result = response.json()
         # result = json.dumps(result, ensure_ascii=False)
         result = remove_none_and_empty_values(result)
         # Check if the response contains errors
-        if 'errors' in result:
-            print("Invalid Query: ", result['errors'])
+        if "errors" in result:
+            print("Invalid Query: ", result["errors"])
             return None
-         # Check if the data field is empty
-        elif not result.get('data') or all(not v for v in result['data'].values()):
+        # Check if the data field is empty
+        elif not result.get("data") or all(not v for v in result["data"].values()):
             print("No data returned")
             return None
         else:
             return result
-    except requests.exceptions.JSONDecodeError as e:
+    except requests.exceptions.JSONDecodeError:
         print("JSONDecodeError: Could not decode the response as JSON")
         return None
 
@@ -63,60 +73,172 @@ class GraphQLTool(BaseTool):
     def __init__(self, tool_config, endpoint_url):
         super().__init__(tool_config)
         self.endpoint_url = endpoint_url
-        self.query_schema = tool_config['query_schema']
-        self.parameters = tool_config['parameter']['properties']
+        self.query_schema = tool_config["query_schema"]
+        self.parameters = tool_config["parameter"]["properties"]
         self.default_size = 5
 
     def run(self, arguments):
         arguments = copy.deepcopy(arguments)
-        if 'size' in self.parameters and 'size' not in arguments:
-            arguments['size'] = 5
-        return execute_query(endpoint_url=self.endpoint_url, query=self.query_schema, variables=arguments)
+        if "size" in self.parameters and "size" not in arguments:
+            arguments["size"] = 5
+        return execute_query(
+            endpoint_url=self.endpoint_url, query=self.query_schema, variables=arguments
+        )
 
 
+@register_tool("OpenTarget")
 class OpentargetTool(GraphQLTool):
     def __init__(self, tool_config):
-        endpoint_url = 'https://api.platform.opentargets.org/api/v4/graphql'
+        endpoint_url = "https://api.platform.opentargets.org/api/v4/graphql"
         super().__init__(tool_config, endpoint_url)
-        
+
     def run(self, arguments):
-        for each_arg, arg_value in arguments.items(): # opentarget api cannot handle '-' in the arguments
-            if isinstance(arg_value, str): 
-                if '-' in arg_value:
-                    arguments[each_arg] = arg_value.replace('-', ' ')
-        return super().run(arguments)
-            
+        # First try without modifying '-'
+        result = super().run(arguments)
+
+        # If no results, try with '-' replaced by ' '
+        if result is None:
+            if "drugName" in arguments and isinstance(arguments["drugName"], str):
+                arguments["drugName"] = arguments["drugName"].split("-")[0]
+            modified_arguments = copy.deepcopy(arguments)
+            for each_arg, arg_value in modified_arguments.items():
+                if isinstance(arg_value, str) and "-" in arg_value:
+                    modified_arguments[each_arg] = arg_value.replace("-", " ")
+            result = super().run(modified_arguments)
+            return result
 
+        return result
+
+
+@register_tool("OpentargetToolDrugNameMatch")
 class OpentargetToolDrugNameMatch(GraphQLTool):
     def __init__(self, tool_config, drug_generic_tool=None):
-        endpoint_url = 'https://api.platform.opentargets.org/api/v4/graphql'
+        endpoint_url = "https://api.platform.opentargets.org/api/v4/graphql"
         self.drug_generic_tool = drug_generic_tool
-        self.possible_drug_name_args = ['drugName']
+        self.possible_drug_name_args = ["drugName"]
         super().__init__(tool_config, endpoint_url)
 
     def run(self, arguments):
         arguments = copy.deepcopy(arguments)
-        results = execute_query(endpoint_url=self.endpoint_url, query=self.query_schema, variables=arguments)
+        results = execute_query(
+            endpoint_url=self.endpoint_url, query=self.query_schema, variables=arguments
+        )
         if results is None:
-            print("No results found for the drug brand name. Trying with the generic name.")
+            print(
+                "No results found for the drug brand name. Trying with the generic name."
+            )
             name_arguments = {}
             for each_args in self.possible_drug_name_args:
                 if each_args in arguments:
-                    name_arguments['drug_name'] = arguments[each_args]
+                    name_arguments["drug_name"] = arguments[each_args]
                     break
-            if len(name_arguments)==0:
+            if len(name_arguments) == 0:
                 print("No drug name found in the arguments.")
                 return None
             drug_name_results = self.drug_generic_tool.run(name_arguments)
-            if drug_name_results is not None and 'openfda.generic_name' in drug_name_results:
-                arguments[each_args] = drug_name_results['openfda.generic_name']
-                print("Found generic name. Trying with the generic name: ", arguments[each_args])
-                results = execute_query(endpoint_url=self.endpoint_url, query=self.query_schema, variables=arguments)
+            if (
+                drug_name_results is not None
+                and "openfda.generic_name" in drug_name_results
+            ):
+                arguments[each_args] = drug_name_results["openfda.generic_name"]
+                print(
+                    "Found generic name. Trying with the generic name: ",
+                    arguments[each_args],
+                )
+                results = execute_query(
+                    endpoint_url=self.endpoint_url,
+                    query=self.query_schema,
+                    variables=arguments,
+                )
         return results
-            
 
 
+@register_tool("OpenTargetGenetics")
 class OpentargetGeneticsTool(GraphQLTool):
     def __init__(self, tool_config):
-        endpoint_url = 'https://api.genetics.opentargets.org/graphql'
+        endpoint_url = "https://api.genetics.opentargets.org/graphql"
         super().__init__(tool_config, endpoint_url)
+
+
+@register_tool("DiseaseTargetScoreTool")
+class DiseaseTargetScoreTool(GraphQLTool):
+    """Tool to extract disease-target association scores from specific data sources"""
+
+    def __init__(self, tool_config, datasource_id=None):
+        endpoint_url = "https://api.platform.opentargets.org/api/v4/graphql"
+        # Get datasource_id from config if not provided as parameter
+        self.datasource_id = datasource_id or tool_config.get("datasource_id")
+        super().__init__(tool_config, endpoint_url)
+
+    def run(self, arguments):
+        """
+        Extract disease-target scores for a specific datasource
+        Arguments should contain: efoId, datasourceId (optional), pageSize (optional)
+        """
+        arguments = copy.deepcopy(arguments)
+        efo_id = arguments.get("efoId")
+        datasource_id = arguments.get("datasourceId", self.datasource_id)
+        page_size = arguments.get("pageSize", 100)
+
+        if not efo_id:
+            return {"error": "efoId is required"}
+        if not datasource_id:
+            return {"error": "datasourceId is required"}
+
+        results = []
+        page_index = 0
+        total_fetched = 0
+        total_count = None
+        disease_info = None
+
+        while True:
+            variables = {"efoId": efo_id, "index": page_index, "size": page_size}
+
+            response_data = execute_query(
+                self.endpoint_url, self.query_schema, variables
+            )
+            if not response_data or "data" not in response_data:
+                break
+
+            disease_data = response_data["data"]["disease"]
+            if not disease_data:
+                break
+
+            if disease_info is None:
+                disease_info = {
+                    "disease_id": disease_data["id"],
+                    "disease_name": disease_data["name"],
+                }
+
+            rows = disease_data["associatedTargets"]["rows"]
+            if total_count is None:
+                total_count = disease_data["associatedTargets"]["count"]
+
+            for row in rows:
+                symbol = row["target"]["approvedSymbol"]
+                target_id = row["target"]["id"]
+                score_entry = next(
+                    (ds for ds in row["datasourceScores"] if ds["id"] == datasource_id),
+                    None,
+                )
+                if score_entry:
+                    results.append(
+                        {
+                            "target_symbol": symbol,
+                            "target_id": target_id,
+                            "datasource": datasource_id,
+                            "score": score_entry["score"],
+                        }
+                    )
+
+            total_fetched += len(rows)
+            if total_fetched >= total_count or len(rows) == 0:
+                break
+            page_index += 1
+
+        return {
+            "disease_info": disease_info,
+            "datasource": datasource_id,
+            "total_targets_with_scores": len(results),
+            "target_scores": results,
+        }

tooluniverse/gwas_tool.py

@@ -0,0 +1,358 @@
+import requests
+from typing import Dict, Any, Optional
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+
+
+class GWASRESTTool(BaseTool):
+    """Base class for GWAS Catalog REST API tools."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.base_url = "https://www.ebi.ac.uk/gwas/rest/api"
+        self.endpoint = ""  # Will be set by subclasses
+
+    def _make_request(
+        self, endpoint: str, params: Optional[Dict] = None
+    ) -> Dict[str, Any]:
+        """Make a request to the GWAS Catalog API."""
+        url = f"{self.base_url}{endpoint}"
+        try:
+            response = requests.get(url, params=params, timeout=30)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            return {"error": f"Request failed: {str(e)}"}
+
+    def _extract_embedded_data(
+        self, data: Dict[str, Any], data_type: str
+    ) -> Dict[str, Any]:
+        """Extract data from the _embedded structure and add metadata."""
+        if "error" in data:
+            return data
+
+        result: Dict[str, Any] = {"data": [], "metadata": {}}
+        metadata: Dict[str, Any] = {}
+
+        # Extract the main data from _embedded
+        if "_embedded" in data and data_type in data["_embedded"]:
+            result["data"] = data["_embedded"][data_type]
+
+        # Extract pagination metadata
+        if "page" in data:
+            metadata["pagination"] = data["page"]
+
+        # Extract links metadata
+        if "_links" in data:
+            metadata["links"] = data["_links"]
+
+        if metadata:
+            result["metadata"] = metadata
+
+        # If no _embedded structure, return the original data
+        if not result["data"]:
+            result["data"] = data
+
+        return result
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Execute the tool with given arguments."""
+        return self._make_request(self.endpoint, arguments)
+
+
+@register_tool("GWASAssociationSearch")
+class GWASAssociationSearch(GWASRESTTool):
+    """Search for GWAS associations by various criteria."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/associations"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Search for associations with optional filters."""
+        params = {}
+
+        # Handle various search parameters based on examples
+        if "efo_trait" in arguments:
+            params["efo_trait"] = arguments["efo_trait"]
+        if "rs_id" in arguments:
+            params["rs_id"] = arguments["rs_id"]
+        if "accession_id" in arguments:
+            params["accession_id"] = arguments["accession_id"]
+        if "sort" in arguments:
+            params["sort"] = arguments["sort"]
+        if "direction" in arguments:
+            params["direction"] = arguments["direction"]
+        if "size" in arguments:
+            params["size"] = arguments["size"]
+        if "page" in arguments:
+            params["page"] = arguments["page"]
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "associations")
+
+
+@register_tool("GWASStudySearch")
+class GWASStudySearch(GWASRESTTool):
+    """Search for GWAS studies by various criteria."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/studies"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Search for studies with optional filters."""
+        params = {}
+
+        if "efo_trait" in arguments:
+            params["efo_trait"] = arguments["efo_trait"]
+        if "disease_trait" in arguments:
+            params["disease_trait"] = arguments["disease_trait"]
+        if "cohort" in arguments:
+            params["cohort"] = arguments["cohort"]
+        if "gxe" in arguments:
+            params["gxe"] = arguments["gxe"]
+        if "full_pvalue_set" in arguments:
+            params["full_pvalue_set"] = arguments["full_pvalue_set"]
+        if "size" in arguments:
+            params["size"] = arguments["size"]
+        if "page" in arguments:
+            params["page"] = arguments["page"]
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "studies")
+
+
+@register_tool("GWASSNPSearch")
+class GWASSNPSearch(GWASRESTTool):
+    """Search for GWAS single nucleotide polymorphisms (SNPs)."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/single-nucleotide-polymorphisms"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Search for SNPs with optional filters."""
+        params = {}
+
+        if "rs_id" in arguments:
+            params["rs_id"] = arguments["rs_id"]
+        if "mapped_gene" in arguments:
+            params["mapped_gene"] = arguments["mapped_gene"]
+        if "size" in arguments:
+            params["size"] = arguments["size"]
+        if "page" in arguments:
+            params["page"] = arguments["page"]
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "snps")
+
+
+# Get by ID tools
+@register_tool("GWASAssociationByID")
+class GWASAssociationByID(GWASRESTTool):
+    """Get a specific GWAS association by its ID."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/associations"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get association by ID."""
+        if "association_id" not in arguments:
+            return {"error": "association_id is required"}
+
+        association_id = arguments["association_id"]
+        return self._make_request(f"{self.endpoint}/{association_id}")
+
+
+@register_tool("GWASStudyByID")
+class GWASStudyByID(GWASRESTTool):
+    """Get a specific GWAS study by its ID."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/studies"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get study by ID."""
+        if "study_id" not in arguments:
+            return {"error": "study_id is required"}
+
+        study_id = arguments["study_id"]
+        return self._make_request(f"{self.endpoint}/{study_id}")
+
+
+@register_tool("GWASSNPByID")
+class GWASSNPByID(GWASRESTTool):
+    """Get a specific GWAS SNP by its rs ID."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/single-nucleotide-polymorphisms"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get SNP by rs ID."""
+        if "rs_id" not in arguments:
+            return {"error": "rs_id is required"}
+
+        rs_id = arguments["rs_id"]
+        return self._make_request(f"{self.endpoint}/{rs_id}")
+
+
+# Specialized search tools based on common use cases from examples
+@register_tool("GWASVariantsForTrait")
+class GWASVariantsForTrait(GWASRESTTool):
+    """Get all variants associated with a specific trait."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/associations"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get variants for a trait with pagination support."""
+        if "efo_trait" not in arguments:
+            return {"error": "efo_trait is required"}
+
+        params = {
+            "efo_trait": arguments["efo_trait"],
+            "size": arguments.get("size", 200),
+            "page": arguments.get("page", 0),
+        }
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "associations")
+
+
+@register_tool("GWASAssociationsForTrait")
+class GWASAssociationsForTrait(GWASRESTTool):
+    """Get all associations for a specific trait, sorted by p-value."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/associations"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get associations for a trait, sorted by significance."""
+        if "efo_trait" not in arguments:
+            return {"error": "efo_trait is required"}
+
+        params = {
+            "efo_trait": arguments["efo_trait"],
+            "sort": "p_value",
+            "direction": "asc",
+            "size": arguments.get("size", 40),
+            "page": arguments.get("page", 0),
+        }
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "associations")
+
+
+@register_tool("GWASAssociationsForSNP")
+class GWASAssociationsForSNP(GWASRESTTool):
+    """Get all associations for a specific SNP."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/associations"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get associations for a SNP."""
+        if "rs_id" not in arguments:
+            return {"error": "rs_id is required"}
+
+        params = {
+            "rs_id": arguments["rs_id"],
+            "size": arguments.get("size", 200),
+            "page": arguments.get("page", 0),
+        }
+
+        if "sort" in arguments:
+            params["sort"] = arguments["sort"]
+        if "direction" in arguments:
+            params["direction"] = arguments["direction"]
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "associations")
+
+
+@register_tool("GWASStudiesForTrait")
+class GWASStudiesForTrait(GWASRESTTool):
+    """Get studies for a specific trait with optional filters."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/studies"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get studies for a trait with optional filters."""
+        if "efo_trait" not in arguments and "disease_trait" not in arguments:
+            return {"error": "efo_trait or disease_trait is required"}
+
+        params = {
+            "size": arguments.get("size", 200),
+            "page": arguments.get("page", 0),
+        }
+
+        if "efo_trait" in arguments:
+            params["efo_trait"] = arguments["efo_trait"]
+        if "disease_trait" in arguments:
+            params["disease_trait"] = arguments["disease_trait"]
+        if "cohort" in arguments:
+            params["cohort"] = arguments["cohort"]
+        if "gxe" in arguments:
+            params["gxe"] = arguments["gxe"]
+        if "full_pvalue_set" in arguments:
+            params["full_pvalue_set"] = arguments["full_pvalue_set"]
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "studies")
+
+
+@register_tool("GWASSNPsForGene")
+class GWASSNPsForGene(GWASRESTTool):
+    """Get SNPs mapped to a specific gene."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/single-nucleotide-polymorphisms"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get SNPs for a gene."""
+        if "mapped_gene" not in arguments:
+            return {"error": "mapped_gene is required"}
+
+        params = {
+            "mapped_gene": arguments["mapped_gene"],
+            "size": arguments.get("size", 10000),
+            "page": arguments.get("page", 0),
+        }
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "snps")
+
+
+@register_tool("GWASAssociationsForStudy")
+class GWASAssociationsForStudy(GWASRESTTool):
+    """Get all associations for a specific study."""
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.endpoint = "/v2/associations"
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Get associations for a study."""
+        if "accession_id" not in arguments:
+            return {"error": "accession_id is required"}
+
+        params = {
+            "accession_id": arguments["accession_id"],
+            "sort": "p_value",
+            "direction": "asc",
+            "size": arguments.get("size", 200),
+            "page": arguments.get("page", 0),
+        }
+
+        data = self._make_request(self.endpoint, params)
+        return self._extract_embedded_data(data, "associations")