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

main.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""
+main.py
+
+Самокорректирующийся агент на LangGraph.
+- AgentState содержит task, result, attempts, status, error, max_attempts.
+- Узлы: execute_task, verify_result, handle_error.
+- Инструмент unreliable_tool с ~30% вероятностью бросает ValueError.
+- verify_result использует LLM-as-judge (OpenAI Ollama) и отвечает строго "success" или "failed".
+- Граф реализует цикл retry до исчерпания max_attempts.
+- При запуске выводит номера попыток и финальный статус.
+"""
+
+import os
+import random
+from typing import TypedDict, Optional
+
+from langgraph.graph import StateGraph, START, END
+from langchain_openai import ChatOpenAI  # требует OPENAI_API_KEY в окружении
+from langchain_core.messages import HumanMessage
+
+
+# ----------------------------------------------------------------------
+# 1. Состояние агента
+# ----------------------------------------------------------------------
+class AgentState(TypedDict):
+    task: str
+    result: str
+    attempts: int
+    status: str          # pending | success | failed | max_attempts
+    error: Optional[str]
+    max_attempts: int
+
+
+# ----------------------------------------------------------------------
+# 2. Ненадёжный инструмент
+# ----------------------------------------------------------------------
+def unreliable_tool(task: str) -> str:
+    """
+    Выполняет простую арифметическую задачу (eval) с вероятностью ~30% бросить ValueError.
+    Для демонстрации retry.
+    """
+    if random.random() < 0.3:
+        raise ValueError("Случайная ошибка инструмента")
+    # Безопасный eval только для простых арифметических выражений.
+    # В реальном коде следует использовать более строгую валидацию.
+    return str(eval(task, {"__builtins__": {}}))
+
+
+# ----------------------------------------------------------------------
+# 3. LLM‑judge (OpenAI). Если ключ не задан – fallback на простую проверку.
+# ----------------------------------------------------------------------
+def get_llm():
+    api_key = os.getenv("OPENAI_API_KEY")
+    if api_key:
+        return ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
+    return None
+
+
+LLM = get_llm()
+
+
+def judge_with_llm(task: str, result: str) -> str:
+    """
+    Запрашивает у LLM оценку результата.
+    Ожидается ответ exatamente "success" или "failed".
+    """
+    if LLM is None:
+        # Fallback: просто сравниваем с правильным ответом через eval
+        try:
+            correct = str(eval(task, {"__builtins__": {}}))
+            return "success" if result.strip() == correct else "failed"
+        except Exception:
+            return "failed"
+
+    prompt = (
+        "You are a judge. Determine if the result correctly answers the task.\n"
+        f"Task: {task}\n"
+        f"Result: {result}\n"
+        "Respond with exactly one word: 'success' if the result is correct, otherwise 'failed'."
+    )
+    msg = [HumanMessage(content=prompt)]
+    response = LLM.invoke(msg)
+    verdict = response.content.strip().lower()
+    # Приводим к одному из допустимых вариантов
+    if "success" in verdict:
+        return "success"
+    return "failed"
+
+
+# ----------------------------------------------------------------------
+# 4. Узлы графа
+# ----------------------------------------------------------------------
+def execute_task(state: AgentState) -> AgentState:
+    """Выполняет задачу через unreliable_tool."""
+    state["attempts"] += 1
+    state["error"] = None
+    try:
+        state["result"] = unreliable_tool(state["task"])
+    except Exception as e:
+        state["result"] = ""
+        state["error"] = str(e)
+    return state
+
+
+def verify_result(state: AgentState) -> AgentState:
+    """LLM‑as‑judge: ставит status = success/failed."""
+    if state["error"] is not None:
+        # Если инструмент уже упал – считаем failed без вызова LLM
+        state["status"] = "failed"
+        return state
+
+    verdict = judge_with_llm(state["task"], state["result"])
+    state["status"] = verdict
+    return state
+
+
+def handle_error(state: AgentState) -> AgentState:
+    """Подготовка к повторной попытке."""
+    # Ставим статус pending, чтобы граф вернулся к execute_task
+    state["status"] = "pending"
+    # Ошибку оставляем для логов, но не очищаем – можно очистить, если нужно
+    return state
+
+
+# ----------------------------------------------------------------------
+# 5. Построение графа
+# ----------------------------------------------------------------------
+def build_graph() -> StateGraph:
+    workflow = StateGraph(AgentState)
+
+    # Регистрация узлов
+    workflow.add_node("execute_task", execute_task)
+    workflow.add_node("verify_result", verify_result)
+    workflow.add_node("handle_error", handle_error)
+
+    # Входная точка
+    workflow.add_edge(START, "execute_task")
+    # После выполнения всегда идём к проверке
+    workflow.add_edge("execute_task", "verify_result")
+
+    # Условные переходы из verify_result
+    def should_continue(state: AgentState) -> str:
+        if state["status"] == "success":
+            return END
+        if state["status"] == "failed" and state["attempts"] < state["max_attempts"]:
+            return "handle_error"
+        # Либо max_attempts исчерпан, либо иной статус – завершаем
+        return END
+
+    workflow.add_conditional_edges(
+        "verify_result",
+        should_continue,
+        {
+            "handle_error": "handle_error",
+            END: END,
+        },
+    )
+
+    # После handle_error повторяем выполнение задачи
+    workflow.add_edge("handle_error", "execute_task")
+
+    return workflow.compile()
+
+
+# ----------------------------------------------------------------------
+# 6. CLI‑запуск
+# ----------------------------------------------------------------------
+def main():
+    # Пример задачи; можно переопределить через аргумент командной строки
+    import sys
+    task = sys.argv[1] if len(sys.argv) > 1 else "Вычисли 2+2"
+
+    initial_state: AgentState = {
+        "task": task,
+        "result": "",
+        "attempts": 0,
+        "status": "pending",
+        "error": None,
+        "max_attempts": 5,  # можно изменить
+    }
+
+    app = build_graph()
+    final_state = app.invoke(initial_state)
+
+    print(f"Задача: {final_state['task']}")
+    print(f"Итоговый статус: {final_state['status']}")
+    print(f"Попыток сделано: {final_state['attempts']}")
+    if final_state["result"]:
+        print(f"Результат: {final_state['result']}")
+    if final_state["error"]:
+        print(f"Последняя ошибка: {final_state['error']}")
+
+
+if __name__ == "__main__":
+    main()