From c22c332549683f6f2326f947e7268bdc20aa5530 Mon Sep 17 00:00:00 2001 From: RomanKurlygin Date: Tue, 26 May 2026 07:13:03 +0000 Subject: [PATCH] add main.py --- main.py | 180 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 180 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..d253993 --- /dev/null +++ b/main.py @@ -0,0 +1,180 @@ +"""Текстовая игра «выбери свою историю»: 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()