cBioPortal MCP Server

""" Studies endpoint module for the cBioPortal MCP server. Contains all study-related endpoint methods: - get_cancer_studies: List cancer studies with pagination - search_studies: Search studies by keyword - get_study_details: Get detailed study information - get_multiple_studies: Fetch multiple studies concurrently """ import asyncio import time from typing import Any, Dict, List, Optional import httpx from .base import BaseEndpoint, handle_api_errors, validate_paginated_params from ..utils.validation import ( validate_page_params, validate_sort_params, validate_study_id, validate_keyword, ) from ..utils.pagination import collect_all_results from ..utils.logging import get_logger logger = get_logger(__name__) class StudiesEndpoints(BaseEndpoint): """Handles all study-related endpoints for the cBioPortal MCP server.""" def __init__(self, api_client): super().__init__(api_client) @handle_api_errors("get cancer studies") @validate_paginated_params async def get_cancer_studies( self, page_number: int = 0, page_size: int = 50, sort_by: Optional[str] = None, direction: str = "ASC", limit: Optional[int] = None, ) -> Dict[str, Any]: """ Get a list of cancer studies in cBioPortal with pagination support. Args: page_number: Page number to retrieve (0-based) page_size: Number of items per page sort_by: Field to sort by direction: Sort direction (ASC or DESC) limit: Maximum number of items to return across all pages (None for no limit) Returns: Dictionary containing list of studies and metadata """ return await self.paginated_request( endpoint="studies", page_number=page_number, page_size=page_size, sort_by=sort_by, direction=direction, limit=limit, data_key="studies" ) @handle_api_errors("search studies") async def search_studies( self, keyword: str, page_number: int = 0, page_size: int = 50, sort_by: Optional[str] = None, direction: str = "ASC", limit: Optional[int] = None, ) -> Dict[str, Any]: """ Search for cancer studies by keyword in their name or description with pagination support. Args: keyword: Keyword to search for page_number: Page number to retrieve (0-based) page_size: Number of items per page sort_by: Field to sort by direction: Sort direction (ASC or DESC) limit: Maximum number of items to return across all pages (None for no limit) Returns: Dictionary containing list of studies and metadata """ # Input Validation validate_keyword(keyword) validate_page_params(page_number, page_size, limit) validate_sort_params(sort_by, direction) try: # Await the asynchronous API call all_studies = await self.api_client.make_api_request("studies") if not isinstance(all_studies, list): # Handle cases where API request might not return a list (e.g., error response) error_message = ( "Unexpected response from API when fetching all studies." ) if isinstance(all_studies, dict) and "error" in all_studies: error_message = all_studies["error"] return { "error": f"Failed to search studies for '{keyword}': {error_message}" } keyword_lower = keyword.lower() matching_studies = [ study for study in all_studies if keyword_lower in study.get("name", "").lower() or keyword_lower in study.get("description", "").lower() ] if sort_by: reverse = direction.upper() == "DESC" matching_studies.sort( key=lambda s: str(s.get(sort_by, "")), reverse=reverse ) total_count = len(matching_studies) start_idx = page_number * page_size end_idx = start_idx + page_size paginated_studies = matching_studies[start_idx:end_idx] if limit and limit > 0 and len(paginated_studies) > limit: paginated_studies = paginated_studies[:limit] has_more = end_idx < total_count return { "studies": paginated_studies, "pagination": { "page": page_number, "page_size": page_size if limit != 0 else total_count, # Adjust page_size if all were requested "total_found": total_count, "has_more": has_more if limit != 0 else False, # No more if all were requested }, } except httpx.HTTPStatusError as exc: return { "error": f"Failed to search studies for '{keyword}': API request failed with status {exc.response.status_code}" } except httpx.RequestError as exc: return { "error": f"Failed to search studies for '{keyword}': Network error - {str(exc)}" } except Exception as e: return { "error": f"Failed to search studies for '{keyword}': An unexpected error occurred - {str(e)}" } @handle_api_errors("get study details") async def get_study_details(self, study_id: str) -> Dict[str, Any]: """ Get detailed information for a specific cancer study. Args: study_id: The ID of the cancer study Returns: Dictionary containing study details """ # Input Validation validate_study_id(study_id) endpoint = f"studies/{study_id}" try: study = await self.api_client.make_api_request(endpoint) return {"study": study} except Exception as e: return {"error": f"Failed to get study details for {study_id}: {str(e)}"} @handle_api_errors("get multiple studies") async def get_multiple_studies(self, study_ids: List[str]) -> Dict: """ Get details for multiple studies concurrently. This method demonstrates the power of async concurrency by fetching multiple studies in parallel, which is much faster than sequential requests. Args: study_ids: List of study IDs to fetch Returns: Dictionary mapping study IDs to their details, with metadata about the operation """ if not study_ids: return { "studies": {}, "metadata": {"count": 0, "errors": 0, "concurrent": True}, } # Create a reusable async function for fetching a single study async def fetch_study(study_id): try: data = await self.api_client.make_api_request(f"studies/{study_id}") return {"study_id": study_id, "data": data, "success": True} except Exception as e: return {"study_id": study_id, "error": str(e), "success": False} # Create tasks for all study IDs and run them concurrently tasks = [fetch_study(study_id) for study_id in study_ids] start_time = time.perf_counter() # Use asyncio.gather to execute all tasks concurrently results = await asyncio.gather(*tasks) end_time = time.perf_counter() # Process results into a structured response studies_dict = {} error_count = 0 for result in results: if result["success"]: studies_dict[result["study_id"]] = result["data"] else: studies_dict[result["study_id"]] = {"error": result["error"]} error_count += 1 return { "studies": studies_dict, "metadata": { "count": len(study_ids), "errors": error_count, "execution_time": round(end_time - start_time, 3), "concurrent": True, }, } @handle_api_errors("get cancer types") @validate_paginated_params async def get_cancer_types( self, page_number: int = 0, page_size: int = 50, sort_by: Optional[str] = None, direction: str = "ASC", limit: Optional[int] = None, ) -> Dict: """ Get a list of all available cancer types in cBioPortal with pagination support. Args: page_number: Page number (0-based) page_size: Number of items per page sort_by: Field to sort by direction: Sort direction (ASC or DESC) limit: Maximum number of items to return across all pages (None for no limit) Returns: Dictionary containing list of cancer types and metadata """ return await self.paginated_request( endpoint="cancer-types", page_number=page_number, page_size=page_size, sort_by=sort_by, direction=direction, limit=limit, data_key="cancer_types" )