diff --git a/agent.py b/agent.py index 18af32c..16ee480 100644 --- a/agent.py +++ b/agent.py @@ -1,30 +1,33 @@ """ Self‑correcting LangGraph agent. -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 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 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 -from typing import TypedDict +import time +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 - ``ValueError`` with 30 % probability. + Args: + 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 " - "one word: success or failed.\n\nResult: {result}\nAnswer:" # no extra formatting + "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." ) -# --------------------------------------------------------------------------- -# 4. Node functions -# --------------------------------------------------------------------------- +async def verify_result(state: AgentState) -> Dict[str, str]: + """LLM judge that returns status. -async def execute_task(state: AgentState) -> AgentState: - """Execute the task using the unreliable tool. - - On success, store the result. On failure, capture the exception. + 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. + """ + task = state["task"] 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 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 + result = unreliable_tool(task) + return {"result": result, "error": None} + except Exception as e: + return {"result": "", "error": str(e)} # --------------------------------------------------------------------------- -# 5. Graph construction +# 6. Build the graph # --------------------------------------------------------------------------- +builder = StateGraph(AgentState) -def build_graph(max_attempts: int = 5) -> StateGraph[AgentState]: - graph = StateGraph(AgentState) - graph.add_node("execute_task", execute_task) - graph.add_node("verify_result", verify_result) - graph.add_node("handle_error", handle_error) +# Register nodes +builder.add_node("execute_task", execute_task) +builder.add_node("verify_result", verify_result) +builder.add_node("handle_error", handle_error) - # 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") +# Define edges +builder.set_entry_point("execute_task") - # Conditional router based on status after verification. - def router(state: AgentState) -> str: - return state["status"] +# After executing, verify +builder.add_edge("execute_task", "verify_result") - graph.add_conditional_edges( - "verify_result", - router, - { - "success": END, - "failed": "handle_error", - "max_attempts": END, - }, - ) +# After verification +builder.add_conditional_edges( + "verify_result", + lambda x: x["status"], + { + "success": END, + "failed": "handle_error", + "max_attempts": END, # safety, though not expected here + }, +) - graph.set_entry_point("execute_task") - return graph +# 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()) # --------------------------------------------------------------------------- -# 6. CLI driver +# 7. CLI # --------------------------------------------------------------------------- - -async def main(): - if len(sys.argv) > 1: - task = " ".join(sys.argv[1:]) - else: - task = input("Введите задачу: ") - +def main() -> None: + print("Self‑correcting LangGraph agent demo") + task = input("Enter a task: ") max_attempts = 5 - graph = build_graph(max_attempts) - app = graph.compile() - # Initial state - state: AgentState = { + initial_state: AgentState = { "task": task, "result": "", "attempts": 0, @@ -170,28 +168,23 @@ async def main(): "max_attempts": max_attempts, } - # 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 + # Run the graph + result = graph.invoke(initial_state) - # 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']}") + # 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}") if __name__ == "__main__": - import asyncio - - asyncio.run(main()) + main() +" \ No newline at end of file