Экзамен: Самокорректирующийся агент: main.py

main.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 """
-main.py – запуск LangGraph‑агента с самопроверкой и retry.
+main.py – запуск LangGraph‑агента с самопроверкой и повторными попытками.
 
 Требования:
     - Python 3.10+
@@ -10,7 +10,7 @@ main.py – запуск LangGraph‑агента с самопроверкой
 """
 
 import random
-from typing import TypedDict, Dict, Any
+from typing import TypedDict, Dict
 
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
@@ -22,59 +22,47 @@ from langgraph.prebuilt import InMemorySaver
 # --------------------------------------------------------------------------- #
 
 class AgentState(TypedDict):
-    task: str                # исходная задача
-    result: str              # результат выполнения инструмента
-    attempts: int            # количество попыток
-    status: str              # pending | success | failed | max_attempts
-    error: str | None        # сообщение об ошибке (если есть)
-    max_attempts: int        # лимит попыток
+    task: str          # исходная задача
+    result: str        # результат выполнения инструмента
+    attempts: int      # число попыток
+    status: str        # pending | success | failed | max_attempts
+    error: str | None  # сообщение об ошибке, если есть
+    max_attempts: int  # лимит повторов
 
 
 # --------------------------------------------------------------------------- #
-# 2. Инструмент – unreliable_tool
+# 2. Инструмент – «не надёжный» калькулятор
 # --------------------------------------------------------------------------- #
 
-def unreliable_tool(input_text: str) -> str:
+def unreliable_tool(input_str: str) -> str:
     """
     Случайно бросает ValueError с вероятностью ~30%.
-    Иначе возвращает результат выполнения простого арифметического выражения.
+    Иначе возвращает результат простого арифметического выражения.
     """
     if random.random() < 0.3:
-        raise ValueError("Инструмент случайно упал")
-    # Простейший парсер: ищем числа и оператор +, -, *, /
+        raise ValueError("Инструмент временно недоступен")
     try:
-        left, op, right = input_text.split()
-        left_val = float(left)
-        right_val = float(right)
-        if op == '+':
-            return str(left_val + right_val)
-        elif op == '-':
-            return str(left_val - right_val)
-        elif op == '*':
-            return str(left_val * right_val)
-        elif op == '/':
-            return str(left_val / right_val)
-    except Exception:
-        pass
-    # Если не удалось распарсить – просто возвращаем строку
-    return f"результат: {input_text}"
+        # eval безопасный для простых выражений (здесь только числа и +)
+        return str(eval(input_str))
+    except Exception as exc:
+        raise ValueError(f"Невозможно вычислить: {exc}") from exc
 
 
 # --------------------------------------------------------------------------- #
-# 3. Узлы графа
+# 3. Узлы
 # --------------------------------------------------------------------------- #
 
-def execute_task(state: AgentState) -> Dict[str, Any]:
+def execute_task(state: AgentState) -> Dict[str, object]:
     """
     Выполняет задачу через unreliable_tool.
-    При ошибке сохраняется сообщение об ошибке и статус «failed».
+    При ошибке сохраняет сообщение об ошибке и помечает статус как failed.
     """
     try:
         result = unreliable_tool(state["task"])
         return {
             "result": result,
             "error": None,
-            "status": "pending",  # переход к verify_result
+            "status": "pending",  # результат получен, но ещё не проверён
         }
     except Exception as exc:
         return {
@@ -84,41 +72,39 @@ def execute_task(state: AgentState) -> Dict[str, Any]:
         }
 
 
-def verify_result(state: AgentState) -> Dict[str, Any]:
+def verify_result(state: AgentState) -> Dict[str, object]:
     """
-    LLM‑оценка результата. Модель должна ответить «success» или «failed».
-    Если результат пустой – сразу считается «failed».
+    LLM‑проверка результата. Модель должна вернуть ровно 'success' или 'failed'.
     """
-    if not state["result"]:
-        return {"status": "failed"}
-
     llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
     prompt = (
         f"Задача: {state['task']}\n"
-        f"Результат: {state['result']}\n\n"
-        "Оцените результат. Ответьте только одним словом: success или failed."
+        f"Результат: {state['result']}\n"
+        "Оцените результат. Ответьте только одним словом: 'success' или 'failed'."
     )
     response = llm.invoke(prompt).content.strip().lower()
-    if response.startswith("success"):
-        return {"status": "success"}
-    else:
-        return {"status": "failed"}
+    if response not in {"success", "failed"}:
+        # Если модель дала непонятный ответ – считаем это ошибкой
+        return {
+            "status": "failed",
+            "error": f"LLM дал неожиданный ответ: {response}",
+        }
+    return {"status": response}
 
 
-def handle_error(state: AgentState) -> Dict[str, Any]:
+def handle_error(state: AgentState) -> Dict[str, object]:
     """
-    Увеличивает счётчик попыток и готовит повторный запуск.
-    Если достигнут лимит – статус «max_attempts».
+    Увеличиваем счётчик попыток и готовим к повтору.
+    Если лимит превышен – помечаем как max_attempts.
     """
     attempts = state["attempts"] + 1
     if attempts >= state["max_attempts"]:
-        return {"attempts": attempts, "status": "max_attempts"}
-    else:
-        return {"attempts": attempts, "status": "pending"}
+        return {"status": "max_attempts", "attempts": attempts}
+    return {"status": "pending", "attempts": attempts}
 
 
 # --------------------------------------------------------------------------- #
-# 4. Построение графа
+# 4. Создание графа
 # --------------------------------------------------------------------------- #
 
 builder = StateGraph(AgentState)
@@ -131,34 +117,22 @@ builder.add_node("handle_error", handle_error)
 # Определяем переходы
 builder.set_entry_point("execute_task")
 
-builder.add_conditional_edges(
-    "execute_task",
-    lambda x: x["status"],
-    {
-        "pending": "verify_result",
-        "failed": "handle_error",
-    },
-)
-
+builder.add_edge("execute_task", "verify_result")
 builder.add_conditional_edges(
     "verify_result",
     lambda x: x["status"],
     {
         "success": END,
         "failed": "handle_error",
-    },
-)
-
-builder.add_conditional_edges(
-    "handle_error",
-    lambda x: x["status"],
-    {
-        "pending": "execute_task",
         "max_attempts": END,
     },
 )
 
-graph = builder.compile(saver=InMemorySaver())
+builder.add_edge("handle_error", "execute_task")
+
+# Сохраняем состояние в памяти (для отладки)
+memory = InMemorySaver()
+graph = builder.compile(persist_to_db=memory)
 
 
 # --------------------------------------------------------------------------- #
@@ -166,18 +140,11 @@ graph = builder.compile(saver=InMemorySaver())
 # --------------------------------------------------------------------------- #
 
 def main() -> None:
-    """
-    CLI‑интерфейс: пользователь вводит задачу, например «2 + 2».
-    Выводятся номера попыток и финальный статус.
-    """
-    task = input("Введите задачу (пример: 2 + 2): ").strip()
-    if not task:
-        print("Задача не может быть пустой.")
-        return
+    # Пример задачи: простое арифметическое выражение
+    task_text = "2 + 2"
 
-    # Инициализируем состояние
     initial_state: AgentState = {
-        "task": task,
+        "task": task_text,
         "result": "",
         "attempts": 0,
         "status": "pending",
@@ -185,28 +152,25 @@ def main() -> None:
         "max_attempts": 5,
     }
 
-    print("\nЗапуск агента...\n")
-
     # Запускаем граф
-    for step in graph.stream(initial_state, stream_mode="values"):
-        state = step.values
-        attempts = state["attempts"] + 1 if state["status"] != "pending" else state["attempts"]
-        print(f"Попытка {attempts}:")
-        if state.get("error"):
-            print(f"  Ошибка инструмента: {state['error']}")
-        else:
-            print(f"  Результат: {state['result']}")
+    final_state = graph.invoke(initial_state)
 
-        if state["status"] == "success":
-            print("\nИтог: success за", attempts, "попыток")
-            break
-        elif state["status"] == "max_attempts":
-            print("\nИтог: max_attempts – задача не решена после 5 попыток")
-            break
-
-    else:
-        # Если цикл завершился без явного выхода (должно быть)
-        print("\nНеизвестный конец выполнения.")
+    # Выводим результаты
+    print(f"Задача: {task_text}")
+    for i in range(final_state["attempts"] + 1):
+        state_snapshot = memory.get(i)
+        if state_snapshot is None:
+            continue
+        attempt_num = state_snapshot["attempts"]
+        status = state_snapshot["status"]
+        error_msg = state_snapshot.get("error")
+        result = state_snapshot.get("result", "")
+        print(
+            f"Попытка {attempt_num + 1}: "
+            f"{'Error: ' + error_msg if error_msg else 'Результат: ' + result} → "
+            f"verify: {status}"
+        )
+    print(f"\nИтог: {final_state['status']} за {final_state['attempts'] + 1} попыток")
 
 
 if __name__ == "__main__":