205 lines
7.2 KiB
Python
205 lines
7.2 KiB
Python
"""
|
||
Main entry point for the self‑correcting LangGraph agent.
|
||
|
||
The program demonstrates a simple task – evaluating an arithmetic expression –
|
||
and shows how the agent retries until the LLM judge confirms success or the
|
||
maximum number of attempts is reached.
|
||
|
||
Usage:
|
||
python main.py "2+2"
|
||
|
||
Three examples are printed in the README and can be run directly from this file.
|
||
"""
|
||
|
||
import os
|
||
import random
|
||
from typing import TypedDict, Dict, Any
|
||
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage
|
||
from rich.console import Console
|
||
from rich.table import Table
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
MAX_ATTEMPTS = 5
|
||
LLM_MODEL = "openai/gpt-oss-20b:free"
|
||
BASE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1"
|
||
API_KEY_ENV = "JOURNAL_MCP_PAT"
|
||
|
||
console = Console()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LLM instance – BroJS provider
|
||
# ---------------------------------------------------------------------------
|
||
llm = ChatOpenAI(
|
||
model=LLM_MODEL,
|
||
base_url=BASE_URL,
|
||
api_key=os.getenv(API_KEY_ENV),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# State definition
|
||
# ---------------------------------------------------------------------------
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str | None
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None
|
||
max_attempts: int
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tool – unreliable arithmetic evaluator
|
||
# ---------------------------------------------------------------------------
|
||
def unreliable_tool(expr: str) -> str:
|
||
"""Evaluate a simple arithmetic expression.
|
||
|
||
With 30 % probability it raises ValueError to simulate an external
|
||
failure. The function is intentionally minimal – the goal is to
|
||
demonstrate retry logic, not complex parsing.
|
||
"""
|
||
if random.random() < 0.3:
|
||
raise ValueError("Simulated evaluation error")
|
||
try:
|
||
# Safe eval: only arithmetic operators are allowed.
|
||
result = eval(expr, {"__builtins__": None}, {})
|
||
except Exception as exc:
|
||
raise ValueError(f"Invalid expression: {expr}") from exc
|
||
return str(result)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node implementations
|
||
# ---------------------------------------------------------------------------
|
||
async def execute_task(state: AgentState) -> AgentState:
|
||
"""Run the unreliable tool and store its output.
|
||
|
||
The node updates ``result`` and ``error`` fields. If the tool raises an
|
||
exception, the error message is captured and ``status`` is set to
|
||
``failed`` – this will trigger a retry.
|
||
"""
|
||
expr = state["task"]
|
||
try:
|
||
result = unreliable_tool(expr)
|
||
state.update(result=result, error=None, status="pending")
|
||
except Exception as exc: # pragma: no cover – exercised via retries
|
||
state.update(result=None, error=str(exc), status="failed")
|
||
return state
|
||
|
||
async def verify_result(state: AgentState) -> AgentState:
|
||
"""Ask the LLM to judge whether the result is correct.
|
||
|
||
The prompt explicitly asks for a single word answer – ``success`` or
|
||
``failed``. Any other response is treated as failure.
|
||
"""
|
||
if state["result"] is None:
|
||
# No result – treat as failed to trigger retry logic.
|
||
return state
|
||
|
||
prompt = (
|
||
f"You are a judge evaluating the correctness of an arithmetic\n"
|
||
f"expression: {state['task']}\n"
|
||
f"Result produced by the agent: {state['result']}\n"
|
||
"Is this result correct? Respond with only one word: success or failed."
|
||
)
|
||
msg = HumanMessage(content=prompt)
|
||
response = await llm.ainvoke([msg])
|
||
verdict = response.content.strip().lower()
|
||
if verdict == "success":
|
||
state.update(status="success")
|
||
else:
|
||
state.update(status="failed")
|
||
return state
|
||
|
||
async def handle_error(state: AgentState) -> AgentState:
|
||
"""Increment attempt counter and decide whether to retry.
|
||
|
||
If the maximum number of attempts is reached, set status to
|
||
``max_attempts``. Otherwise reset error and result for a fresh run.
|
||
"""
|
||
state["attempts"] += 1
|
||
if state["attempts"] >= state["max_attempts"]:
|
||
state.update(status="max_attempts", error=None, result=None)
|
||
else:
|
||
# Prepare for retry: clear previous result and error.
|
||
state.update(error=None, result=None, status="pending")
|
||
return state
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Graph construction
|
||
# ---------------------------------------------------------------------------
|
||
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: {
|
||
"success": END,
|
||
"failed": "handle_error",
|
||
"max_attempts": END,
|
||
}[x["status"]],
|
||
)
|
||
builder.add_edge("handle_error", "execute_task")
|
||
|
||
graph = builder.compile(checkpointer=MemorySaver())
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI helper – run a single task and print progress table
|
||
# ---------------------------------------------------------------------------
|
||
async def run_single(expr: str) -> None:
|
||
initial_state: AgentState = {
|
||
"task": expr,
|
||
"result": None,
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": MAX_ATTEMPTS,
|
||
}
|
||
|
||
state = initial_state
|
||
table = Table(title=f"Self‑correcting agent – {expr}")
|
||
table.add_column("Attempt", justify="right")
|
||
table.add_column("Result")
|
||
table.add_column("Status")
|
||
table.add_column("Error")
|
||
|
||
while state["status"] not in ("success", "max_attempts"):
|
||
state = await graph.ainvoke(state)
|
||
attempt_no = state["attempts"] + 1 if state["status"] == "failed" else state["attempts"]
|
||
table.add_row(
|
||
str(attempt_no),
|
||
str(state.get("result", "")),
|
||
state["status"],
|
||
state.get("error", "") or "",
|
||
)
|
||
|
||
console.print(table)
|
||
if state["status"] == "success":
|
||
console.print(f"✅ Success in {state['attempts'] + 1} attempt(s). Result: {state['result']}\n")
|
||
else:
|
||
console.print("❌ Max attempts reached without success.\n")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main entry point – run three examples
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__": # pragma: no cover – manual execution only
|
||
import sys
|
||
|
||
if len(sys.argv) > 1:
|
||
expressions = [" ".join(sys.argv[1:])]
|
||
else:
|
||
expressions = ["2+2", "10/3", "5*6-7"]
|
||
|
||
for expr in expressions:
|
||
import asyncio
|
||
|
||
asyncio.run(run_single(expr))
|