196 lines
7.4 KiB
Python
196 lines
7.4 KiB
Python
#!/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() |