Update agent.py
This commit is contained in:
@@ -1,20 +1,11 @@
|
||||
"""Self‑correcting LangGraph agent.
|
||||
"""
|
||||
Self‑correcting LangGraph agent.
|
||||
|
||||
The agent receives a *task* string. It executes the task via an
|
||||
`unreliable_tool` that sometimes raises a ValueError. After execution
|
||||
it asks the LLM (OpenAI or Ollama) to judge whether the *result* is
|
||||
correct. If the judge says ``failed`` the agent retries until
|
||||
``max_attempts`` is reached.
|
||||
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 implementation uses LangGraph's low‑level API: a StateGraph
|
||||
with three nodes – execute_task, verify_result, handle_error – and a
|
||||
simple loop.
|
||||
|
||||
Run the script with:
|
||||
|
||||
python agent.py
|
||||
|
||||
It will ask for a task, then show the attempts and final status.
|
||||
The implementation uses LangGraph 1.x and LangChain 1.x.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -23,7 +14,11 @@ import random
|
||||
import sys
|
||||
from typing import TypedDict
|
||||
|
||||
# LangChain imports
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage, AIMessage
|
||||
|
||||
# LangGraph imports
|
||||
from langgraph.graph import StateGraph, END
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -32,164 +27,171 @@ from langgraph.graph import StateGraph, END
|
||||
|
||||
class AgentState(TypedDict):
|
||||
task: str
|
||||
result: str | None
|
||||
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
|
||||
# 2. Unreliable tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def unreliable_tool(task: str) -> str:
|
||||
"""Simulate a tool that sometimes fails.
|
||||
"""Simulate a tool that fails 30 % of the time.
|
||||
|
||||
The function simply returns ``task`` reversed (as a dummy result) but
|
||||
raises a ValueError 30 % of the time.
|
||||
The tool simply returns the string ``f"Result of {task}"`` but raises a
|
||||
``ValueError`` with 30 % probability.
|
||||
"""
|
||||
if random.random() < 0.3:
|
||||
raise ValueError("Simulated tool failure")
|
||||
return task[::-1] # dummy "computation"
|
||||
return f"Result of {task}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. LLM judge – asks for "success" or "failed"
|
||||
# 3. LLM judge
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
llm = ChatOpenAI(temperature=0, model="gpt-4o-mini") # or use Ollama
|
||||
# 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
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Node functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def execute_task(state: AgentState) -> AgentState:
|
||||
"""Execute the task using the unreliable tool.
|
||||
|
||||
On success, store the result. On failure, capture the exception.
|
||||
"""
|
||||
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 whether the result is correct.
|
||||
"""Ask the LLM to judge the result.
|
||||
|
||||
The prompt forces the model to answer only "success" or "failed".
|
||||
The LLM must return either "success" or "failed".
|
||||
"""
|
||||
if state["result"] is None:
|
||||
# Should not happen – guard
|
||||
if state["error"]:
|
||||
# If the tool raised an exception, we consider it a failure.
|
||||
state["status"] = "failed"
|
||||
return state
|
||||
|
||||
prompt = (
|
||||
f"Task: {state['task']}\n"
|
||||
f"Result: {state['result']}\n"
|
||||
"Is this result correct? Respond with only 'success' or 'failed'."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
verdict = response.content.strip().lower()
|
||||
# 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
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Execute task node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def execute_task(state: AgentState) -> AgentState:
|
||||
"""Run the unreliable tool and capture errors."""
|
||||
try:
|
||||
result = unreliable_tool(state["task"])
|
||||
state["result"] = result
|
||||
state["error"] = None
|
||||
except Exception as exc: # catch ValueError
|
||||
state["result"] = None
|
||||
state["error"] = str(exc)
|
||||
return state
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Handle error / retry node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def handle_error(state: AgentState) -> AgentState:
|
||||
"""Increment attempt counter and decide whether to retry."""
|
||||
"""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 next try
|
||||
state["result"] = None
|
||||
# Reset result and error for the next attempt.
|
||||
state["result"] = ""
|
||||
state["error"] = None
|
||||
return state
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. Build the graph
|
||||
# 5. Graph construction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
graph = 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)
|
||||
|
||||
# Add nodes
|
||||
graph.add_node("execute_task", execute_task)
|
||||
graph.add_node("verify_result", verify_result)
|
||||
graph.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
|
||||
# Start → execute_task
|
||||
graph.set_entry_point("execute_task")
|
||||
# Conditional router based on status after verification.
|
||||
def router(state: AgentState) -> str:
|
||||
return state["status"]
|
||||
|
||||
# After execution, go to verification
|
||||
graph.add_edge("execute_task", "verify_result")
|
||||
graph.add_conditional_edges(
|
||||
"verify_result",
|
||||
router,
|
||||
{
|
||||
"success": END,
|
||||
"failed": "handle_error",
|
||||
"max_attempts": END,
|
||||
},
|
||||
)
|
||||
|
||||
# Verification outcomes
|
||||
# success → END
|
||||
# failed → check attempts
|
||||
# max_attempts → END
|
||||
|
||||
# We use a conditional edge on the status field
|
||||
|
||||
def verify_cond(state: AgentState):
|
||||
return state["status"]
|
||||
|
||||
# Map status to next node
|
||||
graph.add_conditional_edges(
|
||||
"verify_result",
|
||||
verify_cond,
|
||||
{
|
||||
"success": END,
|
||||
"failed": "handle_error",
|
||||
"max_attempts": END,
|
||||
},
|
||||
)
|
||||
|
||||
# From handle_error back to execute_task
|
||||
graph.add_edge("handle_error", "execute_task")
|
||||
|
||||
# Compile the graph into a runnable chain
|
||||
agent = graph.compile()
|
||||
graph.set_entry_point("execute_task")
|
||||
return graph
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. CLI entry point
|
||||
# 6. CLI driver
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
print("Self‑correcting LangGraph agent demo")
|
||||
task = input("Enter a task: ")
|
||||
if not task:
|
||||
print("No task provided. Exiting.")
|
||||
sys.exit(0)
|
||||
async def main():
|
||||
if len(sys.argv) > 1:
|
||||
task = " ".join(sys.argv[1:])
|
||||
else:
|
||||
task = input("Введите задачу: ")
|
||||
|
||||
max_attempts = 5
|
||||
graph = build_graph(max_attempts)
|
||||
app = graph.compile()
|
||||
|
||||
# Initial state
|
||||
state: AgentState = {
|
||||
"task": task,
|
||||
"result": None,
|
||||
"attempts": 1,
|
||||
"result": "",
|
||||
"attempts": 0,
|
||||
"status": "pending",
|
||||
"error": None,
|
||||
"max_attempts": 5,
|
||||
"max_attempts": max_attempts,
|
||||
}
|
||||
|
||||
# Run the chain
|
||||
final_state = agent.invoke(state)
|
||||
# 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
|
||||
|
||||
print("\n--- Result ---")
|
||||
print(f"Task: {final_state['task']}")
|
||||
print(f"Attempts: {final_state['attempts']}")
|
||||
print(f"Status: {final_state['status']}")
|
||||
if final_state['result']:
|
||||
print(f"Result: {final_state['result']}")
|
||||
if final_state['error']:
|
||||
print(f"Last error: {final_state['error']}")
|
||||
# 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']}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
"
|
||||
import asyncio
|
||||
|
||||
asyncio.run(main())
|
||||
|
||||
Reference in New Issue
Block a user