Update app.py

app.py

@@ -27,6 +27,9 @@ import pubchempy as pcp
 import gradio as gr
 from PIL import Image
 
+from chem_nodes import *
+from agent_nodes import *
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 hf = HuggingFacePipeline.from_model_id(
@@ -36,3 +39,97 @@ hf = HuggingFacePipeline.from_model_id(
 
 chat_model = ChatHuggingFace(llm=hf)
 
+builder = StateGraph(State)
+builder.add_node("first_node", first_node)
+builder.add_node("retry_node", retry_node)
+builder.add_node("smiles_node", smiles_node)
+builder.add_node("name_node", name_node)
+builder.add_node("similars_node", similars_node)
+builder.add_node("loop_node", loop_node)
+builder.add_node("parser_node", parser_node)
+builder.add_node("reflect_node", reflect_node)
+
+builder.add_edge(START, "first_node")
+builder.add_conditional_edges("first_node", get_chemtool, {
+    "smiles_tool": "smiles_node",
+    "name_tool": "name_node",
+    "similars_tool": "similars_node",
+    None: "parser_node"})
+
+builder.add_conditional_edges("retry_node", get_chemtool, {
+    "smiles_tool": "smiles_node",
+    "name_tool": "name_node",
+    "similars_tool": "similars_node",
+    None: "parser_node"})
+
+builder.add_edge("smiles_node", "loop_node")
+builder.add_edge("name_node", "loop_node")
+builder.add_edge("similars_node", "loop_node")
+
+builder.add_conditional_edges("loop_node", get_chemtool, {
+    "smiles_tool": "smiles_node",
+    "name_tool": "name_node",
+    "similars_tool": "similars_node",
+    "loop_again": "first_node",
+    None: "parser_node"})
+
+builder.add_conditional_edges("parser_node", loop_or_not, {
+    True: "retry_node",
+    False: "reflect_node"})
+
+builder.add_edge("reflect_node", END)
+
+graph = builder.compile()
+
+def MoleculeAgent(smiles, name, task):
+
+  #if Similars_image.png exists, remove it
+  if os.path.exists('Similars_image.png'):
+    os.remove('Similars_image.png')
+
+  input = {
+    "messages": [
+        HumanMessage(f'query_smiles: {smiles}, query_task: {task}, query_name: {name}')
+    ]
+  }
+  #print(input)
+
+  replies = []
+  for c in graph.stream(input): #, stream_mode='updates'):
+    m = re.findall(r'[a-z]+\_node', str(c))
+    if len(m) != 0:
+      reply = c[str(m[0])]['messages']
+      if 'assistant' in str(reply):
+        reply = str(reply).split("<|assistant|>")[-1].split('#')[0].strip()
+        replies.append(reply)
+  #check if image exists
+  if os.path.exists('Similars_image.png'):
+    img_loc = 'Similars_image.png'
+    img = Image.open(img_loc)
+  #else create a dummy blank image
+  else:
+    img = Image.new('RGB', (250, 250), color = (255, 255, 255))
+
+  return replies[-1], img
+
+with gr.Blocks(fill_height=True) as forest:
+  gr.Markdown('''
+              # Molecule Agent - fetches names and SMILES from PubChem or finds
+              similar molecules.
+              ''')
+
+  name, smiles = None, None
+  with gr.Row():
+    with gr.Column():
+      smiles = gr.Textbox(label="Molecule SMILES of interest (optional): ", placeholder='none')
+      name = gr.Textbox(label="Molecule Name of interest (optional): ", placeholder='none')
+      task = gr.Textbox(label="Task for Agent: ")
+      calc_btn = gr.Button(value = "Submit to Agent")
+    with gr.Column():
+      props = gr.Textbox(label="Agent results: ", lines=20 )
+      pic = gr.Image(label="Molecule")
+
+
+      calc_btn.click(MoleculeAgent, inputs = [smiles, name, task], outputs = [props, pic])
+
+forest.launch(debug=True, share=True)