Update app.py

app.py

@@ -29,7 +29,6 @@ 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"
 
@@ -40,6 +39,252 @@ hf = HuggingFacePipeline.from_model_id(
 
 chat_model = ChatHuggingFace(llm=hf)
 
+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
+  '''
+  query_smiles = None
+  state["query_smiles"] = query_smiles
+  query_task = None
+  state["query_task"] = query_task
+  query_name = None
+  state["query_name"] = query_name
+  query_reference = None
+  state["query_reference"] = query_reference
+  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_name' in part:
+      query_name = part.split(':')[1]
+      if query_name.lower() == 'none':
+        query_name = None
+      state["query_name"] = query_name
+    if 'query_reference' in part:
+      query_reference = part.split(':')[1]
+      state["query_reference"] = query_reference
+
+  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_NAME: {query_name}.\n \
+Tools: \n \
+smiles_tool: queries Pubchem for the smiles string of the molecule based on the name.\n \
+name_tool: queries Pubchem for the NAME of the molecule based on the smiles string.\n \
+similars_tool: queries Pubchem for similar molecules based on the smiles string or name and returns 20 results. \
+returns the names, SMILES strings, molecular weights and logP values for the similar molecules. \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:
+    if tool_choices[0].strip().lower() == 'none':
+      tool_choice = (None, None)
+    else:
+      tool_choice = (tool_choices[0].strip().lower(), None)
+  elif len(tool_choices) == 2:
+    if tool_choices[0].strip().lower() == 'none':
+      tool_choice = (None, tool_choices[1].strip().lower())
+    elif tool_choices[1].strip().lower() == 'none':
+      tool_choice = (tool_choices[0].strip().lower(), None)
+    else:
+      tool_choice = (tool_choices[0].strip().lower(), tool_choices[1].strip().lower())
+  else:
+    tool_choice = None
+
+  state["tool_choice"] = tool_choice
+  state["which_tool"] = 0
+  print(f"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_name = state["query_name"]
+
+  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. TYou 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 \
+QUERY_TASK: {query_task}.\n \
+The information provided by the user is:\n \
+QUERY_SMILES: {query_smiles}.\n \
+QUERY_NAME: {query_name}.\n \
+Tools: \n \
+smiles_tool: queries Pubchem for the smiles string of the molecule based on the name as input.\n \
+name_tool: queries Pubchem for the NAME (IUPAC) of the molecule based on the smiles string as input. \
+Also returns a short list of common names for the molecule. \n \
+similars_tool: queries Pubchem for similar molecules based on the smiles string or name as input and returns 20 results. \
+Returns the names, SMILES strings, molecular weights and logP values for the similar molecules. \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:
+    if tool_choices[0].strip().lower() == 'none':
+      tool_choice = (None, None)
+    else:
+      tool_choice = (tool_choices[0].strip().lower(), None)
+  elif len(tool_choices) == 2:
+    if tool_choices[0].strip().lower() == 'none':
+      tool_choice = (None, tool_choices[1].strip().lower())
+    elif tool_choices[1].strip().lower() == 'none':
+      tool_choice = (tool_choices[0].strip().lower(), None)
+    else:
+      tool_choice = (tool_choices[0].strip().lower(), tool_choices[1].strip().lower())
+  elif 'none' in tool_choices[0].strip().lower():
+    tool_choice = None
+  else:
+    tool_choice = 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"]
+
+  check_prompt = f'Determine if there is enough CONTEXT below to answer the original \
+QUERY TASK. If there is, respond with "PROCEED #" . If there is not enough information \
+to answer the QUERY TASK, respond with "LOOP #" \n \
+CONTEXT: {props_string}.\n \
+QUERY_TASK: {query_task}.\n'
+
+  res = chat_model.invoke(check_prompt)
+  # print('*'*50)
+  # print(res)
+  # print('*'*50)
+  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
+
+    prompt = f'Using the CONTEXT below, answer the original query, which \
+was to answer the QUERY_TASK. End your answer with a "#" \
+QUERY_TASK: {query_task}.\n \
+CONTEXT: {props_string}.\n '
+
+    res = chat_model.invoke(prompt)
+    return {"messages": res}
+
+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 get_chemtool(state):
+  '''
+  '''
+  which_tool = state["which_tool"]
+  tool_choice = state["tool_choice"]
+  #print(tool_choice)
+  if tool_choice == None:
+    return None
+  if which_tool == 0 or which_tool == 1:
+    current_tool = tool_choice[which_tool]
+    if current_tool == "smiles_tool" and ("query_name" not in state.keys()):
+      current_tool = "name_tool"
+      print("Switching from smiles tool to name tool")
+    elif current_tool == "name_tool" and ("query_smiles" not in state.keys()):
+      current_tool = "smiles_tool"
+      print("Switching from name tool to smiles tool")
+
+  elif which_tool > 1:
+    current_tool = None
+
+  return current_tool
+
+def loop_or_not(state):
+  '''
+  '''
+  print(f"Loop? {state['loop_again']}")
+  if state["loop_again"] == "loop_again":
+    return True
+  else:
+    return False
+
+def pretty_print(answer):
+  final = str(answer['messages'][-1]).split('<|assistant|>')[-1].split('#')[0].strip("n").strip('\\').strip('n').strip('\\')
+  for i in range(0,len(final),100):
+    print(final[i:i+100])
+
+def print_short(answer):
+  for i in range(0,len(answer),100):
+    print(answer[i:i+100])
+
 builder = StateGraph(State)
 builder.add_node("first_node", first_node)
 builder.add_node("retry_node", retry_node)