Update agent.py

agent.py

@@ -1,29 +1,23 @@
 """
 Self‑correcting LangGraph agent.
 
-The agent:
+The agent takes a user task, executes it with an unreliable tool, asks an LLM to judge the result, and retries until success or a maximum number of attempts.
-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.
 
-Run with:
+Usage:
     python agent.py
 
-Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`.
+The agent will prompt for a task and print the outcome.
 """
 
-from __future__ import annotations
-
 import os
 import random
 import time
-from typing import Dict, TypedDict
+from typing import TypedDict, Any
 
-from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, END, START
-from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import ToolNode
 from langgraph.checkpoint.memory import InMemorySaver
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
 
 # ---------------------------------------------------------------------------
 # 1. State definition
@@ -32,158 +26,168 @@ class AgentState(TypedDict):
     task: str
     result: str
     attempts: int
-    status: str  # pending | success | failed | max_attempts
+    status: str  # "pending" | "success" | "failed" | "max_attempts"
     error: str | None
     max_attempts: int
 
 # ---------------------------------------------------------------------------
-# 2. Unreliable tool
+# 2. Unreliable tool – 30% chance of raising ValueError
 # ---------------------------------------------------------------------------
+
 def unreliable_tool(task: str) -> str:
-    """Simulates a tool that fails 30% of the time.
+    """Simulate a tool that fails 30% of the time.
 
     Args:
         task: The task string.
 
     Returns:
-        A string result (here we just echo the task for demo).
+        A fabricated result string.
 
     Raises:
         ValueError: Simulated failure.
     """
     if random.random() < 0.3:
         raise ValueError("Simulated tool failure")
-    # Simulate some work
+    # Simulate some processing time
     time.sleep(0.5)
     return f"Result for task: {task}"
 
 # ---------------------------------------------------------------------------
-# 3. LLM judge node
+# 3. Nodes
 # ---------------------------------------------------------------------------
-llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
 
-JUDGE_PROMPT = (
+def execute_task(state: AgentState) -> AgentState:
-    "You are a judge that evaluates the result of a task. "
+    """Execute the task using the unreliable tool.
-    "Given the task and the result, reply with either 'success' or 'failed'. "
-    "Do not add any other text."
-)
 
-async def verify_result(state: AgentState) -> Dict[str, str]:
+    Updates ``result`` and ``status``.
-    """LLM judge that returns status.
-
-    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(task)
-        return {"result": result, "error": None}
+        state["result"] = result
+        state["status"] = "pending"
+        state["error"] = None
     except Exception as e:
-        return {"result": "", "error": str(e)}
+        state["result"] = ""
+        state["status"] = "failed"
+        state["error"] = str(e)
+    return state
+
+# LLM for judging the result
+llm = ChatOpenAI(temperature=0, model="gpt-4o-mini")
+
+
+def verify_result(state: AgentState) -> AgentState:
+    """Ask the LLM to judge whether the result is correct.
+
+    The LLM must respond with only "success" or "failed".
+    """
+    task = state["task"]
+    result = state["result"]
+    # If tool failed, we skip LLM and mark as failed
+    if state["status"] == "failed":
+        return state
+
+    prompt = (
+        f"You are a judge. Given the task: {task}\n"
+        f"And the result: {result}\n"
+        "Decide if the result is correct. Respond with only "success" or "failed"."
+    )
+    try:
+        msg = llm([HumanMessage(content=prompt)])
+        verdict = msg.content.strip().lower()
+        if verdict.startswith("success"):
+            state["status"] = "success"
+        else:
+            state["status"] = "failed"
+    except Exception as e:
+        state["status"] = "failed"
+        state["error"] = f"LLM error: {e}"
+    return state
+
+
+def handle_error(state: AgentState) -> AgentState:
+    """Increment attempts and prepare for a retry."""
+    state["attempts"] += 1
+    # If we hit max attempts, set status accordingly
+    if state["attempts"] >= state["max_attempts"]:
+        state["status"] = "max_attempts"
+    else:
+        state["status"] = "pending"
+    return state
 
 # ---------------------------------------------------------------------------
-# 6. Build the graph
+# 4. Build the graph
 # ---------------------------------------------------------------------------
-builder = StateGraph(AgentState)
 
-# Register nodes
+graph = StateGraph(AgentState)
-builder.add_node("execute_task", execute_task)
-builder.add_node("verify_result", verify_result)
-builder.add_node("handle_error", handle_error)
 
-# Define edges
+graph.add_node("execute_task", execute_task)
-builder.set_entry_point("execute_task")
 
-# After executing, verify
+graph.add_node("verify_result", verify_result)
-builder.add_edge("execute_task", "verify_result")
+
+graph.add_node("handle_error", handle_error)
+
+# Entry point
+
+graph.set_entry_point("execute_task")
+
+# Conditional transitions
+
+def verify_cond(state: AgentState) -> str:
+    return state["status"]
 
 # After verification
-builder.add_conditional_edges(
+#   success -> END
+#   failed  -> handle_error (if attempts < max)
+#   max_attempts -> END
+
+graph.add_conditional_edges(
     "verify_result",
-    lambda x: x["status"],
+    verify_cond,
     {
         "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,
     },
 )
 
-# Compile graph
+# After error handling, go back to execute_task
-graph = builder.compile(checkpointer=InMemorySaver())
+
+graph.add_edge("handle_error", "execute_task")
+
+# Build the graph
+flow = graph.compile(checkpointer=InMemorySaver())
 
 # ---------------------------------------------------------------------------
-# 7. CLI
+# 5. CLI
 # ---------------------------------------------------------------------------
-def main() -> None:
+
+def main():
     print("Self‑correcting LangGraph agent demo")
-    task = input("Enter a task: ")
+    while True:
-    max_attempts = 5
+        task = input("Enter a task (or 'exit' to quit): ")
-
+        if task.strip().lower() == "exit":
-    initial_state: AgentState = {
+            break
+        # Initialize state
+        state: AgentState = {
             "task": task,
             "result": "",
             "attempts": 0,
             "status": "pending",
             "error": None,
-        "max_attempts": max_attempts,
+            "max_attempts": 5,
         }
-
+        # Run the flow
-    # Run the graph
+        result = flow(state)
-    result = graph.invoke(initial_state)
+        final_state = result["states"][-1]
-
-    # 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"Status: {final_state['status']}")
-    print(f"Attempts: {attempts}")
+        print(f"Attempts: {final_state['attempts']}")
-    if error:
+        print(f"Result: {final_state['result']}")
-        print(f"Last error: {error}")
+        if final_state["error"]:
-    print(f"Result: {final_result}")
+            print(f"Error: {final_state['error']}")
-
+        print("\n")
 
 if __name__ == "__main__":
     main()