167 lines
4.8 KiB
Python
167 lines
4.8 KiB
Python
"""
|
||
LangGraph Agent that plans and executes a task step by step.
|
||
|
||
Usage:
|
||
python agent.py "Compare Python and JavaScript"
|
||
|
||
Dependencies:
|
||
- langgraph
|
||
- langchain-openai
|
||
- langchain-ollama (optional)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import traceback
|
||
from typing import TypedDict
|
||
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
# ----- State definition (from snippet) -----
|
||
class PlanningState(TypedDict):
|
||
task: str
|
||
plan: list[str] | None
|
||
current_step: int
|
||
results: list[str]
|
||
|
||
# ----- LLM configuration -----
|
||
# Prefer OpenAI if API key is set; otherwise fallback to Ollama
|
||
if os.getenv("OPENAI_API_KEY"):
|
||
llm = ChatOpenAI(
|
||
model="gpt-4o-mini",
|
||
temperature=0,
|
||
)
|
||
else:
|
||
# Fall back to local Ollama model if available
|
||
llm = ChatOpenAI(
|
||
model=os.getenv("CHAT_MODEL", "llama3"),
|
||
base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
|
||
api_key="ollama",
|
||
temperature=0,
|
||
)
|
||
|
||
# ----- Planning node -----
|
||
|
||
def planning(state: PlanningState) -> PlanningState:
|
||
"""Prompt LLM to break the task into 3–6 numbered steps."""
|
||
try:
|
||
prompt = (
|
||
f"Given the task '{state['task']}', break it into 3-6 numbered steps. "
|
||
"Return only the numbered list or a JSON array of steps."
|
||
)
|
||
raw = llm.invoke(prompt)
|
||
text = raw if isinstance(raw, str) else raw.content
|
||
except Exception as e:
|
||
raise RuntimeError(f"LLM failed in planning node: {e}")
|
||
|
||
# Parse plan – first try JSON, then regex
|
||
plan: list[str] | None = None
|
||
try:
|
||
plan = json.loads(text)
|
||
if not isinstance(plan, list):
|
||
plan = None
|
||
except Exception:
|
||
plan = None
|
||
|
||
if plan is None:
|
||
pattern = r"^\s*\d+\.\s+(.*)$"
|
||
plan = [m.group(1).strip() for m in re.finditer(pattern, text, re.MULTILINE)]
|
||
|
||
if not plan or not (3 <= len(plan) <= 6):
|
||
raise ValueError(
|
||
f"Planning node returned invalid plan: {plan}. Expected 3-6 steps."
|
||
)
|
||
|
||
state["plan"] = plan
|
||
state["current_step"] = 0
|
||
state["results"] = []
|
||
return state
|
||
|
||
# ----- Execution node -----
|
||
|
||
def execution(state: PlanningState) -> PlanningState:
|
||
"""Execute a single step from the plan and record the result."""
|
||
step_idx = state["current_step"]
|
||
if state["plan"] is None or step_idx >= len(state["plan"]):
|
||
return state
|
||
step = state["plan"][step_idx]
|
||
try:
|
||
prompt = f"Execute step: {step}\nProvide a concise result."
|
||
raw = llm.invoke(prompt)
|
||
result = raw if isinstance(raw, str) else raw.content
|
||
except Exception as e:
|
||
result = f"Error executing step: {e}"
|
||
state["results"].append(result.strip())
|
||
state["current_step"] = step_idx + 1
|
||
return state
|
||
|
||
# ----- Graph construction -----
|
||
# Graph diagram (from snippet):
|
||
# START → planning → execution → should_continue
|
||
# ↑____________| (execute)
|
||
# finish → END
|
||
|
||
builder = StateGraph(PlanningState)
|
||
builder.add_node("planning", planning)
|
||
builder.add_node("execution", execution)
|
||
|
||
# Conditional edges to loop execution until all steps processed
|
||
builder.add_conditional_edges(
|
||
"planning",
|
||
lambda state: "execution" if state.get("plan") else "END"
|
||
)
|
||
builder.add_conditional_edges(
|
||
"execution",
|
||
lambda state: "execution"
|
||
if state.get("current_step", 0) < len(state.get("plan", []))
|
||
else "END"
|
||
)
|
||
|
||
builder.set_entry_point("planning")
|
||
|
||
# Compile graph with in-memory checkpointing
|
||
graph = builder.compile(checkpointer=InMemorySaver())
|
||
|
||
# ----- Demo -----
|
||
|
||
def run_demo(task: str) -> None:
|
||
initial_state: PlanningState = {
|
||
"task": task,
|
||
"plan": None,
|
||
"current_step": 0,
|
||
"results": [],
|
||
}
|
||
try:
|
||
result = graph.invoke(initial_state)
|
||
except Exception:
|
||
traceback.print_exc()
|
||
sys.exit(1)
|
||
|
||
plan = result.get("plan", [])
|
||
print("\n===== PLAN =====")
|
||
for idx, step in enumerate(plan, 1):
|
||
print(f"{idx}. {step}")
|
||
print("\n===== RESULTS =====")
|
||
for idx, res in enumerate(result.get("results", []), 1):
|
||
print(f"[Step {idx}] {res}\n")
|
||
print("===== SUMMARY =====")
|
||
try:
|
||
summary_prompt = f"Given the collected results: {result.get('results', [])}, produce a concise summary of the task outcome."
|
||
summary = llm.invoke(summary_prompt)
|
||
print(summary if isinstance(summary, str) else summary.content)
|
||
except Exception as e:
|
||
print(f"Error generating summary: {e}")
|
||
|
||
if __name__ == "__main__":
|
||
if len(sys.argv) < 2:
|
||
print("Usage: python agent.py '<task>'")
|
||
sys.exit(1)
|
||
task = " ".join(sys.argv[1:])
|
||
run_demo(task)
|