текстовая игра на основе llm + interrupt: solution.py

solution.py

@@ -0,0 +1,133 @@
+import os
+from typing import TypedDict, List, Optional
+
+import questionary
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import MemorySaver
+
+# ------------------------------------------------------------------
+# 1. Определяем состояние графа
+# ------------------------------------------------------------------
+class StoryState(TypedDict):
+    # Тема/завязка истории
+    theme: Optional[str]
+    # Текст сцены, сгенерированный LLM
+    scene: Optional[str]
+    # Список вариантов выбора
+    options: List[str]
+    # Выбранный пользователем вариант
+    choice: Optional[str]
+    # Концовка, сгенерированная LLM
+    ending: Optional[str]
+
+
+# ------------------------------------------------------------------
+# 2. Инициализируем LLM
+# ------------------------------------------------------------------
+# Убедитесь, что переменная окружения OPENAI_API_KEY установлена
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
+
+
+# ------------------------------------------------------------------
+# 3. Узлы графа
+# ------------------------------------------------------------------
+def generate_scene(state: StoryState) -> StoryState:
+    """
+    Генерируем начальную сцену и варианты выбора.
+    """
+    theme = state.get("theme") or "Вы просыпаетесь в загадочном лесу."
+    prompt = f"""
+    Напиши короткую сцену (не более 150 слов) для истории, где главный герой {theme}
+    и предложи три варианта действий, которые может предпринять герой.
+    Выведи варианты в виде списка, каждый с номером 1-3.
+    """
+    response = llm.invoke(prompt)
+    # Разделяем на сцену и варианты
+    lines = response.content.strip().split("\n")
+    scene_lines = []
+    options = []
+    for line in lines:
+        if line.strip().startswith(("1.", "2.", "3.")):
+            options.append(line.strip())
+        else:
+            scene_lines.append(line.strip())
+    state["scene"] = " ".join(scene_lines)
+    state["options"] = options
+    return state
+
+
+def interrupt_choice(state: StoryState) -> StoryState:
+    """
+    Пауза и запрос выбора у пользователя.
+    """
+    print("\n--- Сцена ---")
+    print(state["scene"])
+    print("\n--- Выберите действие ---")
+    for opt in state["options"]:
+        print(opt)
+    choice = questionary.select(
+        "Выберите вариант:",
+        choices=[opt.split(".", 1)[1].strip() for opt in state["options"]],
+    ).ask()
+    state["choice"] = choice
+    return state
+
+
+def generate_ending(state: StoryState) -> StoryState:
+    """
+    Генерируем короткую концовку на основе выбранного варианта.
+    """
+    prompt = f"""
+    Герой выбрал: "{state['choice']}".
+    Напиши короткую концовку (не более 100 слов), которая логично следует из выбранного действия.
+    """
+    response = llm.invoke(prompt)
+    state["ending"] = response.content.strip()
+    return state
+
+
+# ------------------------------------------------------------------
+# 4. Создаём граф
+# ------------------------------------------------------------------
+builder = StateGraph(StoryState)
+
+# Добавляем узлы
+builder.add_node("scene", generate_scene)
+builder.add_node("choice", interrupt_choice)
+builder.add_node("ending", generate_ending)
+
+# Определяем переходы
+builder.set_entry_point("scene")
+builder.add_edge("scene", "choice")
+builder.add_edge("choice", "ending")
+builder.add_edge("ending", END)
+
+# Сохраняем состояние в памяти (необязательно, но удобно)
+checkpoint = MemorySaver()
+graph = builder.compile(checkpointer=checkpoint)
+
+
+# ------------------------------------------------------------------
+# 5. Запуск
+# ------------------------------------------------------------------
+def main():
+    # Инициализируем состояние
+    init_state: StoryState = {
+        "theme": None,
+        "scene": None,
+        "options": [],
+        "choice": None,
+        "ending": None,
+    }
+
+    # Запускаем граф
+    result = graph.invoke(init_state)
+
+    # Выводим финальный результат
+    print("\n=== Концовка ===")
+    print(result["ending"])
+
+
+if __name__ == "__main__":
+    main()