From a9f9f88b59581f0db28bace8e1833a3c8b1c9f6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 28 May 2026 18:13:27 +0000 Subject: [PATCH] add main.py --- main.py | 204 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 204 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..3b80a7b --- /dev/null +++ b/main.py @@ -0,0 +1,204 @@ +""" +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))