<|channel|>final code<|message|>import os
from typing import TypedDict, List

import questionary
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command


# ---------- 1. Состояние графа ----------
class GraphState(TypedDict):
    """Состояние графа."""
    # Текст истории, генерируемый LLM
    story: str
    # Ответ пользователя после прерывания
    human_value: str | None
    # Внутренний токен для идентификации потока (необязательно)
    thread_id: str


# ---------- 2. Узел с прерыванием ----------
def generate_story(state: GraphState) -> GraphState:
    """
    Генерирует начало истории и предлагает варианты действий.
    После генерации вызывается interrupt, чтобы пользователь выбрал вариант.
    """
    # Если история уже есть, просто возвращаем состояние
    if state.get("story"):
        return state

    llm = ChatOpenAI(temperature=0.7)
    prompt = (
        "Создай короткую интригующую историю о герое, который стоит перед выбором.\n"
        "В конце предложи 3 варианта действий героя (A, B, C)."
    )
    response = llm.invoke(prompt)
    story_text = response.content.strip()

    # Сохраняем сгенерированную часть истории
    state["story"] = story_text

    # Подготавливаем payload для прерывания
    interrupt_payload = {
        "type": "choice",
        "question": "Какое действие выберет герой?",
        "allow_responds": ["A", "B", "C"],
    }

    # Вызываем прерывание: выполнение остановится до возобновления
    interrupt(interrupt_payload)

    # После возобновления сюда вернётся объект с ответом пользователя
    return state


def finish_story(state: GraphState) -> GraphState:
    """
    Завершающий узел, который получает выбор пользователя и дописывает концовку.
    """
    llm = ChatOpenAI(temperature=0.7)
    choice = state["human_value"]
    prompt = (
        f"История:\n{state['story']}\n\n"
        f"Герой выбрал вариант {choice}. Напиши короткую концовку истории."
    )
    response = llm.invoke(prompt)
    state["story"] += "\n\n" + response.content.strip()
    return state


# ---------- 3. Сборка графа ----------
builder = StateGraph(GraphState)

# Добавляем узлы
builder.add_node("generate", generate_story)
builder.add_node("finish", finish_story)

# Определяем переходы
builder.set_entry_point("generate")
builder.add_edge(START, "generate")
builder.add_edge("generate", "finish")

# Компилируем граф с чекпоинтером
graph = builder.compile(checkpointer=InMemorySaver())


# ---------- 4. Цикл запуска ----------
def main() -> None:
    # Уникальный идентификатор потока (можно генерировать случайным образом)
    thread_id = "story_thread_1"
    config = {"configurable": {"thread_id": thread_id}}

    # Запускаем граф до первого прерывания
    stream = graph.stream(Command(state={}), config=config)

    for chunk in stream:
        # Если в чанке есть прерывание, обрабатываем его
        if "__interrupt__" in chunk:
            interrupt_obj = chunk["__interrupt__"][0].value  # payload dict
            print("\n=== Прерывание ===")
            print(f"{interrupt_obj['question']}")
            answer = questionary.select(
                "Выберите вариант:",
                choices=interrupt_obj["allow_responds"],
            ).ask()
            if not answer:
                raise SystemExit("Отмена пользователем")

            # Добавляем ответ в payload и возобновляем граф
            interrupt_obj["answer"] = answer
            stream = graph.stream(Command(resume=interrupt_obj), config=config)
            continue

        # Выводим обычные сообщения (если нужно)
        if "story" in chunk:
            print("\nИстория:")
            print(chunk["story"])

    # После завершения выводим итоговую историю
    final_state = stream.final_state()
    print("\n=== Итоговая история ===")
    print(final_state["story"])


if __name__ == "__main__":
    main()