190 lines
5.5 KiB
Python
190 lines
5.5 KiB
Python
"""
|
||
Self‑correcting LangGraph agent.
|
||
|
||
The agent:
|
||
1. Takes a user task.
|
||
2. Executes it via an unreliable tool.
|
||
3. Asks an LLM to judge the result (success / failed).
|
||
4. Retries until success or max_attempts.
|
||
|
||
Run with:
|
||
python agent.py
|
||
|
||
Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import random
|
||
import time
|
||
from typing import Dict, TypedDict
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.prebuilt import ToolNode
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. State definition
|
||
# ---------------------------------------------------------------------------
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None
|
||
max_attempts: int
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. Unreliable tool
|
||
# ---------------------------------------------------------------------------
|
||
def unreliable_tool(task: str) -> str:
|
||
"""Simulates a tool that fails 30% of the time.
|
||
|
||
Args:
|
||
task: The task string.
|
||
|
||
Returns:
|
||
A string result (here we just echo the task for demo).
|
||
|
||
Raises:
|
||
ValueError: Simulated failure.
|
||
"""
|
||
if random.random() < 0.3:
|
||
raise ValueError("Simulated tool failure")
|
||
# Simulate some work
|
||
time.sleep(0.5)
|
||
return f"Result for task: {task}"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. LLM judge node
|
||
# ---------------------------------------------------------------------------
|
||
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
||
|
||
JUDGE_PROMPT = (
|
||
"You are a judge that evaluates the result of a task. "
|
||
"Given the task and the result, reply with either 'success' or 'failed'. "
|
||
"Do not add any other text."
|
||
)
|
||
|
||
async def verify_result(state: AgentState) -> Dict[str, str]:
|
||
"""LLM judge that returns status.
|
||
|
||
Returns a dict with key 'status' set to 'success' or 'failed'.
|
||
"""
|
||
task = state["task"]
|
||
result = state["result"]
|
||
prompt = f"Task: {task}\nResult: {result}\n{JUDGE_PROMPT}"
|
||
response = await llm.ainvoke(prompt)
|
||
status = response.content.strip().lower()
|
||
if status not in {"success", "failed"}:
|
||
# Fallback to failed if LLM is uncertain
|
||
status = "failed"
|
||
return {"status": status}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4. Handle error / retry node
|
||
# ---------------------------------------------------------------------------
|
||
async def handle_error(state: AgentState) -> Dict[str, int | str | None]:
|
||
"""Increment attempts and decide whether to retry or stop.
|
||
|
||
Returns updated attempts and status.
|
||
"""
|
||
attempts = state["attempts"] + 1
|
||
max_attempts = state["max_attempts"]
|
||
status = "max_attempts" if attempts >= max_attempts else "pending"
|
||
return {"attempts": attempts, "status": status}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5. Execute task node
|
||
# ---------------------------------------------------------------------------
|
||
async def execute_task(state: AgentState) -> Dict[str, str | None]:
|
||
"""Runs the unreliable tool and captures result or error.
|
||
|
||
Returns updated result and error.
|
||
"""
|
||
task = state["task"]
|
||
try:
|
||
result = unreliable_tool(task)
|
||
return {"result": result, "error": None}
|
||
except Exception as e:
|
||
return {"result": "", "error": str(e)}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 6. Build the graph
|
||
# ---------------------------------------------------------------------------
|
||
builder = StateGraph(AgentState)
|
||
|
||
# Register nodes
|
||
builder.add_node("execute_task", execute_task)
|
||
builder.add_node("verify_result", verify_result)
|
||
builder.add_node("handle_error", handle_error)
|
||
|
||
# Define edges
|
||
builder.set_entry_point("execute_task")
|
||
|
||
# After executing, verify
|
||
builder.add_edge("execute_task", "verify_result")
|
||
|
||
# After verification
|
||
builder.add_conditional_edges(
|
||
"verify_result",
|
||
lambda x: x["status"],
|
||
{
|
||
"success": END,
|
||
"failed": "handle_error",
|
||
"max_attempts": END, # safety, though not expected here
|
||
},
|
||
)
|
||
|
||
# After error handling, either retry or end
|
||
builder.add_conditional_edges(
|
||
"handle_error",
|
||
lambda x: x["status"],
|
||
{
|
||
"pending": "execute_task",
|
||
"max_attempts": END,
|
||
},
|
||
)
|
||
|
||
# Compile graph
|
||
graph = builder.compile(checkpointer=InMemorySaver())
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 7. CLI
|
||
# ---------------------------------------------------------------------------
|
||
def main() -> None:
|
||
print("Self‑correcting LangGraph agent demo")
|
||
task = input("Enter a task: ")
|
||
max_attempts = 5
|
||
|
||
initial_state: AgentState = {
|
||
"task": task,
|
||
"result": "",
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": max_attempts,
|
||
}
|
||
|
||
# Run the graph
|
||
result = graph.invoke(initial_state)
|
||
|
||
# Extract final status
|
||
final_status = result["status"]
|
||
attempts = result["attempts"]
|
||
final_result = result["result"]
|
||
error = result["error"]
|
||
|
||
print("\n--- Result ---")
|
||
print(f"Status: {final_status}")
|
||
print(f"Attempts: {attempts}")
|
||
if error:
|
||
print(f"Last error: {error}")
|
||
print(f"Result: {final_result}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
" |