177 lines
6.3 KiB
Python
177 lines
6.3 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
main.py – запуск LangGraph‑агента с самопроверкой и повторными попытками.
|
||
|
||
Требования:
|
||
- Python 3.10+
|
||
- pip install langgraph langchain-openai
|
||
"""
|
||
|
||
import random
|
||
from typing import TypedDict, Dict
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.prebuilt import InMemorySaver
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 1. Состояние графа
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
class AgentState(TypedDict):
|
||
task: str # исходная задача
|
||
result: str # результат выполнения инструмента
|
||
attempts: int # число попыток
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None # сообщение об ошибке, если есть
|
||
max_attempts: int # лимит повторов
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 2. Инструмент – «не надёжный» калькулятор
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def unreliable_tool(input_str: str) -> str:
|
||
"""
|
||
Случайно бросает ValueError с вероятностью ~30%.
|
||
Иначе возвращает результат простого арифметического выражения.
|
||
"""
|
||
if random.random() < 0.3:
|
||
raise ValueError("Инструмент временно недоступен")
|
||
try:
|
||
# eval безопасный для простых выражений (здесь только числа и +)
|
||
return str(eval(input_str))
|
||
except Exception as exc:
|
||
raise ValueError(f"Невозможно вычислить: {exc}") from exc
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 3. Узлы
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def execute_task(state: AgentState) -> Dict[str, object]:
|
||
"""
|
||
Выполняет задачу через unreliable_tool.
|
||
При ошибке сохраняет сообщение об ошибке и помечает статус как failed.
|
||
"""
|
||
try:
|
||
result = unreliable_tool(state["task"])
|
||
return {
|
||
"result": result,
|
||
"error": None,
|
||
"status": "pending", # результат получен, но ещё не проверён
|
||
}
|
||
except Exception as exc:
|
||
return {
|
||
"result": "",
|
||
"error": str(exc),
|
||
"status": "failed",
|
||
}
|
||
|
||
|
||
def verify_result(state: AgentState) -> Dict[str, object]:
|
||
"""
|
||
LLM‑проверка результата. Модель должна вернуть ровно 'success' или 'failed'.
|
||
"""
|
||
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
||
prompt = (
|
||
f"Задача: {state['task']}\n"
|
||
f"Результат: {state['result']}\n"
|
||
"Оцените результат. Ответьте только одним словом: 'success' или 'failed'."
|
||
)
|
||
response = llm.invoke(prompt).content.strip().lower()
|
||
if response not in {"success", "failed"}:
|
||
# Если модель дала непонятный ответ – считаем это ошибкой
|
||
return {
|
||
"status": "failed",
|
||
"error": f"LLM дал неожиданный ответ: {response}",
|
||
}
|
||
return {"status": response}
|
||
|
||
|
||
def handle_error(state: AgentState) -> Dict[str, object]:
|
||
"""
|
||
Увеличиваем счётчик попыток и готовим к повтору.
|
||
Если лимит превышен – помечаем как max_attempts.
|
||
"""
|
||
attempts = state["attempts"] + 1
|
||
if attempts >= state["max_attempts"]:
|
||
return {"status": "max_attempts", "attempts": attempts}
|
||
return {"status": "pending", "attempts": attempts}
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 4. Создание графа
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
builder = StateGraph(AgentState)
|
||
|
||
# Добавляем узлы
|
||
builder.add_node("execute_task", execute_task)
|
||
builder.add_node("verify_result", verify_result)
|
||
builder.add_node("handle_error", handle_error)
|
||
|
||
# Определяем переходы
|
||
builder.set_entry_point("execute_task")
|
||
|
||
builder.add_edge("execute_task", "verify_result")
|
||
builder.add_conditional_edges(
|
||
"verify_result",
|
||
lambda x: x["status"],
|
||
{
|
||
"success": END,
|
||
"failed": "handle_error",
|
||
"max_attempts": END,
|
||
},
|
||
)
|
||
|
||
builder.add_edge("handle_error", "execute_task")
|
||
|
||
# Сохраняем состояние в памяти (для отладки)
|
||
memory = InMemorySaver()
|
||
graph = builder.compile(persist_to_db=memory)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 5. Запуск
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def main() -> None:
|
||
# Пример задачи: простое арифметическое выражение
|
||
task_text = "2 + 2"
|
||
|
||
initial_state: AgentState = {
|
||
"task": task_text,
|
||
"result": "",
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": 5,
|
||
}
|
||
|
||
# Запускаем граф
|
||
final_state = graph.invoke(initial_state)
|
||
|
||
# Выводим результаты
|
||
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__":
|
||
main() |