update main.py

main.py

@@ -15,7 +15,7 @@ Requirements:
 from __future__ import annotations
 
 import os
-from typing import TypedDict, List, Optional
+from typing import List
 
 from langgraph.graph import StateGraph, START, END
 from langchain_core.messages import HumanMessage, SystemMessage
@@ -23,17 +23,11 @@ from langchain_openai import ChatOpenAI
 from rich.console import Console
 from rich.table import Table
 
-# ---------------------------------------------------------------------------
-# 1. State definition
-# ---------------------------------------------------------------------------
-class PlanningState(TypedDict):
-    task: str
-    plan: List[str] | None
-    current_step: int
-    results: List[str]
+# Import shared state definition
+from models import PlanningState
 
 # ---------------------------------------------------------------------------
-# 2. LLM configuration – BroJS endpoint
+# 1. LLM configuration – BroJS endpoint
 # ---------------------------------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
@@ -43,7 +37,7 @@ llm = ChatOpenAI(
 )
 
 # ---------------------------------------------------------------------------
-# 3. Planning node – split task into steps
+# 2. Planning node – split task into steps
 # ---------------------------------------------------------------------------
 def planning(state: PlanningState) -> PlanningState:
     """Ask the LLM to produce a numbered list of steps.
@@ -64,7 +58,6 @@ def planning(state: PlanningState) -> PlanningState:
         data = json.loads(text)
         plan: List[str] = data.get("plan", [])
     except Exception:
-        # Fallback – split on newlines and strip numbering
         plan = []
         for line in text.splitlines():
             line = line.strip()
@@ -82,19 +75,13 @@ def planning(state: PlanningState) -> PlanningState:
     }
 
 # ---------------------------------------------------------------------------
-# 4. Execution node – perform one step (here we just echo the step)
+# 3. Execution node – perform one step (here we just echo the step)
 # ---------------------------------------------------------------------------
 def execution(state: PlanningState) -> PlanningState:
-    """Execute a single step.
-
-    In a real assignment you would replace this with calls to tools or other logic.
-    For demonstration, we simply record the step text as the result.
-    """
     idx = state["current_step"]
-    if idx >= len(state["plan"]):
+    if idx >= len(state.get("plan", [])):
         return state
     step_text = state["plan"][idx]
-    # Simulate execution – in practice you might call a tool here.
     result = f"Executed: {step_text}"
     new_results = state["results"].copy()
     new_results.append(result)
@@ -105,7 +92,7 @@ def execution(state: PlanningState) -> PlanningState:
     }
 
 # ---------------------------------------------------------------------------
-# 5. Decision node – should we continue?
+# 4. Decision node – should we continue?
 # ---------------------------------------------------------------------------
 def should_continue(state: PlanningState) -> str:
     if state["current_step"] >= len(state.get("plan", [])):
@@ -113,7 +100,7 @@ def should_continue(state: PlanningState) -> str:
     return "execute"
 
 # ---------------------------------------------------------------------------
-# 6. Build the graph
+# 5. Build the graph
 # ---------------------------------------------------------------------------
 graph = StateGraph(PlanningState)
 graph.add_node("planning", planning)
@@ -131,21 +118,16 @@ graph.add_conditional_edges(
 agent = graph.compile()
 
 # ---------------------------------------------------------------------------
-# 7. Helper to run the agent and pretty‑print results
+# 6. Helper to run the agent and pretty‑print results
 # ---------------------------------------------------------------------------
 def run_agent(task: str) -> None:
     console = Console()
     state: PlanningState = {"task": task, "plan": None, "current_step": 0, "results": []}
     config = {"configurable": {"thread_id": f"{task[:8]}"}}
 
-    # Run the agent – we capture intermediate states via a callback.
-    results: List[PlanningState] = []
     for event in agent.stream(state, config):
-        if isinstance(event, dict) and "messages" in event:
-            continue  # ignore final message
-        results.append(event)
+        pass  # we only need final state
 
-    # Print plan
     console.print("\n[bold underline]Task:[/]", task)
     if state.get("plan"):
         table = Table(title="Plan", show_header=False, box=None)
@@ -153,7 +135,6 @@ def run_agent(task: str) -> None:
             table.add_row(f"{i}. {step}")
         console.print(table)
 
-    # Print execution results
     console.print("\n[bold underline]Execution Results:[/]")
     for r in state.get("results", []):
         console.print(r)
@@ -161,7 +142,7 @@ def run_agent(task: str) -> None:
     console.print("\n[green]Finished.[/]\n")
 
 # ---------------------------------------------------------------------------
-# 8. Demo examples
+# 7. Demo examples
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
     examples = [