diff --git a/solution.py b/solution.py new file mode 100644 index 0000000..7ac78ce --- /dev/null +++ b/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() \ No newline at end of file