overwrite main.py

main.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 import os
 from typing import List
 
-from langgraph.graph import StateGraph, START, END
+from langgraph.graph import StateGraph, END
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_openai import ChatOpenAI
 from rich.console import Console
@@ -40,10 +40,6 @@ llm = ChatOpenAI(
 # 2. Planning node – split task into steps
 # ---------------------------------------------------------------------------
 def planning(state: PlanningState) -> PlanningState:
-    """Ask the LLM to produce a numbered list of steps.
-
-    The prompt forces JSON output for reliable parsing.
-    """
     system = SystemMessage(
         content="You are a helpful assistant that splits a task into clear, actionable steps. Return a JSON array of strings under the key `plan`."
     )
@@ -52,7 +48,6 @@ def planning(state: PlanningState) -> PlanningState:
     response = llm.invoke([system, user])
     text = response.content.strip()
 
-    # Try to parse JSON; if fails, fall back to simple split by lines.
     try:
         import json
         data = json.loads(text)
@@ -63,7 +58,6 @@ def planning(state: PlanningState) -> PlanningState:
             line = line.strip()
             if not line:
                 continue
-            # Remove leading numbers like "1. " or "- "
             if line[0].isdigit() and (len(line) > 2 and line[1] in ".-"):
                 line = line.split("", 1)[1]
             plan.append(line)
@@ -125,8 +119,8 @@ def run_agent(task: str) -> None:
     state: PlanningState = {"task": task, "plan": None, "current_step": 0, "results": []}
     config = {"configurable": {"thread_id": f"{task[:8]}"}}
 
-    for event in agent.stream(state, config):
-        pass  # we only need final state
+    final_state = agent.invoke(state, config)
+    state = final_state
 
     console.print("\n[bold underline]Task:[/]", task)
     if state.get("plan"):