115 lines
3.2 KiB
Python
115 lines
3.2 KiB
Python
"""
|
||
Self‑correcting LangGraph agent demo.
|
||
|
||
Run with:
|
||
python main.py
|
||
|
||
Requires:
|
||
pip install langgraph langchain-openai
|
||
"""
|
||
import random
|
||
from typing import TypedDict, Dict
|
||
|
||
# ---------- State definition ----------
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str | None
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None
|
||
max_attempts: int
|
||
|
||
# ---------- Unreliable tool ----------
|
||
class UnreliableTool:
|
||
def run(self, input_: str) -> str:
|
||
if random.random() < 0.3:
|
||
raise ValueError("Simulated tool failure")
|
||
# simple eval for demo purposes
|
||
try:
|
||
return str(eval(input_))
|
||
except Exception as e:
|
||
raise ValueError(f"Eval error: {e}")
|
||
|
||
# ---------- LangGraph imports ----------
|
||
from langgraph.graph import StateGraph, END
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
# ---------- Nodes ----------
|
||
def execute_task(state: AgentState) -> AgentState:
|
||
tool = UnreliableTool()
|
||
try:
|
||
result = tool.run(state["task"])
|
||
state["result"] = result
|
||
state["error"] = None
|
||
except Exception as e:
|
||
state["result"] = None
|
||
state["error"] = str(e)
|
||
return state
|
||
|
||
def verify_result(state: AgentState) -> AgentState:
|
||
# Ask LLM to judge success or failed based on result and error
|
||
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
|
||
prompt = (
|
||
f"Task: {state['task']}\n"
|
||
f"Result: {state.get('result')}\n"
|
||
f"Error: {state.get('error')}\n"
|
||
"Respond with only 'success' or 'failed'."
|
||
)
|
||
resp = llm.invoke(prompt)
|
||
verdict = resp.content.strip().lower()
|
||
if verdict not in ("success", "failed"):
|
||
verdict = "failed"
|
||
state["verdict"] = verdict
|
||
return state
|
||
|
||
def handle_error(state: AgentState) -> AgentState:
|
||
state["attempts"] += 1
|
||
if state["attempts"] >= state["max_attempts"]:
|
||
state["status"] = "max_attempts"
|
||
else:
|
||
state["status"] = "pending"
|
||
return state
|
||
|
||
# ---------- Graph ----------
|
||
builder = StateGraph(AgentState)
|
||
builder.add_node("execute_task", execute_task)
|
||
builder.add_node("verify_result", verify_result)
|
||
builder.add_node("handle_error", handle_error)
|
||
|
||
builder.set_entry_point("execute_task")
|
||
builder.add_edge("execute_task", "verify_result")
|
||
# From verify_result, branch on verdict
|
||
builder.add_conditional_edges(
|
||
"verify_result",
|
||
lambda state: state.get("verdict"),
|
||
{
|
||
"success": END,
|
||
"failed": "handle_error",
|
||
},
|
||
)
|
||
builder.add_edge("handle_error", "execute_task")
|
||
|
||
graph = builder.compile()
|
||
|
||
# ---------- Demo runner ----------
|
||
if __name__ == "__main__":
|
||
initial_state: AgentState = {
|
||
"task": "2+2",
|
||
"result": None,
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": 5,
|
||
}
|
||
for attempt in range(1, initial_state["max_attempts"] + 1):
|
||
print(f"Attempt {attempt}:")
|
||
result = graph.invoke(initial_state)
|
||
if result.get("verdict") == "success":
|
||
print(f"Success: {result['result']}")
|
||
break
|
||
else:
|
||
print(f"Failed, error: {result.get('error')}")
|
||
else:
|
||
print("Reached max attempts without success.")
|
||
"" |