diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..1688a33 --- /dev/null +++ b/agent.py @@ -0,0 +1,131 @@ +"""LangGraph agent with retry until success or max_attempts. + +This module implements: +- AgentState TypedDict +- An unreliable tool that fails 30% of the time +- Nodes: execute_task, verify_result, handle_error +- StateGraph with cycle per plan +- CLI to run a single task +""" + +from typing import TypedDict, Literal, Any +import random +import os +import sys + +# LLM imports +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, START, END +from langgraph.checkpoint.memory import MemorySaver + +# 1. Состояние графа +class AgentState(TypedDict): + task: str + result: str + attempts: int + status: Literal["pending", "success", "failed", "max_attempts"] + error: str | None + max_attempts: int + +# 2. Tool: unreliable_tool + +def unreliable_tool(input: str) -> str: + """Simulate an unreliable tool that fails 30% of the time.""" + if random.random() < 0.3: + raise ValueError("Simulated tool failure") + # Example task: compute sum of two numbers in the input string + try: + parts = input.split() + nums = [int(p) for p in parts if p.isdigit()] + return str(sum(nums)) + except Exception as e: + raise ValueError(f"Tool error: {e}") + +# 3. LLM for verification +llm = ChatOpenAI( + model="gpt-4o-mini", # placeholder, user can set via env + temperature=0, + base_url=os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"), + api_key=os.getenv("OPENAI_API_KEY"), +) + +# 4. Node: execute_task +async def execute_task(state: AgentState) -> AgentState: + try: + result = unreliable_tool(state["task"]) + state["result"] = result + state["error"] = None + except Exception as e: + state["result"] = "" + state["error"] = str(e) + return state + +# 5. Node: verify_result +async def verify_result(state: AgentState) -> AgentState: + # Use LLM to decide success or failed + prompt = f"Given the result '{state['result']}', is this a correct answer? Respond with 'success' or 'failed'." + try: + resp = await llm.ainvoke(prompt) + # Extract keyword + verdict = "failed" + if "success" in resp.content.lower(): + verdict = "success" + except Exception as e: + verdict = "failed" + state["error"] = f"LLM error: {e}" + state["status"] = verdict + return state + +# 6. Node: handle_error +async def handle_error(state: AgentState) -> AgentState: + state["attempts"] += 1 + if state["attempts"] >= state["max_attempts"]: + state["status"] = "max_attempts" + else: + state["status"] = "failed" + return state + +# 7. Graph definition +builder = StateGraph(AgentState) +builder.add_node("execute_task", execute_task) +builder.add_node("verify_result", verify_result) +builder.add_node("handle_error", handle_error) + +builder.set_entry_point("execute_task") +builder.add_edge("execute_task", "verify_result") +builder.add_conditional_edges( + "verify_result", + lambda state: state["status"], + { + "success": END, + "failed": "handle_error", + "max_attempts": END, + }, +) +builder.add_edge("handle_error", "execute_task") + +graph = builder.compile(checkpointer=MemorySaver()) + +# 8. CLI +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Run a single task with retry until success or max attempts.") + parser.add_argument("task", type=str, help="Task string to feed to unreliable_tool") + parser.add_argument("--max", type=int, default=5, help="Maximum retry attempts") + args = parser.parse_args() + + initial_state: AgentState = { + "task": args.task, + "result": "", + "attempts": 0, + "status": "pending", + "error": None, + "max_attempts": args.max, + } + + result = graph.invoke(initial_state) + + print("\nFinal state:") + for k, v in result.items(): + print(f"{k}: {v}")