add main

main.py

@@ -0,0 +1,80 @@
+"""
+LangGraph planning agent example.
+
+Run with:
+    python main.py "Compare Python and JavaScript"
+"""
+import json, sys
+from typing import TypedDict, List
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import create_chat_agent
+from langchain_openai import ChatOpenAI
+
+# 1. Define state
+class PlanningState(TypedDict):
+    task: str
+    plan: List[str] | None
+    current_step: int
+    results: List[str]
+
+# 2. LLM for planning and execution
+llm = ChatOpenAI(temperature=0)
+
+# 3. Planning node – split task into steps
+async def planning(state: PlanningState) -> PlanningState:
+    prompt = (
+        "You are a helpful assistant that splits a user task into a numbered list of 3-6 concrete steps.
+        Return only the JSON array of strings, e.g. ["Step 1", "Step 2"] without any explanation.")
+    response = await llm.agenerate([{"role": "user", "content": f"{prompt}\nTask: {state['task']}"}])
+    text = response.generations[0][0].text.strip()
+    try:
+        plan = json.loads(text)
+        if not isinstance(plan, list): raise ValueError
+    except Exception as e:
+        # fallback simple split by lines
+        plan = [line.strip() for line in text.splitlines() if line.strip()]
+    return {**state, "plan": plan, "current_step": 0, "results": []}
+
+# 4. Execution node – run one step
+async def execution(state: PlanningState) -> PlanningState:
+    step = state['plan'][state['current_step']]
+    prompt = f"Execute the following step and return only the result string: {step}"
+    response = await llm.agenerate([{"role": "user", "content": prompt}])
+    result = response.generations[0][0].text.strip()
+    new_results = state['results'] + [result]
+    return {**state, "current_step": state['current_step'] + 1, "results": new_results}
+
+# 5. Condition node – decide to continue or finish
+def should_continue(state: PlanningState) -> str:
+    if state['current_step'] >= len(state['plan']):
+        return END
+    return "execute"
+
+# 6. Build graph
+workflow = StateGraph(PlanningState)
+workflow.add_node("planning", planning)
+workflow.add_node("execution", execution)
+workflow.set_entry_point("planning")
+workflow.add_conditional_edges("planning", lambda _: "execute")
+workflow.add_edge("execution", "should_continue")
+workflow.add_conditional_edges("should_continue", should_continue)
+graph = workflow.compile()
+
+# 7. Run demo
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python main.py '<task>'")
+        sys.exit(1)
+    task = sys.argv[1]
+    result = graph.invoke({"task": task})
+    plan = result['plan']
+    results = result['results']
+    print(f"\nTask: {task}\n")
+    print("Plan:\n")
+    for i, step in enumerate(plan, 1):
+        print(f"{i}. {step}")
+    print("\nResults:\n")
+    for i, res in enumerate(results, 1):
+        print(f"[Step {i}] {res}\n")
+    final = "\nFinal answer: " + "\n".join(results)
+    print(final)