Update app.py

app.py

@@ -769,93 +769,93 @@ Your goal: Provide the EXACT answer in the EXACT format requested.
 **REMEMBER:** One step at a time. Use tools. Format JSON correctly.
 """
 
-        # --- *** UPDATED: Initialize Google Gemini LLM *** ---
-        print("Initializing Google Gemini LLM...")
-        try:
-            # --- Initialize ONE Powerful LLM for all tasks ---
-            self.planner_llm = ChatGoogleGenerativeAI(
-                model="gemini-2.0-flash-exp",  # ← Same model, different config
-                google_api_key=GEMINI_API_KEY,
-                temperature=0,
-                convert_system_message_to_human=True,
-                max_output_tokens=1024  # Planner needs less tokens
-            )
-            print("✅ Planner LLM (Gemini 2.0 Flash) initialized.")
+    # --- *** UPDATED: Initialize Google Gemini LLM *** ---
+    print("Initializing Google Gemini LLM...")
+    try:
+        # --- Initialize ONE Powerful LLM for all tasks ---
+        self.planner_llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash-exp",  # ← Same model, different config
+            google_api_key=GEMINI_API_KEY,
+            temperature=0,
+            convert_system_message_to_human=True,
+            max_output_tokens=1024  # Planner needs less tokens
+        )
+        print("✅ Planner LLM (Gemini 2.0 Flash) initialized.")
 
-        except Exception as e:
-            print(f"❌ Error initializing Gemini: {e}")
-            raise
-        # --- *** END UPDATE *** ---
+    except Exception as e:
+        print(f"❌ Error initializing Gemini: {e}")
+        raise
+    # --- *** END UPDATE *** ---
+    
+    # --- Node 1: The Agent ---
+    def agent_node(state: AgentState):
+        current_turn = state.get('turn', 0) + 1
+        print(f"\n{'='*60}")
+        print(f"AGENT TURN {current_turn}/{MAX_TURNS}")
+        print('='*60)
         
-        # --- Node 1: The Agent ---
-        def agent_node(state: AgentState):
-            current_turn = state.get('turn', 0) + 1
-            print(f"\n{'='*60}")
-            print(f"AGENT TURN {current_turn}/{MAX_TURNS}")
-            print('='*60)
-            
-            # Note: Max turns is also checked in should_continue, but good to have here
-            if current_turn > MAX_TURNS:
-                return {"messages": [SystemMessage(content="Max turns reached.")]}
+        # Note: Max turns is also checked in should_continue, but good to have here
+        if current_turn > MAX_TURNS:
+            return {"messages": [SystemMessage(content="Max turns reached.")]}
 
-            max_retries = 3
-            ai_message = None
-            for attempt in range(max_retries):
-                try:
-                    # Call the single, powerful LLM
-                    ai_message = self.llm_with_tools.invoke(state["messages"])
-                    break
-                except Exception as e:
-                    print(f"⚠️ LLM attempt {attempt+1}/{max_retries} failed: {e}")
-                    if attempt == max_retries - 1:
-                        ai_message = AIMessage(
-                            content=f"Error: LLM failed after {max_retries} attempts: {e}"
-                        )
-                    time.sleep(2 ** attempt)
-            
-            # --- Fallback Parsing Logic ---
-            if not ai_message.tool_calls and isinstance(ai_message.content, str) and ai_message.content.strip():
-                parsed_tool_calls = parse_tool_call_from_string(ai_message.content, self.tools)
-                if parsed_tool_calls:
-                    print("🔧 Fallback SUCCESS: Rebuilding tool call(s).")
-                    ai_message.tool_calls = parsed_tool_calls
-                    ai_message.content = "" # Clear the text content
-                else:
-                    print(f"⚠️ Fallback FAILED: Could not parse any tool call from content:\n{ai_message.content[:200]}...")
-            
-            if ai_message.tool_calls:
-                 print(f"🔧 Agent Tool Call: {ai_message.tool_calls[0]['name']}")
-            else:
-                print(f"💭 Agent Reasoning: {ai_message.content[:200]}...")
-            
-            return {"messages": [ai_message], "turn": current_turn}
-
-        # --- Tool Node ---
-        tool_node = ToolNode(self.tools)
-        
-        # --- Build Graph ---
-        print("Building Single-Agent graph...")
-        graph_builder = StateGraph(AgentState)
-        
-        graph_builder.add_node("agent", agent_node)
-        graph_builder.add_node("tools", tool_node)
-        
-        graph_builder.add_edge(START, "agent")
+        max_retries = 3
+        ai_message = None
+        for attempt in range(max_retries):
+            try:
+                # Call the single, powerful LLM
+                ai_message = self.llm_with_tools.invoke(state["messages"])
+                break
+            except Exception as e:
+                print(f"⚠️ LLM attempt {attempt+1}/{max_retries} failed: {e}")
+                if attempt == max_retries - 1:
+                    ai_message = AIMessage(
+                        content=f"Error: LLM failed after {max_retries} attempts: {e}"
+                    )
+                time.sleep(2 ** attempt)
         
-        graph_builder.add_conditional_edges(
-            "agent",
-            should_continue, # Use the reverted conditional function
-            {
-                "tools": "tools",
-                "agent": "agent", # For loop prevention
-                END: END
-            }
-        )
+        # --- Fallback Parsing Logic ---
+        if not ai_message.tool_calls and isinstance(ai_message.content, str) and ai_message.content.strip():
+            parsed_tool_calls = parse_tool_call_from_string(ai_message.content, self.tools)
+            if parsed_tool_calls:
+                print("🔧 Fallback SUCCESS: Rebuilding tool call(s).")
+                ai_message.tool_calls = parsed_tool_calls
+                ai_message.content = "" # Clear the text content
+            else:
+                print(f"⚠️ Fallback FAILED: Could not parse any tool call from content:\n{ai_message.content[:200]}...")
         
-        graph_builder.add_edge("tools", "agent") # Loop back to agent
+        if ai_message.tool_calls:
+             print(f"🔧 Agent Tool Call: {ai_message.tool_calls[0]['name']}")
+        else:
+            print(f"💭 Agent Reasoning: {ai_message.content[:200]}...")
         
-        self.graph = graph_builder.compile()
-        print("✅ Single-Agent graph compiled successfully.")
+        return {"messages": [ai_message], "turn": current_turn}
+
+    # --- Tool Node ---
+    tool_node = ToolNode(self.tools)
+    
+    # --- Build Graph ---
+    print("Building Single-Agent graph...")
+    graph_builder = StateGraph(AgentState)
+    
+    graph_builder.add_node("agent", agent_node)
+    graph_builder.add_node("tools", tool_node)
+    
+    graph_builder.add_edge(START, "agent")
+    
+    graph_builder.add_conditional_edges(
+        "agent",
+        should_continue, # Use the reverted conditional function
+        {
+            "tools": "tools",
+            "agent": "agent", # For loop prevention
+            END: END
+        }
+    )
+    
+    graph_builder.add_edge("tools", "agent") # Loop back to agent
+    
+    self.graph = graph_builder.compile()
+    print("✅ Single-Agent graph compiled successfully.")
 
     def __call__(self, question: str) -> str:
         print(f"\n--- Starting Agent Run for Question ---")