Update agent.py

agent.py

@@ -1,30 +1,33 @@
 """
 Self‑correcting LangGraph agent.
 
-The agent receives a natural‑language task, executes it via an unreliable tool,
+The agent:
-then asks an LLM to judge whether the result is correct. If the judge says
+1. Takes a user task.
-"failed" the agent retries until success or a maximum number of attempts.
+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 implementation uses LangGraph 1.x and LangChain 1.x.
+Run with:
+    python agent.py
+
+Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`.
 """
 
 from __future__ import annotations
 
+import os
 import random
-import sys
+import time
-from typing import TypedDict
+from typing import Dict, TypedDict
 
-# LangChain imports
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, AIMessage
-
-# LangGraph imports
 from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode
+from langgraph.checkpoint.memory import InMemorySaver
 
 # ---------------------------------------------------------------------------
 # 1. State definition
 # ---------------------------------------------------------------------------
-
 class AgentState(TypedDict):
     task: str
     result: str
@@ -36,132 +39,127 @@ class AgentState(TypedDict):
 # ---------------------------------------------------------------------------
 # 2. Unreliable tool
 # ---------------------------------------------------------------------------
-
 def unreliable_tool(task: str) -> str:
-    """Simulate a tool that fails 30 % of the time.
+    """Simulates a tool that fails 30% of the time.
 
-    The tool simply returns the string ``f"Result of {task}"`` but raises a
+    Args:
-    ``ValueError`` with 30 % probability.
+        task: The task string.
+
+    Returns:
+        A string result (here we just echo the task for demo).
+
+    Raises:
+        ValueError: Simulated failure.
     """
     if random.random() < 0.3:
         raise ValueError("Simulated tool failure")
-    return f"Result of {task}"
+    # Simulate some work
+    time.sleep(0.5)
+    return f"Result for task: {task}"
 
 # ---------------------------------------------------------------------------
-# 3. LLM judge
+# 3. LLM judge node
 # ---------------------------------------------------------------------------
-
-# Create a lightweight LLM instance.  The user must set the OPENAI_API_KEY
-# environment variable or provide a key in the code.
 llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
 
-# The judge prompt asks the model to answer only "success" or "failed".
 JUDGE_PROMPT = (
-    "You are a strict judge. Given the following result of a task, answer only "
+    "You are a judge that evaluates the result of a task. "
-    "one word: success or failed.\n\nResult: {result}\nAnswer:"  # no extra formatting
+    "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]:
-# 4. Node functions
+    """LLM judge that returns status.
-# ---------------------------------------------------------------------------
 
-async def execute_task(state: AgentState) -> AgentState:
+    Returns a dict with key 'status' set to 'success' or 'failed'.
-    """Execute the task using the unreliable tool.
-
-    On success, store the result.  On failure, capture the exception.
     """
+    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(state["task"])
+        result = unreliable_tool(task)
-        state["result"] = result
+        return {"result": result, "error": None}
-        state["error"] = None
+    except Exception as e:
-    except Exception as exc:
+        return {"result": "", "error": str(e)}
-        state["result"] = ""
-        state["error"] = str(exc)
-    return state
-
-async def verify_result(state: AgentState) -> AgentState:
-    """Ask the LLM to judge the result.
-
-    The LLM must return either "success" or "failed".
-    """
-    if state["error"]:
-        # If the tool raised an exception, we consider it a failure.
-        state["status"] = "failed"
-        return state
-
-    # Build the prompt with the result.
-    prompt = JUDGE_PROMPT.format(result=state["result"])
-    messages = [HumanMessage(content=prompt)]
-    ai_msg: AIMessage = await llm.ainvoke(messages)
-    verdict = ai_msg.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.
-
-    If the maximum number of attempts is reached, set status to
-    "max_attempts".
-    """
-    state["attempts"] += 1
-    if state["attempts"] >= state["max_attempts"]:
-        state["status"] = "max_attempts"
-    else:
-        # Reset result and error for the next attempt.
-        state["result"] = ""
-        state["error"] = None
-    return state
 
 # ---------------------------------------------------------------------------
-# 5. Graph construction
+# 6. Build the graph
 # ---------------------------------------------------------------------------
+builder = StateGraph(AgentState)
 
-def build_graph(max_attempts: int = 5) -> StateGraph[AgentState]:
+# Register nodes
-    graph = StateGraph(AgentState)
+builder.add_node("execute_task", execute_task)
-    graph.add_node("execute_task", execute_task)
+builder.add_node("verify_result", verify_result)
-    graph.add_node("verify_result", verify_result)
+builder.add_node("handle_error", handle_error)
-    graph.add_node("handle_error", handle_error)
 
-    # Define the flow: execute → verify → (success → END | failed → handle_error → execute)
+# Define edges
-    graph.add_edge("execute_task", "verify_result")
+builder.set_entry_point("execute_task")
-    graph.add_edge("handle_error", "execute_task")
 
-    # Conditional router based on status after verification.
+# After executing, verify
-    def router(state: AgentState) -> str:
+builder.add_edge("execute_task", "verify_result")
-        return state["status"]
 
-    graph.add_conditional_edges(
+# After verification
+builder.add_conditional_edges(
     "verify_result",
-        router,
+    lambda x: x["status"],
     {
         "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,
     },
 )
 
-    graph.set_entry_point("execute_task")
+# Compile graph
-    return graph
+graph = builder.compile(checkpointer=InMemorySaver())
 
 # ---------------------------------------------------------------------------
-# 6. CLI driver
+# 7. CLI
 # ---------------------------------------------------------------------------
-
+def main() -> None:
-async def main():
+    print("Self‑correcting LangGraph agent demo")
-    if len(sys.argv) > 1:
+    task = input("Enter a task: ")
-        task = " ".join(sys.argv[1:])
-    else:
-        task = input("Введите задачу: ")
-
     max_attempts = 5
-    graph = build_graph(max_attempts)
-    app = graph.compile()
 
-    # Initial state
+    initial_state: AgentState = {
-    state: AgentState = {
         "task": task,
         "result": "",
         "attempts": 0,
@@ -170,28 +168,23 @@ async def main():
         "max_attempts": max_attempts,
     }
 
-    # Run the graph until it ends.
+    # Run the graph
-    async for partial_state in app.stream(state):
+    result = graph.invoke(initial_state)
-        # Print progress when attempts change.
-        if partial_state["attempts"] != state["attempts"]:
-            print(f"Попытка {partial_state['attempts']}:", end=" ")
-            if partial_state["error"]:
-                print(f"Error → {partial_state['error']}")
-            else:
-                print(f"результат {partial_state['result']}")
-        state = partial_state
 
-    # Final status
+    # Extract final status
-    print("\nИтог:")
+    final_status = result["status"]
-    if state["status"] == "success":
+    attempts = result["attempts"]
-        print(f"Успех за {state['attempts']} попыток. Результат: {state['result']}")
+    final_result = result["result"]
-    elif state["status"] == "max_attempts":
+    error = result["error"]
-        print(f"Не удалось за {state['attempts']} попыток. Последняя ошибка: {state['error']}")
+
-    else:
+    print("\n--- Result ---")
-        print(f"Не удалось. Последняя ошибка: {state['error']}")
+    print(f"Status: {final_status}")
+    print(f"Attempts: {attempts}")
+    if error:
+        print(f"Last error: {error}")
+    print(f"Result: {final_result}")
 
 
 if __name__ == "__main__":
-    import asyncio
+    main()
-
+"
-    asyncio.run(main())