add main.py
This commit is contained in:
@@ -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))
|
||||||
Reference in New Issue
Block a user