diff --git a/main.py b/main.py new file mode 100644 index 0000000..1090c8a --- /dev/null +++ b/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() \ No newline at end of file