Self‑correcting LangGraph agent: update main.py

main.py

@@ -1,16 +1,25 @@
 """
-Self‑correcting LangGraph agent demo.
+Self‑correcting LangGraph agent.
 
 Run with:
-    python main.py
+    python main.py "Вычисли 2+2"
 
-Requires:
-    pip install langgraph langchain-openai
+The agent will execute the task, verify the result via an LLM judge, and retry up to max_attempts.
 """
-import random
-from typing import TypedDict, Dict
+import os
+from typing import TypedDict, Dict, Any
+
+from langgraph.graph import StateGraph
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+
+MAX_ATTEMPTS = 5
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    raise RuntimeError("Please set OPENAI_API_KEY environment variable.")
+
+llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
 
-# ---------- State definition ----------
 class AgentState(TypedDict):
     task: str
     result: str | None
@@ -19,97 +28,75 @@ class AgentState(TypedDict):
     error: str | None
     max_attempts: int
 
-# ---------- Unreliable tool ----------
-class UnreliableTool:
-    def run(self, input_: str) -> str:
-        if random.random() < 0.3:
-            raise ValueError("Simulated tool failure")
-        # simple eval for demo purposes
-        try:
-            return str(eval(input_))
-        except Exception as e:
-            raise ValueError(f"Eval error: {e}")
-
-# ---------- LangGraph imports ----------
-from langgraph.graph import StateGraph, END
-
-from langchain_openai import ChatOpenAI
-
-# ---------- Nodes ----------
-def execute_task(state: AgentState) -> AgentState:
-    tool = UnreliableTool()
+async def execute_task(state: AgentState) -> Dict[str, Any]:
     try:
-        result = tool.run(state["task"])
-        state["result"] = result
-        state["error"] = None
+        import random
+        if random.random() < 0.3:
+            raise ValueError("Simulated tool error")
+        result = f"Result for: {state['task']}"
+        return {"result": result, "error": None, "status": "pending"}
     except Exception as e:
-        state["result"] = None
-        state["error"] = str(e)
-    return state
+        return {"result": None, "error": str(e), "status": "failed"}
 
-def verify_result(state: AgentState) -> AgentState:
-    # Ask LLM to judge success or failed based on result and error
-    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+async def verify_result(state: AgentState) -> Dict[str, Any]:
+    if state.get("error"):
+        return {"status": "failed"}
     prompt = (
-        f"Task: {state['task']}\n"
-        f"Result: {state.get('result')}\n"
-        f"Error: {state.get('error')}\n"
-        "Respond with only 'success' or 'failed'."
+        f"You are a judge. The task was: {state['task']}\n"
+        f"The result returned by the tool is: {state['result']}\n"
+        "Determine if this result satisfies the task. Respond with only one word: 'success' or 'failed'."
     )
-    resp = llm.invoke(prompt)
-    verdict = resp.content.strip().lower()
-    if verdict not in ("success", "failed"):
-        verdict = "failed"
-    state["verdict"] = verdict
-    return state
+    response = await llm.agenerate([HumanMessage(content=prompt)])
+    verdict = response.generations[0][0].content.strip().lower()
+    return {"status": verdict if verdict in {"success", "failed"} else "failed"}
 
-def handle_error(state: AgentState) -> AgentState:
-    state["attempts"] += 1
-    if state["attempts"] >= state["max_attempts"]:
-        state["status"] = "max_attempts"
-    else:
-        state["status"] = "pending"
-    return state
+async def handle_error(state: AgentState) -> Dict[str, Any]:
+    new_attempts = state["attempts"] + 1
+    if new_attempts >= state["max_attempts"]:
+        return {"attempts": new_attempts, "status": "max_attempts"}
+    return {"attempts": new_attempts, "status": "pending", "error": None, "result": None}
 
-# ---------- Graph ----------
 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")
-# From verify_result, branch on verdict
 builder.add_conditional_edges(
     "verify_result",
-    lambda state: state.get("verdict"),
-    {
-        "success": END,
-        "failed": "handle_error",
-    },
+    lambda x: x["status"],
+    {"success": "END_success", "failed": "handle_error"},
 )
 builder.add_edge("handle_error", "execute_task")
-
+builder.add_conditional_edges(
+    "handle_error",
+    lambda x: x["status"],
+    {"max_attempts": "END_max"},
+)
+builder.set_finish_nodes(["END_success", "END_max"])
 graph = builder.compile()
 
-# ---------- Demo runner ----------
-if __name__ == "__main__":
+async def run_task(task: str) -> AgentState:
     initial_state: AgentState = {
-        "task": "2+2",
+        "task": task,
         "result": None,
         "attempts": 0,
         "status": "pending",
         "error": None,
-        "max_attempts": 5,
+        "max_attempts": MAX_ATTEMPTS,
     }
-    for attempt in range(1, initial_state["max_attempts"] + 1):
-        print(f"Attempt {attempt}:")
-        result = graph.invoke(initial_state)
-        if result.get("verdict") == "success":
-            print(f"Success: {result['result']}")
-            break
-        else:
-            print(f"Failed, error: {result.get('error')}")
+    return await graph.ainvoke(initial_state)
+
+if __name__ == "__main__":
+    import sys, asyncio
+    if len(sys.argv) < 2:
+        print("Usage: python main.py '<task>'")
+        sys.exit(1)
+    task_text = sys.argv[1]
+    final = asyncio.run(run_task(task_text))
+    print(f"Task: {final['task']}")
+    print(f"Attempts: {final['attempts']}\n")
+    if final.get("status") == "success":
+        print(f"Result: {final['result']}\nStatus: success")
     else:
-        print("Reached max attempts without success.")
-""
+        print(f"Failed after {final['attempts']} attempts. Status: {final.get('status')}\nError: {final.get('error')}")