From 1d27b02481fad6bc1f33246b004747d23b62b5d4 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=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Wed, 3 Jun 2026 12:44:34 +0000 Subject: [PATCH] Add main.py --- main.py | 177 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..421f10b --- /dev/null +++ b/main.py @@ -0,0 +1,177 @@ +""" +Self‑correcting LangGraph agent demo. + +Run: + python main.py + +Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`. +""" + +import random +from typing import TypedDict, Optional + +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, END, START + +# --------------------------------------------------------------------------- +# 1. State definition +# --------------------------------------------------------------------------- +class AgentState(TypedDict): + task: str + result: str + attempts: int + status: str # pending | success | failed | max_attempts + error: Optional[str] + max_attempts: int + +# --------------------------------------------------------------------------- +# 2. Unreliable tool +# --------------------------------------------------------------------------- + +def unreliable_addition(a: int, b: int) -> str: + """Adds two numbers but fails with 30 % probability. + + Returns the sum as a string. + """ + if random.random() < 0.3: + raise ValueError("Random failure in unreliable tool") + return str(a + b) + +# --------------------------------------------------------------------------- +# 3. LLM initialization +# --------------------------------------------------------------------------- +# The OpenAI key must be set in the environment. +llm = ChatOpenAI(temperature=0) + +# --------------------------------------------------------------------------- +# 4. Node implementations +# --------------------------------------------------------------------------- + +def execute_task(state: AgentState) -> AgentState: + """Executes the task using the unreliable tool. + + The task is expected to be a string of the form " + ". + """ + # Increment attempt counter at the start of each execution + state["attempts"] += 1 + + task = state["task"].strip() + try: + a_str, b_str = task.split("+") + a = int(a_str.strip()) + b = int(b_str.strip()) + result = unreliable_addition(a, b) + state["result"] = result + state["error"] = None + except Exception as e: + state["result"] = "" + state["error"] = str(e) + return state + +def verify_result(state: AgentState) -> AgentState: + """LLM checks if the result is correct. + + The LLM is asked to output only "success" or "failed". + """ + task = state["task"].strip() + result = state["result"].strip() + # Compute expected result for comparison + try: + a_str, b_str = task.split("+") + expected = str(int(a_str.strip()) + int(b_str.strip())) + except Exception: + expected = "" + + prompt = ( + f"Task: {task}\n" + f"Result: {result}\n" + f"Expected: {expected}\n" + "Is the result correct? Respond with only 'success' or 'failed'." + ) + response = llm.invoke(prompt).content.strip().lower() + if response == "success": + state["status"] = "success" + else: + state["status"] = "failed" + return state + +def handle_error(state: AgentState) -> AgentState: + """Decide whether to retry or stop based on the number of attempts.""" + if state["attempts"] >= state["max_attempts"]: + state["status"] = "max_attempts" + else: + state["status"] = "pending" + return state + +# --------------------------------------------------------------------------- +# 5. Graph construction +# --------------------------------------------------------------------------- + +graph = StateGraph(AgentState) + +graph.add_node("execute_task", execute_task) + +graph.add_node("verify_result", verify_result) + +graph.add_node("handle_error", handle_error) + +# Define edges +# Start -> execute_task +graph.add_edge(START, "execute_task") +# execute_task -> verify_result +graph.add_edge("execute_task", "verify_result") +# verify_result -> END on success, else -> handle_error on failed + +graph.add_conditional_edges( + "verify_result", + lambda state: "END" if state["status"] == "success" else "handle_error" if state["status"] == "failed" else "END", +) +# handle_error -> execute_task if pending, else END + +graph.add_conditional_edges( + "handle_error", + lambda state: "execute_task" if state["status"] == "pending" else "END", +) + +# Compile the graph +app = graph.compile() + +# --------------------------------------------------------------------------- +# 6. Demo execution +# --------------------------------------------------------------------------- + +def main(): + task = "2 + 2" + initial_state: AgentState = { + "task": task, + "result": "", + "attempts": 0, + "status": "pending", + "error": None, + "max_attempts": 5, + } + + print(f"Task: {task}") + + # Run the graph + for event in app.stream(initial_state): + state = event["state"] + attempts = state["attempts"] + if state["status"] == "pending": + print(f"Попытка {attempts}: выполняется...") + elif state["status"] == "failed": + print(f"Попытка {attempts}: ошибка -> {state['error']}") + elif state["status"] == "success": + print(f"Попытка {attempts}: результат {state['result']} → verify: success") + elif state["status"] == "max_attempts": + print(f"Попытка {attempts}: достигнут лимит попыток → verify: failed") + + final_state = app.get_state() + print("\nИтог:") + if final_state["status"] == "success": + print(f"success за {final_state['attempts']} попыток") + else: + print(f"failed после {final_state['attempts']} попыток") + +if __name__ == "__main__": + main() \ No newline at end of file