Spaces:
Running
Running
| """ | |
| ChemGraph Loop — a guarded, natural-language API around the real ChemGraph agent. | |
| A visitor asks a chemistry question in plain language ("what's the IR spectrum of | |
| water?"). An LLM intent router identifies the target molecule (by any name / | |
| synonym / formula) + the task; the real ChemGraph LLM agent then plans the | |
| workflow and tool calls for that intent, runs (with its deterministic workflow | |
| guard), and we return the full agentic trace + a 3D structure + the task-specific | |
| result so the site can render the loop. | |
| Two tiers, by physics cost on a free CPU: | |
| • energy / dipole → single-point, ~10-25s → run LIVE. | |
| • vibrations·IR / thermochemistry → finite-difference Hessian, ~1-3 min → | |
| served from PRECOMPUTED real agent runs (flagged cached:true). | |
| Public-safe guards: molecule + task allow-list, per-call timeout, hourly rate | |
| limit. OPENAI_API_KEY comes from a Space secret. | |
| """ | |
| import asyncio | |
| import glob | |
| import json | |
| import os | |
| import time | |
| from fastapi import FastAPI, Request | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse | |
| from cg_extract import extract | |
| # --- molecule allow-list + synonyms/formulae (only a fallback; the LLM router | |
| # below is the primary, agentic intent path) ------------------------------ | |
| MOLECULES = { | |
| "water": ["water", "h2o", "dihydrogen monoxide", "dihydrogen oxide", "oxidane", "aqua", "hydrogen oxide"], | |
| "methane": ["methane", "ch4", "marsh gas", "carbane", "natural gas", "methyl hydride"], | |
| "ammonia": ["ammonia", "nh3", "azane", "hydrogen nitride", "spirit of hartshorn"], | |
| "methanol": ["methanol", "ch3oh", "ch4o", "methyl alcohol", "wood alcohol", "carbinol", "meoh"], | |
| "ethanol": ["ethanol", "c2h5oh", "ch3ch2oh", "c2h6o", "ethyl alcohol", "grain alcohol", "drinking alcohol", "etoh"], | |
| "carbon dioxide": ["carbon dioxide", "co2", "carbonic anhydride", "carbonic acid gas", "dry ice"], | |
| "benzene": ["benzene", "c6h6", "benzol", "cyclohexatriene"], | |
| "acetic acid": ["acetic acid", "ch3cooh", "ch3co2h", "c2h4o2", "ethanoic acid", "glacial acetic acid", "vinegar acid"], | |
| "formaldehyde": ["formaldehyde", "ch2o", "hcho", "methanal", "formalin", "methylene oxide"], | |
| "hydrogen peroxide": ["hydrogen peroxide", "h2o2", "dihydrogen dioxide", "hydrogen dioxide", "peroxide"], | |
| } | |
| CALC_LABEL = {"emt": "EMT", "tblite": "TBLite"} | |
| # unicode sub/superscripts → ascii digits, so "CO₂" / "H²O" resolve like "CO2" | |
| _DIGITS = {"₀": "0", "₁": "1", "₂": "2", "₃": "3", "₄": "4", "₅": "5", "₆": "6", "₇": "7", "₈": "8", "₉": "9", | |
| "⁰": "0", "¹": "1", "²": "2", "³": "3", "⁴": "4", "⁵": "5", "⁶": "6", "⁷": "7", "⁸": "8", "⁹": "9"} | |
| def normalize(text: str) -> str: | |
| return "".join(_DIGITS.get(c, c) for c in (text or "")).lower() | |
| TASKS = ("energy", "dipole", "ir", "thermo") | |
| LIVE_TASKS = {"energy", "dipole"} | |
| CACHED_TASKS = {"ir", "thermo"} | |
| MODEL = os.environ.get("CHEMGRAPH_MODEL", "gpt-4o-mini") | |
| INTENT_MODEL = os.environ.get("CHEMGRAPH_INTENT_MODEL", MODEL) | |
| RUN_TIMEOUT = int(os.environ.get("CHEMGRAPH_TIMEOUT", "95")) | |
| RATE_MAX = int(os.environ.get("CHEMGRAPH_RATE_MAX", "40")) | |
| # --- load precomputed heavy runs (real agent, cached) ------------------------- | |
| _PRECOMPUTED = {} | |
| _HERE = os.path.dirname(os.path.abspath(__file__)) | |
| for fp in glob.glob(os.path.join(_HERE, "precomputed", "*.json")): | |
| try: | |
| base = os.path.basename(fp)[:-5] # e.g. carbon_dioxide__ir | |
| slug, task = base.rsplit("__", 1) | |
| with open(fp, "r", encoding="utf-8") as f: | |
| _PRECOMPUTED[(slug, task)] = json.load(f) | |
| except Exception: | |
| pass | |
| def _avail(task: str) -> list: | |
| return sorted({slug.replace("_", " ") for (slug, t) in _PRECOMPUTED if t == task}) | |
| SPECTRA_MOLECULES = sorted({slug.replace("_", " ") for (slug, _t) in _PRECOMPUTED}) | |
| TASK_PHRASE = {"energy": "energy", "dipole": "dipole moment", | |
| "ir": "IR spectrum", "thermo": "thermochemistry"} | |
| app = FastAPI(title="ChemGraph Loop") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=[ | |
| "https://sciencesloop.com", "https://www.sciencesloop.com", | |
| "http://localhost:4321", "http://127.0.0.1:4321", | |
| ], | |
| allow_origin_regex=r"https://[a-z0-9-]+\.vercel\.app", | |
| allow_methods=["GET", "POST", "OPTIONS"], | |
| allow_headers=["*"], | |
| ) | |
| _hits: dict[str, list[float]] = {} | |
| def _rate_ok(ip: str) -> bool: | |
| now = time.time() | |
| w = _hits.setdefault(ip, []) | |
| w[:] = [t for t in w if now - t < 3600] | |
| if len(w) >= RATE_MAX: | |
| return False | |
| w.append(now) | |
| return True | |
| # --- PRIMARY: agentic LLM intent router -------------------------------------- | |
| # An LLM reads the plain-language question and identifies the target molecule | |
| # (by any name/synonym/formula) + the task. The ChemGraph agent then plans the | |
| # actual workflow and tool calls for that intent. This is intentionally not a | |
| # rule/keyword engine — the model handles the open-ended naming space. | |
| _client = None | |
| def _openai(): | |
| global _client | |
| if _client is None: | |
| from openai import OpenAI | |
| _client = OpenAI() | |
| return _client | |
| _INTENT_SYS = ( | |
| "You are the intent router for a computational-chemistry agent that is limited " | |
| "to a fixed set of small molecules and four tasks. Read the user's natural-language " | |
| "question, identify the target molecule and the task, and reply with ONLY a compact " | |
| "JSON object. Do not add commentary." | |
| ) | |
| def llm_parse_query(text: str) -> dict: | |
| mols = ", ".join(sorted(MOLECULES.keys())) | |
| user = ( | |
| f"Allowed molecules (canonical names): {mols}.\n" | |
| "Map the user's molecule to one of these ONLY if it is the SAME compound, by any " | |
| "name — common name, trivial/trade name, IUPAC name, or chemical formula. Examples: " | |
| "'H2O' / 'dihydrogen monoxide' / 'aqua' -> water; 'CO2' / 'dry ice' -> carbon dioxide; " | |
| "'EtOH' / 'ethyl alcohol' / 'grain alcohol' -> ethanol; 'NH3' / 'azane' -> ammonia; " | |
| "'benzol' -> benzene.\n" | |
| "CRITICAL: a substituted or derivative molecule is a DIFFERENT compound — return " | |
| "null, do NOT map it to the parent. e.g. dimethoxybenzene / toluene / nitrobenzene / " | |
| "phenol / aniline are NOT benzene; acetaldehyde is NOT formaldehyde; propanol is NOT " | |
| "ethanol. Never match just because the name contains an allowed molecule's name. If " | |
| "the molecule is not EXACTLY one of the allowed ones, use null.\n" | |
| "Tasks: 'energy' = single-point energy; 'dipole' = dipole moment; 'ir' = " | |
| "vibrational frequencies / IR spectrum; 'thermo' = thermochemistry (enthalpy, " | |
| "entropy, Gibbs free energy). Pick the closest task; default to 'energy' if none " | |
| "is implied.\n" | |
| "Calculator: 'emt' only if the user explicitly asks for EMT and the task is " | |
| "energy; otherwise 'tblite'.\n" | |
| "When molecule is null, also give: 'nearest' = the single closest allowed molecule " | |
| "if the user's is a close relative/derivative (e.g. dimethoxybenzene -> benzene, " | |
| "1-propanol -> ethanol, acetone -> null if nothing close), else null; and 'note' = " | |
| "ONE short, friendly sentence naming the user's molecule and why it's outside this " | |
| "small-molecule demo (e.g. \"Caffeine is too large for this small-molecule demo.\").\n" | |
| f'User question: "{text}"\n' | |
| 'Reply as JSON: {"molecule": <canonical name or null>, "task": ' | |
| '"energy|dipole|ir|thermo", "calculator": "emt|tblite", ' | |
| '"nearest": <allowed name or null>, "note": <string or null>}' | |
| ) | |
| resp = _openai().chat.completions.create( | |
| model=INTENT_MODEL, | |
| temperature=0, | |
| response_format={"type": "json_object"}, | |
| messages=[{"role": "system", "content": _INTENT_SYS}, | |
| {"role": "user", "content": user}], | |
| ) | |
| data = json.loads(resp.choices[0].message.content or "{}") | |
| mol = data.get("molecule") | |
| mol = mol.strip().lower() if isinstance(mol, str) else "" | |
| task = data.get("task") or "energy" | |
| task = task.strip().lower() if isinstance(task, str) else "energy" | |
| if mol not in MOLECULES: | |
| nearest = data.get("nearest") | |
| nearest = nearest.strip().lower() if isinstance(nearest, str) else None | |
| return {"error": "no_molecule", | |
| "note": data.get("note") if isinstance(data.get("note"), str) else None, | |
| "nearest": nearest if nearest in MOLECULES else None, | |
| "task": task if task in TASKS else "energy"} | |
| if task not in TASKS: | |
| task = "energy" | |
| # calculator isn't "intent": TBLite (real QM) by default; EMT only for an | |
| # energy query where the user explicitly named it. Deterministic, not model-guessed. | |
| calc = "emt" if (task == "energy" and "emt" in normalize(text)) else "tblite" | |
| return {"molecule": mol, "task": task, "calculator": calc} | |
| def keyword_fallback(text: str) -> dict: | |
| """Non-LLM fallback used ONLY if the intent model call fails. Token-level | |
| alias match (no regex) so e.g. 'ch4' does not match the token 'ch4o'.""" | |
| t = normalize(text) | |
| toks = set(t.replace(",", " ").replace("?", " ").replace("(", " ").replace(")", " ").split()) | |
| molecule = None | |
| for canonical, aliases in MOLECULES.items(): | |
| for a in aliases: | |
| hit = (a in t) if " " in a else (a in toks) | |
| if hit: | |
| molecule = canonical | |
| break | |
| if molecule: | |
| break | |
| if molecule is None: | |
| return {"error": "no_molecule"} | |
| task = "energy" | |
| for name, kws in (("thermo", ("thermo", "enthalpy", "entropy", "gibbs")), | |
| ("ir", ("ir", "infrared", "vibration", "vibrational", "frequency", "frequencies", "spectrum", "spectra")), | |
| ("dipole", ("dipole", "polarity"))): | |
| if toks & set(kws): | |
| task = name | |
| break | |
| calc = "emt" if ("emt" in toks and task == "energy") else "tblite" | |
| return {"molecule": molecule, "task": task, "calculator": calc} | |
| def canonical_query(molecule: str, task: str, calc: str) -> str: | |
| label = CALC_LABEL[calc] | |
| if task == "energy": | |
| return f"Compute the single point energy of {molecule} using the {label} calculator." | |
| if task == "dipole": | |
| return f"Compute the dipole moment of {molecule} using the {label} calculator." | |
| if task == "ir": | |
| return f"Compute the IR spectrum of {molecule} using the {label} calculator." | |
| if task == "thermo": | |
| return f"Compute the thermochemistry (enthalpy, entropy, Gibbs free energy) of {molecule} using the {label} calculator at 298.15 K." | |
| return f"Compute the single point energy of {molecule} using the {label} calculator." | |
| def health(): | |
| return { | |
| "ok": True, "service": "chemgraph-loop", "model": MODEL, | |
| "molecules": sorted(MOLECULES.keys()), | |
| "tasks": ["energy", "dipole", "ir", "thermo"], | |
| "live_tasks": sorted(LIVE_TASKS), | |
| "ir_molecules": _avail("ir"), | |
| "thermo_molecules": _avail("thermo"), | |
| } | |
| async def run(req: Request): | |
| ip = (req.headers.get("x-forwarded-for") or (req.client.host if req.client else "?") or "?").split(",")[0].strip() | |
| if not _rate_ok(ip): | |
| return JSONResponse({"error": "Too many requests — try again later."}, status_code=429) | |
| try: | |
| body = await req.json() | |
| except Exception: | |
| body = {} | |
| # Accept NL {query}, or legacy {molecule, calculator}. | |
| if body.get("query"): | |
| user_text = str(body["query"]) | |
| # PRIMARY: LLM intent router; keyword fallback only if the model call fails. | |
| try: | |
| parsed = llm_parse_query(user_text) | |
| except Exception: | |
| parsed = keyword_fallback(user_text) | |
| else: | |
| mol = str(body.get("molecule", "water")).strip().lower() | |
| calc = str(body.get("calculator", "emt")).strip().lower() | |
| parsed = {"molecule": mol if mol in MOLECULES else None, | |
| "task": "energy", "calculator": calc if calc in CALC_LABEL else "emt"} | |
| if parsed["molecule"] is None: | |
| parsed = {"error": "no_molecule"} | |
| user_text = None | |
| # ---- CLARIFICATION NODE: molecule not in the demo's small-molecule set ---- | |
| if parsed.get("error") == "no_molecule": | |
| note = parsed.get("note") or "That doesn't look like one of the small molecules in this demo." | |
| nearest = parsed.get("nearest") | |
| want = parsed.get("task", "energy") | |
| suggestions = [] | |
| if nearest: | |
| note += f" Did you mean {nearest}?" | |
| if want in CACHED_TASKS and nearest in _avail(want): | |
| suggestions.append({"label": f"{TASK_PHRASE[want]} of {nearest}", | |
| "query": f"{TASK_PHRASE[want]} of {nearest}"}) | |
| suggestions += [{"label": f"energy of {nearest}", "query": f"energy of {nearest}"}, | |
| {"label": f"dipole of {nearest}", "query": f"dipole of {nearest}"}] | |
| else: | |
| suggestions = [{"label": "IR spectrum of water", "query": "IR spectrum of water"}, | |
| {"label": "dipole of ammonia", "query": "dipole of ammonia"}, | |
| {"label": "energy of benzene", "query": "energy of benzene"}] | |
| return JSONResponse({"clarify": True, "message": note, | |
| "molecules": sorted(MOLECULES.keys()), "suggestions": suggestions}) | |
| molecule, task, calc = parsed["molecule"], parsed["task"], parsed["calculator"] | |
| calc_label = CALC_LABEL[calc] | |
| query = canonical_query(molecule, task, calc) | |
| display_query = user_text or query | |
| # ---- cached heavy tasks (vibrations/IR, thermochemistry) ---- | |
| if task in CACHED_TASKS: | |
| slug = molecule.replace(" ", "_") | |
| payload = _PRECOMPUTED.get((slug, task)) | |
| if payload is None: | |
| # CLARIFICATION NODE: molecule is supported, but this heavy task wasn't | |
| # precomputed for it — offer its live options + where the task IS available. | |
| avail = _avail(task) | |
| msg = (f"I don't have a precomputed {TASK_PHRASE[task]} for {molecule} — those need a " | |
| f"slow vibrational (Hessian) run, so they're prepared ahead of time" | |
| + (f" for {', '.join(avail)}" if avail else "") + ". " | |
| f"But I can run {molecule}'s energy or dipole live right now.") | |
| suggestions = [{"label": f"energy of {molecule}", "query": f"energy of {molecule}"}, | |
| {"label": f"dipole of {molecule}", "query": f"dipole of {molecule}"}] | |
| if avail: | |
| suggestions.append({"label": f"{TASK_PHRASE[task]} of {avail[0]}", | |
| "query": f"{TASK_PHRASE[task]} of {avail[0]}"}) | |
| return JSONResponse({"clarify": True, "message": msg, "suggestions": suggestions}) | |
| out = dict(payload) | |
| out["display_query"] = display_query | |
| return JSONResponse(out) | |
| # ---- live single-point tasks (energy, dipole) ---- | |
| try: | |
| from chemgraph.agent.llm_agent import ChemGraph | |
| agent = ChemGraph(model_name=MODEL, return_option="state", recursion_limit=20) | |
| state = await asyncio.wait_for(agent.run(query), timeout=RUN_TIMEOUT) | |
| except asyncio.TimeoutError: | |
| return JSONResponse({"error": "The workflow took too long — please retry."}, status_code=504) | |
| except Exception as e: # noqa: BLE001 | |
| return JSONResponse({"error": f"Agent error: {str(e)[:200]}"}, status_code=500) | |
| out = extract(state, query, molecule, calc_label, task, cached=False) | |
| out["display_query"] = display_query | |
| return JSONResponse(out) | |