Add agent.py

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}")