127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
"""
|
||
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())
|