81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
"""
|
||
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)
|