131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
"""
|
||
Self‑correcting LangGraph agent.
|
||
|
||
The agent executes a task, verifies the result with an LLM judge, and retries on failure
|
||
until success or a maximum number of attempts is reached.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import TypedDict
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
|
||
# ---------- State ----------
|
||
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None
|
||
max_attempts: int
|
||
|
||
# ---------- Unreliable tool ----------
|
||
|
||
import random
|
||
|
||
class UnreliableTool:
|
||
"""Simulates a tool that fails 30 % of the time."""
|
||
|
||
def __call__(self, input_text: str) -> str:
|
||
if random.random() < 0.3:
|
||
raise ValueError("Simulated tool failure")
|
||
return f"Result for: {input_text}"
|
||
|
||
unreliable_tool = UnreliableTool()
|
||
|
||
# ---------- Nodes ----------
|
||
|
||
async def execute_task(state: AgentState) -> AgentState:
|
||
"""Runs the task using the unreliable tool."""
|
||
try:
|
||
result = unreliable_tool(state["task"])
|
||
state["result"] = result
|
||
state["error"] = None
|
||
except Exception as e:
|
||
state["result"] = ""
|
||
state["error"] = str(e)
|
||
return state
|
||
|
||
async def verify_result(state: AgentState) -> AgentState:
|
||
"""LLM judge that returns 'success' or 'failed' based on the result."""
|
||
llm = ChatOpenAI(temperature=0)
|
||
prompt = (
|
||
"""You are a judge. Evaluate the following result for the task:\n"
|
||
f"Task: {state['task']}\n"
|
||
f"Result: {state['result']}\n"
|
||
"If the result is correct, reply with the word 'success'.\n"
|
||
"If the result is incorrect or missing, reply with the word 'failed'."""
|
||
)
|
||
response = llm.invoke(prompt)
|
||
verdict = response.content.strip().lower()
|
||
if verdict == "success":
|
||
state["status"] = "success"
|
||
else:
|
||
state["status"] = "failed"
|
||
return state
|
||
|
||
async def handle_error(state: AgentState) -> AgentState:
|
||
"""Increment attempts and prepare for retry."""
|
||
state["attempts"] += 1
|
||
state["status"] = "pending"
|
||
return state
|
||
|
||
# ---------- Graph ----------
|
||
|
||
def create_agent_graph(max_attempts: int = 3) -> StateGraph:
|
||
graph = StateGraph(AgentState)
|
||
graph.add_node("execute_task", execute_task)
|
||
graph.add_node("verify_result", verify_result)
|
||
graph.add_node("handle_error", handle_error)
|
||
|
||
# Start state
|
||
graph.set_entry_point("execute_task")
|
||
|
||
# Transitions
|
||
graph.add_conditional_edges(
|
||
"verify_result",
|
||
lambda state: (
|
||
"END" if state["status"] == "success" else
|
||
"handle_error" if state["attempts"] < state["max_attempts"] else
|
||
"END"
|
||
),
|
||
)
|
||
|
||
graph.add_edge("handle_error", "execute_task")
|
||
|
||
# Final states
|
||
graph.add_edge("execute_task", "verify_result")
|
||
|
||
return graph
|
||
|
||
# ---------- Runner ----------
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="Self‑correcting agent demo")
|
||
parser.add_argument("task", type=str, help="Task to execute")
|
||
parser.add_argument("--max_attempts", type=int, default=3, help="Maximum retry attempts")
|
||
args = parser.parse_args()
|
||
|
||
initial_state: AgentState = {
|
||
"task": args.task,
|
||
"result": "",
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": args.max_attempts,
|
||
}
|
||
|
||
graph = create_agent_graph(max_attempts=args.max_attempts)
|
||
memory = InMemorySaver()
|
||
app = graph.compile(checkpointer=memory)
|
||
|
||
final_state = app.invoke(initial_state)
|
||
print("\nFinal state:")
|
||
for k, v in final_state.items():
|
||
print(f"{k}: {v}") |