Update agent.py
This commit is contained in:
@@ -1,29 +1,23 @@
|
||||
"""
|
||||
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.
|
||||
The agent takes a user task, executes it with an unreliable tool, asks an LLM to judge the result, and retries until success or a maximum number of attempts.
|
||||
|
||||
Run with:
|
||||
Usage:
|
||||
python agent.py
|
||||
|
||||
Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`.
|
||||
The agent will prompt for a task and print the outcome.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from typing import Dict, TypedDict
|
||||
from typing import TypedDict, Any
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.graph import StateGraph, END
|
||||
from langgraph.prebuilt import ToolNode
|
||||
from langgraph.graph import StateGraph, END, START
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. State definition
|
||||
@@ -32,158 +26,168 @@ class AgentState(TypedDict):
|
||||
task: str
|
||||
result: str
|
||||
attempts: int
|
||||
status: str # pending | success | failed | max_attempts
|
||||
status: str # "pending" | "success" | "failed" | "max_attempts"
|
||||
error: str | None
|
||||
max_attempts: int
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Unreliable tool
|
||||
# 2. Unreliable tool – 30% chance of raising ValueError
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def unreliable_tool(task: str) -> str:
|
||||
"""Simulates a tool that fails 30% of the time.
|
||||
"""Simulate 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).
|
||||
A fabricated result string.
|
||||
|
||||
Raises:
|
||||
ValueError: Simulated failure.
|
||||
"""
|
||||
if random.random() < 0.3:
|
||||
raise ValueError("Simulated tool failure")
|
||||
# Simulate some work
|
||||
# Simulate some processing time
|
||||
time.sleep(0.5)
|
||||
return f"Result for task: {task}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. LLM judge node
|
||||
# 3. Nodes
|
||||
# ---------------------------------------------------------------------------
|
||||
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."
|
||||
)
|
||||
def execute_task(state: AgentState) -> AgentState:
|
||||
"""Execute the task using the unreliable tool.
|
||||
|
||||
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.
|
||||
Updates ``result`` and ``status``.
|
||||
"""
|
||||
task = state["task"]
|
||||
try:
|
||||
result = unreliable_tool(task)
|
||||
return {"result": result, "error": None}
|
||||
state["result"] = result
|
||||
state["status"] = "pending"
|
||||
state["error"] = None
|
||||
except Exception as e:
|
||||
return {"result": "", "error": str(e)}
|
||||
state["result"] = ""
|
||||
state["status"] = "failed"
|
||||
state["error"] = str(e)
|
||||
return state
|
||||
|
||||
# LLM for judging the result
|
||||
llm = ChatOpenAI(temperature=0, model="gpt-4o-mini")
|
||||
|
||||
|
||||
def verify_result(state: AgentState) -> AgentState:
|
||||
"""Ask the LLM to judge whether the result is correct.
|
||||
|
||||
The LLM must respond with only "success" or "failed".
|
||||
"""
|
||||
task = state["task"]
|
||||
result = state["result"]
|
||||
# If tool failed, we skip LLM and mark as failed
|
||||
if state["status"] == "failed":
|
||||
return state
|
||||
|
||||
prompt = (
|
||||
f"You are a judge. Given the task: {task}\n"
|
||||
f"And the result: {result}\n"
|
||||
"Decide if the result is correct. Respond with only "success" or "failed"."
|
||||
)
|
||||
try:
|
||||
msg = llm([HumanMessage(content=prompt)])
|
||||
verdict = msg.content.strip().lower()
|
||||
if verdict.startswith("success"):
|
||||
state["status"] = "success"
|
||||
else:
|
||||
state["status"] = "failed"
|
||||
except Exception as e:
|
||||
state["status"] = "failed"
|
||||
state["error"] = f"LLM error: {e}"
|
||||
return state
|
||||
|
||||
|
||||
def handle_error(state: AgentState) -> AgentState:
|
||||
"""Increment attempts and prepare for a retry."""
|
||||
state["attempts"] += 1
|
||||
# If we hit max attempts, set status accordingly
|
||||
if state["attempts"] >= state["max_attempts"]:
|
||||
state["status"] = "max_attempts"
|
||||
else:
|
||||
state["status"] = "pending"
|
||||
return state
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. Build the graph
|
||||
# 4. 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)
|
||||
graph = StateGraph(AgentState)
|
||||
|
||||
# Define edges
|
||||
builder.set_entry_point("execute_task")
|
||||
graph.add_node("execute_task", execute_task)
|
||||
|
||||
# After executing, verify
|
||||
builder.add_edge("execute_task", "verify_result")
|
||||
graph.add_node("verify_result", verify_result)
|
||||
|
||||
graph.add_node("handle_error", handle_error)
|
||||
|
||||
# Entry point
|
||||
|
||||
graph.set_entry_point("execute_task")
|
||||
|
||||
# Conditional transitions
|
||||
|
||||
def verify_cond(state: AgentState) -> str:
|
||||
return state["status"]
|
||||
|
||||
# After verification
|
||||
builder.add_conditional_edges(
|
||||
# success -> END
|
||||
# failed -> handle_error (if attempts < max)
|
||||
# max_attempts -> END
|
||||
|
||||
graph.add_conditional_edges(
|
||||
"verify_result",
|
||||
lambda x: x["status"],
|
||||
verify_cond,
|
||||
{
|
||||
"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())
|
||||
# After error handling, go back to execute_task
|
||||
|
||||
graph.add_edge("handle_error", "execute_task")
|
||||
|
||||
# Build the graph
|
||||
flow = graph.compile(checkpointer=InMemorySaver())
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. CLI
|
||||
# 5. CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def main() -> None:
|
||||
|
||||
def main():
|
||||
print("Self‑correcting LangGraph agent demo")
|
||||
task = input("Enter a task: ")
|
||||
max_attempts = 5
|
||||
|
||||
initial_state: AgentState = {
|
||||
while True:
|
||||
task = input("Enter a task (or 'exit' to quit): ")
|
||||
if task.strip().lower() == "exit":
|
||||
break
|
||||
# Initialize state
|
||||
state: AgentState = {
|
||||
"task": task,
|
||||
"result": "",
|
||||
"attempts": 0,
|
||||
"status": "pending",
|
||||
"error": None,
|
||||
"max_attempts": max_attempts,
|
||||
"max_attempts": 5,
|
||||
}
|
||||
|
||||
# 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"]
|
||||
|
||||
# Run the flow
|
||||
result = flow(state)
|
||||
final_state = result["states"][-1]
|
||||
print("\n--- Result ---")
|
||||
print(f"Status: {final_status}")
|
||||
print(f"Attempts: {attempts}")
|
||||
if error:
|
||||
print(f"Last error: {error}")
|
||||
print(f"Result: {final_result}")
|
||||
|
||||
print(f"Status: {final_state['status']}")
|
||||
print(f"Attempts: {final_state['attempts']}")
|
||||
print(f"Result: {final_state['result']}")
|
||||
if final_state["error"]:
|
||||
print(f"Error: {final_state['error']}")
|
||||
print("\n")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user