124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
"""
|
||
Self‑correcting LangGraph agent.
|
||
|
||
Run with:
|
||
python main.py "Вычисли 2+2"
|
||
|
||
The script will keep retrying until the LLM judge says `success` or the maximum number of attempts is reached.
|
||
"""
|
||
import random
|
||
from typing import TypedDict, Dict
|
||
|
||
# LangGraph imports
|
||
from langgraph.graph import StateGraph
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage, AIMessage
|
||
|
||
# ---------- 1. State definition -------------------------------------------
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None
|
||
max_attempts: int
|
||
|
||
# ---------- 2. Tool -------------------------------------------------------
|
||
class UnreliableTool:
|
||
"""Simulates a tool that fails with ~30% probability."""
|
||
|
||
def __call__(self, input_: str) -> str:
|
||
if random.random() < 0.3:
|
||
raise ValueError("Simulated tool failure")
|
||
# Very simple evaluation: try to compute arithmetic expression
|
||
try:
|
||
return str(eval(input_))
|
||
except Exception as e:
|
||
raise ValueError(f"Evaluation error: {e}")
|
||
|
||
unreliable_tool = UnreliableTool()
|
||
|
||
# ---------- 3. Nodes -----------------------------------------------------
|
||
async def execute_task(state: AgentState) -> Dict[str, str]:
|
||
"""Runs the task using the unreliable tool."""
|
||
try:
|
||
result = unreliable_tool(state["task"])
|
||
return {"result": result, "error": None}
|
||
except Exception as e:
|
||
return {"result": "", "error": str(e)}
|
||
|
||
async def verify_result(state: AgentState) -> Dict[str, str]:
|
||
"""LLM judge that decides success or failed."""
|
||
llm = ChatOpenAI(temperature=0)
|
||
# Ask the model to output only 'success' or 'failed'
|
||
prompt = (
|
||
f"Task: {state['task']}\n"
|
||
f"Result: {state['result']}\n"
|
||
"Is this result correct? Respond with either 'success' or 'failed'."
|
||
)
|
||
response = await llm.ainvoke(HumanMessage(content=prompt))
|
||
verdict = response.content.strip().lower()
|
||
if verdict not in {"success", "failed"}:
|
||
# Fallback: treat as failed
|
||
verdict = "failed"
|
||
return {"status": verdict}
|
||
|
||
async def handle_error(state: AgentState) -> Dict[str, str]:
|
||
"""Increment attempts and prepare for retry."""
|
||
new_attempts = state["attempts"] + 1
|
||
if new_attempts >= state["max_attempts"]:
|
||
return {"status": "max_attempts", "attempts": new_attempts}
|
||
return {"attempts": new_attempts, "status": "pending"}
|
||
|
||
# ---------- 4. 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())
|
||
|
||
# ---------- 5. CLI -------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
import sys
|
||
|
||
if len(sys.argv) < 2:
|
||
print("Usage: python main.py '<task>'")
|
||
sys.exit(1)
|
||
|
||
task_input = sys.argv[1]
|
||
initial_state: AgentState = {
|
||
"task": task_input,
|
||
"result": "",
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": 5,
|
||
}
|
||
|
||
result = graph.invoke(initial_state)
|
||
final_status = result["status"]
|
||
attempts = result.get("attempts", 0) + 1 # include last attempt
|
||
print(f"Задача: {task_input}")
|
||
if final_status == "success":
|
||
print(f"Итог: success за {attempts} попытки{'и' if attempts>1 else ''}")
|
||
elif final_status == "max_attempts":
|
||
print(f"Не удалось достичь успеха после {attempts} попыток.")
|
||
else:
|
||
print("Непредвиденный статус", final_status)
|
||
"""
|