Update app.py

app.py

@@ -130,21 +130,27 @@ def similars_node(state: State) -> State:
   print(state['query_name'])
   try:
   #while x == 2:
-    if (state['query_smiles'] != None) and (state['query_smiles'] != '') and (state['query_smiles'] != 'None'):
-      smiles = state["query_smiles"]
-      print('trying with smiles')
-      res = pcp.get_compounds(smiles, "smiles", searchtype="similarity",listkey_count=20)
-      props_string = f'The following molecules are similar to {smiles}: \n'
-    # elif (state['query_name'] != None) and (state['query_name'] != '') and (state['query_name'] != 'None'):
-    else:
+    if (state['query_smiles'] != None):
+      try:
+        smiles = state["query_smiles"]
+        print('trying with smiles')
+        res = pcp.get_compounds(smiles, "smiles", searchtype="similarity",listkey_count=50)
+        props_string = f'The following molecules are similar to {smiles}: \n'
+        print('got similar molecules with smiles')
+      except:
+        print('No SMILES.')
+    elif (state['query_name'] != None):
       try:
         name = state["query_name"]
         print('trying with name', name)
-        res = pcp.get_compounds(name, "name", searchtype="similarity",listkey_count=20)
+        res = pcp.get_compounds(name, "name", searchtype="similarity",listkey_count=50)
         props_string = f'The following molecules are similar to {name}: \n'
+        print('got similar molecules with name')
       except:
-        print('Not enough information to run similars tool')
-        return state
+        print('No Name.')
+    else:
+      print('Not enough information to run similars tool')
+      return state
 
     sub_smiles = []
 
@@ -479,7 +485,7 @@ builder.add_edge("reflect_node", END)
 graph = builder.compile()
 
 @spaces.GPU
-def MoleculeAgent(smiles, name, task):
+def MoleculeAgent(task, smiles, name):
 
   #if Similars_image.png exists, remove it
   if os.path.exists('Similars_image.png'):
@@ -536,6 +542,7 @@ with gr.Blocks(fill_height=True) as forest:
       pic = gr.Image(label="Molecule")
 
 
-      calc_btn.click(MoleculeAgent, inputs = [smiles, name, task], outputs = [props, pic])
+      calc_btn.click(MoleculeAgent, inputs = [task, smiles, name], outputs = [props, pic])
+      task.submit(MoleculeAgent, inputs = [task, smiles, name], outputs = [props, pic])
 
 forest.launch(debug=False, mcp_server=True)