Add main.py

main.py

@@ -0,0 +1,177 @@
+"""
+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()