Update LG_agent.py

LG_agent.py

@@ -43,23 +43,22 @@ def build_graph(max_steps: int = 5):
 
     graph = builder.compile()
 
-    def limited_invoke(input_state):
-        steps = 0
-        state = input_state
-        while steps < max_steps:
-            print(f"🔁 Step {steps + 1}")
-            state = graph.invoke(state)
-            latest_message = state["messages"][-1] if state["messages"] else None
-
-            # Print tool usage if any
-            for m in state["messages"]:
-                if isinstance(m, ToolMessage):
-                    print(f"🔧 Tool Used: {m.tool_call_id} — {m.content}")
-
-            if latest_message and isinstance(latest_message, AIMessage):
-                break
-            steps += 1
-        return state
+    def limited_invoke(state):
+    steps = 0
+    while steps < max_steps:
+        print(f"🔁 Step {steps + 1}")
+        state = graph.invoke(state)
+        latest_message = state["messages"][-1] if state["messages"] else None
+
+        # Print tool usage if any
+        for m in state["messages"]:
+            if hasattr(m, "tool_call_id"):
+                print(f"🔧 Tool Used: {m.tool_call_id} — {m.content}")
+
+        if latest_message and isinstance(latest_message, AIMessage):
+            break
+        steps += 1
+    return state
 
     return limited_invoke