cBioPortal MCP Server

""" Genes endpoint module for the cBioPortal MCP server. Contains all gene-related endpoint methods: - search_genes: Search genes by keyword with pagination - get_genes: Get gene information using batch endpoint - get_multiple_genes: Fetch multiple genes concurrently - get_mutations_in_gene: Get mutations in a specific gene """ import asyncio import time from typing import Dict, List, Optional from ..api_client import APIClient from ..constants import FETCH_ALL_PAGE_SIZE from .base import handle_api_errors from ..utils.validation import ( validate_page_params, validate_sort_params, validate_keyword, ) from ..utils.logging import get_logger logger = get_logger(__name__) class GenesEndpoints: """Handles all gene-related endpoints for the cBioPortal MCP server.""" def __init__(self, api_client: APIClient, config=None): self.api_client = api_client self.config = config @handle_api_errors("search genes") async def search_genes( self, keyword: str, page_number: int = 0, page_size: int = 50, sort_by: Optional[str] = None, direction: str = "ASC", limit: Optional[int] = None, ) -> Dict: """ Search for genes by keyword in their symbol or name with pagination support. """ # Input Validation validate_keyword(keyword) validate_page_params(page_number, page_size, limit) validate_sort_params(sort_by, direction) try: api_call_params = { "keyword": keyword, "pageNumber": page_number, "pageSize": page_size, "direction": direction, } if sort_by: api_call_params["sortBy"] = sort_by if limit == 0: api_call_params["pageSize"] = FETCH_ALL_PAGE_SIZE genes_from_api = await self.api_client.make_api_request( "genes", params=api_call_params ) # Determine if the API might have more data api_might_have_more = len(genes_from_api) == api_call_params["pageSize"] # If 'fetch all' was intended and API returned less than max fetch size, then it's definitely the end. if ( api_call_params["pageSize"] == FETCH_ALL_PAGE_SIZE and len(genes_from_api) < FETCH_ALL_PAGE_SIZE ): api_might_have_more = False # Apply server-side limit if specified (after fetching the page from API) genes_for_response = genes_from_api if limit and limit > 0 and len(genes_from_api) > limit: genes_for_response = genes_from_api[:limit] # total_found in pagination now means number of items in this specific response payload # This makes the count consistent with the actual returned data structure total_items_in_response = len(genes_for_response) return { "genes": genes_for_response, "pagination": { "page": page_number, "page_size": page_size, # Report original requested page_size to client "total_found": total_items_in_response, "has_more": api_might_have_more, }, } except Exception as e: return {"error": f"Failed to search genes: {str(e)}"} @handle_api_errors("get genes") async def get_genes( self, gene_ids: List[str], gene_id_type: str = "ENTREZ_GENE_ID", projection: str = "SUMMARY", ) -> Dict: """ Get information about specific genes by their Hugo symbol or Entrez ID using batch endpoint. Args: gene_ids: List of gene IDs (Entrez IDs or Hugo symbols) gene_id_type: Type of gene ID provided (ENTREZ_GENE_ID or HUGO_GENE_SYMBOL) projection: Level of detail to return (ID, SUMMARY, DETAILED) Returns: Dictionary with gene information """ try: params = {"geneIdType": gene_id_type, "projection": projection} gene_data = await self.api_client.make_api_request( "genes/fetch", method="POST", params=params, json_data=gene_ids ) return {"genes": gene_data} except Exception as e: return {"error": f"Failed to get gene information: {str(e)}"} @handle_api_errors("get multiple genes") async def get_multiple_genes( self, gene_ids: List[str], gene_id_type: str = "ENTREZ_GENE_ID", projection: str = "SUMMARY", ) -> Dict: """ Get information about multiple genes concurrently. This method uses concurrency to fetch multiple genes in parallel, which is much more efficient than sequential requests for large batches. Args: gene_ids: List of gene IDs (Entrez IDs or Hugo symbols) gene_id_type: Type of gene ID provided (ENTREZ_GENE_ID or HUGO_GENE_SYMBOL) projection: Level of detail to return (ID, SUMMARY, DETAILED) Returns: Dictionary with gene information and performance metadata """ overall_start_time = time.time() # Single start time for the whole operation if not gene_ids: return { "genes": {}, "metadata": { "count": 0, "total_requested": 0, "errors": 0, "concurrent": True, "batches": 0, "execution_time": round(time.time() - overall_start_time, 3), }, } # For large gene lists, break into smaller batches for API compatibility batch_size = ( self.config.get("api.batch_size.genes", 100) if self.config else 100 # Default to 100 if no config provided ) gene_batches = [ gene_ids[i : i + batch_size] for i in range(0, len(gene_ids), batch_size) ] async def fetch_gene_batch(batch): try: params = {"geneIdType": gene_id_type, "projection": projection} batch_data = await self.api_client.make_api_request( "genes/fetch", method="POST", params=params, json_data=batch ) return {"data": batch_data, "success": True} except Exception as e: return {"error": str(e), "success": False} # Create tasks for all batches and run them concurrently tasks = [fetch_gene_batch(batch) for batch in gene_batches] batch_results = await asyncio.gather(*tasks) # Process results all_genes = [] error_count = 0 for result in batch_results: if result["success"]: all_genes.extend(result["data"]) else: error_count += 1 # Convert to dictionary for easier lookup genes_dict = {} key_field = ( "hugoGeneSymbol" if gene_id_type == "HUGO_GENE_SYMBOL" else "entrezGeneId" ) for gene in all_genes: gene_key_value = gene.get(key_field) if gene_key_value: genes_dict[str(gene_key_value)] = gene return { "genes": genes_dict, "metadata": { "count": len(genes_dict), "total_requested": len(gene_ids), "errors": error_count, "execution_time": round(time.time() - overall_start_time, 3), "concurrent": True, "batches": len(gene_batches), }, } @handle_api_errors("get mutations in gene") async def get_mutations_in_gene( self, gene_id: str, study_id: str, sample_list_id: str, page_number: int = 0, page_size: int = 50, sort_by: Optional[str] = None, direction: str = "ASC", limit: Optional[int] = None, ) -> Dict: """ Get mutations in a specific gene for a given study and sample list, with pagination support. Uses the /molecular-profiles/{molecularProfileId}/mutations endpoint with GET and query parameters. The molecularProfileId is dynamically determined based on the studyId. """ try: molecular_profiles_response = await self.api_client.make_api_request( f"studies/{study_id}/molecular-profiles" ) if ( isinstance(molecular_profiles_response, dict) and "api_error" in molecular_profiles_response ): return { "error": f"Failed to fetch molecular profiles for study {study_id} to find mutation profile", "details": molecular_profiles_response, } mutation_profile_id = None if isinstance(molecular_profiles_response, list): for profile in molecular_profiles_response: if profile.get("molecularAlterationType") == "MUTATION_EXTENDED": mutation_profile_id = profile.get("molecularProfileId") break if not mutation_profile_id: return { "error": f"No MUTATION_EXTENDED molecular profile found for study {study_id}" } api_call_params = { "studyId": study_id, "sampleListId": sample_list_id, "pageNumber": page_number, "pageSize": page_size, "direction": direction, } if sort_by: api_call_params["sortBy"] = sort_by if limit == 0: api_call_params["pageSize"] = FETCH_ALL_PAGE_SIZE if str(gene_id).isdigit(): api_call_params["entrezGeneId"] = gene_id else: api_call_params["hugoGeneSymbol"] = gene_id endpoint = f"molecular-profiles/{mutation_profile_id}/mutations" mutations_from_api = await self.api_client.make_api_request( endpoint, method="GET", params=api_call_params ) if ( isinstance(mutations_from_api, dict) and "api_error" in mutations_from_api ): return { "error": "API error fetching mutations", "details": mutations_from_api, "request_params": api_call_params, } if not isinstance(mutations_from_api, list): return { "error": "Unexpected API response type for mutations (expected list)", "details": mutations_from_api, "request_params": api_call_params, } api_might_have_more = len(mutations_from_api) == api_call_params["pageSize"] if ( api_call_params["pageSize"] == FETCH_ALL_PAGE_SIZE and len(mutations_from_api) < FETCH_ALL_PAGE_SIZE ): api_might_have_more = False mutations_for_response = mutations_from_api if limit and limit > 0 and len(mutations_from_api) > limit: mutations_for_response = mutations_from_api[:limit] total_items_in_response = len(mutations_for_response) return { "mutations": mutations_for_response, "pagination": { "page": page_number, "page_size": page_size, "total_found": total_items_in_response, "has_more": api_might_have_more, }, } except Exception as e: return { "error": f"An unexpected error occurred in get_mutations_in_gene: {str(e)}" }