add main.py

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))