Add SDF file support and enhance molecule input handling in handle_molecule_input function

app.py

@@ -7,7 +7,7 @@ import zipfile
 
 import pandas as pd
 import streamlit as st
-from CGRtools.files import SMILESRead
+from CGRtools.files import SMILESRead, SDFRead
 from streamlit_ketcher import st_ketcher
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import disable_progress_bars
@@ -262,7 +262,45 @@ def handle_molecule_input():
     # When text_input_changed_callback increments ketcher_render_count,
     # Ketcher's key changes, forcing it to re-initialize with the current st.session_state.shared_smiles.
     ketcher_key = f"ketcher_widget_for_sync_{st.session_state.ketcher_render_count}"
-    smile_code_output_from_ketcher = st_ketcher(st.session_state.shared_smiles, key=ketcher_key)
+    # smile_code_output_from_ketcher = st_ketcher(st.session_state.shared_smiles, key=ketcher_key)
+    col_kethcer, col_upload = st.columns([0.8, 0.2])
+    with col_kethcer:
+        smile_code_output_from_ketcher = st_ketcher(st.session_state.shared_smiles, key=ketcher_key)
+    with col_upload:
+        mol_file = st.file_uploader("Upload a molecule file (e.g., SDF, SMILES)", type=["sdf", "smi"], key="molecule_file_uploader")
+    if mol_file is not None:
+        # print(mol_file)
+        try:
+            from io import StringIO
+            stringio = StringIO(mol_file.getvalue().decode("utf-8"))
+
+            if mol_file.name.endswith(".sdf"):
+                with SDFRead(stringio) as sdf_reader:
+                    for i, molecule in enumerate(sdf_reader):
+                        if molecule is None:
+                            st.warning(f"Skipping invalid molecule at index {i}.")
+                            continue
+                        molecule.implicify_hydrogens()
+                        smiles_uploaded = str(molecule)
+            
+            # Read the uploaded file using SMILESRead
+            if mol_file.name.endswith(".smi"):
+                with SMILESRead(stringio) as sdf_reader:
+                    for i, molecule in enumerate(sdf_reader):
+                        if molecule is None:
+                            st.warning(f"Skipping invalid molecule at index {i}.")
+                            continue
+                        molecule.implicify_hydrogens()
+                        smiles_uploaded = str(molecule)
+        except Exception as e:
+            st.error(f"Error reading the uploaded file: {e}")
+        else:
+            # Priority to file-based SMILES: sync only if upload gave a new value
+            if 'smiles_uploaded' in locals() and smiles_uploaded and smiles_uploaded != st.session_state.shared_smiles:
+                st.session_state.shared_smiles = smiles_uploaded
+                st.session_state.ketcher = smiles_uploaded
+                st.session_state.ketcher_render_count += 1
+                st.rerun()
 
     # Synchronize if Ketcher changed the SMILES (user drew and applied)
     if smile_code_output_from_ketcher != st.session_state.shared_smiles:
@@ -342,6 +380,13 @@ def setup_planning_options():
             placeholder="Type a number...",
             key="c_ucb_input",
         )
+        exp_time = st.number_input(
+            "Expansion time (in seconds)",
+            value=120,
+            placeholder="Type a number...",
+            step=10,
+            key="exp_time_input",
+        )
 
     with col_options_2:
         max_iterations = st.slider(
@@ -380,6 +425,7 @@ def setup_planning_options():
         "max_iterations": max_iterations,
         "max_depth": max_depth,
         "min_mol_size": min_mol_size,
+        "max_time": exp_time,
     }
 
     if st.button("Start retrosynthetic planning", key="submit_planning_button"):
@@ -430,14 +476,15 @@ def setup_planning_options():
 
                     tree_config = TreeConfig(
                         search_strategy=planning_params["search_strategy"],
-                        evaluation_type="rollout",  # This was hardcoded, keeping it.
+                        evaluation_type="rollout",
+                        max_time=planning_params["max_time"],
                         max_iterations=planning_params["max_iterations"],
                         max_depth=planning_params["max_depth"],
                         min_mol_size=planning_params["min_mol_size"],
-                        init_node_value=0.5,  # This was hardcoded
+                        init_node_value=0.5,  
                         ucb_type=planning_params["ucb_type"],
                         c_ucb=planning_params["c_ucb"],
-                        silent=True,  # This was hardcoded
+                        silent=True,  
                     )
 
                     tree = Tree(
@@ -446,7 +493,7 @@ def setup_planning_options():
                         reaction_rules=reaction_rules,
                         building_blocks=building_blocks,
                         expansion_function=policy_function,
-                        evaluation_function=None,  # This was hardcoded
+                        evaluation_function=None,  
                     )
 
                     mcts_progress_text = "Running MCTS iterations..."