""" AlphaFold Adaptive API FastAPI server for the AlphaFold custom GPT HuggingFace Space. Endpoints consumed by GPT Actions + interactive HTML UI at /. """ from __future__ import annotations import base64, uuid, os, traceback from typing import Any from fastapi import FastAPI, Request, HTTPException, Query from fastapi.responses import HTMLResponse, JSONResponse, Response from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from adaptive import get_session, build_adaptive_prompt, update_session_protein, generate_followups from alphafold_client import ( fetch_alphafold_summary, fetch_uniprot_entry, extract_protein_metadata, extract_plddt_scores, fetch_plddt_scores, ) from morphic import generate_morphic_image, generate_confidence_heatmap from models import embed_sequence, predict_disorder, classify_secondary_structure_heuristic # ── App setup ───────────────────────────────────────────────────────────────── app = FastAPI( title="AlphaFold Adaptive API", description="Adaptive protein intelligence with morphic simulations and ESM-2 embeddings.", version="1.0.0", docs_url="/docs", openapi_url="/openapi.json", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) # ── Request / Response models ───────────────────────────────────────────────── class AnalyzeRequest(BaseModel): uniprot_id: str = Field(..., description="UniProt accession number (e.g. P04637)") session_id: str = Field(default="", description="Session ID for adaptive context tracking") query: str = Field(default="", description="User's current question for expertise calibration") class ContextRequest(BaseModel): session_id: str = Field(default="", description="Session ID") query: str = Field(default="", description="Current user query") class SimulateRequest(BaseModel): uniprot_id: str = Field(..., description="UniProt accession number") plddt: float = Field(default=75.0, ge=0, le=100) length: int = Field(default=300, ge=1) function: str = Field(default="unknown") protein_name: str = Field(default="") # ── Health ──────────────────────────────────────────────────────────────────── @app.get("/health") async def health(): return {"status": "ok", "service": "alphafold-adaptive-api"} # ── Structure endpoint ──────────────────────────────────────────────────────── @app.get("/structure/{uniprot_id}") async def get_structure(uniprot_id: str, session_id: str = Query(default="")): """Fetch AlphaFold structure summary and UniProt metadata.""" uid = uniprot_id.strip().upper() try: af_data, up_data = await _fetch_both(uid) except Exception as e: raise HTTPException(status_code=404, detail=f"Could not fetch data for {uid}: {e}") meta = extract_protein_metadata(up_data) plddt_scores = await fetch_plddt_scores(af_data) avg_plddt = sum(plddt_scores) / len(plddt_scores) if plddt_scores else 75.0 # Update session if session_id: sess = get_session(session_id) update_session_protein(sess, uid, meta.get("keywords", [])) return { "uniprot_id": uid, "protein_name": meta["protein_name"], "gene_name": meta["gene_name"], "organism": meta["organism"], "length": meta["length"], "function": meta["function_text"], "function_category": meta["function_category"], "keywords": meta["keywords"], "subcellular_location": meta["subcellular_location"], "disease_associations": meta["disease_associations"], "avg_plddt": round(avg_plddt, 2), "plddt_scores": plddt_scores[:500], # cap for payload size "model_url": af_data.get("pdbUrl") or af_data.get("cifUrl", ""), "pae_image_url": af_data.get("paeImageUrl", ""), "alphafold_model_id": af_data.get("entryId", ""), } @app.post("/analyze") async def analyze_protein(req: AnalyzeRequest): """ Full protein analysis: structure + ESM-2 embeddings + disorder prediction + adaptive context injection for the custom GPT. """ uid = req.uniprot_id.strip().upper() sid = req.session_id or str(uuid.uuid4()) try: af_data, up_data = await _fetch_both(uid) except Exception as e: raise HTTPException(status_code=404, detail=f"Could not fetch data for {uid}: {e}") meta = extract_protein_metadata(up_data) plddt_scores = await fetch_plddt_scores(af_data) avg_plddt = sum(plddt_scores) / len(plddt_scores) if plddt_scores else 75.0 # ESM-2 embedding + heuristics (run synchronously — CPU only) seq = meta.get("sequence", "") embedding = embed_sequence(seq) if seq else [] ss = classify_secondary_structure_heuristic(seq) if seq else {} disorder = predict_disorder(seq)[:50] if seq else [] # first 50 residues # Adaptive context sess = get_session(sid) update_session_protein(sess, uid, meta.get("keywords", [])) adaptive_context = build_adaptive_prompt(sess, req.query) followups = generate_followups(sess) return { "session_id": sid, "uniprot_id": uid, "protein_name": meta["protein_name"], "gene_name": meta["gene_name"], "organism": meta["organism"], "length": meta["length"], "function": meta["function_text"], "function_category": meta["function_category"], "keywords": meta["keywords"], "subcellular_location": meta["subcellular_location"], "disease_associations": meta["disease_associations"], "avg_plddt": round(avg_plddt, 2), "secondary_structure": ss, "disorder_n_term": [round(d, 3) for d in disorder], "embedding_dim": len(embedding), "model_url": af_data.get("pdbUrl") or af_data.get("cifUrl", ""), "pae_image_url": af_data.get("paeImageUrl", ""), "adaptive_context": adaptive_context, "suggested_followups": followups, "simulation_url": f"/simulate/image/{uid}?plddt={avg_plddt:.1f}&length={meta['length']}&function={meta['function_category']}&name={meta['protein_name'][:40]}", "heatmap_url": f"/simulate/heatmap/{uid}", } @app.post("/context") async def get_adaptive_context(req: ContextRequest): """Return the current adaptive system prompt context for a session.""" sid = req.session_id or str(uuid.uuid4()) sess = get_session(sid) prompt = build_adaptive_prompt(sess, req.query) followups = generate_followups(sess) return { "session_id": sid, "adaptive_context": prompt, "expertise_score": sess.expertise_score, "preferred_depth": sess.preferred_depth, "proteins_visited": sess.proteins_visited, "interaction_count": sess.interaction_count, "suggested_followups": followups, } # ── Simulation endpoints ────────────────────────────────────────────────────── @app.post("/simulate") async def simulate_json(req: SimulateRequest): """Return morphic simulation as base64-encoded PNG in JSON.""" png = generate_morphic_image( uniprot_id=req.uniprot_id, plddt=req.plddt, length=req.length, function=req.function, protein_name=req.protein_name, ) return { "uniprot_id": req.uniprot_id, "image_b64": base64.b64encode(png).decode(), "content_type": "image/png", "description": ( f"Gray-Scott reaction-diffusion morphic simulation for {req.uniprot_id}. " f"Pattern self-assembled from: pLDDT={req.plddt:.1f}, " f"length={req.length}aa, function={req.function}." ), } @app.get("/simulate/image/{uniprot_id}") async def simulate_image( uniprot_id: str, plddt: float = Query(default=75.0), length: int = Query(default=300), function: str = Query(default="unknown"), name: str = Query(default=""), ): """Return morphic simulation PNG directly (for img src use).""" png = generate_morphic_image( uniprot_id=uniprot_id, plddt=plddt, length=length, function=function, protein_name=name, ) return Response(content=png, media_type="image/png") @app.get("/simulate/heatmap/{uniprot_id}") async def simulate_heatmap(uniprot_id: str): """Fetch live pLDDT scores and return per-residue heatmap PNG.""" uid = uniprot_id.strip().upper() try: af_data, up_data = await _fetch_both(uid) except Exception as e: raise HTTPException(status_code=404, detail=str(e)) meta = extract_protein_metadata(up_data) scores = await fetch_plddt_scores(af_data) if len(scores) <= 1: scores = [scores[0]] * meta["length"] png = generate_confidence_heatmap( residue_scores=scores, uniprot_id=uid, protein_name=meta["protein_name"], ) return Response(content=png, media_type="image/png") # ── Interactive UI ──────────────────────────────────────────────────────────── @app.get("/", response_class=HTMLResponse) async def ui(): return HTMLResponse(content=_HTML_UI) # ── Helper ──────────────────────────────────────────────────────────────────── async def _fetch_both(uid: str): import asyncio af_task = fetch_alphafold_summary(uid) up_task = fetch_uniprot_entry(uid) af_data, up_data = await asyncio.gather(af_task, up_task) return af_data, up_data # ── Embedded HTML UI ───────────────────────────────────────────────────────── _HTML_UI = """ AlphaFold Adaptive API

AlphaFold Adaptive API

Adaptive intelligence · Morphic simulations · ESM-2 embeddings · Custom GPT integration

Enter a UniProt ID and click Analyze.
API Endpoints
Swagger UI OpenAPI JSON Health

POST /analyze · GET /structure/{id} · POST /simulate · GET /simulate/image/{id} · GET /simulate/heatmap/{id} · POST /context

"""