текстовая игра на основе llm + interrupt: game.py

This commit is contained in:
2026-05-27 11:43:46 +00:00
parent be2665f8c9
commit cd30d4602a
@@ -0,0 +1,146 @@
import uuid
from typing import TypedDict, List, Dict, Any
import questionary
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command
# ---------- 1. Состояние графа ----------
class GameState(TypedDict):
theme: str
story: str
options: List[str]
choice: str
ending: str
# ---------- 2. Узлы ----------
def generate_scene(state: GameState) -> Dict[str, Any]:
"""
Генерируем завязку и варианты действий.
После генерации вызываем interrupt для выбора пользователем.
"""
theme = state["theme"]
llm_prompt = (
f"Тема: {theme}. "
"Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя. "
"Ответь в формате:\n"
"Завязка:\n<текст>\n\nВарианты:\n1. <вариант1>\n2. <вариант2>\n3. <вариант3>"
)
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
response = llm.invoke(llm_prompt).content
# Парсим ответ
parts = response.split("Варианты:")
story_part = parts[0].replace("Завязка:", "").strip()
options_raw = parts[1] if len(parts) > 1 else ""
options = [opt.strip() for opt in options_raw.split("\n") if opt.strip()]
# Убираем номера
options = [opt.split(".", 1)[-1].strip() for opt in options]
state["story"] = story_part
state["options"] = options
# Создаём payload для прерывания
interrupt_payload = {
"type": "choice",
"question": f"{story_part}\n\nЧто делает герой?",
"allow_responds": options,
}
# Возвращаем объект, который будет передан в interrupt
return interrupt(interrupt_payload)
def finish_story(state: GameState) -> GameState:
"""
После выбора пользователя генерируем короткую концовку.
"""
story = state["story"]
choice = state["choice"]
llm_prompt = (
f"Завязка: {story}\n"
f"Выбор игрока: {choice}\n"
"Допиши короткую концовку (2–3 предложения)."
)
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
ending = llm.invoke(llm_prompt).content.strip()
state["ending"] = ending
return state
# ---------- 3. Создание графа ----------
builder = StateGraph(GameState)
builder.add_node("generate_scene", generate_scene)
builder.add_node("finish_story", finish_story)
builder.set_entry_point("generate_scene")
builder.add_edge(START, "generate_scene")
builder.add_edge("generate_scene", "finish_story")
graph = builder.compile(checkpointer=InMemorySaver())
# ---------- 4. Запуск и обработка прерываний ----------
def run_game(theme: str):
thread_id = str(uuid.uuid4())
config = {"configurable": {"thread_id": thread_id}}
# Инициализируем состояние
state: GameState = {
"theme": theme,
"story": "",
"options": [],
"choice": "",
"ending": "",
}
# Запускаем поток
stream = graph.stream(state, config)
for chunk in stream:
if "__interrupt__" in chunk:
interrupt_payload = chunk["__interrupt__"][0].value # dict passed to interrupt()
question = interrupt_payload.get("question", "Выберите вариант:")
options = interrupt_payload.get("allow_responds", [])
# Показываем пользователю
answer = questionary.select(
message=question,
choices=options,
).ask()
if answer is None:
print("\nОтмена выбора. Завершаем игру.")
return
# Добавляем ответ в payload и возобновляем поток
interrupt_payload["choice"] = answer
resume_cmd = Command(resume=interrupt_payload)
stream = graph.stream(resume_cmd, config)
else:
# Выводим обычный вывод LLM (если есть)
if "text" in chunk:
print(chunk["text"], end="")
# После завершения печатаем финальное состояние
final_state: GameState = stream.final_state()
print("\n\n--- Конец истории ---")
print(f"Тема: {final_state['theme']}")
print(f"Завязка: {final_state['story']}")
print(f"Выбор игрока: {final_state['choice']}")
print(f"Концовка: {final_state['ending']}")
if __name__ == "__main__":
theme_input = questionary.text("Введите тему истории:").ask()
if theme_input:
run_game(theme_input)
else:
print("Тема не указана. Выход.")