Update agent.py

agent.py

@@ -1,194 +1,32 @@
-"""
+# Updated section of create_agent to use single conditional
-Self‑correcting LangGraph 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.
+def create_agent() -> StateGraph:
+    """Build and return the LangGraph StateGraph."""
+    graph = StateGraph()
+    graph.add_node("execute_task", execute_task)
+    graph.add_node("verify_result", verify_result)
+    graph.add_node("handle_error", handle_error)
 
-Usage:
+    # Entry point
-    python agent.py
+    graph.set_entry_point("execute_task")
 
-The agent will prompt for a task and print the outcome.
+    # After execution, decide whether to verify or end due to max attempts
-"""
+    def check_max(state: AgentState):
+        return "max_attempts" if state["attempts"] >= state["max_attempts"] else "verify_result"
 
-import os
+    graph.add_conditional_edges("execute_task", check_max, {"max_attempts": END, "verify_result": "verify_result"})
-import random
-import time
-from typing import TypedDict, Any
 
-from langgraph.graph import StateGraph, END, START
+    # After verification, either finish or retry
-from langgraph.checkpoint.memory import InMemorySaver
+    graph.add_conditional_edges(
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
-
-# ---------------------------------------------------------------------------
-# 1. State definition
-# ---------------------------------------------------------------------------
-class AgentState(TypedDict):
-    task: str
-    result: str
-    attempts: int
-    status: str  # "pending" | "success" | "failed" | "max_attempts"
-    error: str | None
-    max_attempts: int
-
-# ---------------------------------------------------------------------------
-# 2. Unreliable tool – 30% chance of raising ValueError
-# ---------------------------------------------------------------------------
-
-def unreliable_tool(task: str) -> str:
-    """Simulate a tool that fails 30% of the time.
-
-    Args:
-        task: The task string.
-
-    Returns:
-        A fabricated result string.
-
-    Raises:
-        ValueError: Simulated failure.
-    """
-    if random.random() < 0.3:
-        raise ValueError("Simulated tool failure")
-    # Simulate some processing time
-    time.sleep(0.5)
-    return f"Result for task: {task}"
-
-# ---------------------------------------------------------------------------
-# 3. Nodes
-# ---------------------------------------------------------------------------
-
-def execute_task(state: AgentState) -> AgentState:
-    """Execute the task using the unreliable tool.
-
-    Updates ``result`` and ``status``.
-    """
-    task = state["task"]
-    try:
-        result = unreliable_tool(task)
-        state["result"] = result
-        state["status"] = "pending"
-        state["error"] = None
-    except Exception as 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
-
-# ---------------------------------------------------------------------------
-# 4. Build the graph
-# ---------------------------------------------------------------------------
-
-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
-
-graph.set_entry_point("execute_task")
-
-# Conditional transitions
-
-def verify_cond(state: AgentState) -> str:
-    return state["status"]
-
-# After verification
-#   success -> END
-#   failed  -> handle_error (if attempts < max)
-#   max_attempts -> END
-
-graph.add_conditional_edges(
         "verify_result",
-    verify_cond,
+        lambda x: x["status"],
         {
             "success": END,
             "failed": "handle_error",
-        "max_attempts": END,
         },
-)
+    )
 
-# After error handling, go back to execute_task
+    # Retry path
+    graph.add_edge("handle_error", "execute_task")
 
-graph.add_edge("handle_error", "execute_task")
+    return graph
-
-# Build the graph
-flow = graph.compile(checkpointer=InMemorySaver())
-
-# ---------------------------------------------------------------------------
-# 5. CLI
-# ---------------------------------------------------------------------------
-
-def main():
-    print("Self‑correcting LangGraph agent demo")
-    while True:
-        task = input("Enter a task (or 'exit' to quit): ")
-        if task.strip().lower() == "exit":
-            break
-        # Initialize state
-        state: AgentState = {
-            "task": task,
-            "result": "",
-            "attempts": 0,
-            "status": "pending",
-            "error": None,
-            "max_attempts": 5,
-        }
-        # Run the flow
-        result = flow(state)
-        final_state = result["states"][-1]
-        print("\n--- Result ---")
-        print(f"Status: {final_state['status']}")
-        print(f"Attempts: {final_state['attempts']}")
-        print(f"Result: {final_state['result']}")
-        if final_state["error"]:
-            print(f"Error: {final_state['error']}")
-        print("\n")
-
-if __name__ == "__main__":
-    main()
-"