Update agent.py
This commit is contained in:
@@ -1,30 +1,33 @@
|
||||
"""
|
||||
Self‑correcting LangGraph agent.
|
||||
|
||||
The agent receives a natural‑language task, executes it via an unreliable tool,
|
||||
then asks an LLM to judge whether the result is correct. If the judge says
|
||||
"failed" the agent retries until success or a maximum number of attempts.
|
||||
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 implementation uses LangGraph 1.x and LangChain 1.x.
|
||||
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 sys
|
||||
from typing import TypedDict
|
||||
import time
|
||||
from typing import Dict, TypedDict
|
||||
|
||||
# LangChain imports
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage, AIMessage
|
||||
|
||||
# LangGraph imports
|
||||
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
|
||||
@@ -36,132 +39,127 @@ class AgentState(TypedDict):
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Unreliable tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def unreliable_tool(task: str) -> str:
|
||||
"""Simulate a tool that fails 30 % of the time.
|
||||
"""Simulates a tool that fails 30% of the time.
|
||||
|
||||
The tool simply returns the string ``f"Result of {task}"`` but raises a
|
||||
``ValueError`` with 30 % probability.
|
||||
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")
|
||||
return f"Result of {task}"
|
||||
# Simulate some work
|
||||
time.sleep(0.5)
|
||||
return f"Result for task: {task}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. LLM judge
|
||||
# 3. LLM judge node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Create a lightweight LLM instance. The user must set the OPENAI_API_KEY
|
||||
# environment variable or provide a key in the code.
|
||||
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
||||
|
||||
# The judge prompt asks the model to answer only "success" or "failed".
|
||||
JUDGE_PROMPT = (
|
||||
"You are a strict judge. Given the following result of a task, answer only "
|
||||
"one word: success or failed.\n\nResult: {result}\nAnswer:" # no extra formatting
|
||||
"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."
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Node functions
|
||||
# ---------------------------------------------------------------------------
|
||||
async def verify_result(state: AgentState) -> Dict[str, str]:
|
||||
"""LLM judge that returns status.
|
||||
|
||||
async def execute_task(state: AgentState) -> AgentState:
|
||||
"""Execute the task using the unreliable tool.
|
||||
|
||||
On success, store the result. On failure, capture the exception.
|
||||
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(state["task"])
|
||||
state["result"] = result
|
||||
state["error"] = None
|
||||
except Exception as exc:
|
||||
state["result"] = ""
|
||||
state["error"] = str(exc)
|
||||
return state
|
||||
|
||||
async def verify_result(state: AgentState) -> AgentState:
|
||||
"""Ask the LLM to judge the result.
|
||||
|
||||
The LLM must return either "success" or "failed".
|
||||
"""
|
||||
if state["error"]:
|
||||
# If the tool raised an exception, we consider it a failure.
|
||||
state["status"] = "failed"
|
||||
return state
|
||||
|
||||
# Build the prompt with the result.
|
||||
prompt = JUDGE_PROMPT.format(result=state["result"])
|
||||
messages = [HumanMessage(content=prompt)]
|
||||
ai_msg: AIMessage = await llm.ainvoke(messages)
|
||||
verdict = ai_msg.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.
|
||||
|
||||
If the maximum number of attempts is reached, set status to
|
||||
"max_attempts".
|
||||
"""
|
||||
state["attempts"] += 1
|
||||
if state["attempts"] >= state["max_attempts"]:
|
||||
state["status"] = "max_attempts"
|
||||
else:
|
||||
# Reset result and error for the next attempt.
|
||||
state["result"] = ""
|
||||
state["error"] = None
|
||||
return state
|
||||
result = unreliable_tool(task)
|
||||
return {"result": result, "error": None}
|
||||
except Exception as e:
|
||||
return {"result": "", "error": str(e)}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Graph construction
|
||||
# 6. Build the graph
|
||||
# ---------------------------------------------------------------------------
|
||||
builder = StateGraph(AgentState)
|
||||
|
||||
def build_graph(max_attempts: int = 5) -> StateGraph[AgentState]:
|
||||
graph = StateGraph(AgentState)
|
||||
graph.add_node("execute_task", execute_task)
|
||||
graph.add_node("verify_result", verify_result)
|
||||
graph.add_node("handle_error", handle_error)
|
||||
# Register nodes
|
||||
builder.add_node("execute_task", execute_task)
|
||||
builder.add_node("verify_result", verify_result)
|
||||
builder.add_node("handle_error", handle_error)
|
||||
|
||||
# Define the flow: execute → verify → (success → END | failed → handle_error → execute)
|
||||
graph.add_edge("execute_task", "verify_result")
|
||||
graph.add_edge("handle_error", "execute_task")
|
||||
# Define edges
|
||||
builder.set_entry_point("execute_task")
|
||||
|
||||
# Conditional router based on status after verification.
|
||||
def router(state: AgentState) -> str:
|
||||
return state["status"]
|
||||
# After executing, verify
|
||||
builder.add_edge("execute_task", "verify_result")
|
||||
|
||||
graph.add_conditional_edges(
|
||||
# After verification
|
||||
builder.add_conditional_edges(
|
||||
"verify_result",
|
||||
router,
|
||||
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,
|
||||
},
|
||||
)
|
||||
|
||||
graph.set_entry_point("execute_task")
|
||||
return graph
|
||||
# Compile graph
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. CLI driver
|
||||
# 7. CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def main():
|
||||
if len(sys.argv) > 1:
|
||||
task = " ".join(sys.argv[1:])
|
||||
else:
|
||||
task = input("Введите задачу: ")
|
||||
|
||||
def main() -> None:
|
||||
print("Self‑correcting LangGraph agent demo")
|
||||
task = input("Enter a task: ")
|
||||
max_attempts = 5
|
||||
graph = build_graph(max_attempts)
|
||||
app = graph.compile()
|
||||
|
||||
# Initial state
|
||||
state: AgentState = {
|
||||
initial_state: AgentState = {
|
||||
"task": task,
|
||||
"result": "",
|
||||
"attempts": 0,
|
||||
@@ -170,28 +168,23 @@ async def main():
|
||||
"max_attempts": max_attempts,
|
||||
}
|
||||
|
||||
# Run the graph until it ends.
|
||||
async for partial_state in app.stream(state):
|
||||
# Print progress when attempts change.
|
||||
if partial_state["attempts"] != state["attempts"]:
|
||||
print(f"Попытка {partial_state['attempts']}:", end=" ")
|
||||
if partial_state["error"]:
|
||||
print(f"Error → {partial_state['error']}")
|
||||
else:
|
||||
print(f"результат {partial_state['result']}")
|
||||
state = partial_state
|
||||
# Run the graph
|
||||
result = graph.invoke(initial_state)
|
||||
|
||||
# Final status
|
||||
print("\nИтог:")
|
||||
if state["status"] == "success":
|
||||
print(f"Успех за {state['attempts']} попыток. Результат: {state['result']}")
|
||||
elif state["status"] == "max_attempts":
|
||||
print(f"Не удалось за {state['attempts']} попыток. Последняя ошибка: {state['error']}")
|
||||
else:
|
||||
print(f"Не удалось. Последняя ошибка: {state['error']}")
|
||||
# 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__":
|
||||
import asyncio
|
||||
|
||||
asyncio.run(main())
|
||||
main()
|
||||
"
|
||||
Reference in New Issue
Block a user