From 396ac7c15d2aea70c6f548eb2c4d80448e13278f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Wed, 3 Jun 2026 12:44:28 +0000 Subject: [PATCH] Update agent.py --- agent.py | 210 +++++++------------------------------------------------ 1 file changed, 24 insertions(+), 186 deletions(-) diff --git a/agent.py b/agent.py index 006c4f9..6e37632 100644 --- a/agent.py +++ b/agent.py @@ -1,194 +1,32 @@ -""" -Self‑correcting LangGraph agent. +# Updated section of create_agent to use single conditional -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: - python agent.py + # Entry point + 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 -import random -import time -from typing import TypedDict, Any + graph.add_conditional_edges("execute_task", check_max, {"max_attempts": END, "verify_result": "verify_result"}) -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 -# --------------------------------------------------------------------------- -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"." + # After verification, either finish or retry + graph.add_conditional_edges( + "verify_result", + lambda x: x["status"], + { + "success": END, + "failed": "handle_error", + }, ) - 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 + # Retry path + graph.add_edge("handle_error", "execute_task") -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, - { - "success": END, - "failed": "handle_error", - "max_attempts": END, - }, -) - -# After error handling, go back to execute_task - -graph.add_edge("handle_error", "execute_task") - -# 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() -" \ No newline at end of file + return graph \ No newline at end of file