132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
"""LangGraph agent with retry until success or max_attempts.
|
|
|
|
This module implements:
|
|
- AgentState TypedDict
|
|
- An unreliable tool that fails 30% of the time
|
|
- Nodes: execute_task, verify_result, handle_error
|
|
- StateGraph with cycle per plan
|
|
- CLI to run a single task
|
|
"""
|
|
|
|
from typing import TypedDict, Literal, Any
|
|
import random
|
|
import os
|
|
import sys
|
|
|
|
# LLM imports
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
|
|
# 1. Состояние графа
|
|
class AgentState(TypedDict):
|
|
task: str
|
|
result: str
|
|
attempts: int
|
|
status: Literal["pending", "success", "failed", "max_attempts"]
|
|
error: str | None
|
|
max_attempts: int
|
|
|
|
# 2. Tool: unreliable_tool
|
|
|
|
def unreliable_tool(input: str) -> str:
|
|
"""Simulate an unreliable tool that fails 30% of the time."""
|
|
if random.random() < 0.3:
|
|
raise ValueError("Simulated tool failure")
|
|
# Example task: compute sum of two numbers in the input string
|
|
try:
|
|
parts = input.split()
|
|
nums = [int(p) for p in parts if p.isdigit()]
|
|
return str(sum(nums))
|
|
except Exception as e:
|
|
raise ValueError(f"Tool error: {e}")
|
|
|
|
# 3. LLM for verification
|
|
llm = ChatOpenAI(
|
|
model="gpt-4o-mini", # placeholder, user can set via env
|
|
temperature=0,
|
|
base_url=os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
)
|
|
|
|
# 4. Node: execute_task
|
|
async def execute_task(state: AgentState) -> AgentState:
|
|
try:
|
|
result = unreliable_tool(state["task"])
|
|
state["result"] = result
|
|
state["error"] = None
|
|
except Exception as e:
|
|
state["result"] = ""
|
|
state["error"] = str(e)
|
|
return state
|
|
|
|
# 5. Node: verify_result
|
|
async def verify_result(state: AgentState) -> AgentState:
|
|
# Use LLM to decide success or failed
|
|
prompt = f"Given the result '{state['result']}', is this a correct answer? Respond with 'success' or 'failed'."
|
|
try:
|
|
resp = await llm.ainvoke(prompt)
|
|
# Extract keyword
|
|
verdict = "failed"
|
|
if "success" in resp.content.lower():
|
|
verdict = "success"
|
|
except Exception as e:
|
|
verdict = "failed"
|
|
state["error"] = f"LLM error: {e}"
|
|
state["status"] = verdict
|
|
return state
|
|
|
|
# 6. Node: handle_error
|
|
async def handle_error(state: AgentState) -> AgentState:
|
|
state["attempts"] += 1
|
|
if state["attempts"] >= state["max_attempts"]:
|
|
state["status"] = "max_attempts"
|
|
else:
|
|
state["status"] = "failed"
|
|
return state
|
|
|
|
# 7. Graph definition
|
|
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")
|
|
builder.add_conditional_edges(
|
|
"verify_result",
|
|
lambda state: state["status"],
|
|
{
|
|
"success": END,
|
|
"failed": "handle_error",
|
|
"max_attempts": END,
|
|
},
|
|
)
|
|
builder.add_edge("handle_error", "execute_task")
|
|
|
|
graph = builder.compile(checkpointer=MemorySaver())
|
|
|
|
# 8. CLI
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Run a single task with retry until success or max attempts.")
|
|
parser.add_argument("task", type=str, help="Task string to feed to unreliable_tool")
|
|
parser.add_argument("--max", type=int, default=5, 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,
|
|
}
|
|
|
|
result = graph.invoke(initial_state)
|
|
|
|
print("\nFinal state:")
|
|
for k, v in result.items():
|
|
print(f"{k}: {v}")
|