181 lines
5.8 KiB
Python
181 lines
5.8 KiB
Python
"""Текстовая игра «выбери свою историю»: LLM + interrupt (LangGraph)."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import sys
|
|
from typing import TypedDict
|
|
|
|
import questionary
|
|
from dotenv import load_dotenv
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
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 Command, interrupt
|
|
|
|
load_dotenv()
|
|
|
|
llm = ChatOpenAI(
|
|
model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
|
|
base_url=os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1"),
|
|
api_key=os.getenv("OPENAI_API_KEY", "fake"),
|
|
temperature=0.8,
|
|
)
|
|
|
|
DEFAULT_THEME = "космический кот"
|
|
|
|
|
|
class StoryState(TypedDict):
|
|
theme: str
|
|
scene_text: str
|
|
choices: list[str]
|
|
user_choice: str
|
|
ending: str
|
|
|
|
|
|
def _call_llm(system: str, user: str) -> str:
|
|
response = llm.invoke([SystemMessage(content=system), HumanMessage(content=user)])
|
|
return str(response.content).strip()
|
|
|
|
|
|
def _parse_setup(raw: str) -> tuple[str, list[str]]:
|
|
"""Разбор ответа LLM: завязка + три варианта."""
|
|
setup = raw
|
|
choices: list[str] = []
|
|
|
|
if "ВАРИАНТЫ:" in raw:
|
|
setup, variants_block = raw.split("ВАРИАНТЫ:", maxsplit=1)
|
|
setup = setup.replace("ЗАВЯЗКА:", "").strip()
|
|
for line in variants_block.strip().splitlines():
|
|
line = re.sub(r"^\d+[\).\]]\s*", "", line.strip())
|
|
if line:
|
|
choices.append(line)
|
|
else:
|
|
numbered = re.findall(r"(?:^|\n)\s*\d+[\).\]]\s*(.+)", raw)
|
|
if numbered:
|
|
choices = [c.strip() for c in numbered[:3]]
|
|
setup = re.split(r"\n\s*1[\).\]]", raw, maxsplit=1)[0].strip()
|
|
|
|
if len(choices) < 3:
|
|
for line in raw.splitlines():
|
|
if line.strip().startswith(("-", "•")):
|
|
choices.append(line.strip().lstrip("-• ").strip())
|
|
choices = choices[:3]
|
|
|
|
while len(choices) < 3:
|
|
choices.append(f"Вариант {len(choices) + 1}")
|
|
|
|
return setup.strip(), choices[:3]
|
|
|
|
|
|
def story_node(state: StoryState) -> StoryState:
|
|
"""Генерация завязки → прерывание (выбор) → концовка от LLM."""
|
|
theme = state.get("theme") or DEFAULT_THEME
|
|
|
|
setup_raw = _call_llm(
|
|
system=(
|
|
"Ты автор интерактивных текстовых игр. Отвечай строго в формате:\n"
|
|
"ЗАВЯЗКА:\n<2-3 предложения>\n\n"
|
|
"ВАРИАНТЫ:\n1) <вариант 1>\n2) <вариант 2>\n3) <вариант 3>"
|
|
),
|
|
user=f"Тема: {theme}. Придумай короткую завязку и ровно 3 варианта поступка героя.",
|
|
)
|
|
scene_text, choices = _parse_setup(setup_raw)
|
|
|
|
print(f"\n[LLM] {scene_text}\n")
|
|
|
|
payload = interrupt(
|
|
{
|
|
"type": "choice",
|
|
"question": f"{scene_text}\n\nЧто делаем?",
|
|
"allow_responds": choices,
|
|
"theme": theme,
|
|
"scene_text": scene_text,
|
|
"choices": choices,
|
|
}
|
|
)
|
|
|
|
user_choice = payload.get("answer", choices[0])
|
|
print(f"\n> {user_choice}\n")
|
|
|
|
ending = _call_llm(
|
|
system="Ты автор интерактивных историй. Допиши короткую концовку (2–3 предложения).",
|
|
user=(
|
|
f"Тема: {theme}\n"
|
|
f"Завязка: {scene_text}\n"
|
|
f"Выбор пользователя: {user_choice}\n"
|
|
"Допиши короткую концовку с учётом выбора."
|
|
),
|
|
)
|
|
|
|
print(f"[LLM] {ending}\n")
|
|
|
|
return {
|
|
"theme": theme,
|
|
"scene_text": scene_text,
|
|
"choices": choices,
|
|
"user_choice": user_choice,
|
|
"ending": ending,
|
|
}
|
|
|
|
|
|
def _ask_user(payload: dict) -> dict:
|
|
question = payload.get("question", "Что делаем?")
|
|
options = payload.get("allow_responds") or ["Вариант 1", "Вариант 2", "Вариант 3"]
|
|
choice = questionary.select(question, choices=options).ask()
|
|
if choice is None:
|
|
choice = options[0]
|
|
updated = dict(payload)
|
|
updated["answer"] = choice
|
|
return updated
|
|
|
|
|
|
def play_story(theme: str) -> StoryState:
|
|
builder = StateGraph(StoryState)
|
|
builder.add_node("story", story_node)
|
|
builder.add_edge(START, "story")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
config = {"configurable": {"thread_id": f"story-{theme}"}}
|
|
initial: StoryState = {
|
|
"theme": theme,
|
|
"scene_text": "",
|
|
"choices": [],
|
|
"user_choice": "",
|
|
"ending": "",
|
|
}
|
|
|
|
final: StoryState = initial
|
|
stream = graph.stream(initial, config)
|
|
|
|
for chunk in stream:
|
|
if "__interrupt__" in chunk:
|
|
payload = dict(chunk["__interrupt__"][0].value)
|
|
resumed = _ask_user(payload)
|
|
resume_stream = graph.stream(Command(resume=resumed), config)
|
|
for resume_chunk in resume_stream:
|
|
if "story" in resume_chunk:
|
|
final = resume_chunk["story"]
|
|
elif "story" in chunk:
|
|
final = chunk["story"]
|
|
|
|
return final
|
|
|
|
|
|
def main() -> None:
|
|
theme = " ".join(sys.argv[1:]).strip() or DEFAULT_THEME
|
|
print(f"Тема: {theme}\n")
|
|
result = play_story(theme)
|
|
|
|
print("=== Итоговая история ===")
|
|
print(f"Тема: {result['theme']}")
|
|
print(f"Завязка: {result['scene_text']}")
|
|
print(f"Выбор: {result['user_choice']}")
|
|
print(f"Концовка: {result['ending']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|