Update app.py

app.py

@@ -120,21 +120,50 @@ class AdvancedAgent:
             SystemMessage(content=system_prompt),
             HumanMessage(content=question)
         ]
+
+
         
         # Run the graph
         try:
             result = self.graph.invoke({"messages": messages})
-            # Extract the last AI message as the answer
-            for msg in reversed(result["messages"]):
-                if isinstance(msg, AIMessage) and not getattr(msg, "tool_call_id", None):
+            
+            # Extract the final AI message as the answer
+            final_messages = result["messages"]
+            
+            # Find the last non-tool AI message
+            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."
+            
+            # Get the last AI message that's not a tool call
+            for msg in reversed(ai_messages):
+                # Check if this is not a tool call message
+                if not hasattr(msg, 'tool_calls') or not msg.tool_calls:
                     return msg.content
             
-            # Fallback if no valid AI message found
-            return "I wasn't able to generate a proper response. Please try again."
+            # If we only have tool call messages, return the content of the last AI message
+            return ai_messages[-1].content if ai_messages 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)}"
 
+
+        
+        # # Run the graph
+        # try:
+        #     result = self.graph.invoke({"messages": messages})
+        #     # Extract the last AI message as the answer
+        #     for msg in reversed(result["messages"]):
+        #         if isinstance(msg, AIMessage) and not getattr(msg, "tool_call_id", None):
+        #             return msg.content
+            
+        #     # Fallback if no valid AI message found
+        #     return "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)}"
+
 # Function to run and submit all questions
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """