Update app.py

app.py

@@ -1,504 +1,504 @@
-import torch
-from typing import Annotated, TypedDict, Literal
-from langchain_community.tools import DuckDuckGoSearchRun
-from langchain_core.tools import tool
-from langgraph.prebuilt import ToolNode, tools_condition
-from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import add_messages
-from langchain_core.messages import SystemMessage, trim_messages, AIMessage, HumanMessage, ToolCall
-
-from langchain_huggingface.llms import HuggingFacePipeline
-from langchain_huggingface import ChatHuggingFace
-from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
-from langchain_core.runnables import chain
-from uuid import uuid4
-import re
-import matplotlib.pyplot as plt
-import spaces
-
-from dockstring import load_target
-from rdkit import Chem
-from rdkit.Chem import AllChem, QED
-from rdkit.Chem import Draw
-from rdkit.Chem.Draw import MolsToGridImage
-import os, re
-import gradio as gr
-from PIL import Image
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-hf = HuggingFacePipeline.from_model_id(
-    model_id= "microsoft/Phi-4-mini-instruct",
-    task="text-generation",
-    pipeline_kwargs = {"max_new_tokens": 500, "temperature": 0.4})
-
-chat_model = ChatHuggingFace(llm=hf)
-
-cpuCount = os.cpu_count()
-print(f"Number of CPUs: {cpuCount}")
-
-class State(TypedDict):
-  '''
-    The state of the agent.
-  '''
-  messages: Annotated[list, add_messages]
-  #for the agent
-  tool_choice: tuple
-  which_tool: int
-  props_string: str
-  similars_img: str
-  loop_again: str
-
-  #for the user input
-  query_smiles: str
-  query_task: str
-  query_protein: str
-
-def docking_node(state: State) -> State:
-  '''
-    Docking tool: uses dockstring to dock the molecule into the protein
-  '''
-  print("docking tool")
-  print('===================================================')
-  current_props_string = state["props_string"]
-  query_protein = state["query_protein"]
-  query_smiles = state["query_smiles"]
-
-  print(f'query_protein: {query_protein}')
-  print(f'query_smiles: {query_smiles}')
-
-  try:
-    target = load_target(query_protein)
-    print("===============================================")
-    print(f"Docking molecule with {cpuCount} cpu cores.")
-    score, aux = target.dock(query_smiles, num_cpus = cpuCount)
-    mol = aux['ligand']
-    print(f"Docking score: {score}")
-    print("===============================================")
-    atoms_list = ""
-    template = mol
-    molH = Chem.AddHs(mol)
-    AllChem.ConstrainedEmbed(molH,template, useTethers=True)
-    xyz_string = f"{molH.GetNumAtoms()}\n\n"
-    for atom in molH.GetAtoms():
-      atoms_list += atom.GetSymbol()
-      pos = molH.GetConformer().GetAtomPosition(atom.GetIdx())
-      xyz_string += f"{atom.GetSymbol()} {pos[0]} {pos[1]} {pos[2]}\n"
-    prop_string = f"Docking score: {score} kcal/mol \n\n"
-    prop_string += f"pose structure: {xyz_string}\n"
-
-  except:
-    print(f"Molecule could not be docked!")
-    prop_string = ''
-  
-
-  current_props_string += prop_string
-  state["props_string"] = current_props_string
-  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
-    to determine which is the best tool to use to answer the QUERY TASK.
-      Input: the initial prompt from the user. should contain only one of more of the following:
-             smiles: the smiles string, task: the query task, path: the path to the file,
-             reference: the reference smiles
-             the value should be separated from the name by a ':' and each field should
-             be separated from the previous one by a ','.
-             All of these values are saved to the state
-      Output: the tool choice
-  '''
-  #for the user input
-  query_smiles = None
-  state["query_smiles"] = query_smiles
-  query_task = None
-  state["query_task"] = query_task
-  query_protein = None
-  state["query_protein"] = query_protein
-  #for the agent
-  state['similars_img'] = None
-  props_string = ""
-  state["props_string"] = props_string
-  state["loop_again"] = None
-
-  raw_input = state["messages"][-1].content
-  #print(raw_input)
-  parts = raw_input.split(',')
-  for part in parts:
-    if 'query_smiles' in part:
-      query_smiles = part.split(':')[1]
-      if query_smiles.lower() == 'none':
-        query_smiles = None
-      state["query_smiles"] = query_smiles
-    if 'query_task' in part:
-      query_task = part.split(':')[1]
-      state["query_task"] = query_task
-    if 'query_protein' in part:
-      query_protein = part.split(':')[1]
-      state["query_protein"] = query_protein
-
-  prompt = f'For the QUERY_TASK given below, determine if one or two of the tools descibed below \
-can complete the task. If so, reply with only the tool names followed by "#". If two tools \
-are required, reply with both tool names separated by a comma and followed by "#". \
-If the tools cannot complete the task, reply with "None #".\n \
-QUERY_TASK: {query_task}.\n \
-The information provided by the user is:\n \
-QUERY_SMILES: {query_smiles}.\n \
-QUERY_PROTEIN: {query_protein}.\n \
-Tools: \n \
-docking_tool: uses dockstring to dock the molecule into the protein, producing a pose structure and a docking score.\n \
-'
-
-  res = chat_model.invoke(prompt)
-  print(res)
-
-  tool_choices = str(res).split('<|assistant|>')[1].split('#')[0].strip()
-  tool_choices = tool_choices.split(',')
-  print(tool_choices)
-
-  if len(tool_choices) == 1:
-    tool1 = tool_choices[0].strip()
-    if tool1.lower() == 'none':
-      tool_choice = (None, None)
-    else:
-      tool_choice = (tool1, None)
-  elif len(tool_choices) == 2:
-    tool1 = tool_choices[0].strip()
-    tool2 = tool_choices[1].strip()
-    if tool1.lower() == 'none' and tool2.lower() == 'none':
-      tool_choice = (None, None)
-    elif tool1.lower() == 'none' and tool2.lower() != 'none':
-      tool_choice = (None, tool2)
-    elif tool2.lower() == 'none' and tool1.lower() != 'none':
-      tool_choice = (tool1, None)
-    else:
-      tool_choice = (tool1, tool2)
-  else:
-    tool_choice = (None, None)
-
-  state["tool_choice"] = tool_choice
-  state["which_tool"] = 0
-  print(f"First Node. The chosen tools are: {tool_choice}")
-
-  return state
-
-def retry_node(state: State) -> State:
-  '''
-    If the previous loop of the agent does not get enough informartion from the 
-    tools to answer the query, this node is called to retry the previous loop.
-      Input: the previous loop of the agent.
-      Output: the tool choice
-  '''
-  query_task = state["query_task"]
-  query_smiles = state["query_smiles"]
-  query_protein = state["query_protein"]
-
-  prompt = f'You were previously given the QUERY_TASK below, and asked to determine if one \
-or two of the tools descibed below could complete the task. The tool choices did not succeed. \
-Please re-examine the tool choices and determine if one or two of the tools descibed below \
-can complete the task. If so, reply with only the tool names followed by "#". If two tools \
-are required, reply with both tool names separated by a comma and followed by "#". \
-If the tools cannot complete the task, reply with "None #".\n \
-The information provided by the user is:\n \
-QUERY_SMILES: {query_smiles}.\n \
-QUERY_PROTEIN: {query_protein}.\n \
-The task is: \
-QUERY_TASK: {query_task}.\n \
-Tool options: \n \
-docking_tool: uses dockstring to dock the molecule into the protein, producing a pose structure and a docking score.\n \
-'
-
-  res = chat_model.invoke(prompt)
-
-  tool_choices = str(res).split('<|assistant|>')[1].split('#')[0].strip()
-  tool_choices = tool_choices.split(',')
-  if len(tool_choices) == 1:
-    tool1 = tool_choices[0].strip()
-    if tool1.lower() == 'none':
-      tool_choice = (None, None)
-    else:
-      tool_choice = (tool1, None)
-  elif len(tool_choices) == 2:
-    tool1 = tool_choices[0].strip()
-    tool2 = tool_choices[1].strip()
-    if tool1.lower() == 'none' and tool2.lower() == 'none':
-      tool_choice = (None, None)
-    elif tool1.lower() == 'none' and tool2.lower() != 'none':
-      tool_choice = (None, tool2)
-    elif tool2.lower() == 'none' and tool1.lower() != 'none':
-      tool_choice = (tool1, None)
-    else:
-      tool_choice = (tool1, tool2)
-  else:
-    tool_choice = (None, None)
-
-  state["tool_choice"] = tool_choice
-  state["which_tool"] = 0
-  print(f"The chosen tools are (Retry): {tool_choice}")
-
-  return state
-
-def loop_node(state: State) -> State:
-  '''
-    This node accepts the tool returns and decides if it needs to call another
-    tool or go on to the parser node.
-      Input: the tool returns.
-      Output: the next node to call.
-  '''
-  return state
-
-def parser_node(state: State) -> State:
-  '''
-    This is the third node in the agent. It receives the output from the tool,
-    puts it into a prompt as CONTEXT, and asks the LLM to answer the original
-    query.
-      Input: the output from the tool.
-      Output: the answer to the original query.
-  '''
-  props_string = state["props_string"]
-  query_task = state["query_task"]
-  tool_choice = state["tool_choice"]
-
-  if type(tool_choice) != tuple and tool_choice == None:
-    state["loop_again"] = "finish_gracefully"
-    return state
-  elif type(tool_choice) == tuple and (tool_choice[0] == None) and (tool_choice[1] == None):
-    state["loop_again"] = "finish_gracefully"
-    return state
-
-  prompt = f'Using the CONTEXT below, answer the original query, which \
-was to answer the QUERY_TASK. End your answer with a "#" \
-CONTEXT: {props_string}.\n \
-QUERY_TASK: {query_task}.\n '
-
-  res = chat_model.invoke(prompt)
-  trial_answer = str(res).split('<|assistant|>')[1]
-  print('parser 1 ', trial_answer)
-  state["messages"] = res
-
-  check_prompt = f'Determine if the TRIAL ANSWER below answers the original \
-QUERY TASK. If it does, respond with "PROCEED #" . If the TRIAL ANSWER did not \
-answer the QUERY TASK, respond with "LOOP #" \n \
-Only loop again if the TRIAL ANSWER did not answer the QUERY TASK. \
-TRIAL ANSWER: {trial_answer}.\n \
-QUERY_TASK: {query_task}.\n'
-
-  res = chat_model.invoke(check_prompt)
-  print('parser, loop again? ', res)
-
-  if str(res).split('<|assistant|>')[1].split('#')[0].strip().lower() == "loop":
-    state["loop_again"] = "loop_again"
-    return state
-  elif str(res).split('<|assistant|>')[1].split('#')[0].strip().lower() == "proceed":
-    state["loop_again"] = None
-    print('trying to break loop')
-  elif "proceed" in str(res).split('<|assistant|>')[1].lower():
-    state["loop_again"] = None
-    print('trying to break loop')
-
-  return state
-
-def reflect_node(state: State) -> State:
-  '''
-    This is the fourth node of the agent. It recieves the LLMs previous answer and
-    tries to improve it.
-      Input: the LLMs last answer.
-      Output: the improved answer.
-  '''
-  previous_answer = state["messages"][-1].content
-  props_string = state["props_string"]
-
-  prompt = f'Look at the PREVIOUS ANSWER below which you provided and the \
-TOOL RESULTS. Write an improved answer based on the PREVIOUS ANSWER and the \
-TOOL RESULTS by adding additional clarifying and enriching information. End \
-your new answer with a "#" \
-PREVIOUS ANSWER: {previous_answer}.\n \
-TOOL RESULTS: {props_string}. '
-
-  res = chat_model.invoke(prompt)
-
-  return {"messages": res}
-
-def graceful_exit_node(state: State) -> State:
-  '''
-    Called when the Agent cannot assign any tools for the task
-  '''
-  props_string = state["props_string"]
-  prompt = f'Summarize the information in the CONTEXT, including any useful chemical information. Start your answer with: \
-Here is what I found: \n \
-CONTEXT: {props_string}'
-
-  res = chat_model.invoke(prompt)
-
-  return {"messages": res}
-    
-
-def get_chemtool(state):
-  '''
-  '''
-  which_tool = state["which_tool"]
-  tool_choice = state["tool_choice"]
-  print('in get_chemtool ',tool_choice)
-  if tool_choice == None:
-    return None
-  if which_tool == 0 or which_tool == 1:
-    current_tool = tool_choice[which_tool]
-  elif which_tool > 1:
-    current_tool = None
-
-  return current_tool
-
-def loop_or_not(state):
-  '''
-  '''
-  print(f"(line 417) Loop? {state['loop_again']}")
-  if state["loop_again"] == "loop_again":
-    return True
-  elif state["loop_again"] == "finish_gracefully":
-    return 'lets_get_outta_here'
-  else:
-    return False
-
-builder = StateGraph(State)
-#for the agent
-builder.add_node("first_node", first_node)
-builder.add_node("retry_node", retry_node)
-builder.add_node("loop_node", loop_node)
-builder.add_node("parser_node", parser_node)
-builder.add_node("reflect_node", reflect_node)
-builder.add_node("graceful_exit_node", graceful_exit_node)
-#for the tools
-builder.add_node("docking_node", docking_node)
-
-builder.add_edge(START, "first_node")
-builder.add_conditional_edges("first_node", get_chemtool, {
-    "docking_tool": "docking_node",
-    None: "parser_node"})
-
-builder.add_conditional_edges("retry_node", get_chemtool, {
-    "docking_tool": "docking_node",
-    None: "parser_node"})
-
-builder.add_edge("docking_node", "loop_node")
-
-builder.add_conditional_edges("loop_node", get_chemtool, {
-    "docking_tool" : "docking_node",
-    "loop_again": "first_node",
-    None: "parser_node"})
-
-builder.add_conditional_edges("parser_node", loop_or_not, {
-    True: "retry_node",
-    'lets_get_outta_here': "graceful_exit_node",
-    False: "reflect_node"})
-
-builder.add_edge("reflect_node", END)
-builder.add_edge("graceful_exit_node", END)
-
-graph = builder.compile()
-
-@spaces.GPU
-def DockAgent(task, smiles, protein): # add variables as needed
-
-  #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_protein: {protein}') # add variables as needed
-    ]
-  }
-  #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:
-      try:
-          reply = c[str(m[0])]['messages']
-          if 'assistant' in str(reply):
-            reply = str(reply).split("<|assistant|>")[-1].split('#')[0].strip()
-            replies.append(reply)
-      except:
-          reply = str(c).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
-
-dudes = ['IGF1R', 'JAK2', 'KIT', 'LCK', 'MAPK14', 'MAPKAPK2', 'MET', 'PTK2', 'PTPN1', 'SRC', 'ABL1', 'AKT1', 'AKT2', 'CDK2', 'CSF1R', 'EGFR', 'KDR', 'MAPK1', 'FGFR1', 'ROCK1', 'MAP2K1', 'PLK1',
-         'HSD11B1', 'PARP1', 'PDE5A', 'PTGS2', 'ACHE', 'MAOB', 'CA2', 'GBA', 'HMGCR', 'NOS1', 'REN', 'DHFR', 'ESR1', 'ESR2', 'NR3C1', 'PGR', 'PPARA', 'PPARD', 'PPARG',
-        'AR','THRB','ADAM17', 'F10', 'F2', 'BACE1', 'CASP3', 'MMP13', 'DPP4', 'ADRB1', 'ADRB2', 'DRD2', 'DRD3','ADORA2A','CYP2C9', 'CYP3A4', 'HSP90AA1']
-
-with gr.Blocks(fill_height=True) as forest:
-  gr.Markdown('''
-              # Docking Agent
-              - uses dockstring to dock a molecule into a protein using only a SMILES string and a protein name
-              - produces a pose structure and a docking score
-              ''')
-
-  with gr.Accordion("ProteinOptions", open=False):
-    gr.Markdown('''
-              # Protein Options
-              ## Kinase
-              ### Highest quality
-              - IGF1R, JAK2, KIT, LCK, MAPK14, MAPKAPK2, MET, PTK2, PTPN1, SRC
-              ### Medium quality
-              - ABL1, AKT1, AKT2, CDK2, CSF1R, EGFR, KDR, MAPK1, FGFR1, ROCK1
-              ### Lower quality
-              - MAP2K1, PLK1
-              ## Enzyme
-              ### Highest quality
-              - HSD11B1, PARP1, PDE5A, PTGS2
-              ### Medium quality
-              - ACHE, MAOB
-              ### Lower quality
-              - CA2, GBA, HMGCR, NOS1, REN, DHFR
-              ## Nuclear Receptor
-              ### Highest quality
-              - ESR1, ESR2, NR3C1, PGR, PPARA, PPARD, PPARG
-              ### Medium quality
-              - AR
-              ### Lower quality
-              - THRB
-              ## Protease
-              ### Higher quality
-              - ADAM17, F10, F2
-              ### Medium quality
-              - BACE1, CASP3, MMP13
-              ### Lower quality
-              - DPP4
-              ## GPCR
-              ### Medium quality
-              - ADRB1, ADRB2, DRD2, DRD3
-              ### Lower quality
-              - ADORA2A
-              ## Cytochrome
-              ### Medium quality
-              - CYP2C9, CYP3A4
-              ## Chaperone
-              ### Lower quality
-              - HSP90AA1
-              ''')
-  with gr.Row():
-    with gr.Column():
-      smiles = gr.Textbox(label="Molecule SMILES of interest (optional): ", placeholder='none')
-      protein = gr.Dropdown(dudes, label="Protein name (see options): ")
-      task = gr.Textbox(label="Task for Agent: ")
-      # add variables as needed
-      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(DockAgent, inputs = [task, smiles, protein], outputs = [props, pic])
-      task.submit(DockAgent, inputs = [task, smiles, protein], outputs = [props, pic])
-
+import torch
+from typing import Annotated, TypedDict, Literal
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.tools import tool
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
+from langchain_core.messages import SystemMessage, trim_messages, AIMessage, HumanMessage, ToolCall
+
+from langchain_huggingface.llms import HuggingFacePipeline
+from langchain_huggingface import ChatHuggingFace
+from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
+from langchain_core.runnables import chain
+from uuid import uuid4
+import re
+import matplotlib.pyplot as plt
+import spaces
+
+from dockstring import load_target
+from rdkit import Chem
+from rdkit.Chem import AllChem, QED
+from rdkit.Chem import Draw
+from rdkit.Chem.Draw import MolsToGridImage
+import os, re
+import gradio as gr
+from PIL import Image
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+hf = HuggingFacePipeline.from_model_id(
+    model_id= "microsoft/Phi-4-mini-instruct",
+    task="text-generation",
+    pipeline_kwargs = {"max_new_tokens": 500, "temperature": 0.4})
+
+chat_model = ChatHuggingFace(llm=hf)
+
+cpuCount = os.cpu_count()
+print(f"Number of CPUs: {cpuCount}")
+
+class State(TypedDict):
+  '''
+    The state of the agent.
+  '''
+  messages: Annotated[list, add_messages]
+  #for the agent
+  tool_choice: tuple
+  which_tool: int
+  props_string: str
+  similars_img: str
+  loop_again: str
+
+  #for the user input
+  query_smiles: str
+  query_task: str
+  query_protein: str
+
+def docking_node(state: State) -> State:
+  '''
+    Docking tool: uses dockstring to dock the molecule into the protein
+  '''
+  print("docking tool")
+  print('===================================================')
+  current_props_string = state["props_string"]
+  query_protein = state["query_protein"].strip()
+  query_smiles = state["query_smiles"].strip()
+
+  print(f'query_protein: {query_protein}')
+  print(f'query_smiles: {query_smiles}')
+
+  try:
+    target = load_target(query_protein)
+    print("===============================================")
+    print(f"Docking molecule with {cpuCount} cpu cores.")
+    score, aux = target.dock(query_smiles, num_cpus = cpuCount)
+    mol = aux['ligand']
+    print(f"Docking score: {score}")
+    print("===============================================")
+    atoms_list = ""
+    template = mol
+    molH = Chem.AddHs(mol)
+    AllChem.ConstrainedEmbed(molH,template, useTethers=True)
+    xyz_string = f"{molH.GetNumAtoms()}\n\n"
+    for atom in molH.GetAtoms():
+      atoms_list += atom.GetSymbol()
+      pos = molH.GetConformer().GetAtomPosition(atom.GetIdx())
+      xyz_string += f"{atom.GetSymbol()} {pos[0]} {pos[1]} {pos[2]}\n"
+    prop_string = f"Docking score: {score} kcal/mol \n\n"
+    prop_string += f"pose structure: {xyz_string}\n"
+
+  except:
+    print(f"Molecule could not be docked!")
+    prop_string = ''
+  
+
+  current_props_string += prop_string
+  state["props_string"] = current_props_string
+  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
+    to determine which is the best tool to use to answer the QUERY TASK.
+      Input: the initial prompt from the user. should contain only one of more of the following:
+             smiles: the smiles string, task: the query task, path: the path to the file,
+             reference: the reference smiles
+             the value should be separated from the name by a ':' and each field should
+             be separated from the previous one by a ','.
+             All of these values are saved to the state
+      Output: the tool choice
+  '''
+  #for the user input
+  query_smiles = None
+  state["query_smiles"] = query_smiles
+  query_task = None
+  state["query_task"] = query_task
+  query_protein = None
+  state["query_protein"] = query_protein
+  #for the agent
+  state['similars_img'] = None
+  props_string = ""
+  state["props_string"] = props_string
+  state["loop_again"] = None
+
+  raw_input = state["messages"][-1].content
+  #print(raw_input)
+  parts = raw_input.split(',')
+  for part in parts:
+    if 'query_smiles' in part:
+      query_smiles = part.split(':')[1]
+      if query_smiles.lower() == 'none':
+        query_smiles = None
+      state["query_smiles"] = query_smiles
+    if 'query_task' in part:
+      query_task = part.split(':')[1]
+      state["query_task"] = query_task
+    if 'query_protein' in part:
+      query_protein = part.split(':')[1]
+      state["query_protein"] = query_protein
+
+  prompt = f'For the QUERY_TASK given below, determine if one or two of the tools descibed below \
+can complete the task. If so, reply with only the tool names followed by "#". If two tools \
+are required, reply with both tool names separated by a comma and followed by "#". \
+If the tools cannot complete the task, reply with "None #".\n \
+QUERY_TASK: {query_task}.\n \
+The information provided by the user is:\n \
+QUERY_SMILES: {query_smiles}.\n \
+QUERY_PROTEIN: {query_protein}.\n \
+Tools: \n \
+docking_tool: uses dockstring to dock the molecule into the protein, producing a pose structure and a docking score.\n \
+'
+
+  res = chat_model.invoke(prompt)
+  print(res)
+
+  tool_choices = str(res).split('<|assistant|>')[1].split('#')[0].strip()
+  tool_choices = tool_choices.split(',')
+  print(tool_choices)
+
+  if len(tool_choices) == 1:
+    tool1 = tool_choices[0].strip()
+    if tool1.lower() == 'none':
+      tool_choice = (None, None)
+    else:
+      tool_choice = (tool1, None)
+  elif len(tool_choices) == 2:
+    tool1 = tool_choices[0].strip()
+    tool2 = tool_choices[1].strip()
+    if tool1.lower() == 'none' and tool2.lower() == 'none':
+      tool_choice = (None, None)
+    elif tool1.lower() == 'none' and tool2.lower() != 'none':
+      tool_choice = (None, tool2)
+    elif tool2.lower() == 'none' and tool1.lower() != 'none':
+      tool_choice = (tool1, None)
+    else:
+      tool_choice = (tool1, tool2)
+  else:
+    tool_choice = (None, None)
+
+  state["tool_choice"] = tool_choice
+  state["which_tool"] = 0
+  print(f"First Node. The chosen tools are: {tool_choice}")
+
+  return state
+
+def retry_node(state: State) -> State:
+  '''
+    If the previous loop of the agent does not get enough informartion from the 
+    tools to answer the query, this node is called to retry the previous loop.
+      Input: the previous loop of the agent.
+      Output: the tool choice
+  '''
+  query_task = state["query_task"]
+  query_smiles = state["query_smiles"]
+  query_protein = state["query_protein"]
+
+  prompt = f'You were previously given the QUERY_TASK below, and asked to determine if one \
+or two of the tools descibed below could complete the task. The tool choices did not succeed. \
+Please re-examine the tool choices and determine if one or two of the tools descibed below \
+can complete the task. If so, reply with only the tool names followed by "#". If two tools \
+are required, reply with both tool names separated by a comma and followed by "#". \
+If the tools cannot complete the task, reply with "None #".\n \
+The information provided by the user is:\n \
+QUERY_SMILES: {query_smiles}.\n \
+QUERY_PROTEIN: {query_protein}.\n \
+The task is: \
+QUERY_TASK: {query_task}.\n \
+Tool options: \n \
+docking_tool: uses dockstring to dock the molecule into the protein, producing a pose structure and a docking score.\n \
+'
+
+  res = chat_model.invoke(prompt)
+
+  tool_choices = str(res).split('<|assistant|>')[1].split('#')[0].strip()
+  tool_choices = tool_choices.split(',')
+  if len(tool_choices) == 1:
+    tool1 = tool_choices[0].strip()
+    if tool1.lower() == 'none':
+      tool_choice = (None, None)
+    else:
+      tool_choice = (tool1, None)
+  elif len(tool_choices) == 2:
+    tool1 = tool_choices[0].strip()
+    tool2 = tool_choices[1].strip()
+    if tool1.lower() == 'none' and tool2.lower() == 'none':
+      tool_choice = (None, None)
+    elif tool1.lower() == 'none' and tool2.lower() != 'none':
+      tool_choice = (None, tool2)
+    elif tool2.lower() == 'none' and tool1.lower() != 'none':
+      tool_choice = (tool1, None)
+    else:
+      tool_choice = (tool1, tool2)
+  else:
+    tool_choice = (None, None)
+
+  state["tool_choice"] = tool_choice
+  state["which_tool"] = 0
+  print(f"The chosen tools are (Retry): {tool_choice}")
+
+  return state
+
+def loop_node(state: State) -> State:
+  '''
+    This node accepts the tool returns and decides if it needs to call another
+    tool or go on to the parser node.
+      Input: the tool returns.
+      Output: the next node to call.
+  '''
+  return state
+
+def parser_node(state: State) -> State:
+  '''
+    This is the third node in the agent. It receives the output from the tool,
+    puts it into a prompt as CONTEXT, and asks the LLM to answer the original
+    query.
+      Input: the output from the tool.
+      Output: the answer to the original query.
+  '''
+  props_string = state["props_string"]
+  query_task = state["query_task"]
+  tool_choice = state["tool_choice"]
+
+  if type(tool_choice) != tuple and tool_choice == None:
+    state["loop_again"] = "finish_gracefully"
+    return state
+  elif type(tool_choice) == tuple and (tool_choice[0] == None) and (tool_choice[1] == None):
+    state["loop_again"] = "finish_gracefully"
+    return state
+
+  prompt = f'Using the CONTEXT below, answer the original query, which \
+was to answer the QUERY_TASK. End your answer with a "#" \
+CONTEXT: {props_string}.\n \
+QUERY_TASK: {query_task}.\n '
+
+  res = chat_model.invoke(prompt)
+  trial_answer = str(res).split('<|assistant|>')[1]
+  print('parser 1 ', trial_answer)
+  state["messages"] = res
+
+  check_prompt = f'Determine if the TRIAL ANSWER below answers the original \
+QUERY TASK. If it does, respond with "PROCEED #" . If the TRIAL ANSWER did not \
+answer the QUERY TASK, respond with "LOOP #" \n \
+Only loop again if the TRIAL ANSWER did not answer the QUERY TASK. \
+TRIAL ANSWER: {trial_answer}.\n \
+QUERY_TASK: {query_task}.\n'
+
+  res = chat_model.invoke(check_prompt)
+  print('parser, loop again? ', res)
+
+  if str(res).split('<|assistant|>')[1].split('#')[0].strip().lower() == "loop":
+    state["loop_again"] = "loop_again"
+    return state
+  elif str(res).split('<|assistant|>')[1].split('#')[0].strip().lower() == "proceed":
+    state["loop_again"] = None
+    print('trying to break loop')
+  elif "proceed" in str(res).split('<|assistant|>')[1].lower():
+    state["loop_again"] = None
+    print('trying to break loop')
+
+  return state
+
+def reflect_node(state: State) -> State:
+  '''
+    This is the fourth node of the agent. It recieves the LLMs previous answer and
+    tries to improve it.
+      Input: the LLMs last answer.
+      Output: the improved answer.
+  '''
+  previous_answer = state["messages"][-1].content
+  props_string = state["props_string"]
+
+  prompt = f'Look at the PREVIOUS ANSWER below which you provided and the \
+TOOL RESULTS. Write an improved answer based on the PREVIOUS ANSWER and the \
+TOOL RESULTS by adding additional clarifying and enriching information. End \
+your new answer with a "#" \
+PREVIOUS ANSWER: {previous_answer}.\n \
+TOOL RESULTS: {props_string}. '
+
+  res = chat_model.invoke(prompt)
+
+  return {"messages": res}
+
+def graceful_exit_node(state: State) -> State:
+  '''
+    Called when the Agent cannot assign any tools for the task
+  '''
+  props_string = state["props_string"]
+  prompt = f'Summarize the information in the CONTEXT, including any useful chemical information. Start your answer with: \
+Here is what I found: \n \
+CONTEXT: {props_string}'
+
+  res = chat_model.invoke(prompt)
+
+  return {"messages": res}
+    
+
+def get_chemtool(state):
+  '''
+  '''
+  which_tool = state["which_tool"]
+  tool_choice = state["tool_choice"]
+  print('in get_chemtool ',tool_choice)
+  if tool_choice == None:
+    return None
+  if which_tool == 0 or which_tool == 1:
+    current_tool = tool_choice[which_tool]
+  elif which_tool > 1:
+    current_tool = None
+
+  return current_tool
+
+def loop_or_not(state):
+  '''
+  '''
+  print(f"(line 417) Loop? {state['loop_again']}")
+  if state["loop_again"] == "loop_again":
+    return True
+  elif state["loop_again"] == "finish_gracefully":
+    return 'lets_get_outta_here'
+  else:
+    return False
+
+builder = StateGraph(State)
+#for the agent
+builder.add_node("first_node", first_node)
+builder.add_node("retry_node", retry_node)
+builder.add_node("loop_node", loop_node)
+builder.add_node("parser_node", parser_node)
+builder.add_node("reflect_node", reflect_node)
+builder.add_node("graceful_exit_node", graceful_exit_node)
+#for the tools
+builder.add_node("docking_node", docking_node)
+
+builder.add_edge(START, "first_node")
+builder.add_conditional_edges("first_node", get_chemtool, {
+    "docking_tool": "docking_node",
+    None: "parser_node"})
+
+builder.add_conditional_edges("retry_node", get_chemtool, {
+    "docking_tool": "docking_node",
+    None: "parser_node"})
+
+builder.add_edge("docking_node", "loop_node")
+
+builder.add_conditional_edges("loop_node", get_chemtool, {
+    "docking_tool" : "docking_node",
+    "loop_again": "first_node",
+    None: "parser_node"})
+
+builder.add_conditional_edges("parser_node", loop_or_not, {
+    True: "retry_node",
+    'lets_get_outta_here': "graceful_exit_node",
+    False: "reflect_node"})
+
+builder.add_edge("reflect_node", END)
+builder.add_edge("graceful_exit_node", END)
+
+graph = builder.compile()
+
+@spaces.GPU
+def DockAgent(task, smiles, protein): # add variables as needed
+
+  #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_protein: {protein}') # add variables as needed
+    ]
+  }
+  #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:
+      try:
+          reply = c[str(m[0])]['messages']
+          if 'assistant' in str(reply):
+            reply = str(reply).split("<|assistant|>")[-1].split('#')[0].strip()
+            replies.append(reply)
+      except:
+          reply = str(c).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
+
+dudes = ['IGF1R', 'JAK2', 'KIT', 'LCK', 'MAPK14', 'MAPKAPK2', 'MET', 'PTK2', 'PTPN1', 'SRC', 'ABL1', 'AKT1', 'AKT2', 'CDK2', 'CSF1R', 'EGFR', 'KDR', 'MAPK1', 'FGFR1', 'ROCK1', 'MAP2K1', 'PLK1',
+         'HSD11B1', 'PARP1', 'PDE5A', 'PTGS2', 'ACHE', 'MAOB', 'CA2', 'GBA', 'HMGCR', 'NOS1', 'REN', 'DHFR', 'ESR1', 'ESR2', 'NR3C1', 'PGR', 'PPARA', 'PPARD', 'PPARG',
+        'AR','THRB','ADAM17', 'F10', 'F2', 'BACE1', 'CASP3', 'MMP13', 'DPP4', 'ADRB1', 'ADRB2', 'DRD2', 'DRD3','ADORA2A','CYP2C9', 'CYP3A4', 'HSP90AA1']
+
+with gr.Blocks(fill_height=True) as forest:
+  gr.Markdown('''
+              # Docking Agent
+              - uses dockstring to dock a molecule into a protein using only a SMILES string and a protein name
+              - produces a pose structure and a docking score
+              ''')
+
+  with gr.Accordion("ProteinOptions", open=False):
+    gr.Markdown('''
+              # Protein Options
+              ## Kinase
+              ### Highest quality
+              - IGF1R, JAK2, KIT, LCK, MAPK14, MAPKAPK2, MET, PTK2, PTPN1, SRC
+              ### Medium quality
+              - ABL1, AKT1, AKT2, CDK2, CSF1R, EGFR, KDR, MAPK1, FGFR1, ROCK1
+              ### Lower quality
+              - MAP2K1, PLK1
+              ## Enzyme
+              ### Highest quality
+              - HSD11B1, PARP1, PDE5A, PTGS2
+              ### Medium quality
+              - ACHE, MAOB
+              ### Lower quality
+              - CA2, GBA, HMGCR, NOS1, REN, DHFR
+              ## Nuclear Receptor
+              ### Highest quality
+              - ESR1, ESR2, NR3C1, PGR, PPARA, PPARD, PPARG
+              ### Medium quality
+              - AR
+              ### Lower quality
+              - THRB
+              ## Protease
+              ### Higher quality
+              - ADAM17, F10, F2
+              ### Medium quality
+              - BACE1, CASP3, MMP13
+              ### Lower quality
+              - DPP4
+              ## GPCR
+              ### Medium quality
+              - ADRB1, ADRB2, DRD2, DRD3
+              ### Lower quality
+              - ADORA2A
+              ## Cytochrome
+              ### Medium quality
+              - CYP2C9, CYP3A4
+              ## Chaperone
+              ### Lower quality
+              - HSP90AA1
+              ''')
+  with gr.Row():
+    with gr.Column():
+      smiles = gr.Textbox(label="Molecule SMILES of interest (optional): ", placeholder='none')
+      protein = gr.Dropdown(dudes, label="Protein name (see options): ")
+      task = gr.Textbox(label="Task for Agent: ")
+      # add variables as needed
+      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(DockAgent, inputs = [task, smiles, protein], outputs = [props, pic])
+      task.submit(DockAgent, inputs = [task, smiles, protein], outputs = [props, pic])
+
 forest.launch(debug=False, mcp_server=True)