diff --git a/solution.py b/solution.py index 22a2e6b..28553af 100644 --- a/solution.py +++ b/solution.py @@ -1,114 +1,184 @@ -<|channel|>final solution.py<|message|>import sys -from typing import Optional +from typing import TypedDict, List -# Библиотеки, которые должны быть установлены в окружении задачи. -# Мы используем langchain для взаимодействия с LLM и rich для красивого вывода. -try: - from langchain.chat_models import ChatOpenAI -except ImportError as exc: # pragma: no cover - raise RuntimeError( - "Не удалось импортировать langchain. Убедитесь, что пакет установлен." - ) 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 +import questionary +from langchain_openai import ChatOpenAI +from langchain.prompts import ChatPromptTemplate +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.constants import START +from langgraph.graph import StateGraph +from langgraph.types import interrupt, Command -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: степень случайности генерации текста - """ - 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 + # Парсим ответ: разделяем завязку и варианты + setup = "" + choices = [] - def _display_welcome(self) -> None: - """Показываем приветственное сообщение.""" - welcome = Text("Добро пожаловать в текстовую игру!", style="bold green") - self.console.print(welcome) + 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) - def _get_player_input(self, prompt: str = "Что вы делаете?") -> Optional[str]: - """ - Получаем ввод от игрока. Если пользователь введёт 'quit', - возвращаем None для завершения игры. - """ - try: - user_input = Prompt.ask(prompt) - except KeyboardInterrupt: # pragma: no cover - self.console.print("\n[red]Игра прервана пользователем.[/red]") - return None + # Защита: если парсинг не дал нужных данных + if not setup: + setup = raw + if len(choices) < 3: + choices = choices + [f"Вариант {i+1}" for i in range(len(choices), 3)] - if user_input.strip().lower() in {"quit", "exit"}: - return None - return user_input + state["setup"] = setup + state["choices"] = choices - def _generate_response(self, player_action: str) -> str: - """ - Отправляем действие игрока в LLM и получаем ответ. + # --- Шаг 2: прерывание --- + interrupt_payload = { + "type": "choice", + "question": setup + "\n\nЧто делаем?", + "allow_responds": choices, + } - :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 + # Выполнение останавливается здесь до Command(resume=...) + resume_value = interrupt(interrupt_payload) - return response.content.strip() + # После resume в resume_value — словарь с полем "answer" + human_choice = resume_value.get("answer", choices[0]) + state["human_choice"] = human_choice - 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 + # --- Шаг 3: генерируем концовку через LangChain --- + ending_response = ending_chain.invoke({ + "setup": setup, + "human_choice": human_choice, + }) + state["ending"] = ending_response.content.strip() - self.console.print(f"[cyan]Вы сказали:[/cyan] {action}") - try: - result = self._generate_response(action) - except Exception as exc: # pragma: no cover - self.console.print(f"[red]Ошибка генерации ответа: {exc}[/red]") - continue + return state - self.console.print("\n[green]Ответ модели:[/green]\n") - self.console.print(result, style="italic") - self.console.print("-" * 40) + +# ---------- 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 не задана, - # langchain выбросит исключение при создании модели. - game = TextAdventureGame() - try: - game.run() - except Exception as exc: # pragma: no cover - sys.stderr.write(f"Произошла ошибка: {exc}\n") - sys.exit(1) \ No newline at end of file + main() \ No newline at end of file