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

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

@@ -1,5 +1,5 @@
-<|channel|>final code<|message|>import os
+import os
-from typing import TypedDict, List
+from typing import TypedDict, List, Optional
 
 import questionary
 from langchain_openai import ChatOpenAI
@@ -10,118 +10,163 @@ from langgraph.types import interrupt, Command
 
 
 # ---------- 1. Состояние графа ----------
+
 class GraphState(TypedDict):
-    """Состояние графа."""
+    """Состояние графа с отдельными полями для каждого этапа."""
-    # Текст истории, генерируемый LLM
+    topic: str               # Тема / завязка (входной параметр)
-    story: str
+    setup: str               # Сгенерированная завязка от LLM
-    # Ответ пользователя после прерывания
+    choices: List[str]       # Список вариантов выбора (3 штуки)
-    human_value: str | None
+    human_choice: str        # Выбранный пользователем вариант
-    # Внутренний токен для идентификации потока (необязательно)
+    ending: str              # Концовка от LLM
-    thread_id: str
 
 
-# ---------- 2. Узел с прерыванием ----------
+# ---------- 2. Узел: генерация сцены + прерывание + концовка ----------
-def generate_story(state: GraphState) -> GraphState:
+
+def story_node(state: GraphState) -> GraphState:
     """
-    Генерирует начало истории и предлагает варианты действий.
+    Шаг 1: генерирует завязку и варианты через LLM.
-    После генерации вызывается interrupt, чтобы пользователь выбрал вариант.
+    Шаг 2: вызывает interrupt() — граф ставится на паузу.
+    Шаг 3: после resume дописывает концовку через LLM.
     """
-    # Если история уже есть, просто возвращаем состояние
-    if state.get("story"):
-        return state
-
     llm = ChatOpenAI(temperature=0.7)
-    prompt = (
+
-        "Создай короткую интригующую историю о герое, который стоит перед выбором.\n"
+    # --- Шаг 1: генерируем завязку и варианты ---
-        "В конце предложи 3 варианта действий героя (A, B, C)."
+    topic = state["topic"]
+    scene_prompt = (
+        f"Тема: {topic}.\n"
+        "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
+        "Ответь строго в формате:\n"
+        "ЗАВЯЗКА: <текст завязки>\n"
+        "ВАРИАНТЫ:\n"
+        "1) <первый вариант>\n"
+        "2) <второй вариант>\n"
+        "3) <третий вариант>"
     )
-    response = llm.invoke(prompt)
+    scene_response = llm.invoke(scene_prompt)
-    story_text = response.content.strip()
+    raw = scene_response.content.strip()
 
-    # Сохраняем сгенерированную часть истории
+    # Парсим ответ: разделяем завязку и варианты
-    state["story"] = story_text
+    setup = ""
+    choices = []
 
-    # Подготавливаем payload для прерывания
+    lines = raw.splitlines()
+    in_choices = False
+    for line in lines:
+        line = line.strip()
+        if line.upper().startswith("ЗАВЯЗКА:"):
+            setup = line[len("ЗАВЯЗКА:"):].strip()
+        elif line.upper().startswith("ВАРИАНТЫ:"):
+            in_choices = True
+        elif in_choices and line:
+            # Убираем нумерацию вида "1)", "2)", "3)" или "1.", "2.", "3."
+            for prefix in ("1)", "2)", "3)", "1.", "2.", "3."):
+                if line.startswith(prefix):
+                    line = line[len(prefix):].strip()
+                    break
+            if line:
+                choices.append(line)
+
+    # Защита: если парсинг не дал нужных данных — берём весь текст как завязку
+    if not setup:
+        setup = raw
+    if len(choices) < 3:
+        choices = choices + [f"Вариант {i+1}" for i in range(len(choices), 3)]
+
+    # Сохраняем в состояние
+    state["setup"] = setup
+    state["choices"] = choices
+
+    # --- Шаг 2: прерывание --- 
     interrupt_payload = {
         "type": "choice",
-        "question": "Какое действие выберет герой?",
+        "question": setup + "\n\nЧто делаем?",
-        "allow_responds": ["A", "B", "C"],
+        "allow_responds": choices,
     }
 
-    # Вызываем прерывание: выполнение остановится до возобновления
+    # Выполнение останавливается здесь до Command(resume=...)
-    interrupt(interrupt_payload)
+    resume_value = interrupt(interrupt_payload)
 
-    # После возобновления сюда вернётся объект с ответом пользователя
+    # После resume в resume_value — словарь с полем "answer"
-    return state
+    human_choice = resume_value.get("answer", choices[0])
+    state["human_choice"] = human_choice
 
-
+    # --- Шаг 3: генерируем концовку ---
-def finish_story(state: GraphState) -> GraphState:
+    ending_prompt = (
-    """
+        f"Завязка: {setup}\n"
-    Завершающий узел, который получает выбор пользователя и дописывает концовку.
+        f"Выбор пользователя: {human_choice}\n\n"
-    """
+        "Допиши короткую концовку истории (2–3 предложения)."
-    llm = ChatOpenAI(temperature=0.7)
-    choice = state["human_value"]
-    prompt = (
-        f"История:\n{state['story']}\n\n"
-        f"Герой выбрал вариант {choice}. Напиши короткую концовку истории."
     )
-    response = llm.invoke(prompt)
+    ending_response = llm.invoke(ending_prompt)
-    state["story"] += "\n\n" + response.content.strip()
+    state["ending"] = ending_response.content.strip()
+
     return state
 
 
 # ---------- 3. Сборка графа ----------
+
 builder = StateGraph(GraphState)
+builder.add_node("story", story_node)
+builder.add_edge(START, "story")
 
-# Добавляем узлы
-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. Цикл запуска ----------
+# ---------- 4. Цикл запуска с обработкой прерывания ----------
+
 def main() -> None:
-    # Уникальный идентификатор потока (можно генерировать случайным образом)
+    topic = input("Введите тему истории (например, «космический кот»): ").strip()
+    if not topic:
+        topic = "космический кот"
+
+    print(f"\nТема: {topic}\n")
+
     thread_id = "story_thread_1"
     config = {"configurable": {"thread_id": thread_id}}
 
-    # Запускаем граф до первого прерывания
+    initial_state: GraphState = {
-    stream = graph.stream(Command(state={}), config=config)
+        "topic": topic,
+        "setup": "",
+        "choices": [],
+        "human_choice": "",
+        "ending": "",
+    }
 
-    for chunk in stream:
+    # --- Первый запуск: до прерывания ---
-        # Если в чанке есть прерывание, обрабатываем его
+    for chunk in graph.stream(initial_state, config=config):
         if "__interrupt__" in chunk:
-            interrupt_obj = chunk["__interrupt__"][0].value  # payload dict
+            interrupt_obj = chunk["__interrupt__"][0].value  # наш payload
-            print("\n=== Прерывание ===")
+
-            print(f"{interrupt_obj['question']}")
+            print(f"\n[LLM] {interrupt_obj['question']}\n")
+
+            # Показываем варианты через questionary
             answer = questionary.select(
-                "Выберите вариант:",
+                "Выберите действие:",
                 choices=interrupt_obj["allow_responds"],
             ).ask()
-            if not answer:
+
-                raise SystemExit("Отмена пользователем")
+            if answer is None:
+                raise SystemExit("Отмена пользователем.")
 
             # Добавляем ответ в payload и возобновляем граф
             interrupt_obj["answer"] = answer
-            stream = graph.stream(Command(resume=interrupt_obj), config=config)
-            continue
 
-        # Выводим обычные сообщения (если нужно)
+            # --- Возобновление через Command(resume=...) ---
-        if "story" in chunk:
+            for resumed_chunk in graph.stream(
-            print("\nИстория:")
+                Command(resume=interrupt_obj), config=config
-            print(chunk["story"])
+            ):
+                # После resume граф дописывает концовку — просто ждём завершения
+                pass
 
-    # После завершения выводим итоговую историю
+    # Получаем финальное состояние
-    final_state = stream.final_state()
+    final_state = graph.get_state(config).values
-    print("\n=== Итоговая история ===")
+
-    print(final_state["story"])
+    print("\n" + "=" * 50)
+    print("=== Итоговая история ===")
+    print("=" * 50)
+    print(f"\nЗавязка:\n{final_state['setup']}")
+    print(f"\nВыбор: {final_state['human_choice']}")
+    print(f"\n[LLM] {final_state['ending']}")
+    print("\n" + "=" * 50)
 
 
 if __name__ == "__main__":