Update app.py

app.py

@@ -65,8 +65,14 @@ class AdvancedAgent:
         builder = StateGraph(MessagesState)
         builder.add_node("assistant", assistant)
         builder.add_node("tools", tools_node)
+        builder.add_edge("tools", "assistant")
         builder.add_edge("assistant", "tools")
         builder.set_entry_point("assistant")
+        builder.add_conditional_edges(
+            "tools",
+            lambda x: "assistant",
+            {"assistant": "assistant"}
+        )
         builder.add_conditional_edges(
             "assistant",
             tools_condition,
@@ -84,13 +90,19 @@ class AdvancedAgent:
             # Initial state must be a dict with "messages" key!
             result = self.graph.invoke({"messages": messages})
             final_messages = result["messages"]
+            
+            # Extract only AI messages from the final state
             ai_messages = [msg for msg in final_messages if isinstance(msg, AIMessage)]
+            
             if not ai_messages:
                 return "I wasn't able to generate a proper response. Please try again."
-            for msg in reversed(ai_messages):
-                if not hasattr(msg, 'tool_calls') or not msg.tool_calls:
-                    return msg.content
-            return ai_messages[-1].content if ai_messages else "I wasn't able to generate a proper response. Please try again."
+            
+            # Get the last AI message - this should be the final answer
+            last_message = ai_messages[-1]
+            
+            # Return the content of the last AI message
+            return last_message.content if last_message.content else "I wasn't able to generate a proper response. Please try again."
+            
         except Exception as e:
             print(f"Error running agent graph: {e}")
             return f"Sorry, I encountered an error while processing your question: {str(e)}"