Update app.py

app.py

@@ -520,7 +520,21 @@ def should_continue(state: AgentState):
         print("--- Condition: Tools called, routing to tools node. ---")
         return "tools"
 
-    # 4. Loop back to agent (reasoning/planning step)
+    # 4. NEW LOOP PREVENTION:
+    # Check for consecutive AI messages without tool calls.
+    # This catches "thinking" loops or raw answer dribbling (like "58").
+    if isinstance(last_message, AIMessage) and not last_message.tool_calls:
+        # Check if the message *before* this one was ALSO an AIMessage.
+        # We need at least 3 messages total (System, Human, AI-Turn1-Plan)
+        # for this check to be valid, so we check len > 2.
+        if len(state['messages']) > 2 and isinstance(state['messages'][-2], AIMessage):
+            print(f"--- Condition: Detected 2+ consecutive AI messages (Turn {current_turn}). Ending to prevent loop. ---")
+            state['messages'].append(
+                SystemMessage(content=f"SYSTEM: Agent stuck in a loop (consecutive non-tool-call AI messages). Ending execution.")
+            )
+            return END
+
+    # 5. Default: Loop back to agent (e.g., after Turn 1 plan)
     print(f"--- Condition: No tool call (Turn {current_turn}). Continuing to agent. ---")
     return "agent"
 
@@ -614,7 +628,7 @@ Your goal: Provide the EXACT answer in the EXACT format requested.
             chat_llm = ChatGroq(
                 temperature=0,  # Maximum determinism
                 groq_api_key=GROQ_API_KEY,
-                model_name="openai/gpt-oss-120b",  # Best reasoning model
+                model_name="qwen/qwen3-32b",  # Best reasoning model
                 max_tokens=4096,
                 timeout=60
             )
@@ -627,10 +641,33 @@ Your goal: Provide the EXACT answer in the EXACT format requested.
         print("✅ Tools bound to LLM")
         
         # --- Agent Node ---
-        def agent_node(state: AgentState):
-            current_turn = state.get('turn', 0) + 1
+def agent_node(state: AgentState):
+            
+            # --- Turn Counter Logic ---
+            # We need to check if this is a retry of a failed turn (e.g., Turn 1 violation)
+            # We identify a retry if the *last* message was our "Protocol Violation" message
+            last_msg = state['messages'][-1]
+            is_a_retry = False
+            if isinstance(last_msg, SystemMessage) and "Protocol Violation" in last_msg.content:
+                is_a_retry = True
+            
+            # Get the state's current turn number
+            current_turn = state.get('turn', 0)
+            
+            # If this is NOT a retry, increment the turn.
+            # If it IS a retry, we *stay on the same turn number*
+            if not is_a_retry:
+                current_turn += 1
+            
+            # Handle the very first run (where state['turn'] is 0)
+            if current_turn == 0:
+                current_turn = 1
+            # --- End Turn Counter Logic ---
+
             print(f"\n{'='*60}")
             print(f"AGENT TURN {current_turn}/{MAX_TURNS}")
+            if is_a_retry:
+                print("--- (Re-trying after protocol violation) ---")
             print('='*60)
             
             messages_to_send = state["messages"]
@@ -651,27 +688,29 @@ Your goal: Provide the EXACT answer in the EXACT format requested.
                         )
                         return {"messages": [error_msg], "turn": current_turn}
                     time.sleep(2 ** attempt)  # Exponential backoff
+            
             # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
-            # --- ADD THIS RULE ENFORCEMENT BLOCK ---
+            # --- (FIX #1) RULE ENFORCEMENT BLOCK ---
             #
-            # Force compliance with the "Plan on Turn 1" rule.
-            # If the agent tries to call a tool on turn 1, we
-            # strip the tool call and send a correction message.
+            # If it's Turn 1 AND the agent tried to call tools, we reject it
+            # and force it to re-do Turn 1.
             if current_turn == 1 and ai_message.tool_calls:
                 print("⚠️ AGENT VIOLATION: Tried to call tools on Turn 1. Forcing replan.")
                 
                 # Strip the illegal tool call
                 ai_message.tool_calls = []
                 
-                # Create a correction message
+                # Create the correction message that forces the plan
                 correction_message = SystemMessage(
                     content="SYSTEM: Protocol Violation. Your FIRST turn MUST be a plan with NO tool calls. "
                             "You are not allowed to call any tools on your first turn. "
                             "Re-read the protocol and provide your 2-3 sentence plan now."
                 )
                 
-                # Return the (now harmless) AI message + the correction
-                return {"messages": [ai_message, correction_message], "turn": current_turn}
+                # Return the messages.
+                # Critically, we set the state's turn counter back to 1.
+                # This ensures the *next* run of this node is *still* Turn 1.
+                return {"messages": [ai_message, correction_message], "turn": 1}
             # --- END OF RULE ENFORCEMENT BLOCK ---
 # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
             # --- FIX #2: REPLACE THE FALLBACK PARSING BLOCK ---