Update app.py

app.py

@@ -138,10 +138,39 @@ def find_file(path: str) -> Optional[Path]:
     
     return None
 
+
 # =============================================================================
 # PLANNING & REFLECTION TOOLS
 # =============================================================================
 
+class ThinkInput(BaseModel):
+    reasoning: str = Field(description="Your step-by-step reasoning for a logic puzzle (keep under 200 chars)")
+
+@tool(args_schema=ThinkInput)
+def think_through_logic(reasoning: str) -> str:
+    """
+    Use this to work through logic puzzles, riddles, or reasoning problems.
+    
+    Call this when:
+    - The question is a riddle or brain teaser
+    - You need to reason through a logical problem
+    - No external information is needed, just thinking
+    
+    After thinking through the logic, use calculator if math is involved,
+    then validate_answer and final_answer_tool.
+    
+    NOTE: Keep reasoning summary brief (under 200 chars).
+    """
+    print(f"🧠 Thinking through logic: {reasoning[:100]}...")
+    
+    return f"""✅ Logic reasoning recorded: {reasoning}
+
+Now:
+1. If there's any math to calculate, use calculator()
+2. Once you have the answer, call validate_answer()
+3. Then call final_answer_tool() with just the answer"""
+
+
 class PlanInput(BaseModel):
     question: str = Field(description="Brief summary of the task (keep under 100 chars)")
 
@@ -273,6 +302,7 @@ def validate_answer(proposed_answer: str, original_question: str) -> str:
     return "✅ VALIDATION PASSED: Answer looks good! Proceed with final_answer_tool now."
 
 
+# =============================================================================
 # =============================================================================
 # CORE TOOLS
 # =============================================================================
@@ -822,41 +852,6 @@ def parse_tool_call_from_string(content: str, tools: List) -> List[ToolCall]:
     return []
 
 
-# =============================================================================
-# CONDITIONAL EDGE FUNCTION
-# =============================================================================
-def should_continue(state: AgentState):
-    """Decide whether to continue, call tools, or end."""
-    last_message = state['messages'][-1]
-    current_turn = state.get('turn', 0)
-
-    # Check for final_answer_tool
-    if isinstance(last_message, AIMessage) and last_message.tool_calls:
-        for tool_call in last_message.tool_calls:
-            if tool_call.get("name") == "final_answer_tool":
-                print("--- Condition: final_answer_tool called, ending. ---")
-                return END
-
-    # Check turn limit
-    if current_turn >= MAX_TURNS:
-        print(f"--- Condition: Max turns ({MAX_TURNS}) reached. Ending. ---")
-        return END
-
-    # Route to tools if tool calls exist
-    if isinstance(last_message, AIMessage) and last_message.tool_calls:
-        print("--- Condition: Tools called, routing to tools node. ---")
-        return "tools"
-    
-    # Loop prevention
-    if len(state['messages']) > 2 and isinstance(last_message, AIMessage) and isinstance(state['messages'][-2], AIMessage):
-        print(f"--- Condition: Detected 2+ consecutive AI messages (Turn {current_turn}). Ending to prevent loop. ---")
-        return END
-
-    # Loop back to agent
-    print(f"--- Condition: No tool call (Turn {current_turn}). Continuing to agent. ---")
-    return "agent"
-
-
 # =============================================================================
 # ENHANCED AGENT CLASS WITH PLANNING & REFLECTION
 # =============================================================================
@@ -896,31 +891,58 @@ class PlanningReflectionAgent:
 🎯 YOUR MISSION: Provide the EXACT answer in the EXACT format requested.
 
 ═══════════════════════════════════════════════════════════════
-📋 MANDATORY PROTOCOL - FOLLOW THIS RELIGIOUSLY:
+📋 QUESTION TYPES & STRATEGIES:
+═══════════════════════════════════════════════════════════════
+
+**TYPE 1: LOGIC PUZZLES / RIDDLES** (No tools needed)
+- Riddles, brain teasers, logical reasoning problems
+- Strategy: Think through the logic, use calculator for any math
+- Example: "If all but 30 of 200 coins are face-up, make equal face-down piles"
+  → This is pure logic. Think it through, then use final_answer_tool
+
+**TYPE 2: FACTUAL QUESTIONS** (Need web search)
+- Who, what, when, where questions about real world
+- Strategy: search_tool → scrape_and_retrieve if needed
+- Example: "What was Einstein's birthplace population in 1900?"
+
+**TYPE 3: DATA ANALYSIS** (Need files + code)
+- Questions about CSV, Excel, or other data files
+- Strategy: list_directory → read_file → code_interpreter
+- Example: "What's the average of column X in data.csv?"
+
+**TYPE 4: CALCULATIONS** (Need calculator/code)
+- Math problems, computations
+- Strategy: calculator for simple math, code_interpreter for complex
+- Example: "What is 127 × 83 + sqrt(144)?"
+
+═══════════════════════════════════════════════════════════════
+📋 MANDATORY PROTOCOL:
 ═══════════════════════════════════════════════════════════════
 
-**PHASE 1: PLANNING (For complex/multi-step questions)**
-├─ 1. Call create_plan() to think through your approach
-├─ 2. Identify what information you need
-└─ 3. Determine the sequence of steps
-
-**PHASE 2: EXECUTION (One step at a time)**
-├─ 1. Take ONE action per turn
-├─ 2. Use the RIGHT tool for each task:
-│     • Simple math → calculator()
-│     • Complex data → code_interpreter()
-│     • Web info → search_tool()
-│     • Specific page → scrape_and_retrieve()
-│     • Files → read_file()
+**PHASE 1: IDENTIFY QUESTION TYPE**
+├─ Is this a logic puzzle? → Think through it, use calculator if needed
+├─ Need real-world facts? → Use search/scrape tools
+├─ Need to analyze files? → Use file/code tools
+└─ Just math? → Use calculator
+
+**PHASE 2: FOR TOOL-BASED QUESTIONS**
+├─ 1. Call create_plan() for multi-step questions
+├─ 2. Execute ONE step at a time
 ├─ 3. After EACH tool, evaluate the result
 └─ 4. Ask: "Do I have enough to answer now?"
 
-**PHASE 3: REFLECTION (If stuck)**
+**PHASE 3: FOR LOGIC PUZZLES**
+├─ 1. Think through the logic step-by-step
+├─ 2. Use calculator ONLY if there's arithmetic
+├─ 3. Once you've solved it, call validate_answer()
+└─ 4. Then call final_answer_tool()
+
+**PHASE 4: REFLECTION (If stuck)**
 ├─ If no progress after 3-5 turns → call reflect_on_progress()
 ├─ If tools keep failing → try different approach
 └─ If going in circles → step back and reconsider
 
-**PHASE 4: VALIDATION & SUBMISSION**
+**PHASE 5: VALIDATION & SUBMISSION**
 ├─ 1. When you have the answer → call validate_answer()
 ├─ 2. If validation passes → call final_answer_tool()
 └─ 3. If validation fails → fix the issue first
@@ -929,21 +951,28 @@ class PlanningReflectionAgent:
 🎓 EXAMPLES - LEARN FROM THESE:
 ═══════════════════════════════════════════════════════════════
 
-**Example 1: Simple Math**
+**Example 1: Logic Puzzle (NO TOOLS EXCEPT CALCULATOR/FINAL)**
+Q: If you have 200 coins with 30 face-down, and divide into 2 piles with equal face-down...
+Turn 1: Think through: If I take 30 coins and flip them all, one pile has X face-down...
+Turn 2: calculator("30") → 30
+Turn 3: validate_answer("30", original_q) → ✅ Pass
+Turn 4: final_answer_tool("30")
+
+**Example 2: Simple Math**
 Q: What is 127 × 83?
 Turn 1: calculator("127 * 83") → 10541
 Turn 2: validate_answer("10541", "What is 127 × 83?") → ✅ Pass
 Turn 3: final_answer_tool("10541")
 
-**Example 2: Multi-step Research**
+**Example 3: Multi-step Research**
 Q: What was the population of Einstein's birthplace in 1900?
-Turn 1: create_plan("What was the population of Einstein's birthplace in 1900?")
+Turn 1: create_plan("Brief: Einstein birthplace pop 1900")
 Turn 2: search_tool("Albert Einstein birthplace") → Ulm, Germany
 Turn 3: search_tool("Ulm Germany population 1900") → approximately 50,000
 Turn 4: validate_answer("50000", "What was the population...") → ✅ Pass
 Turn 5: final_answer_tool("50000")
 
-**Example 3: File + Calculation**
+**Example 4: File + Calculation**
 Q: What's the average of the 'score' column in data.csv?
 Turn 1: list_directory(".") → [files shown]
 Turn 2: read_file("data.csv") → [content]
@@ -952,11 +981,11 @@ Turn 3: code_interpreter("import pandas as pd; df = pd.read_csv('data.csv'); pri
 Turn 4: validate_answer("78.5", "What's the average...") → ✅ Pass
 Turn 5: final_answer_tool("78.5")
 
-**Example 4: Getting Unstuck**
+**Example 5: Getting Unstuck**
 Q: What's the GDP of the 2016 Olympics host?
 Turn 1: search_tool("2016 Olympics") → [general info, no clear answer]
 Turn 2: search_tool("Olympics 2016 location") → [still unclear]
-Turn 3: reflect_on_progress("Tried searching but not getting clear host country")
+Turn 3: reflect_on_progress("Searching but not getting host country")
         → Try: "2016 Summer Olympics host country"
 Turn 4: search_tool("2016 Summer Olympics host country") → Brazil
 Turn 5: search_tool("Brazil GDP 2016") → $1.796 trillion
@@ -967,13 +996,13 @@ Turn 7: final_answer_tool("1.796 trillion")
 ⚠️ CRITICAL RULES - NEVER VIOLATE THESE:
 ═══════════════════════════════════════════════════════════════
 
-1. **NO GUESSING**: Always use tools. Never use your own knowledge.
-2. **ONE STEP AT A TIME**: Don't try to do multiple things in one turn.
-3. **EXACT FORMAT**: Answer must be EXACTLY what was asked for.
-4. **NO FLUFF**: Never add "The answer is" or explanations in final answer.
-5. **ALWAYS VALIDATE**: Call validate_answer() before final_answer_tool().
-6. **PLAN COMPLEX TASKS**: Multi-step questions need create_plan() first.
-7. **REFLECT WHEN STUCK**: If no progress after 5 turns, call reflect_on_progress().
+1. **IDENTIFY QUESTION TYPE FIRST**: Logic puzzle vs. factual vs. data vs. math
+2. **LOGIC PUZZLES**: Don't use search/file tools. Just think + validate + final_answer
+3. **ONE STEP AT A TIME**: Don't try to do multiple things in one turn
+4. **EXACT FORMAT**: Answer must be EXACTLY what was asked for
+5. **NO FLUFF**: Never add "The answer is" or explanations in final answer
+6. **ALWAYS VALIDATE**: Call validate_answer() before final_answer_tool()
+7. **DON'T LOOP**: If 2 consecutive turns produce no tool calls, you're stuck - call a tool!
 
 ═══════════════════════════════════════════════════════════════
 📚 AVAILABLE TOOLS:
@@ -982,7 +1011,10 @@ Turn 7: final_answer_tool("1.796 trillion")
 {tool_descriptions}
 
 ═══════════════════════════════════════════════════════════════
-🎯 REMEMBER: Quality over speed. Think carefully, plan ahead, execute methodically.
+🎯 REMEMBER: 
+- Logic puzzles: Think → Calculator (if needed) → Validate → Final Answer
+- Factual questions: Plan → Search → Validate → Final Answer
+- Always call a tool - never just output reasoning text!
 ═══════════════════════════════════════════════════════════════
 """
 
@@ -1145,7 +1177,7 @@ Turn 7: final_answer_tool("1.796 trillion")
         
         self.graph = graph_builder.compile()
         print("✅ Planning & Reflection Agent graph compiled successfully.")
-        
+
 
     def __call__(self, question: str) -> str:
         print(f"\n--- Starting Agent Run for Question ---")