Update app.py

app.py

@@ -1,67 +1,16 @@
 from fastapi import FastAPI, Query, Path, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import RedirectResponse, JSONResponse, StreamingResponse, FileResponse
-import httpx, asyncio, time, os, hashlib, json, io, gzip, math
+from fastapi.responses import RedirectResponse, JSONResponse, StreamingResponse
+import httpx, asyncio, time, hashlib, json, io, gzip, math
 from typing import Dict, Any, Tuple, Optional, List
 
-# === Block 1: entity resolver (drug vs gene) + safe helpers ===
-
-KNOWN_DRUG_TARGETS = {
-    "bupropion": ["SLC6A3","SLC6A2","CHRNA4","CHRNB2","CHRNA3","CHRNB4"],
-    "bupropion hcl": ["SLC6A3","SLC6A2","CHRNA4","CHRNB2","CHRNA3","CHRNB4"],
-    "atomoxetine": ["SLC6A2"],
-    "reboxetine": ["SLC6A2"],
-    "methylphenidate": ["SLC6A3","SLC6A2"],
-    "dexmethylphenidate": ["SLC6A3","SLC6A2"],
-    "lisdexamfetamine": ["SLC6A3","SLC6A2"],
-    "dextroamphetamine": ["SLC6A3","SLC6A2"],
-}
-
-def _norm(s: str) -> str:
-    return (s or "").strip().lower()
-
-async def resolve_mode(receptor: str) -> dict:
-    """
-    Decide if 'receptor' is a gene (via MyGene) or a drug (fallback map).
-    Returns {"mode":"gene"|"drug","canonical":..., "genes":[...]}
-    """
-    rq = (receptor or "").strip()
-    if not rq:
-        return {"mode":"gene","canonical":rq,"genes":[]}
-    # Try exact gene symbol via MyGene
-    try:
-        async with httpx.AsyncClient(headers=UA, timeout=15) as client:
-            r = await client.get("https://mygene.info/v3/query",
-                                 params={"q": f"symbol:{rq}", "fields": "symbol", "size": 1})
-            if r.status_code == 200:
-                js = r.json() or {}
-                hits = js.get("hits") or []
-                if hits and hits[0].get("symbol"):
-                    sym = hits[0]["symbol"]
-                    return {"mode":"gene","canonical":sym,"genes":[sym]}
-    except Exception:
-        pass
-    # Not a sure gene: treat as drug
-    dn = _norm(rq)
-    genes = KNOWN_DRUG_TARGETS.get(dn, [])
-    return {"mode":"drug","canonical":rq,"genes":genes}
-
-# REPLACES the old version
-def collect_gene_symbols_from_string(edges: List[dict], focus: str) -> List[str]:
-    return collect_gene_symbols_from_string_safe(edges, focus)
-
-    """
-    STRING can fail or return dicts. Be defensive and only read from lists of dicts.
-    """
-    genes = set()
-
-
+# ------------------ App constants ------------------
 APP_NAME = "neuro-mechanism-backend"
-CALLER_ID = "neuro-mech-backend-demo"   # shows in STRING logs / rate fairness
+CALLER_ID = "neuro-mech-backend-demo"  # polite ID for STRING
 UA = {"User-Agent": f"{APP_NAME}/1.2 (HF Space)"}
 
+# ------------------ FastAPI app ------------------
 app = FastAPI(title=APP_NAME)
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True,
@@ -96,7 +45,7 @@ def endpoints():
         ]
     })
 
-# ----------------- tiny in-memory TTL cache -----------------
+# ------------------ tiny in-memory TTL cache ------------------
 class TTLCache:
     def __init__(self, max_items=512):
         self.store: Dict[str, Tuple[float, Any]] = {}
@@ -125,7 +74,7 @@ class TTLCache:
 
 CACHE = TTLCache()
 
-# --------------- polite throttling for STRING ----------------
+# ------------------ polite throttling for STRING ------------------
 _last_string_call = 0.0
 async def throttle_string():
     """Be nice to STRING; ~1 req/sec is a good courtesy."""
@@ -136,7 +85,7 @@ async def throttle_string():
         await asyncio.sleep(wait)
     _last_string_call = time.time()
 
-# ----------------- Helpers -----------------
+# ------------------ small helpers ------------------
 async def get_json_cached(url: str, params: Optional[dict], ttl: int):
     try:
         return await CACHE.get(url, params, ttl)
@@ -154,9 +103,10 @@ def gz_json_bytes(obj: Any) -> bytes:
         gz.write(b)
     return bio.getvalue()
 
-# ----------------- External API wrappers -----------------
+# ------------------ External API wrappers ------------------
 @app.get("/lit/eupmc")
 async def europe_pmc_search(query: str, pageSize: int = 5):
+    # Europe PMC returns a 'hitCount' field in the response for quick counts.
     url = "https://www.ebi.ac.uk/europepmc/webservices/rest/search"
     params = {"query": query, "format": "json", "pageSize": pageSize}
     return await get_json_cached(url, params, ttl=600)
@@ -204,12 +154,52 @@ async def gpcrdb_protein(entry: str):
 @app.get("/string/network")
 async def string_network(identifiers: str, species: int = 9606, limit: int = 50):
     await throttle_string()
+    # STRING API supports a 'caller_identity' parameter – good practice to include it.
     url = "https://string-db.org/api/json/network"
     params = {"identifiers": identifiers, "species": species, "caller_identity": CALLER_ID, "limit": limit}
     return await get_json_cached(url, params, ttl=3600)
 
-# ----------------- Synonyms (regions/genes/phenotypes) --------------
-# Simple built-in expansions + OLS/MyGene lookups.
+# ------------------ Entity resolver (drug vs gene) ------------------
+KNOWN_DRUG_TARGETS = {
+    "bupropion": ["SLC6A3","SLC6A2","CHRNA4","CHRNB2","CHRNA3","CHRNB4"],
+    "bupropion hcl": ["SLC6A3","SLC6A2","CHRNA4","CHRNB2","CHRNA3","CHRNB4"],
+    "atomoxetine": ["SLC6A2"],
+    "reboxetine": ["SLC6A2"],
+    "methylphenidate": ["SLC6A3","SLC6A2"],
+    "dexmethylphenidate": ["SLC6A3","SLC6A2"],
+    "lisdexamfetamine": ["SLC6A3","SLC6A2"],
+    "dextroamphetamine": ["SLC6A3","SLC6A2"],
+}
+def _norm(s: str) -> str:
+    return (s or "").strip().lower()
+
+async def resolve_mode(receptor: str) -> dict:
+    """
+    Decide if 'receptor' is a gene (via MyGene) or a drug (fallback map).
+    Returns {"mode":"gene"|"drug","canonical":..., "genes":[...]}
+    """
+    rq = (receptor or "").strip()
+    if not rq:
+        return {"mode":"gene","canonical":rq,"genes":[]}
+    # Try exact gene symbol first (MyGene)
+    try:
+        async with httpx.AsyncClient(headers=UA, timeout=15) as client:
+            r = await client.get("https://mygene.info/v3/query",
+                                 params={"q": f"symbol:{rq}", "fields": "symbol", "size": 1})
+            if r.status_code == 200:
+                js = r.json() or {}
+                hits = js.get("hits") or []
+                if hits and hits[0].get("symbol"):
+                    sym = hits[0]["symbol"]
+                    return {"mode":"gene","canonical":sym,"genes":[sym]}
+    except Exception:
+        pass
+    # Not a sure gene → treat as drug
+    dn = _norm(rq)
+    genes = KNOWN_DRUG_TARGETS.get(dn, [])
+    return {"mode":"drug","canonical":rq,"genes":genes}
+
+# ------------------ Synonyms (regions/genes/phenotypes) ------------------
 REGION_SEED_SYNONYMS = {
     "prefrontal cortex": ["PFC","mPFC","vmPFC","dlPFC","dorsolateral prefrontal cortex","ventromedial prefrontal cortex"],
     "anterior cingulate cortex": ["ACC","dACC","pgACC","sgACC","subgenual cingulate"],
@@ -225,7 +215,6 @@ REGION_SEED_SYNONYMS = {
 }
 
 async def ols4_synonyms(term: str, ontology: Optional[str] = None) -> List[str]:
-    # OLS4 generic search (best-effort parse)
     url = "https://www.ebi.ac.uk/ols4/api/search"
     params = {"q": term, "rows": 20}
     if ontology:
@@ -233,7 +222,7 @@ async def ols4_synonyms(term: str, ontology: Optional[str] = None) -> List[str]:
     data = await get_json_cached(url, params, ttl=86400)
     syns = []
     try:
-        docs = data.get("response", {}).get("docs", []) or data.get("response", {}).get("docs", [])
+        docs = data.get("response", {}).get("docs", []) or []
         for d in docs:
             if "synonym" in d:
                 syns.extend(d.get("synonym", []))
@@ -241,18 +230,14 @@ async def ols4_synonyms(term: str, ontology: Optional[str] = None) -> List[str]:
                 syns.append(d["label"])
     except Exception:
         pass
-    # Dedup & lowercase normalize
-    out = []
-    seen = set()
+    out, seen = [], set()
     for s in syns:
         s2 = s.strip()
-        if s2.lower() not in seen:
-            out.append(s2)
-            seen.add(s2.lower())
+        if s2 and s2.lower() not in seen:
+            out.append(s2); seen.add(s2.lower())
     return out[:50]
 
 async def mygene_synonyms(symbol: str) -> List[str]:
-    # MyGene.info gene synonyms/aliases
     url = "https://mygene.info/v3/query"
     params = {"q": symbol, "fields": "symbol,name,alias,other_names", "size": 5}
     data = await get_json_cached(url, params, ttl=86400)
@@ -265,7 +250,6 @@ async def mygene_synonyms(symbol: str) -> List[str]:
                 if k in hit and isinstance(hit[k], list): syns.extend(hit[k])
     except Exception:
         pass
-    # unique
     out, seen = [], set()
     for s in syns:
         s2 = str(s).strip()
@@ -284,11 +268,10 @@ async def util_synonyms(term: str, kind: str = Query("region", enum=["region","g
         mg = await mygene_synonyms(term_norm)
         return {"term": term_norm, "kind": kind, "synonyms": sorted(set([term_norm] + mg))}
     else:
-        # phenotype via OLS (HPO)
         ols = await ols4_synonyms(term_norm, ontology="hp")
         return {"term": term_norm, "kind": kind, "synonyms": sorted(set([term_norm] + ols))}
 
-# ----------------- Regions heuristic (improved) -----------------
+# ------------------ Regions heuristic ------------------
 REGION_TERMS_DEFAULT = [
     "prefrontal cortex","anterior cingulate cortex","nucleus accumbens","ventral striatum",
     "dorsal striatum","caudate","putamen","amygdala","hippocampus","thalamus","hypothalamus",
@@ -296,6 +279,7 @@ REGION_TERMS_DEFAULT = [
 ]
 
 async def eupmc_hitcount(q: str) -> int:
+    # Europe PMC responses include <hitCount> / "hitCount" for fast counts.
     url = "https://www.ebi.ac.uk/europepmc/webservices/rest/search"
     params = {"query": q, "format": "json", "pageSize": 0}
     data = await get_json_cached(url, params, ttl=1800)
@@ -304,12 +288,12 @@ async def eupmc_hitcount(q: str) -> int:
     except Exception:
         return 0
 
-# safe parser for STRING results
-def collect_gene_symbols_from_string(edges: list, focus: str) -> list:
+def collect_gene_symbols_from_string(edges: Any, focus: str) -> List[str]:
+    """Defensive parse of STRING results (may not always be a list)."""
     genes = set()
-    f = (focus or "").upper()
     if not isinstance(edges, list):
         return []
+    f = (focus or "").upper()
     for e in edges:
         if not isinstance(e, dict):
             continue
@@ -319,7 +303,6 @@ def collect_gene_symbols_from_string(edges: list, focus: str) -> list:
                 genes.add(g)
     return list(genes)
 
-
 @app.get("/heuristics/regions_from_string")
 async def regions_from_string(
     receptor: str = Query(..., description="gene or drug"),
@@ -330,107 +313,75 @@ async def regions_from_string(
     symptom: Optional[str] = None
 ):
     """
-    Robust region scoring that works for **genes** and **drugs**.
-    - gene: (region_syns) AND (receptor OR STRING neighbors OR gene_syns)
-    - drug: (region_syns) AND (drug)
-    Falls back tier-by-tier and never 500s.
+    Robust region scoring that works for genes AND drugs.
+    gene: (region_syns) AND (receptor OR STRING neighbors OR gene_syns)
+    drug: (region_syns) AND (drug)
     """
     mode = await resolve_mode(receptor)
 
-    # --- candidate regions ---
-    REGION_TERMS_DEFAULT = [
-        "prefrontal cortex","anterior cingulate cortex","nucleus accumbens","ventral striatum",
-        "dorsal striatum","caudate","putamen","amygdala","hippocampus","thalamus","hypothalamus",
-        "insula","ventral tegmental area","substantia nigra","cerebellum"
-    ]
+    # candidates
     region_list = [r.strip() for r in (regions.split(",") if regions else REGION_TERMS_DEFAULT) if r.strip()]
 
-    # --- collect region synonyms ---
+    # region synonyms
     region_syns_map = {}
     if use_synonyms:
-        syn_results = []
         for rgn in region_list:
             try:
                 syn = await util_synonyms(rgn, "region")
-                syn_results.append((rgn, syn.get("synonyms", [])[:10] or [rgn]))
+                region_syns_map[rgn] = syn.get("synonyms", [])[:10] or [rgn]
             except Exception:
-                syn_results.append((rgn, [rgn]))
-        region_syns_map = dict(syn_results)
+                region_syns_map[rgn] = [rgn]
     else:
         region_syns_map = {rgn: [rgn] for rgn in region_list}
 
-    # --- build RHS terms depending on mode ---
-    rhs_terms = []
-    neighbors = []
-    gene_syns = []
+    # RHS terms
+    rhs_terms: List[str] = []
+    neighbors: List[str] = []
+    gene_syns: List[str] = []
 
     if mode["mode"] == "gene":
-        # STRING neighbors (defensive)
         try:
             sdata = await string_network(receptor, species=species, limit=limit)
         except Exception:
             sdata = []
-        neighbors = collect_gene_symbols_from_string_safe(sdata, receptor)
-
-        # limited gene synonyms via MyGene
+        neighbors = collect_gene_symbols_from_string(sdata, receptor)
         try:
             gs = await mygene_synonyms(receptor)
             gene_syns = gs[:20]
         except Exception:
             gene_syns = []
-
         rhs_terms = list({t for t in [receptor] + neighbors[:25] + gene_syns[:25] if t})
     else:
-        # drug path — no STRING
+        # drug path — skip STRING
         rhs_terms = [mode["canonical"]]
 
-    # --- query Europe PMC (hitCount) with tiered fallbacks ---
-    # (Europe PMC REST 'search' returns a 'hitCount' for fast counts.) 
-    # Docs: europepmc.org/RestfulWebService
+    # Europe PMC hit counts (fast)
     results = []
     symptom_clause = f" AND ({symptom})" if symptom else ""
 
-    async def hitcount(q: str) -> int:
-        try:
-            url = "https://www.ebi.ac.uk/europepmc/webservices/rest/search"
-            params = {"query": q, "format": "json", "pageSize": 0}
-            data = await get_json_cached(url, params, ttl=1800)
-            return int(data.get("hitCount", 0))
-        except Exception:
-            return 0
-
     for region in region_list:
         syns = region_syns_map.get(region, [region])
         lhs = " OR ".join(syns)
-        # Tier 1
         if mode["mode"] == "gene":
             rhs = " OR ".join(rhs_terms) if rhs_terms else receptor
             q1 = f"({lhs}) AND ({rhs}){symptom_clause}"
         else:
-            # drug
             q1 = f"({lhs}) AND ({mode['canonical']}){symptom_clause}"
 
-        h1 = await hitcount(q1)
+        h1 = await eupmc_hitcount(q1)
         hits, tier = h1, "T1"
 
         if h1 == 0 and mode["mode"] == "gene":
-            # Tier 2: region_syns AND receptor
             q2 = f"({lhs}) AND ({receptor}){symptom_clause}"
-            h2 = await hitcount(q2)
+            h2 = await eupmc_hitcount(q2)
             hits, tier = h2, "T2"
             if h2 == 0:
-                # Tier 3: label AND receptor
                 q3 = f"({region}) AND ({receptor}){symptom_clause}"
-                h3 = await hitcount(q3)
+                h3 = await eupmc_hitcount(q3)
                 hits, tier = h3, "T3"
 
         score = math.log10(hits + 1.0)
-        results.append({
-            "region": region,
-            "hits": hits,
-            "tier": tier,
-            "weighted_score": round(score, 4)
-        })
+        results.append({"region": region, "hits": hits, "tier": tier, "weighted_score": round(score, 4)})
 
     results.sort(key=lambda x: x["weighted_score"], reverse=True)
     return {
@@ -440,80 +391,7 @@ async def regions_from_string(
         "notes": f"mode={mode['mode']} genes={mode['genes']}"
     }
 
-    """
-    Rank brain regions by co-mention with (receptor OR STRING neighbors OR synonyms), with fallbacks.
-    Tiered search:
-      T1: (region_syns) AND (receptor OR neighbors OR gene_syns)
-      T2: (region_syns) AND (receptor)
-      T3: (region) AND (receptor)
-    Unquoted broad matches are used to avoid exact-phrase misses.
-    """
-    # 1) STRING neighbors
-    edges = await string_network(receptor, species=species, limit=limit)
-    neighbors = collect_gene_symbols_from_string(edges, receptor)
-
-    # 2) synonyms
-    region_list = [r.strip() for r in (regions.split(",") if regions else REGION_TERMS_DEFAULT) if r.strip()]
-    region_syns_map: Dict[str, List[str]] = {}
-    if use_synonyms:
-        syn_tasks = [util_synonyms(r, "region") for r in region_list]
-        # run as local function calls (not HTTP)
-        syn_results = await asyncio.gather(*[t if asyncio.iscoroutine(t) else asyncio.create_task(t) for t in syn_tasks])
-        for r, syn in zip(region_list, syn_results):
-            region_syns_map[r] = syn.get("synonyms", [])[:10] or [r]
-        # gene synonyms for top neighbors (cap 20)
-        gene_syns: List[str] = []
-        for g in neighbors[:20]:
-            gs = await util_synonyms(g, "gene")
-            gene_syns.extend(gs.get("synonyms", [])[:5])
-        gene_syns = list({s for s in gene_syns if s})
-    else:
-        for r in region_list:
-            region_syns_map[r] = [r]
-        gene_syns = []
-
-    # 3) Europe PMC hits per region, tiered
-    results = []
-    # build RHS (receptor OR neighbors OR gene_syns)
-    rhs_terms = [receptor] + neighbors[:25] + gene_syns[:25]
-    rhs = " OR ".join({t for t in rhs_terms if t})
-
-    for region in region_list:
-        syns = region_syns_map.get(region, [region])
-        lhs = " OR ".join(syns)
-        symptom_clause = f" AND ({symptom})" if symptom else ""
-
-        # T1
-        q1 = f"({lhs}) AND ({rhs}){symptom_clause}"
-        hc1 = await eupmc_hitcount(q1)
-        score = math.log10(hc1 + 1.0)
-        if hc1 == 0:
-            # T2
-            q2 = f"({lhs}) AND ({receptor}){symptom_clause}"
-            hc2 = await eupmc_hitcount(q2)
-            score = math.log10(hc2 + 1.0)
-            if hc2 == 0:
-                # T3
-                q3 = f"({region}) AND ({receptor}){symptom_clause}"
-                hc3 = await eupmc_hitcount(q3)
-                score = math.log10(hc3 + 1.0)
-                results.append({"region": region, "hits": hc3, "tier": "T3", "weighted_score": round(score, 4)})
-            else:
-                results.append({"region": region, "hits": hc2, "tier": "T2", "weighted_score": round(score, 4)})
-        else:
-            results.append({"region": region, "hits": hc1, "tier": "T1", "weighted_score": round(score, 4)})
-
-    results.sort(key=lambda x: x["weighted_score"], reverse=True)
-    return {
-        "focus": receptor,
-        "neighbors_considered": neighbors[:25],
-        "regions_ranked": results,
-        "notes": "Heuristic uses STRING neighbors + Europe PMC co-mentions with synonyms and fallbacks."
-    }
-
-# ----------------- Manifest / Section / Download -----------------
-
-# ephemeral in-memory store of assembled sections (by job_id)
+# ------------------ Manifest / Section / Download ------------------
 JOBS: Dict[str, Dict[str, Any]] = {}
 
 @app.get("/mechanism_graph_manifest")
@@ -524,24 +402,18 @@ async def mechanism_graph_manifest(
     string_limit: int = 50,
     lit_page_size: int = 10
 ):
-    """
-    Returns a job_id and the list of available sections with approximate sizes.
-    """
+    """Return a job_id + available sections (lightweight counts only)."""
     jid = job_key(receptor, symptom)
 
-    # Pre-compute lightweight counts; store minimal context for later sections
-    # STRING count
     sdata = await string_network(receptor, species=species, limit=string_limit)
     s_count = len(sdata) if isinstance(sdata, list) else 0
 
-    # Literature hitCount
     ldata = await europe_pmc_search(f"{receptor} AND {symptom}", pageSize=0)
     try:
         lit_hits = int(ldata.get("hitCount", 0))
     except Exception:
         lit_hits = 0
 
-    # Regions heuristic preview (no synonyms parameter here; section can recalc)
     rdata = await regions_from_string(receptor=receptor, species=species, limit=40, regions=None, use_synonyms=True, symptom=symptom)
     r_count = len(rdata.get("regions_ranked", [])) if isinstance(rdata, dict) else 0
 
@@ -551,7 +423,6 @@ async def mechanism_graph_manifest(
             "receptor": receptor, "symptom": symptom,
             "counts": {"string_edges": s_count, "literature_hits": lit_hits, "regions": r_count}
         }
-        # other sections are created lazily below
     }
 
     sections = [
@@ -560,7 +431,6 @@ async def mechanism_graph_manifest(
         {"name": "literature", "approx_size": f"{lit_hits} hits (pageSize={lit_page_size})"},
         {"name": "regions",    "approx_size": f"{r_count} entries"}
     ]
-
     return {"job_id": jid, "sections": sections}
 
 @app.get("/mechanism_graph/{section}")
@@ -573,10 +443,7 @@ async def mechanism_graph_section(
     lit_page_size: int = 10,
     job_id: Optional[str] = Query(None, description="optional; use manifest if you want stable ids")
 ):
-    """
-    Returns one section. If job_id is missing or unknown, builds on the fly.
-    """
-    # pull context from job if available
+    """Return one section; builds on-the-fly if no job_id."""
     ctx = None
     if job_id and job_id in JOBS:
         ctx = JOBS[job_id].get("_meta", {})
@@ -592,7 +459,6 @@ async def mechanism_graph_section(
             jid = job_key(receptor, symptom or "")
             JOBS.setdefault(jid, {"_meta": {"receptor": receptor, "symptom": symptom or "", "species": species}})
             job_id = jid
-        # ensure overview exists
         if "overview" not in JOBS[job_id]:
             sdata = await string_network(receptor, species=species, limit=string_limit)
             s_count = len(sdata) if isinstance(sdata, list) else 0
@@ -623,13 +489,11 @@ async def mechanism_graph_section(
 
 @app.get("/download/{job_id}/{section}")
 async def download_section(job_id: str, section: str):
-    """
-    Gzipped JSON download of a section; if section not built yet, tries to return what's there.
-    """
+    """Gzipped JSON download of a section; returns what's there."""
     data = JOBS.get(job_id, {}).get(section) or JOBS.get(job_id, {}).get("_meta")
     if not data:
         raise HTTPException(status_code=404, detail="job/section not found")
     gz = gz_json_bytes({"job_id": job_id, "section": section, "data": data})
     return StreamingResponse(io.BytesIO(gz),
                              media_type="application/gzip",
-                             headers={"Content-Disposition": f'attachment; filename="{job_id}_{section}.json.gz"'})
+                             headers={"Content-Disposition": f'attachment; filename="{job_id}_{section}.json.gz"'} )