Update agent.py

agent.py

@@ -1,32 +1,131 @@
-# Updated section of create_agent to use single conditional
+"""
+Self‑correcting LangGraph agent.
 
-def create_agent() -> StateGraph:
+The agent executes a task, verifies the result with an LLM judge, and retries on failure
-    """Build and return the LangGraph StateGraph."""
+until success or a maximum number of attempts is reached.
-    graph = StateGraph()
+"""
+
+from __future__ import annotations
+
+from typing import TypedDict
+
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import InMemorySaver
+
+# ---------- State ----------
+
+class AgentState(TypedDict):
+    task: str
+    result: str
+    attempts: int
+    status: str  # pending | success | failed | max_attempts
+    error: str | None
+    max_attempts: int
+
+# ---------- Unreliable tool ----------
+
+import random
+
+class UnreliableTool:
+    """Simulates a tool that fails 30 % of the time."""
+
+    def __call__(self, input_text: str) -> str:
+        if random.random() < 0.3:
+            raise ValueError("Simulated tool failure")
+        return f"Result for: {input_text}"
+
+unreliable_tool = UnreliableTool()
+
+# ---------- Nodes ----------
+
+async def execute_task(state: AgentState) -> AgentState:
+    """Runs the task using the unreliable tool."""
+    try:
+        result = unreliable_tool(state["task"])
+        state["result"] = result
+        state["error"] = None
+    except Exception as e:
+        state["result"] = ""
+        state["error"] = str(e)
+    return state
+
+async def verify_result(state: AgentState) -> AgentState:
+    """LLM judge that returns 'success' or 'failed' based on the result."""
+    llm = ChatOpenAI(temperature=0)
+    prompt = (
+        """You are a judge. Evaluate the following result for the task:\n"
+        f"Task: {state['task']}\n"
+        f"Result: {state['result']}\n"
+        "If the result is correct, reply with the word 'success'.\n"
+        "If the result is incorrect or missing, reply with the word 'failed'."""
+    )
+    response = llm.invoke(prompt)
+    verdict = response.content.strip().lower()
+    if verdict == "success":
+        state["status"] = "success"
+    else:
+        state["status"] = "failed"
+    return state
+
+async def handle_error(state: AgentState) -> AgentState:
+    """Increment attempts and prepare for retry."""
+    state["attempts"] += 1
+    state["status"] = "pending"
+    return state
+
+# ---------- Graph ----------
+
+def create_agent_graph(max_attempts: int = 3) -> StateGraph:
+    graph = StateGraph(AgentState)
     graph.add_node("execute_task", execute_task)
     graph.add_node("verify_result", verify_result)
     graph.add_node("handle_error", handle_error)
 
-    # Entry point
+    # Start state
     graph.set_entry_point("execute_task")
 
-    # After execution, decide whether to verify or end due to max attempts
+    # Transitions
-    def check_max(state: AgentState):
-        return "max_attempts" if state["attempts"] >= state["max_attempts"] else "verify_result"
-
-    graph.add_conditional_edges("execute_task", check_max, {"max_attempts": END, "verify_result": "verify_result"})
-
-    # After verification, either finish or retry
     graph.add_conditional_edges(
         "verify_result",
-        lambda x: x["status"],
+        lambda state: (
-        {
+            "END" if state["status"] == "success" else
-            "success": END,
+            "handle_error" if state["attempts"] < state["max_attempts"] else
-            "failed": "handle_error",
+            "END"
-        },
+        ),
     )
 
-    # Retry path
     graph.add_edge("handle_error", "execute_task")
 
-    return graph
+    # Final states
+    graph.add_edge("execute_task", "verify_result")
+
+    return graph
+
+# ---------- Runner ----------
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Self‑correcting agent demo")
+    parser.add_argument("task", type=str, help="Task to execute")
+    parser.add_argument("--max_attempts", type=int, default=3, help="Maximum retry attempts")
+    args = parser.parse_args()
+
+    initial_state: AgentState = {
+        "task": args.task,
+        "result": "",
+        "attempts": 0,
+        "status": "pending",
+        "error": None,
+        "max_attempts": args.max_attempts,
+    }
+
+    graph = create_agent_graph(max_attempts=args.max_attempts)
+    memory = InMemorySaver()
+    app = graph.compile(checkpointer=memory)
+
+    final_state = app.invoke(initial_state)
+    print("\nFinal state:")
+    for k, v in final_state.items():
+        print(f"{k}: {v}")