Обновить solution.py

solution.py

@@ -1,114 +1,184 @@
-<|channel|>final solution.py<|message|>import sys
+from typing import TypedDict, List
-from typing import Optional
 
-# Библиотеки, которые должны быть установлены в окружении задачи.
+import questionary
-# Мы используем langchain для взаимодействия с LLM и rich для красивого вывода.
+from langchain_openai import ChatOpenAI
-try:
+from langchain.prompts import ChatPromptTemplate
-    from langchain.chat_models import ChatOpenAI
+from langgraph.checkpoint.memory import InMemorySaver
-except ImportError as exc:  # pragma: no cover
+from langgraph.constants import START
-    raise RuntimeError(
+from langgraph.graph import StateGraph
-        "Не удалось импортировать langchain. Убедитесь, что пакет установлен."
+from langgraph.types import interrupt, Command
-    ) from exc
-
-try:
-    from rich.console import Console
-    from rich.prompt import Prompt
-    from rich.text import Text
-except ImportError as exc:  # pragma: no cover
-    raise RuntimeError(
-        "Не удалось импортировать rich. Убедитесь, что пакет установлен."
-    ) from exc
 
 
-class TextAdventureGame:
+# ---------- 1. Состояние графа ----------
+
+class GraphState(TypedDict):
+    """Состояние графа с отдельными полями для каждого этапа."""
+    topic: str          # Тема (входной параметр)
+    setup: str          # Сгенерированная завязка от LLM
+    choices: List[str]  # Список вариантов выбора (3 штуки)
+    human_choice: str   # Выбранный пользователем вариант
+    ending: str         # Концовка от LLM
+
+
+# ---------- LangChain: промпты и цепочки ----------
+
+llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.7)
+
+# Цепочка для генерации завязки и вариантов
+scene_prompt = ChatPromptTemplate.from_template(
+    "Тема: {topic}.\n"
+    "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
+    "Ответь строго в формате:\n"
+    "ЗАВЯЗКА: <текст завязки>\n"
+    "ВАРИАНТЫ:\n"
+    "1) <первый вариант>\n"
+    "2) <второй вариант>\n"
+    "3) <третий вариант>"
+)
+scene_chain = scene_prompt | llm
+
+# Цепочка для генерации концовки
+ending_prompt = ChatPromptTemplate.from_template(
+    "Завязка: {setup}\n"
+    "Выбор пользователя: {human_choice}\n\n"
+    "Допиши короткую концовку истории (2–3 предложения)."
+)
+ending_chain = ending_prompt | llm
+
+
+# ---------- 2. Узел: генерация сцены + прерывание + концовка ----------
+
+def story_node(state: GraphState) -> GraphState:
     """
-    Простая текстовая игра на основе LLM.
+    Шаг 1: генерирует завязку и варианты через LangChain-цепочку.
-    Игрок вводит команды, которые отправляются в модель,
+    Шаг 2: вызывает interrupt() — граф ставится на паузу.
-    а ответ модели выводится пользователю.
+    Шаг 3: после resume дописывает концовку через LangChain-цепочку.
     """
 
-    def __init__(self, model_name: str = "gpt-3.5-turbo", temperature: float = 0.7):
+    # --- Шаг 1: генерируем завязку и варианты через LangChain ---
-        """
+    scene_response = scene_chain.invoke({"topic": state["topic"]})
-        Инициализация игрового процесса.
+    raw = scene_response.content.strip()
 
-        :param model_name: имя LLM из OpenAI
+    # Парсим ответ: разделяем завязку и варианты
-        :param temperature: степень случайности генерации текста
+    setup = ""
-        """
+    choices = []
-        self.console = Console()
-        # Создаём объект модели. Для работы требуется переменная окружения OPENAI_API_KEY.
-        try:
-            self.llm = ChatOpenAI(model=model_name, temperature=temperature)
-        except Exception as exc:  # pragma: no cover
-            raise RuntimeError(
-                "Не удалось создать экземпляр LLM. Проверьте ключ OpenAI."
-            ) from exc
 
-    def _display_welcome(self) -> None:
+    lines = raw.splitlines()
-        """Показываем приветственное сообщение."""
+    in_choices = False
-        welcome = Text("Добро пожаловать в текстовую игру!", style="bold green")
+    for line in lines:
-        self.console.print(welcome)
+        line = line.strip()
-
+        if line.upper().startswith("ЗАВЯЗКА:"):
-    def _get_player_input(self, prompt: str = "Что вы делаете?") -> Optional[str]:
+            setup = line[len("ЗАВЯЗКА:"):].strip()
-        """
+        elif line.upper().startswith("ВАРИАНТЫ:"):
-        Получаем ввод от игрока. Если пользователь введёт 'quit',
+            in_choices = True
-        возвращаем None для завершения игры.
+        elif in_choices and line:
-        """
+            # Убираем нумерацию вида "1)", "2)", "3)" или "1.", "2.", "3."
-        try:
+            for prefix in ("1)", "2)", "3)", "1.", "2.", "3."):
-            user_input = Prompt.ask(prompt)
+                if line.startswith(prefix):
-        except KeyboardInterrupt:  # pragma: no cover
+                    line = line[len(prefix):].strip()
-            self.console.print("\n[red]Игра прервана пользователем.[/red]")
-            return None
-
-        if user_input.strip().lower() in {"quit", "exit"}:
-            return None
-        return user_input
-
-    def _generate_response(self, player_action: str) -> str:
-        """
-        Отправляем действие игрока в LLM и получаем ответ.
-
-        :param player_action: строка с командой игрока
-        :return: текст ответа модели
-        """
-        # Формируем сообщение для чат‑модели.
-        messages = [
-            {"role": "system", "content": "Вы — мудрый рассказчик, описывающий события в мире приключений."},
-            {"role": "user", "content": player_action},
-        ]
-        try:
-            response = self.llm.invoke(messages)
-        except Exception as exc:  # pragma: no cover
-            raise RuntimeError("Ошибка при вызове LLM.") from exc
-
-        return response.content.strip()
-
-    def run(self) -> None:
-        """Основной цикл игры."""
-        self._display_welcome()
-        while True:
-            action = self._get_player_input()
-            if action is None:  # пользователь ввёл quit или прервал
-                self.console.print("[bold magenta]Спасибо за игру![/bold magenta]")
                     break
+            if line:
+                choices.append(line)
 
-            self.console.print(f"[cyan]Вы сказали:[/cyan] {action}")
+    # Защита: если парсинг не дал нужных данных
-            try:
+    if not setup:
-                result = self._generate_response(action)
+        setup = raw
-            except Exception as exc:  # pragma: no cover
+    if len(choices) < 3:
-                self.console.print(f"[red]Ошибка генерации ответа: {exc}[/red]")
+        choices = choices + [f"Вариант {i+1}" for i in range(len(choices), 3)]
-                continue
 
-            self.console.print("\n[green]Ответ модели:[/green]\n")
+    state["setup"] = setup
-            self.console.print(result, style="italic")
+    state["choices"] = choices
-            self.console.print("-" * 40)
+
+    # --- Шаг 2: прерывание ---
+    interrupt_payload = {
+        "type": "choice",
+        "question": setup + "\n\nЧто делаем?",
+        "allow_responds": choices,
+    }
+
+    # Выполнение останавливается здесь до Command(resume=...)
+    resume_value = interrupt(interrupt_payload)
+
+    # После resume в resume_value — словарь с полем "answer"
+    human_choice = resume_value.get("answer", choices[0])
+    state["human_choice"] = human_choice
+
+    # --- Шаг 3: генерируем концовку через LangChain ---
+    ending_response = ending_chain.invoke({
+        "setup": setup,
+        "human_choice": human_choice,
+    })
+    state["ending"] = ending_response.content.strip()
+
+    return state
+
+
+# ---------- 3. Сборка графа ----------
+
+builder = StateGraph(GraphState)
+builder.add_node("story", story_node)
+builder.add_edge(START, "story")
+
+graph = builder.compile(checkpointer=InMemorySaver())
+
+
+# ---------- 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 = {
+        "topic": topic,
+        "setup": "",
+        "choices": [],
+        "human_choice": "",
+        "ending": "",
+    }
+
+    # --- Первый запуск: до прерывания ---
+    for chunk in graph.stream(initial_state, config=config):
+        if "__interrupt__" in chunk:
+            interrupt_obj = chunk["__interrupt__"][0].value  # наш payload
+
+            print(f"\n[LLM] {interrupt_obj['question']}\n")
+
+            # Показываем варианты через questionary
+            answer = questionary.select(
+                "Выберите действие:",
+                choices=interrupt_obj["allow_responds"],
+            ).ask()
+
+            if answer is None:
+                raise SystemExit("Отмена пользователем.")
+
+            # Добавляем ответ в payload и возобновляем граф
+            interrupt_obj["answer"] = answer
+
+            # --- Возобновление через Command(resume=...) ---
+            for resumed_chunk in graph.stream(
+                Command(resume=interrupt_obj), config=config
+            ):
+                # После resume граф дописывает концовку — просто ждём завершения
+                pass
+
+    # Получаем финальное состояние
+    final_state = graph.get_state(config).values
+
+    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__":
-    # Запускаем игру. Если переменная окружения OPENAI_API_KEY не задана,
+    main()
-    # langchain выбросит исключение при создании модели.
-    game = TextAdventureGame()
-    try:
-        game.run()
-    except Exception as exc:  # pragma: no cover
-        sys.stderr.write(f"Произошла ошибка: {exc}\n")
-        sys.exit(1)