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

solutions/69b1a07c67bbf488a1177da4_текстовая_игра_на_основе_llm___interrupt/graph.py

@@ -0,0 +1,153 @@
+import os
+from typing import TypedDict, List, Dict, Any
+
+import questionary
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, START, interrupt, Command
+from langgraph.types import InMemorySaver
+
+
+# ------------------------------------------------------------------
+# 1. Состояние графа
+# ------------------------------------------------------------------
+class StoryState(TypedDict):
+    topic: str | None  # тема истории (устанавливается при запуске)
+    intro: str | None   # завязка от LLM
+    options: List[str] | None  # варианты действий героя
+    choice: str | None          # выбранный пользователем вариант
+    ending: str | None         # концовка от LLM
+
+
+# ------------------------------------------------------------------
+# 2. Узел генерации сцены и прерывание
+# ------------------------------------------------------------------
+llm = ChatOpenAI(temperature=0.7, model="gpt-4o-mini")  # можно заменить на любой доступный LLM
+
+def generate_scene(state: StoryState) -> Dict[str, Any]:
+    """
+    Шаг 1 – генерация завязки и вариантов.
+    Шаг 2 – вызов interrupt с вопросом и вариантами.
+    После возобновления – второй запрос к LLM для концовки.
+    """
+    # Если intro ещё не сформирован – генерируем его
+    if state["intro"] is None:
+        prompt = (
+            f"Тема: {state['topic']}\n"
+            "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
+            "Ответь в формате:\n"
+            "Завязка:\n<текст>\n\nВарианты:\n1. <вариант1>\n2. <вариант2>\n3. <вариант3>"
+        )
+        response = llm.invoke(prompt)
+        text = response.content.strip()
+
+        # Разбираем ответ
+        parts = text.split("Варианты:")
+        intro_part = parts[0].replace("Завязка:", "").strip()
+        options_part = parts[1] if len(parts) > 1 else ""
+        options = [opt.strip() for opt in options_part.splitlines() if opt.strip().startswith(tuple(str(i) + "." for i in range(1, 4)))]
+        # Если парсинг не дал ровно 3 варианта – берём любые строки
+        if len(options) != 3:
+            options = [opt.strip() for opt in options_part.splitlines() if opt.strip()][:3]
+
+        state["intro"] = intro_part
+        state["options"] = options
+
+    # Шаг 2 – прерывание для выбора пользователя
+    interrupt_payload = {
+        "type": "choice",
+        "question": f"{state['intro']}\n\nЧто делаем?",
+        "choices": state["options"],
+    }
+    return interrupt(interrupt_payload)
+
+
+def resume_scene(state: StoryState, payload: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    После возобновления – получаем выбор пользователя из payload,
+    сохраняем его и генерируем концовку.
+    """
+    # Сохраняем выбор
+    state["choice"] = payload.get("answer", "")
+
+    # Генерация концовки
+    prompt = (
+        f"Завязка: {state['intro']}\n"
+        f"Выбор пользователя: {state['choice']}\n\n"
+        "Допиши короткую концовку (2–3 предложения)."
+    )
+    response = llm.invoke(prompt)
+    state["ending"] = response.content.strip()
+
+    return {"story_state": state}
+
+
+# ------------------------------------------------------------------
+# 3. Сборка графа
+# ------------------------------------------------------------------
+def build_graph() -> StateGraph:
+    graph = StateGraph(StoryState)
+
+    # Узел генерации сцены (с прерыванием)
+    graph.add_node("generate_scene", generate_scene)
+
+    # Узел завершения – после resume
+    graph.add_node("resume_scene", resume_scene)
+
+    # Переходы
+    graph.set_entry_point("generate_scene")
+    graph.add_edge(START, "generate_scene")
+    graph.add_conditional_edges(
+        "generate_scene",
+        lambda x: "__interrupt__" in x,
+        {
+            "__interrupt__": "resume_scene"
+        },
+    )
+    graph.add_edge("resume_scene", START)  # завершение цикла
+
+    return graph
+
+
+# ------------------------------------------------------------------
+# 4. Запуск с обработкой прерываний
+# ------------------------------------------------------------------
+def run_story(topic: str):
+    graph = build_graph()
+    saver = InMemorySaver()  # чекпоинтер в памяти
+    config = {"configurable": {"thread_id": "story_thread"}}
+
+    # Инициализируем состояние
+    state: StoryState = {
+        "topic": topic,
+        "intro": None,
+        "options": None,
+        "choice": None,
+        "ending": None,
+    }
+
+    # Запускаем поток
+    stream = graph.stream(state, config=config, checkpoint=saver)
+
+    for chunk in stream:
+        if "__interrupt__" in chunk:
+            interrupt_data = chunk["__interrupt__"][0].value  # dict с вопросом и вариантами
+            answer = questionary.select(
+                interrupt_data["question"],
+                choices=interrupt_data["choices"]
+            ).ask()
+            # Возобновляем граф
+            resume_payload = {"answer": answer}
+            stream = graph.stream(Command(resume=resume_payload), config=config, checkpoint=saver)
+        else:
+            # Выводим любой текст (LLM ответ или финальная история)
+            if "story_state" in chunk and chunk["story_state"]["ending"]:
+                print("\n[LLM] " + chunk["story_state"]["ending"])
+            elif "story_state" in chunk and chunk["story_state"]["intro"]:
+                # При первом запуске выводим завязку
+                print("\n[LLM] " + chunk["story_state"]["intro"])
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("OPENAI_API_KEY", "")
+    topic = questionary.text("Введите тему истории:").ask()
+    run_story(topic)