alphafold.py

from typing import Any, Dict, Optional
import httpx
from fastmcp import FastMCP

# Initialize the MCP server
mcp = FastMCP("alphafold")

# Constants
ALPHAFOLD_API_BASE = "https://alphafold.ebi.ac.uk/api"
USER_AGENT = "AlphaFold-MCP/1.0"

async def make_alphafold_request(endpoint: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Make a request to the AlphaFold API."""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                f"{ALPHAFOLD_API_BASE}/{endpoint}",
                params=params,
                headers={"User-Agent": USER_AGENT},
                timeout=30.0
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPError as e:
            return {"error": f"HTTP error occurred: {str(e)}"}
        except Exception as e:
            return {"error": f"An error occurred: {str(e)}"}

@mcp.tool()
async def get_prediction(uniprot_id: str) -> Dict[str, Any]:
    """Get AlphaFold prediction for a given UniProt ID."""
    result = await make_alphafold_request(f"prediction/{uniprot_id}")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def search_protein(query: str) -> Dict[str, Any]:
    """Search for proteins in the AlphaFold database."""
    result = await make_alphafold_request("search", params={"query": query})
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_organism_predictions(tax_id: str) -> Dict[str, Any]:
    """Get all predictions for a specific organism using its NCBI Taxonomy ID."""
    result = await make_alphafold_request(f"organism/{tax_id}")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_entry(uniprot_id: str) -> Dict[str, Any]:
    """Get detailed information about a specific entry."""
    result = await make_alphafold_request(f"entry/{uniprot_id}")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_pdb(uniprot_id: str) -> Dict[str, Any]:
    """Get the PDB file for a specific entry."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/pdb")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_cif(uniprot_id: str) -> Dict[str, Any]:
    """Get the mmCIF file for a specific entry."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/cif")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_json(uniprot_id: str) -> Dict[str, Any]:
    """Get the JSON representation of a specific entry."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/json")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_pae(uniprot_id: str) -> Dict[str, Any]:
    """Get the predicted aligned error (PAE) data."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/pae")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_confidence(uniprot_id: str) -> Dict[str, Any]:
    """Get the confidence scores."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/confidence")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_coverage(uniprot_id: str) -> Dict[str, Any]:
    """Get the coverage information."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/coverage")
    if "error" in result:
        return {"error": result["error"]}
    return result

@mcp.tool()
async def get_sequence(uniprot_id: str) -> Dict[str, Any]:
    """Get the sequence information."""
    result = await make_alphafold_request(f"entry/{uniprot_id}/sequence")
    if "error" in result:
        return {"error": result["error"]}
    return result

if __name__ == "__main__":
    mcp.run(transport="stdio")