133 lines
4.7 KiB
Python
133 lines
4.7 KiB
Python
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() |