import os, json, pickle, pandas as pd, numpy as np from sklearn.metrics.pairwise import cosine_similarity from langchain_openai import AzureOpenAIEmbeddings from tqdm.auto import tqdm from typing import List, Dict, Union, Set, Optional, Any import requests from Bio.Blast import NCBIWWW, NCBIXML from Bio.Seq import Seq from Bio import Entrez # Removed biomni dependency to avoid environment variable loading side effects import traceback import time from smolagents import tool, OpenAIServerModel OPENROUTER_API_KEY_STRING = "" # Use absolute path for schema database SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) AGENTS_STELLA_DIR = os.path.dirname(SCRIPT_DIR) SCHEMA_DB_PATH = os.path.join(AGENTS_STELLA_DIR, "resource", "schema_db") def parse_hpo_obo(obo_file_path: str) -> Dict[str, str]: """ Simple HPO OBO file parser to extract term IDs and names. Args: obo_file_path (str): Path to the HPO OBO file Returns: Dict[str, str]: Dictionary mapping HPO IDs to names """ hpo_dict = {} # Handle relative paths by making them absolute relative to script location if not os.path.isabs(obo_file_path): obo_file_path = os.path.join(SCRIPT_DIR, obo_file_path) try: if not os.path.exists(obo_file_path): print(f"Warning: HPO OBO file not found at {obo_file_path}") return hpo_dict with open(obo_file_path, 'r', encoding='utf-8') as f: current_term = {} in_term_block = False for line in f: line = line.strip() if line == '[Term]': in_term_block = True current_term = {} elif line == '' and in_term_block: # End of term block if 'id' in current_term and 'name' in current_term: hpo_dict[current_term['id']] = current_term['name'] in_term_block = False elif in_term_block and ':' in line: key, value = line.split(':', 1) key = key.strip() value = value.strip() if key == 'id': current_term['id'] = value elif key == 'name': current_term['name'] = value # Handle last term if file doesn't end with empty line if in_term_block and 'id' in current_term and 'name' in current_term: hpo_dict[current_term['id']] = current_term['name'] except Exception as e: print(f"Error parsing HPO OBO file: {e}") return hpo_dict gemini_model = OpenAIServerModel( model_id="google/gemini-2.5-pro", api_base="https://openrouter.ai/api/v1", api_key=OPENROUTER_API_KEY_STRING, temperature=0.1, # Lower temperature for more consistent analysis ) def _query_gemini_for_api(prompt, schema, system_template, model=None): """ Helper function to query Gemini for generating API calls based on natural language prompts. Args: prompt (str): Natural language query to process schema (dict): API schema to include in the system prompt system_template (str): Template string for the system prompt (should have {schema} placeholder) model: Gemini model instance to use (defaults to global gemini_model) Returns: dict: Dictionary with 'success', 'data' (if successful), 'error' (if failed), and optional 'raw_response' """ # Use global gemini_model if none provided model = gemini_model try: if schema is not None: # Format the system prompt with the schema schema_json = json.dumps(schema, indent=2) system_prompt = system_template.format(schema=schema_json) else: system_prompt = system_template # Combine system prompt and user prompt for Gemini full_prompt = f"{system_prompt}\n\nUser query: {prompt}" # Create messages in the correct format for OpenAIServerModel messages = [{"role": "user", "content": full_prompt}] response = model(messages) # Extract content from ChatMessage response if hasattr(response, 'content'): gemini_text = response.content.strip() else: gemini_text = str(response).strip() # Find JSON boundaries (in case Gemini adds explanations) json_start = gemini_text.find('{') json_end = gemini_text.rfind('}') + 1 if json_start >= 0 and json_end > json_start: json_text = gemini_text[json_start:json_end] result = json.loads(json_text) else: # If no JSON found, try the whole response result = json.loads(gemini_text) return { "success": True, "data": result, "raw_response": gemini_text } except (json.JSONDecodeError, KeyError, IndexError) as e: return { "success": False, "error": f"Failed to parse Gemini's response: {str(e)}", "raw_response": gemini_text if 'gemini_text' in locals() else "No content found" } except Exception as e: return { "success": False, "error": f"Error querying Gemini: {str(e)}" } # Function to map HPO terms to names def get_hpo_names(hpo_terms: List[str], ) -> List[str]: """ Retrieve the names of given HPO terms. Args: hpo_terms (List[str]): A list of HPO terms (e.g., ['HP:0001250']). Returns: List[str]: A list of corresponding HPO term names. """ hp_dict = parse_hpo_obo(os.path.join(AGENTS_STELLA_DIR, 'resource', 'hp.obo')) hpo_names = [] for term in hpo_terms: name = hp_dict.get(term, f"Unknown term: {term}") hpo_names.append(name) return hpo_names def _query_rest_api(endpoint, method="GET", params=None, headers=None, json_data=None, description=None): """ General helper function to query REST APIs with consistent error handling. Args: endpoint (str): Full URL endpoint to query method (str): HTTP method ("GET" or "POST") params (dict, optional): Query parameters to include in the URL headers (dict, optional): HTTP headers for the request json_data (dict, optional): JSON data for POST requests description (str, optional): Description of this query for error messages Returns: dict: Dictionary containing the result or error information """ # Set default headers if not provided if headers is None: headers = {"Accept": "application/json"} # Set default description if not provided if description is None: description = f"{method} request to {endpoint}" url_error = None try: # Make the API request if method.upper() == "GET": response = requests.get(endpoint, params=params, headers=headers) elif method.upper() == "POST": response = requests.post(endpoint, params=params, headers=headers, json=json_data) else: return {"error": f"Unsupported HTTP method: {method}"} url_error = str(response.text) response.raise_for_status() # Try to parse JSON response try: result = response.json() except ValueError: # Return raw text if not JSON result = {"raw_text": response.text} return { "success": True, "query_info": { "endpoint": endpoint, "method": method, "description": description }, "result": result } except requests.exceptions.RequestException as e: error_msg = str(e) response_text = "" # Try to get more detailed error info from response if hasattr(e, 'response') and e.response: try: error_json = e.response.json() if 'messages' in error_json: error_msg = "; ".join(error_json['messages']) elif 'message' in error_json: error_msg = error_json['message'] elif 'error' in error_json: error_msg = error_json['error'] elif 'detail' in error_json: error_msg = error_json['detail'] except: response_text = e.response.text return { "success": False, "error": f"API error: {error_msg}", "query_info": { "endpoint": endpoint, "method": method, "description": description }, "response_url_error": url_error, "response_text": response_text } except Exception as e: return { "success": False, "error": f"Error: {str(e)}", "query_info": { "endpoint": endpoint, "method": method, "description": description } } def _query_ncbi_database( database: str, search_term: str, result_formatter = None, max_results: int = 3, ) -> Dict[str, Any]: """ Core function to query NCBI databases using Claude for query interpretation and NCBI eutils. Args: database (str): NCBI database to query (e.g., "clinvar", "gds", "geoprofiles") result_formatter (callable): Function to format results from the database api_key (str): Anthropic API key. If None, will look for ANTHROPIC_API_KEY environment variable model (str): Anthropic model to use max_results (int): Maximum number of results to return verbose (bool): Whether to return verbose results Returns: dict: Dictionary containing both the structured query and the results """ # Query NCBI API using the structured search term esearch_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi" esearch_params = { "db": database, "term": search_term, "retmode": "json", "retmax": 100, "usehistory": "y" # Use history server to store results } # Get IDs of matching entries search_response = _query_rest_api( endpoint=esearch_url, method="GET", params=esearch_params, description="NCBI ESearch API query" ) if not search_response["success"]: return search_response search_data = search_response["result"] # If we have results, fetch the details if "esearchresult" in search_data and int(search_data["esearchresult"]["count"]) > 0: # Extract WebEnv and query_key from the search results webenv = search_data["esearchresult"].get("webenv", "") query_key = search_data["esearchresult"].get("querykey", "") # Use WebEnv and query_key if available if webenv and query_key: # Get details using eSummary esummary_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi" esummary_params = { "db": database, "query_key": query_key, "WebEnv": webenv, "retmode": "json", "retmax": max_results } details_response = _query_rest_api( endpoint=esummary_url, method="GET", params=esummary_params, description="NCBI ESummary API query" ) if not details_response["success"]: return details_response results = details_response["result"] else: # Fall back to direct ID fetch id_list = search_data["esearchresult"]["idlist"][:max_results] # Get details for each ID esummary_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi" esummary_params = { "db": database, "id": ",".join(id_list), "retmode": "json" } details_response = _query_rest_api( endpoint=esummary_url, method="GET", params=esummary_params, description="NCBI ESummary API query" ) if not details_response["success"]: return details_response results = details_response["result"] # Format results using the provided formatter if result_formatter: formatted_results = result_formatter(results) else: formatted_results = results # Return the combined information return { "database": database, "query_interpretation": search_term, "total_results": int(search_data["esearchresult"]["count"]), "formatted_results": formatted_results } else: return { "database": database, "query_interpretation": search_term, "total_results": 0, "formatted_results": [] } def _format_query_results(result, options=None): """ A general-purpose formatter for query function results to reduce output size. Args: result (dict): The original API response dictionary options (dict, optional): Formatting options including: - max_items (int): Maximum number of items to include in lists (default: 5) - max_depth (int): Maximum depth to traverse in nested dictionaries (default: 2) - include_keys (list): Only include these top-level keys (overrides exclude_keys) - exclude_keys (list): Exclude these keys from the output - summarize_lists (bool): Whether to summarize long lists (default: True) - truncate_strings (int): Maximum length for string values (default: 100) Returns: dict: A condensed version of the input results """ def _format_value(value, depth, options): """ Recursively format a value based on its type and formatting options. Args: value: The value to format depth (int): Current recursion depth options (dict): Formatting options Returns: Formatted value """ # Base case: reached max depth if depth >= options['max_depth'] and (isinstance(value, dict) or isinstance(value, list)): if isinstance(value, dict): return { '_summary': f'Nested dictionary with {len(value)} keys', '_keys': list(value.keys())[:options['max_items']] } else: # list return _summarize_list(value, options) # Process based on type if isinstance(value, dict): return _format_dict(value, depth, options) elif isinstance(value, list): return _format_list(value, depth, options) elif isinstance(value, str) and len(value) > options['truncate_strings']: return value[:options['truncate_strings']] + "... (truncated)" else: return value def _format_dict(d, depth, options): """Format a dictionary according to options.""" result = {} # Filter keys based on include/exclude options keys_to_process = d.keys() if depth == 0 and options['include_keys']: # Only apply at top level keys_to_process = [k for k in keys_to_process if k in options['include_keys']] elif depth == 0 and options['exclude_keys']: # Only apply at top level keys_to_process = [k for k in keys_to_process if k not in options['exclude_keys']] # Process each key for key in keys_to_process: result[key] = _format_value(d[key], depth + 1, options) return result def _format_list(lst, depth, options): """Format a list according to options.""" if options['summarize_lists'] and len(lst) > options['max_items']: return _summarize_list(lst, options) result = [] for i, item in enumerate(lst): if i >= options['max_items']: remaining = len(lst) - options['max_items'] result.append(f"... {remaining} more items (omitted)") break result.append(_format_value(item, depth + 1, options)) return result def _summarize_list(lst, options): """Create a summary for a list.""" if not lst: return [] # Sample a few items sample = lst[:min(3, len(lst))] sample_formatted = [_format_value(item, options['max_depth'], options) for item in sample] # For homogeneous lists, provide type info if len(lst) > 0: item_type = type(lst[0]).__name__ homogeneous = all(isinstance(item, type(lst[0])) for item in lst) type_info = f"all {item_type}" if homogeneous else "mixed types" else: type_info = "empty" return { '_summary': f"List with {len(lst)} items ({type_info})", '_sample': sample_formatted } if options is None: options = {} # Default options default_options = { 'max_items': 5, 'max_depth': 20, 'include_keys': None, 'exclude_keys': ['raw_response', 'debug_info', 'request_details'], 'summarize_lists': True, 'truncate_strings': 100 } # Merge provided options with defaults for key, value in default_options.items(): if key not in options: options[key] = value # Filter and format the result formatted = _format_value(result, 0, options) return formatted @tool def query_uniprot(prompt: str = None, endpoint: str = None, max_results: int = 5) -> dict: """ Query the UniProt REST API using either natural language or a direct endpoint. Args: prompt: Natural language query about proteins (e.g., "Find information about human insulin") endpoint: Full or partial UniProt API endpoint URL to query directly max_results: Maximum number of results to return Returns: Dictionary containing the query information and the UniProt API results Examples: - Natural language: query_uniprot(prompt="Find information about human insulin protein") - Direct endpoint: query_uniprot(endpoint="https://rest.uniprot.org/uniprotkb/P01308") """ # Base URL for UniProt API base_url = "https://rest.uniprot.org" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load UniProt schema schema_path = os.path.join(SCHEMA_DB_PATH, "uniprot.pkl") with open(schema_path, "rb") as f: uniprot_schema = pickle.load(f) # Create system prompt template system_template = """ You are a protein biology expert specialized in using the UniProt REST API. Based on the user's natural language request, determine the appropriate UniProt REST API endpoint and parameters. UNIPROT REST API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including base URL, dataset, endpoint type, and parameters) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - Base URL is "https://rest.uniprot.org" - Search in reviewed (Swiss-Prot) entries first before using non-reviewed (TrEMBL) entries - Assume organism is human unless otherwise specified. Human taxonomy ID is 9606 - Use gene_exact: for exact gene name searches - Use specific query fields like accession:, gene:, organism_id: in search queries - Use quotes for terms with spaces: organism_name:"Homo sapiens" Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=uniprot_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Use provided endpoint directly if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = f"Direct query to provided endpoint" # Use the common REST API helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) return api_result @tool def query_alphafold( uniprot_id: str, endpoint: str = "prediction", residue_range: str = None, download: bool = False, output_dir: str = None, file_format: str = "pdb", model_version: str = "v4", model_number: int = 1, ) -> dict: """ Query the AlphaFold Database API for protein structure predictions. Args: uniprot_id: UniProt accession ID (e.g., "P12345") endpoint: Specific AlphaFold API endpoint to query: "prediction", "summary", or "annotations" residue_range: Specific residue range in format "start-end" (e.g., "1-100") download: Whether to download structure files output_dir: Directory to save downloaded files (default: current directory) file_format: Format of the structure file to download - "pdb" or "cif" model_version: AlphaFold model version - "v4" (latest) or "v3", "v2", "v1" model_number: Model number (1-5, with 1 being the highest confidence model) Returns: Dictionary containing both the query information and the AlphaFold results Examples: - Basic query: query_alphafold(uniprot_id="P53_HUMAN") - Download structure: query_alphafold(uniprot_id="P53_HUMAN", download=True, output_dir="./structures") - Get annotations: query_alphafold(uniprot_id="P53_HUMAN", endpoint="annotations") """ # Base URL for AlphaFold API base_url = "https://alphafold.ebi.ac.uk/api" # Ensure we have a UniProt ID if not uniprot_id: return {"error": "UniProt ID is required"} # Validate endpoint valid_endpoints = ["prediction", "summary", "annotations"] if endpoint not in valid_endpoints: return {"error": f"Invalid endpoint. Must be one of: {', '.join(valid_endpoints)}"} # Construct the API URL based on endpoint if endpoint == "prediction": url = f"{base_url}/prediction/{uniprot_id}" elif endpoint == "summary": url = f"{base_url}/uniprot/summary/{uniprot_id}.json" elif endpoint == "annotations": if residue_range: url = f"{base_url}/annotations/{uniprot_id}/{residue_range}" else: url = f"{base_url}/annotations/{uniprot_id}" try: # Make the API request response = requests.get(url) response.raise_for_status() # Parse the response as JSON result = response.json() # Handle download request if specified download_info = None if download: # Ensure output directory exists if not output_dir: output_dir = "." os.makedirs(output_dir, exist_ok=True) # Generate standard AlphaFold filename file_ext = file_format.lower() filename = f"AF-{uniprot_id}-F{model_number}-model_{model_version}.{file_ext}" file_path = os.path.join(output_dir, filename) # Construct download URL download_url = f"https://alphafold.ebi.ac.uk/files/{filename}" # Download the file download_response = requests.get(download_url) if download_response.status_code == 200: with open(file_path, 'wb') as f: f.write(download_response.content) download_info = { "success": True, "file_path": file_path, "url": download_url } else: download_info = { "success": False, "error": f"Failed to download file (status code: {download_response.status_code})", "url": download_url } # Return the query information and results response_data = { "query_info": { "uniprot_id": uniprot_id, "endpoint": endpoint, "residue_range": residue_range, "url": url }, "result": result } if download_info: response_data["download"] = download_info return response_data except requests.exceptions.RequestException as e: error_msg = str(e) response_text = "" # Try to get more detailed error info from response if hasattr(e, 'response') and e.response: try: error_json = e.response.json() if 'message' in error_json: error_msg = error_json['message'] except: response_text = e.response.text return { "error": f"AlphaFold API error: {error_msg}", "query_info": { "uniprot_id": uniprot_id, "endpoint": endpoint, "residue_range": residue_range, "url": url }, "response_text": response_text } except Exception as e: return { "error": f"Error: {str(e)}", "query_info": { "uniprot_id": uniprot_id, "endpoint": endpoint, "residue_range": residue_range } } @tool def query_interpro(prompt: str = None, endpoint: str = None, max_results: int = 3) -> dict: """ Query the InterPro REST API using natural language or a direct endpoint. Args: prompt: Natural language query about protein domains or families endpoint: Direct endpoint path or full URL (e.g., "/entry/interpro/IPR023411") max_results: Maximum number of results to return per page Returns: Dictionary containing both the query information and the InterPro API results Examples: - Natural language: query_interpro("Find information about kinase domains in InterPro") - Direct endpoint: query_interpro(endpoint="/entry/interpro/IPR023411") """ # Base URL for InterPro API base_url = "https://www.ebi.ac.uk/interpro/api" # Default parameters format = "json" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load InterPro schema schema_path = os.path.join(SCHEMA_DB_PATH, "interpro.pkl") with open(schema_path, "rb") as f: interpro_schema = pickle.load(f) # Create system prompt template system_template = """ You are a protein domain expert specialized in using the InterPro REST API. Based on the user's natural language request, determine the appropriate InterPro REST API endpoint. INTERPRO REST API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://www.ebi.ac.uk/interpro/api") 2. "description": A brief description of what the query is doing SPECIAL NOTES: - Path components for data types: entry, protein, structure, set, taxonomy, proteome - Common sources: interpro, pfam, cdd, uniprot, pdb - Protein subtypes can be "reviewed" or "unreviewed" - For specific entries, use lowercase accessions (e.g., "ipr000001" instead of "IPR000001") - Endpoints can be hierarchical like "/entry/interpro/protein/uniprot/P04637" Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=interpro_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Extract the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint # If it's just a path, add the base URL if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = "Direct query to provided endpoint" # Add pagination parameters params = {"page": 1, "page_size": max_results} # Add format parameter if not json if format and format != "json": params["format"] = format # Make the API request api_result = _query_rest_api( endpoint=endpoint, method="GET", params=params, description=description ) return api_result @tool def query_pdb(prompt: str = None, query: dict = None, max_results: int = 3) -> dict: """ Query the RCSB PDB database using natural language or a direct structured query. Args: prompt: Natural language query about protein structures query: Direct structured query in RCSB Search API format (overrides prompt) max_results: Maximum number of results to return Returns: Dictionary containing the structured query, search results, and identifiers Examples: - Natural language: query_pdb("Find structures of human insulin") - Direct query: query_pdb(query={"query": {"type": "terminal", "service": "full_text", "parameters": {"value": "insulin"}}, "return_type": "entry"}) """ # Default parameters return_type = "entry" search_service = "full_text" # Generate search query from natural language if prompt is provided and query is not if prompt and not query: # Load schema from pickle file schema_path = os.path.join(SCHEMA_DB_PATH, "pdb.pkl") with open(schema_path, "rb") as f: schema = pickle.load(f) # Create system prompt template system_template = """ You are a structural biology expert that creates precise RCSB PDB Search API queries based on natural language requests. SEARCH API SCHEMA: {schema} IMPORTANT GUIDELINES: 1. Choose the appropriate search_service based on the query: - Use "text" for attribute-specific searches (REQUIRES attribute, operator, and value) - Use "full_text" for general keyword searches across multiple fields - Use appropriate specialized services for sequence, structure, motif searches 2. For "text" searches, you MUST specify: - attribute: The specific field to search (use common_attributes from schema) - operator: The comparison method (exact_match, contains_words, less_or_equal, etc.) - value: The search term or value 3. For "full_text" searches, only specify: - value: The search term(s) 4. For combined searches, use "group" nodes with logical_operator ("and" or "or") 5. Always specify the appropriate return_type based on what the user is looking for Generate a well-formed Search API query JSON object. Return ONLY the JSON with no additional explanation. """ # Query Gemini to generate the search query gemini_result = _query_gemini_for_api( prompt=prompt, schema=schema, system_template=system_template ) if not gemini_result["success"]: return { "error": gemini_result["error"], "gemini_response": gemini_result.get("raw_response", "No response") } # Get the query from Gemini's response query_json = gemini_result["data"] else: # Use provided query directly query_json = query if query else { "query": { "type": "terminal", "service": search_service, "parameters": { "value": prompt } }, "return_type": return_type } # Ensure return_type is set if "return_type" not in query_json: query_json["return_type"] = return_type # Add request options for pagination, but avoid conflicts with return_all_hits if "request_options" not in query_json: query_json["request_options"] = {} # Only add pagination if return_all_hits is not set if "return_all_hits" in query_json["request_options"] and query_json["request_options"]["return_all_hits"]: # Remove return_all_hits and use pagination instead for limited results query_json["request_options"]["return_all_hits"] = False if "paginate" not in query_json["request_options"]: query_json["request_options"]["paginate"] = { "start": 0, "rows": max_results } # Use query_rest_api to execute the search search_url = "https://search.rcsb.org/rcsbsearch/v2/query" api_result = _query_rest_api( endpoint=search_url, method="POST", json_data=query_json, description="PDB Search API query" ) return api_result @tool def query_pdb_identifiers(identifiers: List[str], return_type: str = "entry", download: bool = False, attributes: List[str] = None) -> dict: """ Retrieve detailed data and/or download files for PDB identifiers. Args: identifiers: List of PDB identifiers (from query_pdb) return_type: Type of results: "entry", "assembly", "polymer_entity", etc. download: Whether to download PDB structure files attributes: List of specific attributes to retrieve Returns: Dictionary containing the detailed data and file paths if downloaded Example: - Search and then get details: results = query_pdb("Find structures of human insulin") details = get_pdb_details(results["identifiers"], download=True) """ if not identifiers: return {"error": "No identifiers provided"} try: # Fetch detailed data using Data API detailed_results = [] for identifier in identifiers: try: # Determine the appropriate endpoint based on return_type and identifier format if return_type == "entry": data_url = f"https://data.rcsb.org/rest/v1/core/entry/{identifier}" elif return_type == "polymer_entity": entry_id, entity_id = identifier.split('_') data_url = f"https://data.rcsb.org/rest/v1/core/polymer_entity/{entry_id}/{entity_id}" elif return_type == "nonpolymer_entity": entry_id, entity_id = identifier.split('_') data_url = f"https://data.rcsb.org/rest/v1/core/nonpolymer_entity/{entry_id}/{entity_id}" elif return_type == "polymer_instance": entry_id, asym_id = identifier.split('.') data_url = f"https://data.rcsb.org/rest/v1/core/polymer_entity_instance/{entry_id}/{asym_id}" elif return_type == "assembly": entry_id, assembly_id = identifier.split('-') data_url = f"https://data.rcsb.org/rest/v1/core/assembly/{entry_id}/{assembly_id}" elif return_type == "mol_definition": data_url = f"https://data.rcsb.org/rest/v1/core/chem_comp/{identifier}" # Fetch data data_response = requests.get(data_url) data_response.raise_for_status() entity_data = data_response.json() # Filter attributes if specified if attributes: filtered_data = {} for attr in attributes: parts = attr.split('.') current = entity_data try: for part in parts[:-1]: current = current[part] filtered_data[attr] = current[parts[-1]] except (KeyError, TypeError): filtered_data[attr] = None entity_data = filtered_data detailed_results.append({ "identifier": identifier, "data": entity_data }) except Exception as e: detailed_results.append({ "identifier": identifier, "error": str(e) }) # Download structure files if requested if download: for identifier in identifiers: if '_' in identifier or '.' in identifier or '-' in identifier: # For non-entry identifiers, extract the PDB ID if '_' in identifier: pdb_id = identifier.split('_')[0] elif '.' in identifier: pdb_id = identifier.split('.')[0] elif '-' in identifier: pdb_id = identifier.split('-')[0] else: pdb_id = identifier try: # Download PDB file pdb_url = f"https://files.rcsb.org/download/{pdb_id}.pdb" pdb_response = requests.get(pdb_url) if pdb_response.status_code == 200: # Create data directory if it doesn't exist data_dir = os.path.join(os.path.dirname(__file__), "data", "pdb") os.makedirs(data_dir, exist_ok=True) # Save PDB file pdb_file_path = os.path.join(data_dir, f"{pdb_id}.pdb") with open(pdb_file_path, 'wb') as pdb_file: pdb_file.write(pdb_response.content) # Add download information to results for result in detailed_results: if result["identifier"] == identifier or result["identifier"].startswith(pdb_id): result["pdb_file_path"] = pdb_file_path except Exception as e: for result in detailed_results: if result["identifier"] == identifier or result["identifier"].startswith(pdb_id): result["download_error"] = str(e) return { "detailed_results": detailed_results } except Exception as e: return { "error": f"Error retrieving PDB details: {str(e)}" } @tool def query_kegg(prompt: str, endpoint: str = None, verbose: bool = True) -> dict: """ Take a natural language prompt and convert it to a structured KEGG API query. Args: prompt: Natural language query about KEGG data (e.g., "Find human pathways related to glycolysis") endpoint: Direct KEGG API endpoint to query verbose: Whether to return detailed results Returns: Dictionary containing both the structured query and the KEGG results """ base_url = "https://rest.kegg.jp" if not prompt and not endpoint: return {"error": "Either a prompt or an endpoint must be provided"} if prompt: # Load schema from pickle file schema_path = os.path.join(SCHEMA_DB_PATH, "kegg.pkl") with open(schema_path, "rb") as f: kegg_schema = pickle.load(f) # Create system prompt template system_template = """ You are a bioinformatics expert that helps convert natural language queries into KEGG API requests. Based on the user's natural language request, you will generate a structured query for the KEGG API. The KEGG API has the following general form: https://rest.kegg.jp//[/[/ ...]] Where can be one of: info, list, find, get, conv, link, ddi Here is the schema of available operations, databases, and other details: {schema} Output only a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://rest.kegg.jp") 2. "description": A brief description of what the query is doing IMPORTANT: Your response must ONLY contain a JSON object with the required fields. EXAMPLES OF CORRECT OUTPUTS: - For "Find information about glycolysis pathway": {{"full_url": "https://rest.kegg.jp/info/pathway/hsa00010", "description": "Finding information about the glycolysis pathway"}} - For "Get information about the human BRCA1 gene": {{"full_url": "https://rest.kegg.jp/get/hsa:672", "description": "Retrieving information about BRCA1 gene in human"}} - For "List all human pathways": {{"full_url": "https://rest.kegg.jp/list/pathway/hsa", "description": "Listing all human-specific pathways"}} - For "Convert NCBI gene ID 672 to KEGG ID": {{"full_url": "https://rest.kegg.jp/conv/genes/ncbi-geneid:672", "description": "Converting NCBI Gene ID 672 to KEGG gene identifier"}} """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=kegg_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Extract the query info from Gemini's response query_info = gemini_result["data"] endpoint = query_info["full_url"] description = query_info["description"] if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } if endpoint: if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = "Direct query to KEGG API" # Execute the KEGG API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_stringdb(prompt: str = None, endpoint: str = None, api_key: str = None, download_image: bool = False, output_dir: str = None, verbose: bool = True) -> dict: """ Query the STRING protein interaction database using natural language or direct endpoint. Args: prompt: Natural language query about protein interactions endpoint: Full URL to query directly (overrides prompt) api_key: Anthropic API key for processing model: Model to use for natural language processing download_image: Whether to download image results output_dir: Directory to save downloaded files verbose: Whether to return detailed results Returns: Dictionary containing the query results or error information Examples: - Natural language: query_stringdb("Show protein interactions for BRCA1 and BRCA2 in humans") - Direct endpoint: query_stringdb(endpoint="https://string-db.org/api/json/network?identifiers=BRCA1,BRCA2&species=9606") """ # Base URL for STRING API base_url = "https://version-12-0.string-db.org/api" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load STRING schema schema_path = os.path.join(SCHEMA_DB_PATH, "stringdb.pkl") with open(schema_path, "rb") as f: stringdb_schema = pickle.load(f) # Create system prompt template system_template = """ You are a protein interaction expert specialized in using the STRING database API. Based on the user's natural language request, determine the appropriate STRING API endpoint and parameters. STRING API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including all parameters) 2. "description": A brief description of what the query is doing 3. "output_format": The format of the output (json, tsv, image, svg) SPECIAL NOTES: - Common species IDs: 9606 (human), 10090 (mouse), 7227 (fruit fly), 4932 (yeast) - For protein identifiers, use either gene names (e.g., "BRCA1") or UniProt IDs (e.g., "P38398") - The "required_score" parameter accepts values from 0 to 1000 (higher means more stringent) - Add "caller_identity=bioagentos_api" as a parameter Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=stringdb_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") output_format = query_info.get("output_format", "json") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Use direct endpoint if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = "Direct query to STRING API" output_format = "json" # Try to determine output format from URL if "image" in endpoint or "svg" in endpoint: output_format = "image" # Check if we're dealing with an image request is_image = output_format in ["image", "highres_image", "svg"] if is_image: if download_image: # For images, we need to handle the download manually try: response = requests.get(endpoint, stream=True) response.raise_for_status() # Create output directory if needed if not output_dir: output_dir = "." os.makedirs(output_dir, exist_ok=True) # Generate filename based on endpoint endpoint_parts = endpoint.split("/") filename = f"string_{endpoint_parts[-2]}_{int(time.time())}.{output_format}" file_path = os.path.join(output_dir, filename) # Save the image with open(file_path, 'wb') as f: for chunk in response.iter_content(chunk_size=1024): if chunk: f.write(chunk) return { "success": True, "query_info": { "endpoint": endpoint, "description": description, "output_format": output_format }, "result": { "image_saved": True, "file_path": file_path, "content_type": response.headers.get('Content-Type') } } except Exception as e: return { "success": False, "error": f"Error downloading image: {str(e)}", "query_info": { "endpoint": endpoint, "description": description } } else: # Just report that an image is available but not downloaded return { "success": True, "query_info": { "endpoint": endpoint, "description": description, "output_format": output_format }, "result": { "image_available": True, "download_url": endpoint, "note": "Set download_image=True to save the image" } } # For non-image requests, use the REST API helper api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_paleobiology(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the Paleobiology Database (PBDB) API using natural language or a direct endpoint. Args: prompt (str, required): Natural language query about fossil records endpoint (str, optional): API endpoint name or full URL api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_paleobiology("Find fossil records of Tyrannosaurus rex") - Direct endpoint: query_paleobiology(endpoint="data1.2/taxa/list.json?name=Tyrannosaurus") """ # Base URL for PBDB API base_url = "https://paleobiodb.org/data1.2" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load PBDB schema schema_path = os.path.join(SCHEMA_DB_PATH, "paleobiology.pkl") with open(schema_path, "rb") as f: pbdb_schema = pickle.load(f) # Create system prompt template system_template = """ You are a paleobiology expert specialized in using the Paleobiology Database (PBDB) API. Based on the user's natural language request, determine the appropriate PBDB API endpoint and parameters. PBDB API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://paleobiodb.org/data1.2" and format extension) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - For taxonomic queries, be specific about taxonomic ranks and names - For geographic queries, use standard country/continent names or coordinate bounding boxes - For time interval queries, use standard geological time names (e.g., "Cretaceous", "Maastrichtian") - Use appropriate format extension (.json, .txt, .csv, .tsv) based on the query - If appropriate, use "vocab=pbdb" (default) or "vocab=com" (compact) parameter in the URL - For detailed occurrence data, include "show=paleoloc,phylo" in the parameters Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=pbdb_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if not endpoint.startswith("http"): # Add base URL if it's just a path if not endpoint.startswith('/'): endpoint = f"{base_url}/{endpoint}" else: endpoint = f"{base_url}{endpoint}" description = "Direct query to PBDB API" # Check if we're dealing with an image request is_image = endpoint.endswith('.png') if is_image: # For image queries, we need special handling try: response = requests.get(endpoint) response.raise_for_status() # Return image metadata without the binary data return { "success": True, "query_info": { "endpoint": endpoint, "description": description, "format": "png" }, "result": { "content_type": response.headers.get('Content-Type'), "size_bytes": len(response.content), "note": "Binary image data not included in response" } } except Exception as e: return { "success": False, "error": f"Error retrieving image: {str(e)}", "query_info": { "endpoint": endpoint, "description": description } } # For non-image requests, use the REST API helper api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_jaspar(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the JASPAR REST API using natural language or a direct endpoint. Args: prompt (str, required): Natural language query about transcription factor binding profiles endpoint (str, optional): API endpoint path (e.g., "/matrix/MA0002.2/") or full URL api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_jaspar("Find all transcription factor matrices for human") - Direct endpoint: query_jaspar(endpoint="/matrix/MA0002.2/") """ # Base URL for JASPAR API base_url = "https://jaspar.elixir.no/api/v1" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load JASPAR schema schema_path = os.path.join(SCHEMA_DB_PATH, "jaspar.pkl") with open(schema_path, "rb") as f: jaspar_schema = pickle.load(f) # Create system prompt template system_template = """ You are a transcription factor binding site expert specialized in using the JASPAR REST API. Based on the user's natural language request, determine the appropriate JASPAR REST API endpoint and parameters. JASPAR REST API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://jaspar.elixir.no/api/v1" and any parameters) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - Common taxonomic groups include: vertebrates, plants, fungi, insects, nematodes, urochordates - Common collections include: CORE, UNVALIDATED, PENDING, etc. - Matrix IDs follow the format MA####.# (e.g., MA0002.2) - For inferring matrices from sequences, provide the protein sequence directly in the path Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=jaspar_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if not endpoint.startswith("http"): # Clean up endpoint format if not endpoint.startswith("/"): endpoint = "/" + endpoint # Ensure endpoint ends with / if not endpoint.endswith("/"): endpoint = endpoint + "/" # Add base URL endpoint = f"{base_url}{endpoint}" description = "Direct query to JASPAR API" # Execute the JASPAR API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_worms(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the World Register of Marine Species (WoRMS) REST API using natural language or a direct endpoint. Args: prompt (str, required): Natural language query about marine species endpoint (str, optional): Full URL or endpoint specification api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_worms("Find information about the blue whale") - Direct endpoint: query_worms(endpoint="https://www.marinespecies.org/rest/AphiaRecordByName/Balaenoptera%20musculus") """ # Base URL for WoRMS API base_url = "https://www.marinespecies.org/rest" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load WoRMS schema schema_path = os.path.join(SCHEMA_DB_PATH, "worms.pkl") with open(schema_path, "rb") as f: worms_schema = pickle.load(f) # Create system prompt template system_template = """ You are a marine biology expert specialized in using the World Register of Marine Species (WoRMS) API. Based on the user's natural language request, determine the appropriate WoRMS API endpoint and parameters. WORMS API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://www.marinespecies.org/rest" and any path/query parameters) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - For taxonomic searches, be precise with scientific names and use proper capitalization - For fuzzy matching, include "fuzzy=true" in the URL query parameters - When searching by name, prefer "AphiaRecordByName" for exact matches and "AphiaRecordsByName" for broader results - AphiaID is the main identifier in WoRMS (e.g., Blue Whale is 137087) - For multiple IDs or names, use the appropriate POST endpoint Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=worms_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL and details from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if not endpoint.startswith("http"): # Add base URL if it's just a path if not endpoint.startswith('/'): endpoint = f"{base_url}/{endpoint}" else: endpoint = f"{base_url}{endpoint}" description = "Direct query to WoRMS API" # Execute the WoRMS API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method='GET', description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_cbioportal(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the cBioPortal REST API using natural language or a direct endpoint. Args: prompt (str, required): Natural language query about cancer genomics data endpoint (str, optional): API endpoint path (e.g., "/studies/brca_tcga/patients") or full URL api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_cbioportal("Find mutations in BRCA1 for breast cancer") - Direct endpoint: query_cbioportal(endpoint="/studies/brca_tcga/molecular-profiles") """ # Base URL for cBioPortal API base_url = "https://www.cbioportal.org/api" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load cBioPortal schema schema_path = os.path.join(SCHEMA_DB_PATH, "cbioportal.pkl") with open(schema_path, "rb") as f: cbioportal_schema = pickle.load(f) # Create system prompt template system_template = """ You are a cancer genomics expert specialized in using the cBioPortal REST API. Based on the user's natural language request, determine the appropriate cBioPortal REST API endpoint and parameters. CBIOPORTAL REST API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://www.cbioportal.org/api" and any parameters) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - For gene queries, use either Hugo symbol (e.g., "BRCA1") or Entrez ID (e.g., 672) - For pagination, include parameters "pageNumber" and "pageSize" if needed - For mutation data queries, always include appropriate sample identifiers - Common studies include: "brca_tcga" (breast cancer), "gbm_tcga" (glioblastoma), "luad_tcga" (lung adenocarcinoma) - For molecular profiles, common IDs follow pattern: "[study]_[data_type]" (e.g., "brca_tcga_mutations") - Consider including "projection=DETAILED" for more comprehensive results when appropriate Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=cbioportal_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if not endpoint.startswith("http"): # Clean up endpoint format if not endpoint.startswith("/"): endpoint = "/" + endpoint # Add base URL endpoint = f"{base_url}{endpoint}" description = "Direct query to cBioPortal API" # Execute the cBioPortal API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_clinvar(prompt: str = None, search_term: str = None, max_results: int = 3) -> dict: """ Take a natural language prompt and convert it to a structured ClinVar query. Args: prompt: Natural language query about genetic variants (e.g., "Find pathogenic BRCA1 variants") search_term: Direct search term for ClinVar max_results: Maximum number of results to return Returns: Dictionary containing both the structured query and the ClinVar results """ if not prompt and not search_term: return {"error": "Either a prompt or an endpoint must be provided"} if prompt: # Load ClinVar schema schema_path = os.path.join(SCHEMA_DB_PATH, "clinvar.pkl") with open(schema_path, "rb") as f: clinvar_schema = pickle.load(f) # ClinVar system prompt template system_prompt_template = """ You are a genetics research assistant that helps convert natural language queries into structured ClinVar search queries. Based on the user's natural language request, you will generate a structured search for the ClinVar database. Output only a JSON object with the following fields: 1. "search_term": The exact search query to use with the ClinVar API IMPORTANT: Your response must ONLY contain a JSON object with the search term field. Your "search_term" MUST strictly follow these ClinVar search syntax rules/tags: {schema} For combining terms: Use AND, OR, NOT (must be capitalized) For complex logic: Use parentheses For terms with multiple words: use double quotes escaped with a backslash or underscore (e.g. breast_cancer[dis] or \"breast cancer\"[dis]) Example: "BRCA1[gene] AND (pathogenic[clinsig] OR likely_pathogenic[clinsig])" EXAMPLES OF CORRECT QUERIES: - For "pathogenic BRCA1 variants": "BRCA1[gene] AND clinsig_pathogenic[prop]" - For "Specific RS": "rs6025[rsid]" - For "Combined search with multiple criteria": "BRCA1[gene] AND origin_germline[prop]" - For "Find variants in a specific genomic region": "17[chr] AND 43000000:44000000[chrpos37]" - If query asks for pathogenicity of a variant, it's asking for all possible germline classifications of the variant, so just [gene] AND [variant] is needed """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=clinvar_schema, system_template=system_prompt_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] search_term = query_info.get("search_term", "") if not search_term: return { "error": "Failed to generate a valid search term from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } return _query_ncbi_database( database="clinvar", search_term=search_term, max_results=max_results, ) @tool def query_geo(prompt: str = None, search_term: str = None, max_results: int = 3) -> dict: """ Query the NCBI Gene Expression Omnibus (GEO) using natural language or a direct search term. Args: prompt: Natural language query about RNA-seq, microarray, or other expression data search_term: Direct search term in GEO syntax max_results: Maximum number of results to return Returns: Dictionary containing the query results or error information Examples: - Natural language: query_geo("Find RNA-seq datasets for breast cancer") - Direct search: query_geo(search_term="RNA-seq AND breast cancer AND gse[ETYP]") """ if not prompt and not search_term: return {"error": "Either a prompt or a search term must be provided"} database = "gds" # Default database if prompt: # Load GEO schema schema_path = os.path.join(SCHEMA_DB_PATH, "geo.pkl") with open(schema_path, "rb") as f: geo_schema = pickle.load(f) # Create system prompt template system_template = """ You are a bioinformatics research assistant that helps convert natural language queries into structured GEO (Gene Expression Omnibus) search queries. Based on the user's natural language request, you will generate a structured search for the GEO database. Output only a JSON object with the following fields: 1. "search_term": The exact search query to use with the GEO API 2. "database": The specific GEO database to search (either "gds" for GEO DataSets or "geoprofiles" for GEO Profiles) IMPORTANT: Your response must ONLY contain a JSON object with the required fields. Your "search_term" MUST strictly follow these GEO search syntax rules/tags: {schema} For combining terms: Use AND, OR, NOT (must be capitalized) For complex logic: Use parentheses For terms with multiple words: use double quotes or underscore (e.g. "breast cancer"[Title]) Date ranges use colon format: 2015/01:2020/12[PDAT] Choose the appropriate database based on the user's query: - gds: GEO DataSets (contains Series, Datasets, Platforms, Samples metadata) - geoprofiles: GEO Profiles (contains gene expression data) If database isn't clearly specified, default to "gds" as it contains most common experiment metadata. EXAMPLES OF CORRECT OUTPUTS: - For "RNA-seq data in breast cancer": {"search_term": "RNA-seq AND breast cancer AND gse[ETYP]", "database": "gds"} - For "Mouse microarray data from 2020": {"search_term": "Mus musculus[ORGN] AND 2020[PDAT] AND microarray AND gse[ETYP]", "database": "gds"} - For "Expression profiles of TP53 in lung cancer": {"search_term": "TP53[Gene Symbol] AND lung cancer", "database": "geoprofiles"} """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=geo_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the search term and database from Gemini's response query_info = gemini_result["data"] search_term = query_info.get("search_term", "") database = query_info.get("database", "gds") if not search_term: return { "error": "Failed to generate a valid search term from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } # Execute the GEO query using the helper function result = _query_ncbi_database( database=database, search_term=search_term, max_results=max_results, ) return result @tool def query_dbsnp(prompt: str = None, search_term: str = None, max_results: int = 3) -> dict: """ Query the NCBI dbSNP database using natural language or a direct search term. Args: prompt: Natural language query about genetic variants/SNPs search_term: Direct search term in dbSNP syntax max_results: Maximum number of results to return Returns: Dictionary containing the query results or error information Examples: - Natural language: query_dbsnp("Find pathogenic variants in BRCA1") - Direct search: query_dbsnp(search_term="BRCA1[Gene Name] AND pathogenic[Clinical Significance]") """ if not prompt and not search_term: return {"error": "Either a prompt or a search term must be provided"} if prompt: # Load dbSNP schema schema_path = os.path.join(SCHEMA_DB_PATH, "dbsnp.pkl") with open(schema_path, "rb") as f: dbsnp_schema = pickle.load(f) # Create system prompt template system_template = """ You are a genetics research assistant that helps convert natural language queries into structured dbSNP search queries. Based on the user's natural language request, you will generate a structured search for the dbSNP database. Output only a JSON object with the following fields: 1. "search_term": The exact search query to use with the dbSNP API IMPORTANT: Your response must ONLY contain a JSON object with the search term field. Your "search_term" MUST strictly follow these dbSNP search syntax rules/tags: {schema} For combining terms: Use AND, OR, NOT (must be capitalized) For complex logic: Use parentheses For terms with multiple words: use double quotes (e.g. "breast cancer"[Disease Name]) EXAMPLES OF CORRECT QUERIES: - For "pathogenic variants in BRCA1": "BRCA1[Gene Name] AND pathogenic[Clinical Significance]" - For "specific SNP rs6025": "rs6025[rs]" - For "SNPs in a genomic region": "17[Chromosome] AND 41196312:41277500[Base Position]" - For "common SNPs in EGFR": "EGFR[Gene Name] AND common[COMMON]" """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=dbsnp_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the search term from Gemini's response query_info = gemini_result["data"] search_term = query_info.get("search_term", "") if not search_term: return { "error": "Failed to generate a valid search term from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } # Execute the dbSNP query using the helper function result = _query_ncbi_database( database="snp", search_term=search_term, max_results=max_results, ) return result @tool def query_ucsc(prompt: str = None, endpoint: str = None, verbose: bool = True) -> dict: """ Query the UCSC Genome Browser API using natural language or a direct endpoint. Args: prompt: Natural language query about genomic data endpoint: Full URL or endpoint specification with parameters verbose: Whether to return detailed results Returns: Dictionary containing the query results or error information Examples: - Natural language: query_ucsc("Get DNA sequence of chromosome M positions 1-100 in human genome") - Direct endpoint: query_ucsc(endpoint="https://api.genome.ucsc.edu/getData/sequence?genome=hg38&chrom=chrM&start=1&end=100") """ # Base URL for UCSC API base_url = "https://api.genome.ucsc.edu" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load UCSC schema schema_path = os.path.join(SCHEMA_DB_PATH, "ucsc.pkl") with open(schema_path, "rb") as f: ucsc_schema = pickle.load(f) # Create system prompt template system_template = """ You are a genomics expert specialized in using the UCSC Genome Browser API. Based on the user's natural language request, determine the appropriate UCSC Genome Browser API endpoint and parameters. UCSC GENOME BROWSER API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "full_url": The complete URL to query (including the base URL "https://api.genome.ucsc.edu" and all parameters) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - For chromosome names, always include the "chr" prefix (e.g., "chr1", "chrX", "chrM") - Genomic positions are 0-based (first base is position 0) - For "start" and "end" parameters, both must be provided together - The "maxItemsOutput" parameter can be used to limit the amount of data returned - Common genomes include: "hg38" (human), "mm39" (mouse), "danRer11" (zebrafish) - For sequence data, use "getData/sequence" endpoint - For chromosome listings, use "list/chromosomes" endpoint - For available genomes, use "list/ucscGenomes" endpoint Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=ucsc_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the full URL from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("full_url", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if not endpoint.startswith("http"): # Add base URL if it's just a path endpoint = f"{base_url}/{endpoint}" description = "Direct query to UCSC Genome Browser API" # Execute the UCSC API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) # Format the results if successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_ensembl(prompt: str = None, endpoint: str = None, verbose: bool = True) -> dict: """ Query the Ensembl REST API using natural language or a direct endpoint. Args: prompt: Natural language query about genomic data endpoint: Direct API endpoint to query (e.g., "lookup/symbol/human/BRCA2") or full URL verbose: Whether to return detailed results Returns: Dictionary containing the query results or error information Examples: - Natural language: query_ensembl("Get information about the human BRCA2 gene") - Direct endpoint: query_ensembl(endpoint="lookup/symbol/homo_sapiens/BRCA2") """ print("IN QUERY ENSEMBL") print("PROMPT: ", prompt) print("ENDPOINT: ", endpoint) # Base URL for Ensembl API base_url = "https://rest.ensembl.org" # Ensure we have either a prompt or an endpoint if not prompt and not endpoint: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load Ensembl schema schema_path = os.path.join(SCHEMA_DB_PATH, "ensembl.pkl") with open(schema_path, "rb") as f: ensembl_schema = pickle.load(f) # Create system prompt template system_template = """ You are a genomics and bioinformatics expert specialized in using the Ensembl REST API. Based on the user's natural language request, determine the appropriate Ensembl REST API endpoint and parameters. ENSEMBL REST API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The API endpoint to query (e.g., "lookup/symbol/homo_sapiens/BRCA2") 2. "params": An object containing query parameters specific to the endpoint 3. "description": A brief description of what the query is doing SPECIAL NOTES: - Chromosome region queries have a maximum length of 4900000 bp inclusive, so bp of start and end should be 4900000 bp apart. If the user's query exceeds this limit, Ensembl will return an error. - For symbol lookups, the format is "lookup/symbol/[species]/[symbol]" - To find the coordinates of a band on a chromosome, use /info/assembly/homo_sapiens/[chromosome] with parameters "band":1 - To find the overlapping genes of a genomic region, use /overlap/region/homo_sapiens/[chromosome]:[start]-[end] - For sequence queries, specify the sequence type in parameters (genomic, cdna, cds, protein) - For converting rsID to hg38 genomic coordinates, use the "GET id/variation/[species]/[rsid]" endpoint - Many endpoints support "content-type" parameter for format specification (application/json, text/xml) Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=ensembl_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") params = query_info.get("params", {}) description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if endpoint.startswith("http"): # If a full URL is provided, extract the endpoint part if endpoint.startswith(base_url): endpoint = endpoint[len(base_url):].lstrip('/') params = {} description = "Direct query to Ensembl API" # Remove leading slash if present if endpoint.startswith("/"): endpoint = endpoint[1:] # Prepare headers for JSON response headers = { "Content-Type": "application/json", "Accept": "application/json" } # Construct the URL url = f"{base_url}/{endpoint}" # Execute the Ensembl API request using the helper function api_result = _query_rest_api( endpoint=url, method="GET", params=params, headers=headers, description=description ) # Format the results if successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_opentarget_genetics(prompt: str = None, query: str = None, variables: dict = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the OpenTargets Genetics API using natural language or a direct GraphQL query. Args: prompt (str, required): Natural language query about genetic targets and variants query (str, optional): Direct GraphQL query string variables (dict, optional): Variables for the GraphQL query api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_opentarget("Get information about variant 1_154453788_C_T") - Direct query: query_opentarget(query="query variantInfo($variantId: String!) {...}", variables={"variantId": "1_154453788_C_T"}) """ # Constants and initialization OPENTARGETS_URL = "https://api.genetics.opentargets.org/graphql" # Ensure we have either a prompt or a query if prompt is None and query is None: return {"error": "Either a prompt or a GraphQL query must be provided"} # If using prompt, parse with Claude if prompt: # Load OpenTargets schema schema_path = os.path.join(SCHEMA_DB_PATH, "opentarget_genetics.pkl") with open(schema_path, "rb") as f: opentarget_schema = pickle.load(f) # Create system prompt template system_template = """ You are an expert in translating natural language requests into GraphQL queries for the OpenTargets Genetics API. Here is a schema of the main types and queries available in the OpenTargets Genetics API: {schema} Translate the user's natural language request into a valid GraphQL query for this API. Return only a JSON object with two fields: 1. "query": The complete GraphQL query string 2. "variables": A JSON object containing the variables needed for the query SPECIAL NOTES: - Variant IDs are typically in the format 'chromosome_position_ref_alt' (e.g., '1_154453788_C_T') - For L2G (locus-to-gene) queries, you need both a variant ID and a study ID - The API can provide variant information, QTLs, PheWAS results, pathogenicity scores, etc. - For mutations by gene, use the approved gene symbol (e.g., "BRCA1") - Always escape special characters, including quotes, in the query string (eg. \" instead of ") Return ONLY the JSON object with no additional text or explanations. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=opentarget_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the query and variables from Gemini's response query_info = gemini_result["data"] query = query_info.get("query", "") if variables is None: # Only use Claude's variables if none provided variables = query_info.get("variables", {}) if not query: return { "error": "Failed to generate a valid GraphQL query from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } # Execute the GraphQL query api_result = _query_rest_api( endpoint=OPENTARGETS_URL, method="POST", json_data={"query": query, "variables": variables or {}}, headers={"Content-Type": "application/json"} ) if not api_result["success"]: return api_result if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_opentarget(prompt: str = None, query: str = None, variables: dict = None, verbose: bool = False) -> dict: """ Query the OpenTargets Platform API using natural language or a direct GraphQL query. Args: prompt: Natural language query about drug targets, diseases, and mechanisms query: Direct GraphQL query string variables: Variables for the GraphQL query verbose: Whether to return detailed results Returns: Dictionary containing the query results or error information Examples: - Natural language: query_opentarget("Find drug targets for Alzheimer's disease") - Direct query: query_opentarget(query="query diseaseAssociations($diseaseId: String!) {...}", variables={"diseaseId": "EFO_0000249"}) """ # Constants and initialization OPENTARGETS_URL = "https://api.platform.opentargets.org/api/v4/graphql" # Ensure we have either a prompt or a query if prompt is None and query is None: return {"error": "Either a prompt or a GraphQL query must be provided"} # If using prompt, parse with Claude if prompt: # Load OpenTargets schema schema_path = os.path.join(SCHEMA_DB_PATH, "opentarget.pkl") with open(schema_path, "rb") as f: opentarget_schema = pickle.load(f) # Create system prompt template system_template = """ You are an expert in translating natural language requests into GraphQL queries for the OpenTargets Platform API. Here is a schema of the main types and queries available in the OpenTargets Platform API: {schema} Translate the user's natural language request into a valid GraphQL query for this API. Return only a JSON object with two fields: 1. "query": The complete GraphQL query string 2. "variables": A JSON object containing the variables needed for the query SPECIAL NOTES: - Disease IDs typically use EFO ontology (e.g., "EFO_0000249" for Alzheimer's disease) - Target IDs typically use Ensembl IDs (e.g., "ENSG00000197386" for ENSG00000197386) - The API can provide information about drug-target associations, disease-target associations, etc. - Always limit results to a reasonable number using "first" parameter (e.g., first: 10) - Always escape special characters, including quotes, in the query string (eg. \\" instead of ") Return ONLY the JSON object with no additional text or explanations. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=opentarget_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the query and variables from Gemini's response query_info = gemini_result["data"] query = query_info.get("query", "") if variables is None: # Only use Claude's variables if none provided variables = query_info.get("variables", {}) if not query: return { "error": "Failed to generate a valid GraphQL query from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } # Execute the GraphQL query api_result = _query_rest_api( endpoint=OPENTARGETS_URL, method="POST", json_data={"query": query, "variables": variables or {}}, headers={"Content-Type": "application/json"}, description="OpenTargets Platform GraphQL query" ) # Format the results if not verbose and successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: api_result["result"] = _format_query_results(api_result["result"]) return api_result @tool def query_gwas_catalog(prompt: str = None, endpoint: str = None, max_results: int = 3) -> dict: """ Query the GWAS Catalog API using natural language or a direct endpoint. Args: prompt: Natural language query about GWAS data endpoint: Full API endpoint to query (e.g., "https://www.ebi.ac.uk/gwas/rest/api/studies?diseaseTraitId=EFO_0001360") max_results: Maximum number of results to return Returns: Dictionary containing the query results or error information Examples: - Natural language: query_gwas_catalog("Find GWAS studies related to Type 2 diabetes") - Direct endpoint: query_gwas_catalog(endpoint="studies", params={"diseaseTraitId": "EFO_0001360"}) """ # Base URL for GWAS Catalog API base_url = "https://www.ebi.ac.uk/gwas/rest/api" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load GWAS Catalog schema schema_path = os.path.join(SCHEMA_DB_PATH, "gwas_catalog.pkl") with open(schema_path, "rb") as f: gwas_schema = pickle.load(f) # Create system prompt template system_template = """ You are a genomics expert specialized in using the GWAS Catalog API. Based on the user's natural language request, determine the appropriate GWAS Catalog API endpoint and parameters. GWAS CATALOG API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The API endpoint to query (e.g., "studies", "associations") 2. "params": An object containing query parameters specific to the endpoint 3. "description": A brief description of what the query is doing SPECIAL NOTES: - For disease/trait searches, consider using the "EFO" identifiers when possible - Common endpoints include: "studies", "associations", "singleNucleotidePolymorphisms", "efoTraits" - For pagination, use "size" and "page" parameters - For filtering by p-value, use "pvalueMax" parameter - GWAS Catalog uses a HAL-based REST API Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=gwas_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") params = query_info.get("params", {}) description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: if endpoint is None: endpoint = "" # Use root endpoint params = {"size": max_results} description = f"Direct query to {endpoint}" # Remove leading slash if present if endpoint.startswith("/"): endpoint = endpoint[1:] # Construct the URL url = f"{base_url}/{endpoint}" # Execute the GWAS Catalog API request using the helper function api_result = _query_rest_api( endpoint=url, method="GET", params=params, description=description ) return api_result @tool def query_gnomad(prompt: str = None, gene_symbol: str = None, verbose: bool = True) -> dict: """ Query gnomAD for variants in a gene using natural language or direct gene symbol. Args: prompt: Natural language query about genetic variants gene_symbol: Gene symbol (e.g., "BRCA1") verbose: Whether to return detailed results Returns: Dictionary containing the query results or error information Examples: - Direct gene: query_gnomad(gene_symbol="BRCA1") - Natural language: query_gnomad(prompt="Find variants in the TP53 gene") """ # Base URL for gnomAD API base_url = "https://gnomad.broadinstitute.org/api" # Ensure we have either a prompt or a gene_symbol if prompt is None and gene_symbol is None: return {"error": "Either a prompt or a gene_symbol must be provided"} # If using prompt, parse with Claude if prompt and not gene_symbol: # Load gnomAD schema schema_path = os.path.join(SCHEMA_DB_PATH, "gnomad.pkl") with open(schema_path, "rb") as f: gnomad_schema = pickle.load(f) # Create system prompt template system_template = """ You are a genomics expert specialized in using the gnomAD GraphQL API. Based on the user's natural language request, extract the gene symbol and relevant parameters and create the gnomAD GraphQL query. GnomAD GraphQL API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "query": The complete GraphQL query string SPECIAL NOTES: - The gene_symbol should be the official gene symbol (e.g., "BRCA1" not "breast cancer gene 1") - If no reference genome is specified, default to GRCh38 - If no dataset is specified, default to gnomad_r4 - Return only a single gene symbol, even if multiple are mentioned - Always escape special characters, including quotes, in the query string (eg. \" instead of ") Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=gnomad_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the gene symbol from Gemini's response query_info = gemini_result["data"] query_str = query_info.get("query", "") if not query_str: return { "error": "Failed to extract a valid query from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Load gnomAD schema for gene_symbol substitution schema_path = os.path.join(SCHEMA_DB_PATH, "gnomad.pkl") with open(schema_path, "rb") as f: gnomad_schema = pickle.load(f) description = f"Query gnomAD for variants in {gene_symbol}" # replace BRCA1 with gene_symbol query_str = gnomad_schema.replace("BRCA1", gene_symbol) api_result = _query_rest_api( endpoint=base_url, method="POST", json_data={"query": query_str}, headers={"Content-Type": "application/json"}, description=description ) if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def blast_sequence(sequence: str, database: str, program: str) -> Union[Dict[str, Union[str, float]], str]: """ Identifies a DNA sequence using NCBI BLAST with improved error handling, timeout management, and debugging Args: sequence: The sequence to identify. If DNA, use database: core_nt, program: blastn; if protein, use database: nr, program: blastp database: The BLAST database to search against program: The BLAST program to use Returns: A dictionary containing the title, e-value, identity percentage, and coverage percentage of the best alignment """ max_attempts = 1 # One initial attempt plus one retry attempts = 0 max_runtime = 600 # 10 minutes in seconds while attempts < max_attempts: try: attempts += 1 query_sequence = Seq(sequence) # Start timer start_time = time.time() # Submit BLAST job print(f"Submitting BLAST job (attempt {attempts}/{max_attempts})...") result_handle = NCBIWWW.qblast(program, database, query_sequence, expect=100, word_size=7, megablast=True) # Parse results with timeout check blast_records = NCBIXML.parse(result_handle) blast_record = None # Try to get the first record with timeout check while time.time() - start_time < max_runtime: try: # Set a short timeout for next operation blast_record = next(blast_records) # Get first record break # Successfully got the record except StopIteration: # No more records return "No BLAST results found" except Exception as e: # Check if we've exceeded the time limit if time.time() - start_time >= max_runtime: if attempts < max_attempts: print("BLAST job timeout exceeded. Resubmitting...") break # Break to retry else: return "BLAST search failed after maximum attempts due to timeout" # Brief pause before trying again time.sleep(1) # Check if we timed out during record retrieval if blast_record is None: if attempts < max_attempts: continue # Retry else: return "BLAST search failed after maximum attempts due to timeout" # Debug information print(f"Number of alignments found: {len(blast_record.alignments)}") if blast_record.alignments: for alignment in blast_record.alignments: print("\nAlignment:") print(f"hit_id: {alignment.hit_id}") print(f"hit_def: {alignment.hit_def}") print(f"accession: {alignment.accession}") for hsp in alignment.hsps: print(f"E-value: {hsp.expect}") print(f"Score: {hsp.score}") print(f"Identities: {hsp.identities}/{hsp.align_length}") return { 'hit_id': alignment.hit_id, 'hit_def': alignment.hit_def, 'accession': alignment.accession, 'e_value': hsp.expect, 'identity': (hsp.identities / float(hsp.align_length)) * 100, 'coverage': len(hsp.query) / len(sequence) * 100 } else: return "No alignments found - sequence might be too short or low complexity" except Exception as e: if attempts < max_attempts: print(f"Error during BLAST search: {str(e)}. Retrying...") time.sleep(2) # Wait briefly before retrying else: return f"Error during BLAST search after maximum attempts: {str(e)}" return "BLAST search failed after maximum attempts" @tool def query_reactome(prompt: str = None, endpoint: str = None, download: bool = False, output_dir: str = None, verbose: bool = True) -> dict: """ Query the Reactome database using natural language or a direct endpoint. Args: prompt: Natural language query about biological pathways endpoint: Direct API endpoint or full URL download: Whether to download pathway diagrams output_dir: Directory to save downloaded files verbose: Whether to return detailed results Returns: Dictionary containing the query results or error information Examples: - Natural language: query_reactome("Find pathways related to DNA repair") - Direct endpoint: query_reactome(endpoint="data/pathways/R-HSA-73894") """ # Base URLs for Reactome APIs content_base_url = "https://reactome.org/ContentService" analysis_base_url = "https://reactome.org/AnalysisService" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # Create output directory if downloading and directory doesn't exist if download and output_dir: os.makedirs(output_dir, exist_ok=True) # If using prompt, parse with Claude if prompt: # Load Reactome schema schema_path = os.path.join(SCHEMA_DB_PATH, "reactome.pkl") with open(schema_path, "rb") as f: reactome_schema = pickle.load(f) # Create system prompt template system_template = """ You are a bioinformatics expert specialized in using the Reactome API. Based on the user's natural language request, determine the appropriate Reactome API endpoint and parameters. REACTOME API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The API endpoint to query (e.g., "data/pathways/PATHWAY_ID", "data/query/GENE_SYMBOL") 2. "base": Which base URL to use ("content" for ContentService or "analysis" for AnalysisService) 3. "params": An object containing query parameters specific to the endpoint 4. "description": A brief description of what the query is doing SPECIAL NOTES: - Reactome has two primary APIs: ContentService (for retrieving specific pathway data) and AnalysisService (for analyzing gene lists) - For pathway queries, use "data/pathways/PATHWAY_ID" with the pathway stable identifier (e.g., R-HSA-73894) - For gene queries, use "data/query/GENE" with official gene symbol (e.g., "BRCA1") - For pathway diagrams, include "download: true" in your response if the query is for pathway visualization - Common human pathway IDs start with "R-HSA-" Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=reactome_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") base = query_info.get("base", "content") # Default to ContentService params = query_info.get("params", {}) description = query_info.get("description", "") should_download = query_info.get("download", download) # Override download if specified if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint if endpoint.startswith("http"): # Full URL already provided if "ContentService" in endpoint: base = "content" elif "AnalysisService" in endpoint: base = "analysis" else: base = "content" # Default else: # Just endpoint provided, assume ContentService by default base = "content" params = {} description = f"Direct query to Reactome {base} API: {endpoint}" should_download = download # Select base URL based on API type base_url = content_base_url if base == "content" else analysis_base_url # Remove leading slash if present if endpoint.startswith("/"): endpoint = endpoint[1:] # Construct the URL if endpoint.startswith("http"): url = endpoint # Full URL already provided else: url = f"{base_url}/{endpoint}" # Execute the Reactome API request using the helper function api_result = _query_rest_api( endpoint=url, method="GET", params=params, description=description ) # Handle downloading pathway diagrams if requested if should_download and api_result.get("success") and "result" in api_result: result = api_result["result"] pathway_id = None # Try to extract pathway ID from result if isinstance(result, dict): pathway_id = result.get("stId") or result.get("dbId") # If we have a pathway ID and output directory, download diagram if pathway_id and output_dir: diagram_url = f"{content_base_url}/data/pathway/{pathway_id}/diagram" try: diagram_response = requests.get(diagram_url) diagram_response.raise_for_status() # Save diagram file diagram_path = os.path.join(output_dir, f"{pathway_id}_diagram.png") with open(diagram_path, "wb") as f: f.write(diagram_response.content) api_result["diagram_path"] = diagram_path except Exception as e: api_result["diagram_error"] = f"Failed to download diagram: {str(e)}" if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: return _format_query_results(api_result["result"]) return api_result @tool def query_regulomedb(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = False) -> dict: """ Query the RegulomeDB database using natural language or direct variant/coordinate specification. Args: prompt (str, required): Natural language query about regulatory elements endpoint (str, optional): Direct endpoint URL or variant/coordinate specification api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_regulomedb("Find regulatory elements for rs35675666") - Direct variant: query_regulomedb(variant="rs35675666") - Coordinates: query_regulomedb(coordinates="chr11:5246919-5246919") """ # Base URL for RegulomeDB API base_url = "https://regulomedb.org/regulome-search/" # Ensure we have either a prompt, variant, or coordinates if prompt is None and endpoint is None: return {"error": "Either a prompt, variant ID, or genomic coordinates must be provided"} # If using prompt, parse with Claude if prompt and not endpoint: # Create system prompt template system_template = """ You are a genomics expert specialized in using the RegulomeDB API. Based on the user's natural language request, extract the variant ID or genomic coordinates they want to query. Your response should be a JSON object with ONLY ONE of the following fields: 1. "endpoint": The API endpoint to query (e.g., "https://regulomedb.org/regulome-search/?regions=chr11:5246919-5246919&genome=GRCh38") SPECIAL NOTES: - RegulomeDB only works with human genome data - Variant IDs should be rsIDs from dbSNP when possible. The endpoint should be in the format https://regulomedb.org/regulome-search/?regions=rsID&genome=GRCh38 - Thumbnails for chip and chromatin should be in the format https://regulomedb.org/regulome-search?regions=chr11:5246919-5246919&genome=GRCh38/thumbnail=chip - Coordinates should be in GRCh37/hg19 format - For single base queries, use the same position for start and end (e.g., "chr11:5246919-5246919") - Chromosome should be specified with "chr" prefix (e.g., "chr11" not just "11") Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=None, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the variant or coordinates from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") if not endpoint: return { "error": "Failed to extract a valid variant ID or coordinates from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: description = f"Query RegulomeDB for {endpoint}" # Construct the request URL endpoint = endpoint # Execute the RegulomeDB API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", headers = {'Accept': 'application/json'} ) # Format the results if not verbose and successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: api_result["result"] = _format_query_results(api_result["result"]) return api_result @tool def query_pride(prompt: str = None, endpoint: str = None, api_key: str = None, max_results: int = 3) -> dict: """ Query the PRIDE (PRoteomics IDEntifications) database using natural language or a direct endpoint. Args: prompt (str, required): Natural language query about proteomics data endpoint (str, optional): The full endpoint to query (e.g., "https://www.ebi.ac.uk/pride/ws/archive/v2/projects?keyword=breast%20cancer") api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results max_results (int): Maximum number of results to return verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_pride("Find proteomics data related to breast cancer") - Direct endpoint: query_pride(endpoint="projects", params={"keyword": "breast cancer"}) """ # Base URL for PRIDE API base_url = "https://www.ebi.ac.uk/pride/ws/archive/v2" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load PRIDE schema schema_path = os.path.join(SCHEMA_DB_PATH, "pride.pkl") with open(schema_path, "rb") as f: pride_schema = pickle.load(f) # Create system prompt template system_template = """ You are a proteomics expert specialized in using the PRIDE API. Based on the user's natural language request, determine the appropriate PRIDE API endpoint and parameters. PRIDE API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The full url endpoint to query 2. "description": A brief description of what the query is doing SPECIAL NOTES: - PRIDE is a repository for proteomics data stored at EBI - Common endpoints include: "projects", "assays", "files", "proteins", "peptideevidences" - For searching projects, you can use parameters like "keyword", "species", "tissue", "disease" - For pagination, use "page" and "pageSize" parameters - Most results include PagingObject and FieldsObject structures Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=pride_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") params = query_info.get("params", {}) description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint params = {"pageSize": max_results, "page": 0} description = f"Direct query to PRIDE {endpoint}" # Remove leading slash if present if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = f"Direct query to provided endpoint" # Execute the PRIDE API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", params=params, description=description ) return api_result @tool def query_gtopdb(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the Guide to PHARMACOLOGY database (GtoPdb) using natural language or a direct endpoint. Args: prompt (str, required): Natural language query about drug targets, ligands, and interactions endpoint (str, optional): Full API endpoint to query (e.g., "https://www.guidetopharmacology.org/services/targets?type=GPCR&name=beta-2") api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_gtopdb("Find ligands that target the beta-2 adrenergic receptor") - Direct endpoint: query_gtopdb(endpoint="targets", params={"type": "GPCR", "name": "beta-2"}) """ # Base URL for GtoPdb API base_url = "https://www.guidetopharmacology.org/services" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load GtoPdb schema schema_path = os.path.join(SCHEMA_DB_PATH, "gtopdb.pkl") with open(schema_path, "rb") as f: gtopdb_schema = pickle.load(f) # Create system prompt template system_template = """ You are a pharmacology expert specialized in using the Guide to PHARMACOLOGY API. Based on the user's natural language request, determine the appropriate GtoPdb API endpoint and parameters. GTOPDB API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The full API endpoint to query 2. "description": A brief description of what the query is doing SPECIAL NOTES: - Main endpoints include: "targets", "ligands", "interactions", "diseases", "refs" - Target types include: "GPCR", "NHR", "LGIC", "VGIC", "OtherIC", "Enzyme", "CatalyticReceptor", "Transporter", "OtherProtein" - Ligand types include: "Synthetic organic", "Metabolite", "Natural product", "Endogenous peptide", "Peptide", "Antibody", "Inorganic", "Approved", "Withdrawn", "Labelled", "INN" - Interaction types include: "Activator", "Agonist", "Allosteric modulator", "Antagonist", "Antibody", "Channel blocker", "Gating inhibitor", "Inhibitor", "Subunit-specific" - For specific target/ligand details, use formats like "targets/{{targetId}}" or "ligands/{{ligandId}}" - For subresources, use formats like "targets/{{targetId}}/interactions" or "ligands/{{ligandId}}/structure" Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=gtopdb_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint description = f"Direct query to GtoPdb {endpoint}" # Remove leading slash if present if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = f"Direct query to provided endpoint" # Execute the GtoPdb API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) # Format the results if not verbose and successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: api_result["result"] = _format_query_results(api_result["result"]) return api_result @tool def region_to_ccre_screen(coord_chrom: str, coord_start: int, coord_end: int, assembly: str = "GRCh38" ) -> str: """ Given starting and ending coordinates, this function retrieves information of intersecting cCREs. Args: coord_chrom: Chromosome of the gene, formatted like 'chr12' coord_start: Starting chromosome coordinate coord_end: Ending chromosome coordinate assembly: Assembly of the genome, formatted like 'GRCh38'. Default is 'GRCh38' Returns: A detailed string explaining the steps and the intersecting cCRE data or any error encountered """ steps = [] try: steps.append(f"Starting cCRE data retrieval for coordinates: {coord_chrom}:{coord_start}-{coord_end} (Assembly: {assembly}).") # Build the URL and request payload url = "https://screen-beta-api.wenglab.org/dataws/cre_table" data = { "assembly": assembly, "coord_chrom": coord_chrom, "coord_start": coord_start, "coord_end": coord_end } steps.append("Sending POST request to API with the following data:") steps.append(str(data)) # Make the request response = requests.post(url, json=data) # Check if the response is successful if not response.ok: raise Exception(f"Request failed with status code {response.status_code}. Response: {response.text}") steps.append("Request executed successfully. Parsing the response...") # Parse the JSON response response_json = response.json() if "errors" in response_json: raise Exception(f"API error: {response_json['errors']}") # Function to reduce and filter response data def reduce_tokens(res_json): # Remove unnecessary fields and round floats res = sorted(res_json["cres"], key=lambda x: x['dnase_zscore'], reverse=True) filtered_res = [] for item in res: new_item = { 'chrom': item['chrom'], 'start': item['start'], 'len': item['len'], 'pct': item['pct'], 'ctcf_zscore': round(item['ctcf_zscore'], 2), 'dnase_zscore': round(item['dnase_zscore'], 2), 'enhancer_zscore': round(item['enhancer_zscore'], 2), 'promoter_zscore': round(item['promoter_zscore'], 2), 'accession': item['info']['accession'], 'isproximal': item['info']['isproximal'], 'concordance': item['info']['concordant'], 'ctcfmax': round(item['info']['ctcfmax'], 2), 'k4me3max': round(item['info']['k4me3max'], 2), 'k27acmax': round(item['info']['k27acmax'], 2) } filtered_res.append(new_item) return filtered_res # Process the response data filtered_data = reduce_tokens(response_json) if not filtered_data: steps.append(f"No intersecting cCREs found for coordinates: {coord_chrom}:{coord_start}-{coord_end}.") return "\n".join(steps + ["No cCRE data available for this genomic region."]) # Format the result into a readable string ccre_data_string = f"Intersecting cCREs for {coord_chrom}:{coord_start}-{coord_end} (Assembly: {assembly}):\n" for i, ccre in enumerate(filtered_data, 1): ccre_data_string += ( f"cCRE {i}:\n" f" Chromosome: {ccre['chrom']}\n" f" Start: {ccre['start']}\n" f" Length: {ccre['len']}\n" f" PCT: {ccre['pct']}\n" f" CTCF Z-score: {ccre['ctcf_zscore']}\n" f" DNase Z-score: {ccre['dnase_zscore']}\n" f" Enhancer Z-score: {ccre['enhancer_zscore']}\n" f" Promoter Z-score: {ccre['promoter_zscore']}\n" f" Accession: {ccre['accession']}\n" f" Is Proximal: {ccre['isproximal']}\n" f" Concordance: {ccre['concordance']}\n" f" CTCFmax: {ccre['ctcfmax']}\n" f" K4me3max: {ccre['k4me3max']}\n" f" K27acmax: {ccre['k27acmax']}\n\n" ) steps.append(f"cCRE data successfully retrieved and formatted for {coord_chrom}:{coord_start}-{coord_end}.") return "\n".join(steps + [ccre_data_string]) except Exception as e: steps.append(f"Exception encountered: {str(e)}") return "\n".join(steps + [f"Error: {str(e)}"]) @tool def get_genes_near_ccre(accession: str, assembly: str, chromosome: str, k: int = 10) -> str: """ Given a cCRE (Candidate cis-Regulatory Element), this function returns a string containing the steps it performs and the k nearest genes sorted by distance. Args: accession: ENCODE Accession ID of query cCRE, e.g., EH38E1516980 assembly: Assembly of the gene, e.g., 'GRCh38' chromosome: Chromosome of the gene, e.g., 'chr12' k: Number of nearby genes to return, sorted by distance. Default is 10 Returns: Steps performed and the result """ steps_log = f"Starting process with accession: {accession}, assembly: {assembly}, chromosome: {chromosome}, k: {k}\n" url = "https://screen-beta-api.wenglab.org/dataws/re_detail/nearbyGenomic" data = { "accession": accession, "assembly": assembly, "coord_chrom": chromosome } steps_log += "Sending POST request to API with given data.\n" response = requests.post(url, json=data) if not response.ok: steps_log += f"API request failed with response: {response.text}\n" return steps_log response_json = response.json() if "errors" in response_json: steps_log += f"API returned errors: {response_json['errors']}\n" return steps_log nearby_genes = response_json.get(accession, {}).get("nearby_genes", []) if not nearby_genes: steps_log += "No nearby genes found for the given accession.\n" return steps_log steps_log += "Successfully retrieved nearby genes. Sorting them by distance.\n" sorted_genes = sorted(nearby_genes, key=lambda x: x['distance'])[:k] steps_log += f"Returning the top {k} nearest genes.\n" steps_log += "Result:\n" for gene in sorted_genes: gene_name = gene.get('name', 'Unknown') distance = gene.get('distance', 'N/A') ensembl_id = gene.get('ensemblid_ver', 'N/A') start = gene.get('start', 'N/A') stop = gene.get('stop', 'N/A') chrom = gene.get('chrom', 'N/A') steps_log += f"Gene: {gene_name}, Distance: {distance}, Ensembl ID: {ensembl_id}, Chromosome: {chrom}, Start: {start}, Stop: {stop}\n" return steps_log @tool def query_remap(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the ReMap database for regulatory elements and transcription factor binding sites. Args: prompt (str, required): Natural language query about transcription factors and binding sites endpoint (str, optional): Full API endpoint to query (e.g., "https://remap.univ-amu.fr/api/v1/catalogue/tf?tf=CTCF") api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_remap("Find CTCF binding sites in chromosome 1") - Direct endpoint: query_remap(endpoint="catalogue/tf", params={"tf": "CTCF"}) """ # Base URL for ReMap API base_url = "https://remap.univ-amu.fr/api/v1" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load ReMap schema schema_path = os.path.join(SCHEMA_DB_PATH, "remap.pkl") with open(schema_path, "rb") as f: remap_schema = pickle.load(f) # Create system prompt template system_template = """ You are a genomics expert specialized in using the ReMap database API. Based on the user's natural language request, determine the appropriate ReMap API endpoint and parameters. REMAP API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The full url endpoint to query 2. "description": A brief description of what the query is doing SPECIAL NOTES: - ReMap is a database of regulatory regions and transcription factor binding sites based on ChIP-seq experiments - Common endpoints include: "catalogue/tf" (transcription factors), "catalogue/biotype" (biotypes), "browse/peaks" (binding sites) - For searching binding sites, you can filter by transcription factor (tf), cell line, biotype, chromosome, etc. - Genomic coordinates should be specified with "chr", "start", and "end" parameters - For limiting results, use "limit" parameter (default is 100) Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=remap_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint description = f"Direct query to ReMap {endpoint}" # Remove leading slash if present if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = f"Direct query to provided endpoint" # Execute the ReMap API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) # Format the results if not verbose and successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: api_result["result"] = _format_query_results(api_result["result"]) return api_result @tool def query_mpd(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the Mouse Phenome Database (MPD) for mouse strain phenotype data. Args: prompt (str, required): Natural language query about mouse phenotypes, strains, or measurements endpoint (str, optional): Full API endpoint to query (e.g., "https://phenomedoc.jax.org/MPD_API/strains") api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_mpd("Find phenotype data for C57BL/6J mice related to blood glucose") - Direct endpoint: query_mpd(endpoint="strains/C57BL/6J/measures") """ # Base URL for MPD API base_url = "https://phenome.jax.org" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load MPD schema schema_path = os.path.join(SCHEMA_DB_PATH, "mpd.pkl") with open(schema_path, "rb") as f: mpd_schema = pickle.load(f) # Create system prompt template system_template = """ You are a mouse genetics expert specialized in using the Mouse Phenome Database (MPD) API. Based on the user's natural language request, determine the appropriate MPD API endpoint and parameters. MPD API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The full url endpoint to query (e.g. https://phenome.jax.org/api/strains) 2. "description": A brief description of what the query is doing SPECIAL NOTES: - The MPD contains phenotype data for diverse strains of laboratory mice - Common endpoints include: "strains" (mouse strains), "measures" (phenotypic measurements), "genes" (gene info) - Use the url to construct the endpoint, not the endpoint name - Common mouse strains include: "C57BL/6J", "DBA/2J", "BALB/cJ", "A/J", "129S1/SvImJ" - Common phenotypic domains include: "behavior", "blood_chemistry", "body_weight", "cardiovascular", "growth", "metabolism" Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=mpd_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint description = f"Direct query to MPD {endpoint}" # Remove leading slash if present if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = f"Direct query to provided endpoint" # Execute the MPD API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", description=description ) # Format the results if not verbose and successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: api_result["result"] = _format_query_results(api_result["result"]) return api_result @tool def query_emdb(prompt: str = None, endpoint: str = None, api_key: str = None, verbose: bool = True) -> dict: """ Query the Electron Microscopy Data Bank (EMDB) for 3D macromolecular structures. Args: prompt (str, required): Natural language query about EM structures and associated data endpoint (str, optional): Full API endpoint to query (e.g., "https://www.ebi.ac.uk/emdb/api/search") api_key (str, optional): Anthropic API key. If None, will use ANTHROPIC_API_KEY env variable verbose (bool): Whether to return detailed results verbose (bool): Whether to return detailed results Returns: dict: Dictionary containing the query results or error information Examples: - Natural language: query_emdb("Find cryo-EM structures of ribosomes at resolution better than 3Å") - Direct endpoint: query_emdb(endpoint="entry/EMD-10000") """ # Base URL for EMDB API base_url = "https://www.ebi.ac.uk/emdb/api" # Ensure we have either a prompt or an endpoint if prompt is None and endpoint is None: return {"error": "Either a prompt or an endpoint must be provided"} # If using prompt, parse with Claude if prompt: # Load EMDB schema schema_path = os.path.join(SCHEMA_DB_PATH, "emdb.pkl") with open(schema_path, "rb") as f: emdb_schema = pickle.load(f) # Create system prompt template system_template = """ You are a structural biology expert specialized in using the Electron Microscopy Data Bank (EMDB) API. Based on the user's natural language request, determine the appropriate EMDB API endpoint and parameters. EMDB API SCHEMA: {schema} Your response should be a JSON object with the following fields: 1. "endpoint": The API endpoint to query (e.g., "search", "entry/EMD-XXXXX") 2. "params": An object containing query parameters specific to the endpoint 3. "description": A brief description of what the query is doing SPECIAL NOTES: - EMDB contains 3D macromolecular structures determined by electron microscopy - Common endpoints include: "search" (search for entries), "entry/EMD-XXXXX" (specific entry details) - For searching, you can filter by resolution, specimen, authors, release date, etc. - Resolution filters should be specified with "resolution_low" and "resolution_high" parameters - For specific entry retrieval, use the format "entry/EMD-XXXXX" where XXXXX is the EMDB ID - Common specimen types include: "ribosome", "virus", "membrane protein", "filament" Return ONLY the JSON object with no additional text. """ # Query Gemini to generate the API call gemini_result = _query_gemini_for_api( prompt=prompt, schema=emdb_schema, system_template=system_template ) if not gemini_result["success"]: return gemini_result # Get the endpoint and parameters from Gemini's response query_info = gemini_result["data"] endpoint = query_info.get("endpoint", "") params = query_info.get("params", {}) description = query_info.get("description", "") if not endpoint: return { "error": "Failed to generate a valid endpoint from the prompt", "gemini_response": gemini_result.get("raw_response", "No response") } else: # Process provided endpoint params = {} description = f"Direct query to EMDB {endpoint}" # Remove leading slash if present if endpoint.startswith("/"): endpoint = f"{base_url}{endpoint}" elif not endpoint.startswith("http"): endpoint = f"{base_url}/{endpoint.lstrip('/')}" description = f"Direct query to provided endpoint" # Execute the EMDB API request using the helper function api_result = _query_rest_api( endpoint=endpoint, method="GET", params=params, description=description ) # Format the results if not verbose and successful if not verbose and "success" in api_result and api_result["success"] and "result" in api_result: api_result["result"] = _format_query_results(api_result["result"]) return api_result