Update app.py

app.py

@@ -69,6 +69,7 @@ class State(TypedDict):
   props_string: str
   loop_again: str
   which_pdbs: int
+  recursion_count: int
   #(Literal["lipinski_tool", "substitution_tool", "pharm_feature_tool"],
   #                   Literal["lipinski_tool", "substitution_tool", "pharm_feature_tool"])
 
@@ -626,6 +627,7 @@ def first_node(state: State) -> State:
   state["props_string"] = props_string
   state["loop_again"] = None
   state['which_pdbs'] = 0
+  state['recursion_count'] = 0
 
   raw_input = state["messages"][-1].content
   parts = raw_input.split(',')
@@ -722,7 +724,7 @@ should be called before this tool.
   else:
     tool_choice = (None, None)
   if (len(tool_choice) == 2) and (tool_choice[1] == tool_choice[0]):
-    tool_choice = (tool1, None)
+    tool_choice = (tool_choice[0], None)
 
   state["tool_choice"] = tool_choice
   state["which_tool"] = 0
@@ -853,7 +855,10 @@ def parser_node(state: State) -> State:
   elif type(tool_choice) == tuple and (tool_choice[0] == None) and (tool_choice[1] == None):
     state["loop_again"] = "finish_gracefully"
     return state
-
+  elif state['recursion_count'] > 20:
+    state["loop_again"] = "finish_gracefully"
+    return state
+      
   prompt = f'Using the CONTEXT below, answer the original query, which \
 was to answer the QUERY_TASK. Remember that novel molecules generated in the CONTEXT \
 were made using a fine-tuned GPT. End your answer with a "#" \
@@ -877,6 +882,7 @@ QUERY_TASK: {query_task}.\n'
 
   if str(res).split('<|assistant|>')[1].split('#')[0].strip().lower() == "loop":
     state["loop_again"] = "loop_again"
+    state['recursion_count'] += 1
     return state
   elif str(res).split('<|assistant|>')[1].split('#')[0].strip().lower() == "proceed":
     state["loop_again"] = None
@@ -885,6 +891,8 @@ QUERY_TASK: {query_task}.\n'
     state["loop_again"] = None
     print('trying to break loop')
 
+  state['recursion_count'] += 1
+
   return state
 
 def reflect_node(state: State) -> State: