Update agent.py

agent.py

@@ -1,20 +1,11 @@
-"""Self‑correcting LangGraph agent.
+"""
+Self‑correcting LangGraph agent.
 
-The agent receives a *task* string. It executes the task via an
-`unreliable_tool` that sometimes raises a ValueError. After execution
-it asks the LLM (OpenAI or Ollama) to judge whether the *result* is
-correct. If the judge says ``failed`` the agent retries until
-``max_attempts`` is reached.
+The agent receives a natural‑language task, executes it via an unreliable tool,
+then asks an LLM to judge whether the result is correct. If the judge says
+"failed" the agent retries until success or a maximum number of attempts.
 
-The implementation uses LangGraph's low‑level API: a StateGraph
-with three nodes – execute_task, verify_result, handle_error – and a
-simple loop.
-
-Run the script with:
-
-    python agent.py
-
-It will ask for a task, then show the attempts and final status.
+The implementation uses LangGraph 1.x and LangChain 1.x.
 """
 
 from __future__ import annotations
@@ -23,7 +14,11 @@ import random
 import sys
 from typing import TypedDict
 
+# LangChain imports
 from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+
+# LangGraph imports
 from langgraph.graph import StateGraph, END
 
 # ---------------------------------------------------------------------------
@@ -32,117 +27,115 @@ from langgraph.graph import StateGraph, END
 
 class AgentState(TypedDict):
     task: str
-    result: str | None
+    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
+# 2. Unreliable tool
 # ---------------------------------------------------------------------------
 
 def unreliable_tool(task: str) -> str:
-    """Simulate a tool that sometimes fails.
+    """Simulate a tool that fails 30 % of the time.
 
-    The function simply returns ``task`` reversed (as a dummy result) but
-    raises a ValueError 30 % of the time.
+    The tool simply returns the string ``f"Result of {task}"`` but raises a
+    ``ValueError`` with 30 % probability.
     """
     if random.random() < 0.3:
         raise ValueError("Simulated tool failure")
-    return task[::-1]  # dummy "computation"
+    return f"Result of {task}"
 
 # ---------------------------------------------------------------------------
-# 3. LLM judge – asks for "success" or "failed"
+# 3. LLM judge
 # ---------------------------------------------------------------------------
 
-llm = ChatOpenAI(temperature=0, model="gpt-4o-mini")  # or use Ollama
+# 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 "
+    "one word: success or failed.\n\nResult: {result}\nAnswer:"  # no extra formatting
+)
+
+# ---------------------------------------------------------------------------
+# 4. Node functions
+# ---------------------------------------------------------------------------
+
+async def execute_task(state: AgentState) -> AgentState:
+    """Execute the task using the unreliable tool.
+
+    On success, store the result.  On failure, capture the exception.
+    """
+    try:
+        result = unreliable_tool(state["task"])
+        state["result"] = result
+        state["error"] = None
+    except Exception as exc:
+        state["result"] = ""
+        state["error"] = str(exc)
+    return state
 
 async def verify_result(state: AgentState) -> AgentState:
-    """Ask the LLM whether the result is correct.
+    """Ask the LLM to judge the result.
 
-    The prompt forces the model to answer only "success" or "failed".
+    The LLM must return either "success" or "failed".
     """
-    if state["result"] is None:
-        # Should not happen – guard
+    if state["error"]:
+        # If the tool raised an exception, we consider it a failure.
         state["status"] = "failed"
         return state
 
-    prompt = (
-        f"Task: {state['task']}\n"
-        f"Result: {state['result']}\n"
-        "Is this result correct? Respond with only 'success' or 'failed'."
-    )
-    response = llm.invoke(prompt)
-    verdict = response.content.strip().lower()
+    # 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
 
-# ---------------------------------------------------------------------------
-# 4. Execute task node
-# ---------------------------------------------------------------------------
-
-async def execute_task(state: AgentState) -> AgentState:
-    """Run the unreliable tool and capture errors."""
-    try:
-        result = unreliable_tool(state["task"])
-        state["result"] = result
-        state["error"] = None
-    except Exception as exc:  # catch ValueError
-        state["result"] = None
-        state["error"] = str(exc)
-    return state
-
-# ---------------------------------------------------------------------------
-# 5. Handle error / retry node
-# ---------------------------------------------------------------------------
-
 async def handle_error(state: AgentState) -> AgentState:
-    """Increment attempt counter and decide whether to retry."""
+    """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 next try
-        state["result"] = None
+        # Reset result and error for the next attempt.
+        state["result"] = ""
         state["error"] = None
     return state
 
 # ---------------------------------------------------------------------------
-# 6. Build the graph
+# 5. Graph construction
 # ---------------------------------------------------------------------------
 
+def build_graph(max_attempts: int = 5) -> StateGraph[AgentState]:
     graph = StateGraph(AgentState)
-
-# Add nodes
     graph.add_node("execute_task", execute_task)
     graph.add_node("verify_result", verify_result)
     graph.add_node("handle_error", handle_error)
 
-# Define edges
-# Start → execute_task
-graph.set_entry_point("execute_task")
-
-# After execution, go to verification
+    # Define the flow: execute → verify → (success → END | failed → handle_error → execute)
     graph.add_edge("execute_task", "verify_result")
+    graph.add_edge("handle_error", "execute_task")
 
-# Verification outcomes
-# success → END
-# failed → check attempts
-# max_attempts → END
-
-# We use a conditional edge on the status field
-
-def verify_cond(state: AgentState):
+    # Conditional router based on status after verification.
+    def router(state: AgentState) -> str:
         return state["status"]
 
-# Map status to next node
     graph.add_conditional_edges(
         "verify_result",
-    verify_cond,
+        router,
         {
             "success": END,
             "failed": "handle_error",
@@ -150,46 +143,55 @@ graph.add_conditional_edges(
         },
     )
 
-# From handle_error back to execute_task
-graph.add_edge("handle_error", "execute_task")
-
-# Compile the graph into a runnable chain
-agent = graph.compile()
+    graph.set_entry_point("execute_task")
+    return graph
 
 # ---------------------------------------------------------------------------
-# 7. CLI entry point
+# 6. CLI driver
 # ---------------------------------------------------------------------------
 
-def main() -> None:
-    print("Self‑correcting LangGraph agent demo")
-    task = input("Enter a task: ")
-    if not task:
-        print("No task provided. Exiting.")
-        sys.exit(0)
+async def main():
+    if len(sys.argv) > 1:
+        task = " ".join(sys.argv[1:])
+    else:
+        task = input("Введите задачу: ")
+
+    max_attempts = 5
+    graph = build_graph(max_attempts)
+    app = graph.compile()
 
     # Initial state
     state: AgentState = {
         "task": task,
-        "result": None,
-        "attempts": 1,
+        "result": "",
+        "attempts": 0,
         "status": "pending",
         "error": None,
-        "max_attempts": 5,
+        "max_attempts": max_attempts,
     }
 
-    # Run the chain
-    final_state = agent.invoke(state)
+    # Run the graph until it ends.
+    async for partial_state in app.stream(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
 
-    print("\n--- Result ---")
-    print(f"Task: {final_state['task']}")
-    print(f"Attempts: {final_state['attempts']}")
-    print(f"Status: {final_state['status']}")
-    if final_state['result']:
-        print(f"Result: {final_state['result']}")
-    if final_state['error']:
-        print(f"Last error: {final_state['error']}")
+    # Final status
+    print("\nИтог:")
+    if state["status"] == "success":
+        print(f"Успех за {state['attempts']} попыток. Результат: {state['result']}")
+    elif state["status"] == "max_attempts":
+        print(f"Не удалось за {state['attempts']} попыток. Последняя ошибка: {state['error']}")
+    else:
+        print(f"Не удалось. Последняя ошибка: {state['error']}")
 
 
 if __name__ == "__main__":
-    main()
-"
+    import asyncio
+
+    asyncio.run(main())