add main.py

main.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Self‑correcting LangGraph agent.
+
+Run with:
+    python main.py "Вычисли 2+2"
+"""
+import os
+import random
+import sys
+from typing import TypedDict, Dict, Any
+
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import InMemorySaver
+
+# ---------- LLM ----------
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.0,
+)
+
+# ---------- State ----------
+class AgentState(TypedDict):
+    task: str
+    result: str
+    attempts: int
+    status: str  # pending | success | failed | max_attempts
+    error: str | None
+    max_attempts: int
+
+# ---------- Tool ----------
+class UnreliableTool:
+    """Tool that fails ~30% of the time."""
+
+    def __call__(self, input: str) -> str:
+        if random.random() < 0.3:
+            raise ValueError("Simulated tool failure")
+        # Simple arithmetic evaluator for demo purposes
+        try:
+            return str(eval(input))
+        except Exception as e:
+            raise ValueError(f"Evaluation error: {e}")
+
+unreliable_tool = UnreliableTool()
+
+# ---------- Nodes ----------
+async def execute_task(state: AgentState) -> Dict[str, Any]:
+    """Run the task using the unreliable tool."""
+    try:
+        result = unreliable_tool(state["task"])
+        return {"result": result, "error": None, "status": "pending"}
+    except Exception as e:
+        return {"result": "", "error": str(e), "status": "failed"}
+
+async def verify_result(state: AgentState) -> Dict[str, Any]:
+    """LLM judge: success or failed."""
+    prompt = (
+        f"Task: {state['task']}\n"
+        f"Result: {state['result']}\n"
+        f"Error: {state['error']}\n"
+        "Is the result correct? Answer with only 'success' or 'failed'."
+    )
+    response = await llm.ainvoke(prompt)
+    verdict = response.content.strip().lower()
+    if verdict not in {"success", "failed"}:
+        verdict = "failed"
+    return {"status": verdict}
+
+async def handle_error(state: AgentState) -> Dict[str, Any]:
+    """Increment attempts and decide whether to retry."""
+    attempts = state["attempts"] + 1
+    if attempts >= state["max_attempts"]:
+        return {"attempts": attempts, "status": "max_attempts"}
+    return {"attempts": attempts, "status": "pending", "error": None, "result": ""}
+
+# ---------- Graph ----------
+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 x: x["status"],
+    {
+        "success": END,
+        "failed": "handle_error",
+        "max_attempts": END,
+    },
+)
+builder.add_edge("handle_error", "execute_task")
+
+graph = builder.compile(checkpointer=InMemorySaver())
+
+# ---------- Runner ----------
+async def run(task: str, max_attempts: int = 5):
+    initial_state: AgentState = {
+        "task": task,
+        "result": "",
+        "attempts": 0,
+        "status": "pending",
+        "error": None,
+        "max_attempts": max_attempts,
+    }
+    async for event in graph.astream_events(initial_state, version="1"):
+        if "update" in event:
+            state = event["update"]
+            if state["status"] == "pending" and state["attempts"] > 0:
+                print(f"Попытка {state['attempts']}: {state['error'] or state['result']}")
+    final_state = graph.get_state()
+    print("\nИтог:", final_state["status"], f"за {final_state['attempts']} попытки(и)")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python main.py '<task>'")
+        sys.exit(1)
+    task = sys.argv[1]
+    import asyncio
+    asyncio.run(run(task))