Spaces:
Sleeping
Sleeping
"""LangGraph node functions for MOF screening pipeline."""
import json
import re
from pathlib import Path
from agent.state import ScreeningState
TOXICITY_KEYS = (
"LC50_Pimephales",
"LC50_Daphnia",
"IGC50_Tetrahymena",
"IBC50_Vibrio",
)
def _append_trace(state: ScreeningState, agent: str, action: str, output: dict) -> list[dict]:
trace = list(state.get("agent_trace", []))
trace.append({
"agent": agent,
"action": action,
"output": output,
})
return trace
def _provider_llm(state: ScreeningState, max_tokens: int = 600):
provider = state.get("llm_provider", "rule_based")
api_key = state.get("llm_api_key")
if provider == "rule_based" or not api_key:
return None
from langchain_openai import ChatOpenAI
if provider == "openai":
return ChatOpenAI(model="gpt-4o-mini", temperature=0.1, max_tokens=max_tokens, api_key=api_key)
if provider == "deepseek":
return ChatOpenAI(
model="deepseek-chat",
temperature=0.1,
max_tokens=max_tokens,
base_url="https://api.deepseek.com/v1",
api_key=api_key,
)
if provider == "qwen":
return ChatOpenAI(
model="qwen-turbo",
temperature=0.1,
max_tokens=max_tokens,
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
api_key=api_key,
)
return None
def _extract_json_object(text: str) -> dict | None:
try:
return json.loads(text)
except Exception:
pass
match = re.search(r"\{.*\}", text, flags=re.S)
if not match:
return None
try:
return json.loads(match.group(0))
except Exception:
return None
def _try_llm_json_agent(state: ScreeningState, agent_name: str, schema_hint: dict) -> dict | None:
try:
from agent.prompts import build_agent_json_prompt
llm = _provider_llm(state)
if llm is None:
return None
response = llm.invoke(build_agent_json_prompt(agent_name, state, schema_hint))
return _extract_json_object(response.content)
except Exception:
return None
def ingest_cif(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
cif_path = state["cif_path"]
try:
p = Path(cif_path)
if not p.exists():
errors.append(f"CIF file not found: {cif_path}")
return {"errors": errors, "warnings": warnings}
mof_id = p.stem
except Exception as e:
errors.append(f"ingest_cif failed: {e}")
mof_id = "unknown"
return {"mof_id": mof_id, "warnings": warnings, "errors": errors}
def six_step_screening_agent(state: ScreeningState) -> dict:
"""Run online screening tools for an uploaded CIF."""
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
if errors:
return {"warnings": warnings, "errors": errors}
try:
from tools.online_screening import ACTIVE_TOOLS, run_online_tools
result = run_online_tools(
state.get("cif_path"),
list(ACTIVE_TOOLS),
)
return {
**result,
"warnings": warnings + list(result.get("warnings", [])),
"errors": errors + list(result.get("errors", [])),
"agent_trace": _append_trace(
state,
"Online Screening Agent",
"ran online CIF screening tools",
{
"gate_status": result.get("gate_status"),
"recommendation": result.get("recommendation"),
},
) + list(result.get("agent_trace", [])),
}
except Exception as e:
errors.append(f"online_screening failed: {e}")
return {
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(
state,
"Online Screening Agent",
"online screening failed",
{"error": str(e)},
),
}
def planning_agent(state: ScreeningState) -> dict:
"""Create a task card before deterministic tools are executed."""
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
mof_id = state.get("mof_id", "unknown")
task_card = {
"objective": "screen MOF for aromatic VOC adsorption and safe-by-design suitability",
"mof_id": mof_id,
"required_tools": [
"compute_descriptor",
"predict_adsorption",
"extract_linker",
"predict_toxicity",
"apply_safety_rules",
],
"required_evidence": [
"SCM descriptor validity",
"benzene and toluene adsorption predictions",
"linker identity and extraction confidence",
"aquatic toxicity endpoints",
"metal safety tier",
"PMT and PFAS pre-filter flags",
],
"decision_policy": {
"adsorption_weight": 0.4,
"safety_weight": 0.3,
"toxicity_weight": 0.3,
"safety_veto": True,
"audit_penalty_enabled": True,
},
}
execution_plan = {
"steps": task_card["required_tools"],
"evidence_mode": "ledger",
"final_review": [
"evidence_audit_agent",
"safety_review_agent",
"reviewer_panel_agent",
"decision_agent",
"report_agent",
],
}
return {
"task_card": task_card,
"execution_plan": execution_plan,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Planning Agent", "created task card", task_card),
}
def tool_execution_agent(state: ScreeningState) -> dict:
"""Translate the task card into an executable tool plan."""
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
task_card = state.get("task_card") or {}
required_tools = task_card.get("required_tools", [])
tool_execution_plan = {
"mode": "deterministic_tools_with_agent_supervision",
"tool_order": required_tools,
"logging_policy": "each tool must write a compact trace entry and evidence ledger record",
"failure_policy": "continue where possible, then route evidence gaps to audit and repair agents",
}
return {
"tool_execution_plan": tool_execution_plan,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(
state,
"Tool Execution Agent",
"converted task card into ordered tool calls",
tool_execution_plan,
),
}
def compute_descriptor(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
if state.get("errors"):
return {"warnings": warnings, "errors": errors}
try:
from tools.descriptors import compute_scm_eigenvalues
import yaml
config_path = Path(__file__).resolve().parent.parent / "configs" / "paths.yaml"
with open(config_path, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
benzene_dim = cfg["adsorption_models"]["benzene_target_dim"]
toluene_dim = cfg["adsorption_models"]["toluene_target_dim"]
result_b = compute_scm_eigenvalues(state["cif_path"], benzene_dim)
result_t = compute_scm_eigenvalues(state["cif_path"], toluene_dim)
scm_meta = {
"benzene_eigenvalues": result_b["eigenvalues"],
"toluene_eigenvalues": result_t["eigenvalues"],
"raw_dim": result_b["raw_dim"],
"benzene_padded": result_b["padded"],
"benzene_truncated": result_b["truncated"],
"toluene_padded": result_t["padded"],
"toluene_truncated": result_t["truncated"],
}
for r in [result_b, result_t]:
if r["applicability_warning"]:
warnings.append(r["applicability_warning"])
trace_output = {
"status": "success",
"raw_dim": scm_meta["raw_dim"],
"benzene_padded": scm_meta["benzene_padded"],
"benzene_truncated": scm_meta["benzene_truncated"],
"toluene_padded": scm_meta["toluene_padded"],
"toluene_truncated": scm_meta["toluene_truncated"],
}
return {
"scm_meta": scm_meta,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "ran compute_descriptor", trace_output),
}
except Exception as e:
errors.append(f"compute_descriptor failed: {e}")
trace_output = {"status": "failed", "error": str(e)}
return {
"scm_meta": None,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "compute_descriptor failed", trace_output),
}
def predict_adsorption(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
if state.get("scm_meta") is None:
errors.append("predict_adsorption skipped: no SCM descriptors available.")
return {"adsorption": None, "warnings": warnings, "errors": errors}
try:
from tools.adsorption import predict_benzene, predict_toluene
meta = state["scm_meta"]
b = predict_benzene(meta["benzene_eigenvalues"])
t = predict_toluene(meta["toluene_eigenvalues"])
adsorption = {
"benzene_uptake_mg_g": b["uptake_mg_g"],
"benzene_model": b["model_version"],
"toluene_uptake_mg_g": t["uptake_mg_g"],
"toluene_model": t["model_version"],
}
for r in [b, t]:
if r["applicability_warning"]:
warnings.append(r["applicability_warning"])
return {
"adsorption": adsorption,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "ran predict_adsorption", adsorption),
}
except Exception as e:
errors.append(f"predict_adsorption failed: {e}")
trace_output = {"status": "failed", "error": str(e)}
return {
"adsorption": None,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "predict_adsorption failed", trace_output),
}
def extract_linker(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
try:
from tools.linker_extraction import extract_linker as _extract
result = _extract(state["cif_path"])
linker = {
"metals": result["metals"],
"linker_smiles": result["linker_smiles"],
"linker_name": result["linker_name"],
"linker_formula": result["linker_formula"],
"extraction_level": result["extraction_level"],
}
if result["extraction_level"] == 3:
warnings.append(result["extraction_note"])
return {
"linker": linker,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "ran extract_linker", linker),
}
except Exception as e:
errors.append(f"extract_linker failed: {e}")
trace_output = {"status": "failed", "error": str(e)}
return {
"linker": None,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "extract_linker failed", trace_output),
}
def predict_toxicity(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
linker = state.get("linker")
smiles = linker.get("linker_smiles") if linker else None
try:
from tools.toxicity import predict_toxicity as _predict
result = _predict(smiles)
trace_output = {
key: result.get(key)
for key in TOXICITY_KEYS
}
trace_output["applicability_domain_flag"] = result.get("applicability_domain_flag")
return {
"toxicity": result,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "ran predict_toxicity", trace_output),
}
except Exception as e:
errors.append(f"predict_toxicity failed: {e}")
trace_output = {"status": "failed", "error": str(e)}
return {
"toxicity": None,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "predict_toxicity failed", trace_output),
}
def apply_safety_rules(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
try:
from tools.safety_rules import check_metal_safety, check_pmt_pre_filter
linker = state.get("linker")
metals = linker.get("metals", []) if linker else []
smiles = linker.get("linker_smiles") if linker else None
metal_result = check_metal_safety(metals) if metals else {
"tier": "unknown", "flagged_metals": [], "details": "No metals found."
}
pmt_result = check_pmt_pre_filter(smiles)
safety = {
"metal_tier": metal_result["tier"],
"metal_flagged": metal_result["flagged_metals"],
"metal_details": metal_result["details"],
"pmt_pass": pmt_result["pmt_pass"],
"pmt_flags": pmt_result["flags"],
"pmt_descriptors": pmt_result["descriptors"],
}
return {
"safety": safety,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "ran apply_safety_rules", safety),
}
except Exception as e:
errors.append(f"apply_safety_rules failed: {e}")
trace_output = {"status": "failed", "error": str(e)}
return {
"safety": None,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Tool Execution Agent", "apply_safety_rules failed", trace_output),
}
def evidence_audit_agent(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
ledger = _build_evidence_ledger(state)
issues = []
scm_meta = state.get("scm_meta") or {}
if scm_meta.get("benzene_padded") or scm_meta.get("toluene_padded"):
issues.append({
"severity": "medium",
"source": "compute_descriptor",
"message": "SCM eigenvalue vector was padded, indicating descriptor dimension mismatch.",
})
if scm_meta.get("benzene_truncated") or scm_meta.get("toluene_truncated"):
issues.append({
"severity": "medium",
"source": "compute_descriptor",
"message": "SCM eigenvalue vector was truncated, indicating descriptor dimension mismatch.",
})
linker = state.get("linker") or {}
if not linker.get("linker_smiles"):
issues.append({
"severity": "high",
"source": "extract_linker",
"message": "No linker SMILES was identified, so linker toxicity evidence is incomplete.",
})
elif linker.get("extraction_level") == 3:
issues.append({
"severity": "medium",
"source": "extract_linker",
"message": "Linker was identified by degraded matching rather than a high-confidence match.",
})
toxicity = state.get("toxicity") or {}
missing_tox = [k for k in TOXICITY_KEYS if toxicity.get(k) is None]
if missing_tox:
issues.append({
"severity": "high" if len(missing_tox) >= 2 else "medium",
"source": "predict_toxicity",
"message": f"Missing toxicity endpoint(s): {', '.join(missing_tox)}.",
})
safety = state.get("safety") or {}
if safety.get("metal_tier") in {"black", "unknown"}:
issues.append({
"severity": "high",
"source": "apply_safety_rules",
"message": f"Metal tier is {safety.get('metal_tier')}.",
})
if not safety.get("pmt_pass", True):
issues.append({
"severity": "high",
"source": "apply_safety_rules",
"message": "PMT pre-filter failed.",
})
if errors:
issues.append({
"severity": "high",
"source": "pipeline",
"message": "One or more upstream pipeline errors occurred.",
})
high_count = sum(1 for issue in issues if issue["severity"] == "high")
medium_count = sum(1 for issue in issues if issue["severity"] == "medium")
confidence_penalty = round(min(0.4, high_count * 0.15 + medium_count * 0.07), 2)
audit_status = "pass" if not issues else "caution"
if high_count >= 2 or errors:
audit_status = "fail"
fallback = {
"audit_status": audit_status,
"issues": issues,
"confidence_penalty": confidence_penalty,
"requires_repair": audit_status != "pass",
}
llm_result = _try_llm_json_agent(state | {"evidence_ledger": ledger}, "Evidence Audit Agent", fallback)
audit_report = llm_result if isinstance(llm_result, dict) else fallback
return {
"evidence_ledger": ledger,
"audit_report": audit_report,
"repair_actions": _repair_actions_from_audit(audit_report),
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Evidence Audit Agent", "audited evidence ledger", audit_report),
}
def repair_agent(state: ScreeningState) -> dict:
"""Create an explicit repair plan when audit finds weak or missing evidence."""
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
audit_report = state.get("audit_report") or {}
issues = audit_report.get("issues", [])
repair_actions = _repair_actions_from_audit(audit_report)
repair_steps = []
for issue in issues:
source = issue.get("source", "unknown")
severity = issue.get("severity", "unknown")
message = issue.get("message", "")
if source == "compute_descriptor":
proposed = "re-parse CIF, verify structure validity, then recompute SCM descriptors"
can_auto_apply = False
elif source == "extract_linker":
proposed = "request curated linker identity or run manual linker validation"
can_auto_apply = False
elif source == "predict_toxicity":
proposed = "rerun endpoint models after linker repair or flag endpoint for experiment"
can_auto_apply = False
elif source == "apply_safety_rules":
proposed = "review metal tier, PMT flags, and potential leaching concern"
can_auto_apply = False
else:
proposed = "resolve pipeline issue before final ranking"
can_auto_apply = False
repair_steps.append({
"source": source,
"severity": severity,
"issue": message,
"proposed_action": proposed,
"can_auto_apply": can_auto_apply,
})
repair_report = {
"repair_status": "not_required" if not repair_steps else "manual_or_external_validation_required",
"repair_steps": repair_steps,
"auto_repair_attempted": False,
"decision_impact": "candidate cannot be promoted to class A until high-severity or required evidence issues are resolved"
if repair_steps else "no repair penalty",
}
return {
"repair_report": repair_report,
"repair_actions": repair_actions,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Repair Agent", "converted audit issues into repair actions", repair_report),
}
def safety_review_agent(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
safety = state.get("safety") or {}
toxicity = state.get("toxicity") or {}
audit_report = state.get("audit_report") or {}
dominant_risks = []
safety_veto = False
metal_tier = safety.get("metal_tier", "unknown")
if metal_tier in {"black", "unknown"}:
dominant_risks.append(f"{metal_tier} metal tier")
safety_veto = metal_tier == "black"
if not safety.get("pmt_pass", True):
dominant_risks.append("PMT pre-filter failure")
dominant_risks.extend([
flag for flag in safety.get("pmt_flags", [])
if "passed" not in str(flag).lower()
])
valid_tox = [toxicity.get(k) for k in TOXICITY_KEYS if toxicity.get(k) is not None]
if valid_tox and sum(valid_tox) / len(valid_tox) >= 4.0:
dominant_risks.append("high predicted aquatic toxicity")
if audit_report.get("audit_status") == "fail":
dominant_risks.append("failed evidence audit")
if safety_veto:
safety_label = "reject"
elif dominant_risks:
safety_label = "caution"
else:
safety_label = "pass"
fallback = {
"safety_label": safety_label,
"dominant_risks": dominant_risks,
"safety_veto": safety_veto,
"required_validation": _required_validation_actions(state, dominant_risks),
}
llm_result = _try_llm_json_agent(state, "Safety Review Agent", fallback)
safety_review = llm_result if isinstance(llm_result, dict) else fallback
return {
"safety_review": safety_review,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Safety Review Agent", "reviewed safety evidence", safety_review),
}
def reviewer_panel_agent(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
adsorption = state.get("adsorption") or {}
toxicity = state.get("toxicity") or {}
audit_report = state.get("audit_report") or {}
safety_review = state.get("safety_review") or {}
benzene = adsorption.get("benzene_uptake_mg_g", 0) or 0
performance_score = round(min(10.0, benzene / 1000 * 10), 1)
safety_score = 8.0
if safety_review.get("safety_label") == "caution":
safety_score = 5.5
elif safety_review.get("safety_label") == "reject":
safety_score = 2.0
valid_tox = [toxicity.get(k) for k in TOXICITY_KEYS if toxicity.get(k) is not None]
if valid_tox:
toxicity_penalty = max(0.0, (sum(valid_tox) / len(valid_tox) - 3.0) * 0.8)
safety_score = round(max(0.0, safety_score - toxicity_penalty), 1)
practicality_score = 7.5
if audit_report.get("requires_repair"):
practicality_score -= 1.5
if (state.get("linker") or {}).get("extraction_level") == 3:
practicality_score -= 1.0
practicality_score = round(max(0.0, practicality_score), 1)
reviewers = [
{
"role": "performance_reviewer",
"score": performance_score,
"comment": "Scores adsorption potential using predicted benzene uptake.",
},
{
"role": "safety_reviewer",
"score": safety_score,
"comment": "Scores metal, PMT, and predicted aquatic toxicity risk.",
},
{
"role": "practicality_reviewer",
"score": practicality_score,
"comment": "Scores evidence completeness and experimental follow-up readiness.",
},
]
panel_score = round(sum(r["score"] for r in reviewers) / len(reviewers), 2)
spread = max(r["score"] for r in reviewers) - min(r["score"] for r in reviewers)
agreement = "high" if spread < 2 else "moderate" if spread < 4 else "low"
fallback = {
"reviewers": reviewers,
"panel_score": panel_score,
"agreement": agreement,
}
llm_result = _try_llm_json_agent(state, "Reviewer Panel Agent", fallback)
panel = llm_result if isinstance(llm_result, dict) else fallback
reviewer_reports = panel.get("reviewers", reviewers)
return {
"reviewer_reports": reviewer_reports,
"decision_record": {"panel_score": panel.get("panel_score", panel_score), "panel_agreement": panel.get("agreement", agreement)},
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Reviewer Panel Agent", "completed multi-perspective review", panel),
}
def decision_agent(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
score, base_recommendation, explanation = _compute_rule_based_score(state)
audit_report = state.get("audit_report") or {}
safety_review = state.get("safety_review") or {}
panel_record = state.get("decision_record") or {}
audit_penalty = float(audit_report.get("confidence_penalty", 0) or 0) * 10
adjusted_score = round(max(0.0, score - audit_penalty), 2)
panel_score = panel_record.get("panel_score")
if isinstance(panel_score, (int, float)):
adjusted_score = round(0.75 * adjusted_score + 0.25 * float(panel_score), 2)
safety_veto = bool(safety_review.get("safety_veto"))
requires_repair = bool(audit_report.get("requires_repair"))
if safety_veto:
decision_class = "D"
agentic_recommendation = "reject_by_safety_rule"
recommendation = "reject"
elif requires_repair:
decision_class = "C"
agentic_recommendation = "repair_evidence_before_ranking"
recommendation = "borderline"
elif adjusted_score >= 7.0:
decision_class = "A"
agentic_recommendation = "recommend_for_experimental_validation"
recommendation = "recommend"
elif adjusted_score >= 5.0:
decision_class = "B"
agentic_recommendation = "validate_before_scaleup"
recommendation = "borderline"
else:
decision_class = "D" if base_recommendation == "reject" else "C"
agentic_recommendation = "reject" if decision_class == "D" else "repair_evidence_before_ranking"
recommendation = "reject" if decision_class == "D" else "borderline"
decision_record = {
**panel_record,
"base_score": score,
"audit_penalty": round(audit_penalty, 2),
"final_score": adjusted_score,
"decision_class": decision_class,
"recommendation": recommendation,
"agentic_recommendation": agentic_recommendation,
"safety_veto": safety_veto,
"blocking_issues": _blocking_issues(state),
"next_actions": _next_actions(state),
"repair_status": (state.get("repair_report") or {}).get("repair_status"),
}
explanation = (
f"{explanation} Agentic decision class {decision_class}: {agentic_recommendation}. "
f"Audit penalty {audit_penalty:.2f}; final agentic score {adjusted_score:.2f}."
)
return {
"final_score": adjusted_score,
"recommendation": recommendation,
"decision_record": decision_record,
"explanation": explanation,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Decision Agent", "fused scores, audit, and safety review", decision_record),
}
def report_agent(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
explanation = state.get("explanation")
llm_report = _try_llm_report(state)
if llm_report:
explanation = llm_report
report_payload = {
"recommendation": state.get("recommendation"),
"final_score": state.get("final_score"),
"decision_record": state.get("decision_record"),
}
return {
"explanation": explanation,
"warnings": warnings,
"errors": errors,
"agent_trace": _append_trace(state, "Report Agent", "generated final report", report_payload),
}
def score_and_explain(state: ScreeningState) -> dict:
"""Backward-compatible single-node entry point for older imports/tests."""
decision = decision_agent(state)
merged = {**state, **decision}
return report_agent(merged)
def _compute_rule_based_score(state: ScreeningState) -> tuple[float, str, str]:
import yaml
config_path = Path(__file__).resolve().parent.parent / "configs" / "thresholds.yaml"
with open(config_path, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
weights = cfg["scoring"]["weights"]
safety_scores = cfg["scoring"]["safety_scores"]
thresholds = cfg["scoring"]["recommendation"]
benzene_ref = cfg["scoring"]["adsorption_normalization"]["benzene_ref_mg_g"]
adsorption = state.get("adsorption") or {}
safety = state.get("safety") or {}
toxicity = state.get("toxicity") or {}
benzene_uptake = adsorption.get("benzene_uptake_mg_g", 0)
adsorption_score = min(benzene_uptake / benzene_ref, 1.0) * 10
tier = safety.get("metal_tier", "unknown")
safety_score = safety_scores.get(tier, safety_scores["unknown"])
if not safety.get("pmt_pass", True):
safety_score = max(safety_score - 3, 0)
tox_values = [
toxicity.get("LC50_Pimephales"),
toxicity.get("LC50_Daphnia"),
toxicity.get("IGC50_Tetrahymena"),
toxicity.get("IBC50_Vibrio"),
]
valid_tox = [v for v in tox_values if v is not None]
if valid_tox:
mean_tox = sum(valid_tox) / len(valid_tox)
tox_score = max(0, min(10, (5.0 - mean_tox) * 2 + 5))
else:
tox_score = 5.0
final_score = (
weights["adsorption"] * adsorption_score
+ weights["safety"] * safety_score
+ weights["toxicity"] * tox_score
)
final_score = round(min(10.0, max(0.0, final_score)), 2)
if final_score >= thresholds["recommend_threshold"]:
recommendation = "recommend"
elif final_score >= thresholds["borderline_threshold"]:
recommendation = "borderline"
else:
recommendation = "reject"
mof_id = state.get("mof_id", "unknown")
linker_data = state.get("linker") or {}
metals_str = ", ".join(linker_data.get("metals", []))
tox_summary = (
f"mean -log toxicity = {sum(valid_tox)/len(valid_tox):.2f}"
if valid_tox else "toxicity data unavailable"
)
explanation = (
f"MOF {mof_id} achieves {benzene_uptake:.1f} mg/g benzene uptake. "
f"Metal node(s) [{metals_str}] are in the {tier} tier. "
f"Predicted aquatic toxicity: {tox_summary}. "
f"Final score: {final_score:.2f} -> {recommendation}."
)
return final_score, recommendation, explanation
def _build_evidence_ledger(state: ScreeningState) -> list[dict]:
scm_meta = state.get("scm_meta")
adsorption = state.get("adsorption")
linker = state.get("linker")
toxicity = state.get("toxicity")
safety = state.get("safety")
return [
{
"tool": "compute_descriptor",
"status": "success" if scm_meta else "failed",
"outputs": {
"raw_dim": (scm_meta or {}).get("raw_dim"),
"benzene_padded": (scm_meta or {}).get("benzene_padded"),
"benzene_truncated": (scm_meta or {}).get("benzene_truncated"),
"toluene_padded": (scm_meta or {}).get("toluene_padded"),
"toluene_truncated": (scm_meta or {}).get("toluene_truncated"),
},
"confidence": "dimension_checked" if scm_meta else "missing",
},
{
"tool": "predict_adsorption",
"status": "success" if adsorption else "failed",
"outputs": adsorption or {},
"confidence": "trained_model" if adsorption else "missing",
},
{
"tool": "extract_linker",
"status": "success" if linker else "failed",
"outputs": linker or {},
"confidence": _linker_confidence(linker),
},
{
"tool": "predict_toxicity",
"status": "success" if toxicity else "failed",
"outputs": toxicity or {},
"confidence": "trained_endpoint_models" if toxicity else "missing",
},
{
"tool": "apply_safety_rules",
"status": "success" if safety else "failed",
"outputs": safety or {},
"confidence": "rule_based" if safety else "missing",
},
]
def _linker_confidence(linker: dict | None) -> str:
if not linker:
return "missing"
level = linker.get("extraction_level")
if level == 1:
return "high"
if level == 2:
return "medium"
if level == 3:
return "low"
return "unknown"
def _repair_actions_from_audit(audit_report: dict) -> list[str]:
actions = []
for issue in audit_report.get("issues", []):
source = issue.get("source", "")
if source == "compute_descriptor":
actions.append("verify CIF parsing and descriptor applicability domain")
elif source == "extract_linker":
actions.append("manually validate linker identity or provide curated linker SMILES")
elif source == "predict_toxicity":
actions.append("rerun or experimentally validate missing toxicity endpoints")
elif source == "apply_safety_rules":
actions.append("review metal tier, PMT flags, and leaching risk")
elif source == "pipeline":
actions.append("resolve upstream pipeline errors before ranking")
return list(dict.fromkeys(actions))
def _required_validation_actions(state: ScreeningState, dominant_risks: list[str]) -> list[str]:
actions = []
if any("metal" in risk for risk in dominant_risks):
actions.append("metal leaching and stability validation")
if any("PMT" in risk for risk in dominant_risks):
actions.append("persistence, mobility, and transformation assessment")
if any("toxicity" in risk for risk in dominant_risks):
actions.append("experimental aquatic toxicity validation")
if state.get("audit_report", {}).get("requires_repair"):
actions.extend(_repair_actions_from_audit(state.get("audit_report", {})))
return list(dict.fromkeys(actions)) or ["standard adsorption and regeneration validation"]
def _blocking_issues(state: ScreeningState) -> list[str]:
issues = []
safety_review = state.get("safety_review") or {}
audit_report = state.get("audit_report") or {}
if safety_review.get("safety_veto"):
issues.append("safety veto triggered")
for issue in audit_report.get("issues", []):
if issue.get("severity") == "high":
issues.append(issue.get("message", "high-severity evidence issue"))
return issues
def _next_actions(state: ScreeningState) -> list[str]:
actions = []
safety_review = state.get("safety_review") or {}
actions.extend(safety_review.get("required_validation", []))
actions.extend(state.get("repair_actions", []))
if not actions:
actions.append("prioritize for experimental adsorption validation")
return list(dict.fromkeys(actions))
def _try_llm_report(state: ScreeningState) -> str | None:
try:
from agent.prompts import build_report_prompt
llm = _provider_llm(state, max_tokens=800)
if llm is None:
return None
response = llm.invoke(build_report_prompt(state))
return response.content
except Exception:
return None
def _try_llm_explanation(
state: ScreeningState, score: float, recommendation: str
) -> str | None:
provider = state.get("llm_provider", "rule_based")
api_key = state.get("llm_api_key")
if provider == "rule_based" or not api_key:
return None
try:
from agent.prompts import build_explanation_prompt
from langchain_openai import ChatOpenAI
prompt = build_explanation_prompt(state, score, recommendation)
if provider == "openai":
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3, max_tokens=300, api_key=api_key)
elif provider == "deepseek":
llm = ChatOpenAI(
model="deepseek-chat",
temperature=0.3,
max_tokens=300,
base_url="https://api.deepseek.com/v1",
api_key=api_key,
)
elif provider == "qwen":
llm = ChatOpenAI(
model="qwen-turbo",
temperature=0.3,
max_tokens=300,
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
api_key=api_key,
)
else:
return None
response = llm.invoke(prompt)
return response.content
except Exception:
return None
"""LangGraph node functions for MOF screening pipeline."""
import json
import re
from pathlib import Path
from agent.state import ScreeningState
TOXICITY_KEYS = (
"LC50_Pimephales",
"LC50_Daphnia",
"IGC50_Tetrahymena",
"IBC50_Vibrio",
)
def _append_trace(state: ScreeningState, agent: str, action: str, output: dict) -> list[dict]:
trace = list(state.get("agent_trace", []))
trace.append({
"agent": agent,
"action": action,
"output": output,
})
return trace
def _provider_llm(state: ScreeningState, max_tokens: int = 600):
provider = state.get("llm_provider", "rule_based")
api_key = state.get("llm_api_key")
if provider == "rule_based" or not api_key:
return None
from langchain_openai import ChatOpenAI
if provider == "openai":
return ChatOpenAI(model="gpt-4o-mini", temperature=0.1, max_tokens=max_tokens, api_key=api_key)
if provider == "deepseek":
return ChatOpenAI(
model="deepseek-chat",
temperature=0.1,
max_tokens=max_tokens,
base_url="https://api.deepseek.com/v1",
api_key=api_key,
)
if provider == "qwen":
return ChatOpenAI(
model="qwen-turbo",
temperature=0.1,
max_tokens=max_tokens,
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
api_key=api_key,
)
return None
def _extract_json_object(text: str) -> dict | None:
try:
return json.loads(text)
except Exception:
pass
match = re.search(r"\{.*\}", text, flags=re.S)
if not match:
return None
try:
return json.loads(match.group(0))
except Exception:
return None
def _try_llm_json_agent(state: ScreeningState, agent_name: str, schema_hint: dict) -> dict | None:
try:
from agent.prompts import build_agent_json_prompt
llm = _provider_llm(state)
if llm is None:
return None
response = llm.invoke(build_agent_json_prompt(agent_name, state, schema_hint))
return _extract_json_object(response.content)
except Exception:
return None
def ingest_cif(state: ScreeningState) -> dict:
warnings = list(state.get("warnings", []))
errors = list(state.get("errors", []))
cif_path = state["cif_path"]
try:
p = Path(cif_path)
if not p.exists():
errors.append(f"CIF file not found: {cif_path}")
return {"errors": errors, "warnings": warnings}
mof_id = p.stem
except Exception as e:
errors.append(f"ingest_cif failed: {e}")
mof_id = "unknown"