graph test

app.py

@@ -57,10 +57,12 @@ Question: {question}
 
 class AnswerState(TypedDict):
     question: str
+    attempt: int
     answer: Optional[str]
     messages: List[Dict[str, Any]]
     is_final_answer: bool
     search_request: Optional[str]
+    search_response: Optional[str]
 
 
 class BasicAgent:
@@ -87,14 +89,17 @@ class BasicAgent:
         self.compiled_graph = self.graph.compile()
 
     def __call__(self, question: str) -> str:
+
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         res = self.compiled_graph.invoke(
             {
                 "question": question,
+                "attempt": 0,
                 "answer": None,
                 "messages": [],
                 "is_final_answer": False,
                 "search_request": None,
+                "search_response": None,
             }
         )
         return res.get("answer", "N/A")
@@ -112,6 +117,12 @@ class BasicAgent:
         }
 
     def try_answer(self, state: AnswerState) -> str:
+        if state["attempt"] > 3:
+            state["answer"] = "Exceeded max number of attempts."
+            return "FINAL_ANSWER"
+        else:
+            state["attempt"] += 1
+
         print("[try_answer] Agent trying to answer")
         prompt = PROMPT.format(question=state["question"])
         state["messages"] = state.get("messages", []) + [
@@ -137,6 +148,16 @@ class BasicAgent:
             print(f"Agent returning answer: {state['answer']}")
             state["answer"] = answer
 
+            return "FINAL_ANSWER"
+        if "TOOL: GoogleSearchAgent" in response.content:
+            request_full = (
+                response.text().split("TOOL: GoogleSearchAgent(")[-1].strip(")")
+            )
+            print(f"Tool invocation request: {request_full}")
+            request = request_full.split("request=")[-1].strip("'")
+            print(f"Search request: {request}")
+            state["search_request"] = request
+
             return "FINAL_ANSWER"
 
         return "FINAL_ANSWER"