177 lines
5.5 KiB
Python
177 lines
5.5 KiB
Python
"""
|
||
Self‑correcting LangGraph agent demo.
|
||
|
||
Run:
|
||
python main.py
|
||
|
||
Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`.
|
||
"""
|
||
|
||
import random
|
||
from typing import TypedDict, Optional
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import StateGraph, END, START
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. State definition
|
||
# ---------------------------------------------------------------------------
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: Optional[str]
|
||
max_attempts: int
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. Unreliable tool
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def unreliable_addition(a: int, b: int) -> str:
|
||
"""Adds two numbers but fails with 30 % probability.
|
||
|
||
Returns the sum as a string.
|
||
"""
|
||
if random.random() < 0.3:
|
||
raise ValueError("Random failure in unreliable tool")
|
||
return str(a + b)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. LLM initialization
|
||
# ---------------------------------------------------------------------------
|
||
# The OpenAI key must be set in the environment.
|
||
llm = ChatOpenAI(temperature=0)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4. Node implementations
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def execute_task(state: AgentState) -> AgentState:
|
||
"""Executes the task using the unreliable tool.
|
||
|
||
The task is expected to be a string of the form "<a> + <b>".
|
||
"""
|
||
# Increment attempt counter at the start of each execution
|
||
state["attempts"] += 1
|
||
|
||
task = state["task"].strip()
|
||
try:
|
||
a_str, b_str = task.split("+")
|
||
a = int(a_str.strip())
|
||
b = int(b_str.strip())
|
||
result = unreliable_addition(a, b)
|
||
state["result"] = result
|
||
state["error"] = None
|
||
except Exception as e:
|
||
state["result"] = ""
|
||
state["error"] = str(e)
|
||
return state
|
||
|
||
def verify_result(state: AgentState) -> AgentState:
|
||
"""LLM checks if the result is correct.
|
||
|
||
The LLM is asked to output only "success" or "failed".
|
||
"""
|
||
task = state["task"].strip()
|
||
result = state["result"].strip()
|
||
# Compute expected result for comparison
|
||
try:
|
||
a_str, b_str = task.split("+")
|
||
expected = str(int(a_str.strip()) + int(b_str.strip()))
|
||
except Exception:
|
||
expected = ""
|
||
|
||
prompt = (
|
||
f"Task: {task}\n"
|
||
f"Result: {result}\n"
|
||
f"Expected: {expected}\n"
|
||
"Is the result correct? Respond with only 'success' or 'failed'."
|
||
)
|
||
response = llm.invoke(prompt).content.strip().lower()
|
||
if response == "success":
|
||
state["status"] = "success"
|
||
else:
|
||
state["status"] = "failed"
|
||
return state
|
||
|
||
def handle_error(state: AgentState) -> AgentState:
|
||
"""Decide whether to retry or stop based on the number of attempts."""
|
||
if state["attempts"] >= state["max_attempts"]:
|
||
state["status"] = "max_attempts"
|
||
else:
|
||
state["status"] = "pending"
|
||
return state
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5. Graph construction
|
||
# ---------------------------------------------------------------------------
|
||
|
||
graph = StateGraph(AgentState)
|
||
|
||
graph.add_node("execute_task", execute_task)
|
||
|
||
graph.add_node("verify_result", verify_result)
|
||
|
||
graph.add_node("handle_error", handle_error)
|
||
|
||
# Define edges
|
||
# Start -> execute_task
|
||
graph.add_edge(START, "execute_task")
|
||
# execute_task -> verify_result
|
||
graph.add_edge("execute_task", "verify_result")
|
||
# verify_result -> END on success, else -> handle_error on failed
|
||
|
||
graph.add_conditional_edges(
|
||
"verify_result",
|
||
lambda state: "END" if state["status"] == "success" else "handle_error" if state["status"] == "failed" else "END",
|
||
)
|
||
# handle_error -> execute_task if pending, else END
|
||
|
||
graph.add_conditional_edges(
|
||
"handle_error",
|
||
lambda state: "execute_task" if state["status"] == "pending" else "END",
|
||
)
|
||
|
||
# Compile the graph
|
||
app = graph.compile()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 6. Demo execution
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main():
|
||
task = "2 + 2"
|
||
initial_state: AgentState = {
|
||
"task": task,
|
||
"result": "",
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": 5,
|
||
}
|
||
|
||
print(f"Task: {task}")
|
||
|
||
# Run the graph
|
||
for event in app.stream(initial_state):
|
||
state = event["state"]
|
||
attempts = state["attempts"]
|
||
if state["status"] == "pending":
|
||
print(f"Попытка {attempts}: выполняется...")
|
||
elif state["status"] == "failed":
|
||
print(f"Попытка {attempts}: ошибка -> {state['error']}")
|
||
elif state["status"] == "success":
|
||
print(f"Попытка {attempts}: результат {state['result']} → verify: success")
|
||
elif state["status"] == "max_attempts":
|
||
print(f"Попытка {attempts}: достигнут лимит попыток → verify: failed")
|
||
|
||
final_state = app.get_state()
|
||
print("\nИтог:")
|
||
if final_state["status"] == "success":
|
||
print(f"success за {final_state['attempts']} попыток")
|
||
else:
|
||
print(f"failed после {final_state['attempts']} попыток")
|
||
|
||
if __name__ == "__main__":
|
||
main() |