Update app.py

app.py

@@ -28,8 +28,6 @@ import pubchempy as pcp
 import gradio as gr
 from PIL import Image
 
-from chem_nodes import *
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 hf = HuggingFacePipeline.from_model_id(
@@ -39,6 +37,141 @@ hf = HuggingFacePipeline.from_model_id(
 
 chat_model = ChatHuggingFace(llm=hf)
 
+class State(TypedDict):
+  '''
+    The state of the agent.
+  '''
+  messages: Annotated[list, add_messages]
+  query_smiles: str
+  query_task: str
+  query_name: str
+  query_reference: str
+  tool_choice: tuple
+  which_tool: int
+  props_string: str
+  similars_img: str
+  loop_again: str
+
+def name_node(state: State) -> State:
+  '''
+    Queries Pubchem for the name of the molecule based on the smiles string.
+      Args:
+        smiles: the input smiles string
+      Returns:
+        name: the name of the molecule
+        props_string: a string of the tool results
+  '''
+  print("name tool")
+  print('===================================================')
+  current_props_string = state["props_string"]
+
+  try:
+    smiles = state["query_smiles"]
+    res = pcp.get_compounds(smiles, "smiles")
+    name = res[0].iupac_name
+    name_string = f'IUPAC molecule name: {name}\n'
+    #print(smiles, name)
+    syn_list = pcp.get_synonyms(res[0].cid)
+    for alt_name in syn_list[0]['Synonym'][:5]:
+      name_string += f'alternative or common name: {alt_name}\n'
+  except:
+    name = "unknown"
+    name_string = 'Could not find name for molecule'
+
+  state["query_name"] = name
+
+  current_props_string += name_string
+  state["props_string"] = current_props_string
+  state["which_tool"] += 1
+  return state
+
+def smiles_node(state: State) -> State:
+  '''
+    Queries Pubchem for the smiles string of the molecule based on the name.
+      Args:
+        smiles: the molecule name
+      Returns:
+        smiles: the smiles string of the molecule
+        props_string: a string of the tool results
+  '''
+  print("smiles tool")
+  print('===================================================')
+  current_props_string = state["props_string"]
+
+  try:
+    name = state["query_name"]
+    res = pcp.get_compounds(name, "name")
+    smiles = res[0].smiles
+    smiles_string = f'molecule SMILES: {smiles}\n'
+  except:
+    smiles = "unknown"
+    smiles_string = 'Could not find smiles for molecule'
+
+  state["query_smiles"] = smiles
+
+  current_props_string += smiles_string
+  state["props_string"] = current_props_string
+  state["which_tool"] += 1
+  return state
+
+def similars_node(state: State) -> State:
+  '''
+    Queries Pubchem for similar molecules based on the smiles string or name
+      Args:
+        smiles: the input smiles string, OR
+        name: the molecule name
+      Returns:
+        props_string: a string of the tool results.
+  '''
+  print("similars tool")
+  print('===================================================')
+  current_props_string = state["props_string"]
+
+  try:
+    if state['query_smiles'] != None:
+      smiles = state["query_smiles"]
+      res = pcp.get_compounds(smiles, "smiles", searchtype="similarity",listkey_count=20)
+    elif state['query_name'] != None:
+      name = state["query_name"]
+      res = pcp.get_compounds(name, "name", searchtype="similarity",listkey_count=20)
+    else:
+      print('Not enough information to run similars tool')
+      return state
+
+    props_string = 'Found Similar compounds: \n'
+    sub_smiles = []
+
+    i = 0
+    for compound in res:
+      if i == 0:
+        print(compound.iupac_name)
+        i+=1
+      sub_smiles.append(compound.smiles)
+      props_string += f'Name: {compound.iupac_name}\n'
+      props_string += f'SMILES: {compound.smiles}\n'
+      props_string += f'Molecular Weight: {compound.molecular_weight}\n'
+      props_string += f'LogP: {compound.xlogp}\n'
+      props_string += '==================='
+
+    sub_mols = [Chem.MolFromSmiles(smile) for smile in sub_smiles]
+    legend = [str(compound.smiles) for compound in res]
+
+    img = Draw.MolsToGridImage(sub_mols, legends=legend, molsPerRow=4, subImgSize=(250, 250))
+    pic = img.data
+
+    filename = "Similars_image"
+    with open(filename+".png",'wb+') as outf:
+        outf.write(pic)
+  except:
+    props_string = 'Could not find similar molecules'
+    filename = None
+
+  current_props_string += props_string
+  state["props_string"] = current_props_string
+  state['similars_img'] = filename
+  state["which_tool"] += 1
+  return state
+
 def first_node(state: State) -> State:
   '''
     The first node of the agent. This node receives the input and asks the LLM