Initial solution for planning agent: add repo/main.py

repo/main.py

@@ -0,0 +1,126 @@
+"""
+LangGraph planning agent example.
+
+This script demonstrates a simple LangGraph agent that:
+1. Takes an input task description.
+2. Uses an LLM to split the task into 3‑6 concrete steps (JSON list).
+3. Executes each step sequentially, collecting results.
+4. Returns a final summary of all results.
+
+Requirements:
+- python >= 3.10
+- langgraph
+- langchain-openai (or langchain-ollama)
+
+Run with:
+    python main.py "Compare Python and JavaScript"
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import TypedDict, List
+
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import create_agent_executor
+from langchain_openai import ChatOpenAI
+
+# ---------- 1. Define the state ---------------------------------
+class PlanningState(TypedDict):
+    task: str
+    plan: List[str] | None
+    current_step: int
+    results: List[str]
+
+# ---------- 2. LLM for planning ---------------------------------
+# The user should set OPENAI_API_KEY in environment.
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+planning_prompt = (
+    '''You are a helpful assistant that plans tasks.
+Given the following task description, break it into 3‑6 concrete steps.
+Return only a JSON array of strings. Example:
+["Step 1: ...", "Step 2: ..."]'''
+)
+
+# ---------- 3. Planning node -----------------------------------
+async def planning(state: PlanningState) -> PlanningState:
+    task = state["task"]
+    # Call LLM to get plan
+    response = await llm.agenerate([planning_prompt + f"\nTask: {task}"])
+    text = response.generations[0][0].text.strip()
+    try:
+        plan = json.loads(text)
+        if not isinstance(plan, list):
+            raise ValueError
+    except Exception:
+        # Fallback: split by newlines
+        plan = [line for line in text.split("\n") if line]
+    return {
+        "task": task,
+        "plan": plan,
+        "current_step": 0,
+        "results": [],
+    }
+
+# ---------- 4. Execution node -----------------------------------
+async def execution(state: PlanningState) -> PlanningState:
+    idx = state["current_step"]
+    step_text = state["plan"][idx]
+    # For demo, just echo the step as result.
+    result = f"Result of {step_text}"
+    new_results = state["results"].copy()
+    new_results.append(result)
+    return {
+        "task": state["task"],
+        "plan": state["plan"],
+        "current_step": idx + 1,
+        "results": new_results,
+    }
+
+# ---------- 5. Condition node -----------------------------------
+def should_continue(state: PlanningState) -> str:
+    if state["current_step"] >= len(state["plan"]):
+        return "finish"
+    return "execute"
+
+# ---------- 6. Build graph -------------------------------------
+builder = StateGraph(PlanningState)
+builder.add_node("planning", planning)
+builder.add_node("execution", execution)
+builder.add_conditional_edges(
+    "planning",
+    lambda _: "execute",
+)
+builder.add_conditional_edges(
+    "execution",
+    should_continue,
+    {
+        "execute": "execution",
+        "finish": END,
+    },
+)
+graph = builder.compile()
+
+# ---------- 7. Demo runner -------------------------------------
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python main.py '<task description>'")
+        sys.exit(1)
+    task_desc = sys.argv[1]
+    # Run graph
+    result = graph.invoke({"task": task_desc})
+    plan = result["plan"]
+    results = result["results"]
+    print("\nTask:", task_desc)
+    print("\nPlan:\n", "\n".join(f"{i+1}. {step}" for i, step in enumerate(plan)))
+    print("\nResults:\n", "\n".join(results))
+    # Final summary via LLM
+    final_prompt = (
+        f"Given the following results: {json.dumps(results)}\nProvide a concise summary."
+    )
+    final_resp = llm.invoke(final_prompt)
+    print("\nFinal Summary:\n", final_resp.content.strip())