graph test

app.py

@@ -74,15 +74,17 @@ class BasicAgent:
         self.model = ChatOpenAI(model="gpt-4o", temperature=0)
         self.graph = StateGraph(AnswerState)
         self.graph.add_node("log_question", self.log_question)
-        self.graph.add_node("try_answer", self.try_answer)
+        self.graph.add_node("invoke_model", self.invoke_model)
+        self.graph.add_node("parse_response", self.parse_response)
         self.graph.add_node("final_answer", self.final_answer)
 
         self.graph.add_edge(START, "log_question")
+        self.graph.add_edge("log_question", "invoke_model")
 
         # Add conditional edges
         self.graph.add_conditional_edges(
-            "log_question",
-            self.try_answer,
+            "invoke_model",
+            self.parse_response,
             {
                 "FINAL_ANSWER": "final_answer",
                 "GOOGLE_SEARCH": "final_answer",
@@ -124,36 +126,24 @@ class BasicAgent:
             "is_final_answer": True,
         }
 
-    def try_answer(self, state: AnswerState) -> str:
-        if state["attempt"] > 3:
+    def parse_response(self, state: AnswerState) -> str:
+        print("[parse_response] parsing last chat response")
+
+        if state["attempt"] > 0:
             state["answer"] = "Exceeded max number of attempts."
             return "FINAL_ANSWER"
-        else:
-            state["attempt"] += 1
-
-        print("[try_answer] Agent trying to answer")
-
-        response = self.model.invoke(state["messages"])
-        print(f"Agent response: {response.content}")
-        state["messages"].append(
-            ChatMessage(
-                role="assistant",
-                content=response.content,
-            )
-        )
 
-        if "FINAL ANSWER:" in response.content:
-            answer = response.text().split("FINAL ANSWER:")[-1].strip()
+        content = self.content2text(state["messages"][-1].content)
+        if "FINAL ANSWER:" in content:
+            answer = content.split("FINAL ANSWER:")[-1].strip()
             print(f"Agent final answer: {answer}")
             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(")")
-            )
+        if "TOOL: GoogleSearchAgent" in content:
+            request_full = content.split("TOOL: GoogleSearchAgent(")[-1].strip(")")
             print(f"Tool invocation request: {request_full}")
             request = request_full.split("request=")[-1].strip("'")
             print(f"Search request: {request}")
@@ -163,6 +153,52 @@ class BasicAgent:
 
         return "FINAL_ANSWER"
 
+        print("[parse_response] Agent parsing response")
+        if state["is_final_answer"]:
+            return "FINAL_ANSWER"
+
+        if state["search_request"]:
+            return "GOOGLE_SEARCH"
+
+        return "FINAL_ANSWER"
+
+    def content2text(self, content) -> str:
+        """Get the text content of the message.
+
+        Returns:
+            The text content of the message.
+        """
+        if isinstance(content, str):
+            return content
+
+        # must be a list
+        blocks = [
+            block
+            for block in content
+            if isinstance(block, str)
+            or (block.get("type") == "text" and isinstance(block.get("text"), str))
+        ]
+        return "".join(
+            block if isinstance(block, str) else block["text"] for block in blocks
+        )
+
+    def invoke_model(self, state: AnswerState) -> Dict[str, Any]:
+        print("[invoke_model] Agent trying to answer")
+
+        response = self.model.invoke(state["messages"])
+        print(f"Agent response: {response.content}")
+        state["messages"].append(
+            ChatMessage(
+                role="assistant",
+                content=response.content,
+            )
+        )
+
+        return {
+            "messages": state["messages"],
+            "attempt": state.get("attempt", 0) + 1,
+        }
+
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """