molecule visualizing, requests->httpx, support multiple molecules

lynxkite-app/web/package-lock.json

@@ -18,11 +18,13 @@
         "@syncedstore/react": "^0.6.0",
         "@types/node": "^22.10.1",
         "@xyflow/react": "^12.3.5",
+        "3dmol": "^2.4.2",
         "axios": "^1.8.2",
         "daisyui": "^4.12.20",
         "echarts": "^5.5.1",
         "fuse.js": "^7.0.0",
         "json-schema-to-typescript": "^15.0.3",
+        "monaco-editor": "^0.52.2",
         "react": "^18.3.1",
         "react-dom": "^18.3.1",
         "react-markdown": "^9.0.1",
@@ -2366,6 +2368,22 @@
         "d3-zoom": "^3.0.0"
       }
     },
+    "node_modules/3dmol": {
+      "version": "2.4.2",
+      "resolved": "https://registry.npmjs.org/3dmol/-/3dmol-2.4.2.tgz",
+      "integrity": "sha512-7R6uFaF5++s9EpO+R0a3bs6igi5idI+sgrVhtozaZlOGHg6oourYkAlAwYxWRaeBGlNcAd+oEIhbhJYhNxGyZA==",
+      "license": "BSD-3-Clause",
+      "dependencies": {
+        "iobuffer": "^5.3.1",
+        "netcdfjs": "^3.0.0",
+        "pako": "^2.1.0",
+        "upng-js": "^2.1.0"
+      },
+      "engines": {
+        "node": ">=16.16.0",
+        "npm": ">=8.11"
+      }
+    },
     "node_modules/abstract-leveldown": {
       "version": "6.2.3",
       "resolved": "https://registry.npmjs.org/abstract-leveldown/-/abstract-leveldown-6.2.3.tgz",
@@ -4187,6 +4205,12 @@
       "integrity": "sha512-0aO8FkhNZlj/ZIbNi7Lxxr12obT7cL1moPfE4tg1LkX7LlLfC6DeX4l2ZEud1ukP9jNQyNnfzQVqwbwmAATY4Q==",
       "license": "MIT"
     },
+    "node_modules/iobuffer": {
+      "version": "5.4.0",
+      "resolved": "https://registry.npmjs.org/iobuffer/-/iobuffer-5.4.0.tgz",
+      "integrity": "sha512-DRebOWuqDvxunfkNJAlc3IzWIPD5xVxwUNbHr7xKB8E6aLJxIPfNX3CoMJghcFjpv6RWQsrcJbghtEwSPoJqMA==",
+      "license": "MIT"
+    },
     "node_modules/is-alphabetical": {
       "version": "2.0.1",
       "resolved": "https://registry.npmjs.org/is-alphabetical/-/is-alphabetical-2.0.1.tgz",
@@ -5465,8 +5489,7 @@
       "version": "0.52.2",
       "resolved": "https://registry.npmjs.org/monaco-editor/-/monaco-editor-0.52.2.tgz",
       "integrity": "sha512-GEQWEZmfkOGLdd3XK8ryrfWz3AIP8YymVXiPHEdewrUq7mh0qrKrfHLNCXcbB6sTnMLnOZ3ztSiKcciFUkIJwQ==",
-      "license": "MIT",
-      "peer": true
+      "license": "MIT"
     },
     "node_modules/ms": {
       "version": "2.1.3",
@@ -5518,6 +5541,15 @@
       "dev": true,
       "license": "MIT"
     },
+    "node_modules/netcdfjs": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmjs.org/netcdfjs/-/netcdfjs-3.0.0.tgz",
+      "integrity": "sha512-LOvT8KkC308qtpUkcBPiCMBtii7ZQCN6LxcVheWgyUeZ6DQWcpSRFV9dcVXLj/2eHZ/bre9tV5HTH4Sf93vrFw==",
+      "license": "MIT",
+      "dependencies": {
+        "iobuffer": "^5.3.2"
+      }
+    },
     "node_modules/no-case": {
       "version": "3.0.4",
       "resolved": "https://registry.npmjs.org/no-case/-/no-case-3.0.4.tgz",
@@ -5649,6 +5681,12 @@
       "integrity": "sha512-g4+387DXDKlZzHkP+9FLt8yKj8+/3tOkPv7DVTJGGRm00RkEWgqbFstX1mXJ4M0VDYhUqsTOiISqNOJnhAu3PQ==",
       "license": "MIT"
     },
+    "node_modules/pako": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/pako/-/pako-2.1.0.tgz",
+      "integrity": "sha512-w+eufiZ1WuJYgPXbV/PO3NCMEc3xqylkKHzp8bxp1uW4qaSNQUkwmLLEc3kKsfz8lpV1F8Ht3U1Cm+9Srog2ug==",
+      "license": "(MIT AND Zlib)"
+    },
     "node_modules/parent-module": {
       "version": "1.0.1",
       "resolved": "https://registry.npmjs.org/parent-module/-/parent-module-1.0.1.tgz",
@@ -7107,6 +7145,21 @@
         "browserslist": ">= 4.21.0"
       }
     },
+    "node_modules/upng-js": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/upng-js/-/upng-js-2.1.0.tgz",
+      "integrity": "sha512-d3xzZzpMP64YkjP5pr8gNyvBt7dLk/uGI67EctzDuVp4lCZyVMo0aJO6l/VDlgbInJYDY6cnClLoBp29eKWI6g==",
+      "license": "MIT",
+      "dependencies": {
+        "pako": "^1.0.5"
+      }
+    },
+    "node_modules/upng-js/node_modules/pako": {
+      "version": "1.0.11",
+      "resolved": "https://registry.npmjs.org/pako/-/pako-1.0.11.tgz",
+      "integrity": "sha512-4hLB8Py4zZce5s4yd9XzopqwVv/yGNhV1Bl8NTmCq1763HeK2+EwVTv+leGeL13Dnh2wfbqowVPXCIO0z4taYw==",
+      "license": "(MIT AND Zlib)"
+    },
     "node_modules/uri-js": {
       "version": "4.4.1",
       "resolved": "https://registry.npmjs.org/uri-js/-/uri-js-4.4.1.tgz",

lynxkite-app/web/package.json

@@ -21,11 +21,13 @@
     "@syncedstore/react": "^0.6.0",
     "@types/node": "^22.10.1",
     "@xyflow/react": "^12.3.5",
+    "3dmol": "^2.4.2",
     "axios": "^1.8.2",
     "daisyui": "^4.12.20",
     "echarts": "^5.5.1",
     "fuse.js": "^7.0.0",
     "json-schema-to-typescript": "^15.0.3",
+    "monaco-editor": "^0.52.2",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "react-markdown": "^9.0.1",

lynxkite-app/web/src/workspace/Workspace.tsx

@@ -36,6 +36,7 @@ import NodeSearch, { type OpsOp, type Catalog, type Catalogs } from "./NodeSearc
 import NodeWithGraphCreationView from "./nodes/GraphCreationNode.tsx";
 import NodeWithImage from "./nodes/NodeWithImage.tsx";
 import NodeWithParams from "./nodes/NodeWithParams";
+import NodeWithPy3Dmol from "./nodes/NodeWithPy3DMol.tsx";
 import NodeWithTableView from "./nodes/NodeWithTableView.tsx";
 import NodeWithVisualization from "./nodes/NodeWithVisualization.tsx";
 
@@ -173,6 +174,7 @@ function LynxKiteFlow() {
       image: NodeWithImage,
       table_view: NodeWithTableView,
       graph_creation_view: NodeWithGraphCreationView,
+      py3dmol: NodeWithPy3Dmol,
     }),
     [],
   );

lynxkite-app/web/src/workspace/nodes/NodeWithPy3DMol.tsx

@@ -0,0 +1,67 @@
+import React, { useEffect } from "react";
+import NodeWithParams from "./NodeWithParams";
+const $3Dmol = await import("3dmol");
+
+const NodeWithPy3Dmol = (props: any) => {
+  const containerRef = React.useRef<HTMLDivElement>(null);
+  const viewerRef = React.useRef<any>(null);
+
+  useEffect(() => {
+    const config = props.data?.display?.value;
+    if (!config || !containerRef.current) return;
+
+    try {
+      // Initialize viewer only once
+      if (!viewerRef.current) {
+        viewerRef.current = $3Dmol.createViewer(containerRef.current, {
+          backgroundColor: "white",
+        });
+      }
+
+      const viewer = viewerRef.current;
+
+      // Clear previous models
+      viewer.clear();
+
+      // Add new model and style it
+      viewer.addModel(config.data, config.format);
+      viewer.setStyle({}, { stick: {} });
+      viewer.zoomTo();
+      viewer.render();
+    } catch (error) {
+      console.error("Error rendering 3D molecule:", error);
+    }
+
+    const resizeObserver = new ResizeObserver(() => {
+      viewerRef.current?.resize();
+    });
+
+    const observed = containerRef.current;
+    resizeObserver.observe(observed);
+
+    return () => {
+      resizeObserver.unobserve(observed);
+      if (viewerRef.current) {
+        viewerRef.current.clear();
+      }
+    };
+  }, [props.data?.display?.value]);
+
+  const nodeStyle = { display: "flex", flexDirection: "column", height: "100%" };
+  const vizStyle = {
+    flex: 1,
+    minHeight: "300px",
+    border: "1px solid #ddd",
+    borderRadius: "4px",
+    overflow: "hidden",
+    position: "relative",
+  };
+
+  return (
+    <NodeWithParams nodeStyle={nodeStyle} collapsed {...props}>
+      <div style={vizStyle} ref={containerRef} />
+    </NodeWithParams>
+  );
+};
+
+export default NodeWithPy3Dmol;

lynxkite-bio/src/lynxkite_bio/nims.py

@@ -3,7 +3,7 @@
 from lynxkite_graph_analytics import Bundle
 from lynxkite.core import ops
 import joblib
-import requests
+import httpx
 import pandas as pd
 import os
 
@@ -15,7 +15,7 @@ op = ops.op_registration(ENV)
 key = os.getenv("NVCF_RUN_KEY")
 
 
-def query_bionemo_nim(
+async def query_bionemo_nim(
     url: str,
     payload: dict,
 ):
@@ -26,17 +26,18 @@ def query_bionemo_nim(
     }
     try:
         print(f"Sending request to {url}")
-        response = requests.post(url, json=payload, headers=headers)
+        async with httpx.AsyncClient(timeout=600) as client:
+            response = await client.post(url, json=payload, headers=headers)
         print(f"Received response from {url}", response.status_code)
         response.raise_for_status()
         return response.json()
-    except requests.exceptions.RequestException as e:
+    except httpx.RequestError as e:
         raise ValueError(f"Query failed: {e}")
 
 
 @op("MSA-search")
 @mem.cache
-def msa_search(
+async def msa_search(
     bundle: Bundle,
     *,
     protein_table: str,
@@ -48,26 +49,34 @@ def msa_search(
     databases: str = "Uniref30_2302,colabfold_envdb_202108",
 ):
     bundle = bundle.copy()
-    response = query_bionemo_nim(
-        url="https://health.api.nvidia.com/v1/biology/colabfold/msa-search/predict",
-        payload={
-            "sequence": bundle.dfs[protein_table][protein_column].iloc[0],
-            "e_value": e_value,
-            "iterations": iterations,
-            "search_type": search_type,
-            "output_alignment_formats": [
-                format for format in output_alignment_formats.split(",")
-            ],
-            "databases": [db for db in databases.split(",")],
-        },
-    )
-    bundle.dfs[protein_table]["alignments"] = [response["alignments"]]
+    bundle.dfs[protein_table]["alignments"] = None
+
+    formats = [format.strip() for format in output_alignment_formats.split(",")]
+    dbs = [db.strip() for db in databases.split(",")]
+
+    for idx, protein_sequence in enumerate(bundle.dfs[protein_table][protein_column]):
+        print(f"Processing protein {idx + 1}/{len(bundle.dfs[protein_table])}")
+
+        response = await query_bionemo_nim(
+            url="https://health.api.nvidia.com/v1/biology/colabfold/msa-search/predict",
+            payload={
+                "sequence": protein_sequence,
+                "e_value": e_value,
+                "iterations": iterations,
+                "search_type": search_type,
+                "output_alignment_formats": formats,
+                "databases": dbs,
+            },
+        )
+
+        bundle.dfs[protein_table].at[idx, "alignments"] = response["alignments"]
+
     return bundle
 
 
 @op("Query OpenFold2")
 @mem.cache
-def query_openfold2(
+async def query_openfold2(
     bundle: Bundle,
     *,
     protein_table: str,
@@ -78,59 +87,57 @@ def query_openfold2(
     relax_prediction: bool = False,
 ):
     bundle = bundle.copy()
-    protein = bundle.dfs[protein_table][protein_column].iloc[0]
-    alignments = bundle.dfs[alignment_table][alignment_column].iloc[0]
-    selected_models = [int(model) for model in selected_models.split(",")]
-    response = query_bionemo_nim(
-        url="https://health.api.nvidia.com/v1/biology/openfold/openfold2/predict-structure-from-msa-and-template",
-        payload={
-            "sequence": protein,
-            "alignments": alignments,
-            "selected_models": selected_models,
-            "relax_prediction": relax_prediction,
-        },
-    )
-    folded_protein = response["structures_in_ranked_order"].pop(0)["structure"]
-    bundle.dfs[protein_table]["folded_protein"] = folded_protein
+
+    bundle.dfs[protein_table]["folded_protein"] = None
+    selected_models_list = [int(model) for model in selected_models.split(",")]
+
+    for idx in range(len(bundle.dfs[protein_table])):
+        print(f"Processing protein {idx + 1}/{len(bundle.dfs[protein_table])}")
+
+        protein = bundle.dfs[protein_table][protein_column].iloc[idx]
+        alignments = bundle.dfs[alignment_table][alignment_column].iloc[idx]
+
+        response = await query_bionemo_nim(
+            url="https://health.api.nvidia.com/v1/biology/openfold/openfold2/predict-structure-from-msa-and-template",
+            payload={
+                "sequence": protein,
+                "alignments": alignments,
+                "selected_models": selected_models_list,
+                "relax_prediction": relax_prediction,
+            },
+        )
+
+        folded_protein = response["structures_in_ranked_order"].pop(0)["structure"]
+        bundle.dfs[protein_table].at[idx, "folded_protein"] = folded_protein
+
+    bundle.dfs["openfold"] = pd.DataFrame()
+
     return bundle
 
 
-@op("View molecules", view="visualization")
-def view_molecules(
+@op("View molecule", view="py3dmol")
+def view_molecule(
     bundle: Bundle,
     *,
     molecule_table: str,
     molecule_column: str,
-    color="spectrum",
+    row_index: int = 0,
 ):
+    molecule_data = bundle.dfs[molecule_table][molecule_column].iloc[row_index]
+
     return {
-        "series": [
-            {
-                "type": "pie",
-                "radius": ["40%", "70%"],
-                "itemStyle": {
-                    "borderRadius": 10,
-                    "borderColor": "#fff",
-                    "borderWidth": 2,
-                },
-                "data": [
-                    {"value": 2, "name": "Hydrogen"},
-                    {"value": 1, "name": "Sulfur"},
-                    {"value": 4, "name": "Oxygen"},
-                ],
-            }
-        ]
+        "data": molecule_data,
+        "format": "pdb"
+        if molecule_data.startswith("ATOM")
+        else "sdf"
+        if molecule_data.startswith("CTfile")
+        else "smiles",
     }
 
 
-@op("Known drug")
-def known_drug(*, drug_name: str):
-    return Bundle()
-
-
 @op("Query GenMol")
 @mem.cache
-def query_genmol(
+async def query_genmol(
     bundle: Bundle,
     *,
     molecule_table: str,
@@ -143,7 +150,7 @@ def query_genmol(
 ):
     bundle = bundle.copy()
 
-    response = query_bionemo_nim(
+    response = await query_bionemo_nim(
         url="https://health.api.nvidia.com/v1/biology/nvidia/genmol/generate",
         payload={
             "smiles": bundle.dfs[molecule_table][molecule_column].iloc[0],
@@ -161,7 +168,7 @@ def query_genmol(
 
 @op("Query DiffDock")
 @mem.cache
-def query_diffdock(
+async def query_diffdock(
     proteins: Bundle,
     ligands: Bundle,
     *,
@@ -174,7 +181,7 @@ def query_diffdock(
     time_divisions=20,
     num_steps=18,
 ):
-    response = query_bionemo_nim(
+    response = await query_bionemo_nim(
         url="https://health.api.nvidia.com/v1/biology/mit/diffdock",
         payload={
             "protein": proteins.dfs[protein_table][protein_column].iloc[0],
@@ -187,17 +194,11 @@ def query_diffdock(
     )
     bundle = Bundle()
     bundle.dfs["diffdock_table"] = pd.DataFrame()
-    bundle.dfs["diffdock_table"]["protein"] = [response["protein"]] * len(
-        response["status"]
-    )
-    bundle.dfs["diffdock_table"]["ligand"] = [response["ligand"]] * len(
-        response["status"]
-    )
+    bundle.dfs["diffdock_table"]["protein"] = [response["protein"]] * len(response["status"])
+    bundle.dfs["diffdock_table"]["ligand"] = [response["ligand"]] * len(response["status"])
     bundle.dfs["diffdock_table"]["trajectory"] = response["trajectory"]
     bundle.dfs["diffdock_table"]["ligand_positions"] = response["ligand_positions"]
-    bundle.dfs["diffdock_table"]["position_confidence"] = response[
-        "position_confidence"
-    ]
+    bundle.dfs["diffdock_table"]["position_confidence"] = response["position_confidence"]
     bundle.dfs["diffdock_table"]["status"] = response["status"]
 
     return bundle

lynxkite-core/src/lynxkite/core/ops.py

@@ -180,6 +180,7 @@ class Op(BaseConfig):
                 "table_view",
                 "graph_creation_view",
                 "image",
+                "py3dmol",
             ]:
                 # If the operation is some kind of visualization, we use the output as the
                 # value to display by default.