Update agent.py
This commit is contained in:
@@ -1,194 +1,32 @@
|
||||
"""
|
||||
Self‑correcting LangGraph agent.
|
||||
# Updated section of create_agent to use single conditional
|
||||
|
||||
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.
|
||||
def create_agent() -> StateGraph:
|
||||
"""Build and return the LangGraph StateGraph."""
|
||||
graph = StateGraph()
|
||||
graph.add_node("execute_task", execute_task)
|
||||
graph.add_node("verify_result", verify_result)
|
||||
graph.add_node("handle_error", handle_error)
|
||||
|
||||
Usage:
|
||||
python agent.py
|
||||
# Entry point
|
||||
graph.set_entry_point("execute_task")
|
||||
|
||||
The agent will prompt for a task and print the outcome.
|
||||
"""
|
||||
# After execution, decide whether to verify or end due to max attempts
|
||||
def check_max(state: AgentState):
|
||||
return "max_attempts" if state["attempts"] >= state["max_attempts"] else "verify_result"
|
||||
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from typing import TypedDict, Any
|
||||
graph.add_conditional_edges("execute_task", check_max, {"max_attempts": END, "verify_result": "verify_result"})
|
||||
|
||||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
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 – 30% chance of raising ValueError
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def unreliable_tool(task: str) -> str:
|
||||
"""Simulate a tool that fails 30% of the time.
|
||||
|
||||
Args:
|
||||
task: The task string.
|
||||
|
||||
Returns:
|
||||
A fabricated result string.
|
||||
|
||||
Raises:
|
||||
ValueError: Simulated failure.
|
||||
"""
|
||||
if random.random() < 0.3:
|
||||
raise ValueError("Simulated tool failure")
|
||||
# Simulate some processing time
|
||||
time.sleep(0.5)
|
||||
return f"Result for task: {task}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Nodes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def execute_task(state: AgentState) -> AgentState:
|
||||
"""Execute the task using the unreliable tool.
|
||||
|
||||
Updates ``result`` and ``status``.
|
||||
"""
|
||||
task = state["task"]
|
||||
try:
|
||||
result = unreliable_tool(task)
|
||||
state["result"] = result
|
||||
state["status"] = "pending"
|
||||
state["error"] = None
|
||||
except Exception as 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
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Build the graph
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
graph = StateGraph(AgentState)
|
||||
|
||||
graph.add_node("execute_task", execute_task)
|
||||
|
||||
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
|
||||
# success -> END
|
||||
# failed -> handle_error (if attempts < max)
|
||||
# max_attempts -> END
|
||||
|
||||
graph.add_conditional_edges(
|
||||
# After verification, either finish or retry
|
||||
graph.add_conditional_edges(
|
||||
"verify_result",
|
||||
verify_cond,
|
||||
lambda x: x["status"],
|
||||
{
|
||||
"success": END,
|
||||
"failed": "handle_error",
|
||||
"max_attempts": END,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# After error handling, go back to execute_task
|
||||
# Retry path
|
||||
graph.add_edge("handle_error", "execute_task")
|
||||
|
||||
graph.add_edge("handle_error", "execute_task")
|
||||
|
||||
# Build the graph
|
||||
flow = graph.compile(checkpointer=InMemorySaver())
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
print("Self‑correcting LangGraph agent demo")
|
||||
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": 5,
|
||||
}
|
||||
# Run the flow
|
||||
result = flow(state)
|
||||
final_state = result["states"][-1]
|
||||
print("\n--- 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()
|
||||
"
|
||||
return graph
|
||||
Reference in New Issue
Block a user