diff --git a/agent.py b/agent.py index 16ee480..006c4f9 100644 --- a/agent.py +++ b/agent.py @@ -1,29 +1,23 @@ """ Self‑correcting LangGraph agent. -The agent: -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. +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. -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 -from langgraph.prebuilt import ToolNode +from langgraph.graph import StateGraph, END, START 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 = ( - "You are a judge that evaluates the result of a task. " - "Given the task and the result, reply with either 'success' or 'failed'. " - "Do not add any other text." -) +def execute_task(state: AgentState) -> AgentState: + """Execute the task using the unreliable tool. -async def verify_result(state: AgentState) -> Dict[str, str]: - """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. + Updates ``result`` and ``status``. """ 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 -builder.add_node("execute_task", execute_task) -builder.add_node("verify_result", verify_result) -builder.add_node("handle_error", handle_error) +graph = StateGraph(AgentState) -# Define edges -builder.set_entry_point("execute_task") +graph.add_node("execute_task", execute_task) -# After executing, verify -builder.add_edge("execute_task", "verify_result") +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 -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 -graph = builder.compile(checkpointer=InMemorySaver()) +# After error handling, go back to execute_task + +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: ") - max_attempts = 5 - - initial_state: AgentState = { - "task": task, - "result": "", - "attempts": 0, - "status": "pending", - "error": None, - "max_attempts": max_attempts, - } - - # Run the graph - result = graph.invoke(initial_state) - - # 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"Attempts: {attempts}") - if error: - print(f"Last error: {error}") - print(f"Result: {final_result}") - + 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()