diff --git a/agent.py b/agent.py index e6e7129..18af32c 100644 --- a/agent.py +++ b/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,164 +27,171 @@ 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 # --------------------------------------------------------------------------- -graph = 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) -# Add nodes -graph.add_node("execute_task", execute_task) -graph.add_node("verify_result", verify_result) -graph.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 -# Start → execute_task -graph.set_entry_point("execute_task") + # Conditional router based on status after verification. + def router(state: AgentState) -> str: + return state["status"] -# After execution, go to verification -graph.add_edge("execute_task", "verify_result") + graph.add_conditional_edges( + "verify_result", + router, + { + "success": END, + "failed": "handle_error", + "max_attempts": END, + }, + ) -# Verification outcomes -# success → END -# failed → check attempts -# max_attempts → END - -# We use a conditional edge on the status field - -def verify_cond(state: AgentState): - return state["status"] - -# Map status to next node -graph.add_conditional_edges( - "verify_result", - verify_cond, - { - "success": END, - "failed": "handle_error", - "max_attempts": END, - }, -) - -# 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() -" \ No newline at end of file + import asyncio + + asyncio.run(main())