From 2a6eb415a40132d45b4a5a330d105bfaac8e9261 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=AD=D0=BC=D0=B8=D0=BB=D1=8C=20=D0=90=D0=BC=D0=B8=D1=80?= =?UTF-8?q?=D0=BE=D0=B2?= Date: Mon, 25 May 2026 22:16:36 +0000 Subject: [PATCH] add main.py --- main.py | 225 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..5aace7f --- /dev/null +++ b/main.py @@ -0,0 +1,225 @@ +"""Интерактивная история: LLM + LangGraph interrupt + questionary.""" +from __future__ import annotations + +import os +import re +import uuid +from typing import TypedDict + +import questionary +from dotenv import load_dotenv + +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 + +THEME_DEFAULT = "космический кот" + +LLM_MODEL = os.getenv("OPENROUTER_MODEL", "openai/gpt-oss-20b:free") +LLM_BASE_URL = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1") + + +class StoryState(TypedDict, total=False): + theme: str + setup: str + scene_text: str + options: list[str] + user_choice: str + ending: str + + +def build_llm() -> ChatOpenAI: + return ChatOpenAI( + model=LLM_MODEL, + base_url=LLM_BASE_URL, + api_key=os.getenv("OPENAI_API_KEY", "fake"), + temperature=0.8, + ) + + +def _parse_scene_response(text: str) -> tuple[str, list[str]]: + """Разбирает ответ LLM на завязку и 3 варианта.""" + options: list[str] = [] + setup = text.strip() + + variants_match = re.search( + r"ВАРИАНТЫ:\s*(.*)", + text, + flags=re.IGNORECASE | re.DOTALL, + ) + if variants_match: + setup = text[: variants_match.start()].strip() + block = variants_match.group(1) + for line in block.splitlines(): + line = line.strip() + if not line: + continue + cleaned = re.sub(r"^\d+[\).\]]\s*", "", line).strip() + if cleaned: + options.append(cleaned) + + if len(options) < 3: + numbered = re.findall(r"^\s*\d+[\).\]]\s*(.+)$", text, flags=re.MULTILINE) + options = [o.strip() for o in numbered if o.strip()] + + if len(options) < 3: + parts = [p.strip() for p in re.split(r"[,;]", text) if p.strip()] + if len(parts) >= 4: + setup = parts[0] + options = parts[1:4] + + while len(options) < 3: + options.append(f"Вариант {len(options) + 1}") + + return setup, options[:3] + + +def generate_scene(state: StoryState) -> dict: + theme = state.get("theme", THEME_DEFAULT) + llm = build_llm() + prompt = ( + f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) " + "и ровно 3 варианта поступка героя.\n" + "Формат ответа строго:\n" + "ЗАВЯЗКА:\n" + "<текст завязки>\n" + "ВАРИАНТЫ:\n" + "1) <вариант 1>\n" + "2) <вариант 2>\n" + "3) <вариант 3>" + ) + response = llm.invoke( + [ + SystemMessage(content="Ты автор интерактивных текстовых историй."), + HumanMessage(content=prompt), + ] + ) + content = str(response.content) + setup, options = _parse_scene_response(content) + return { + "theme": theme, + "setup": setup, + "scene_text": setup, + "options": options, + } + + +def choice_and_ending(state: StoryState) -> dict: + setup = state.get("setup", "") + options = state.get("options", []) + + payload = { + "type": "choice", + "question": f"{setup}\n\nЧто делаем?", + "options": options, + } + resumed = interrupt(payload) + + if isinstance(resumed, dict): + user_choice = ( + resumed.get("user_answer") + or resumed.get("answer") + or resumed.get("choice") + or "" + ) + else: + user_choice = str(resumed) + + llm = build_llm() + ending_prompt = ( + f"Завязка: {setup}\n" + f"Выбор пользователя: {user_choice}\n" + "Допиши короткую концовку (2–3 предложения). Только текст концовки." + ) + ending_response = llm.invoke( + [ + SystemMessage(content="Ты автор интерактивных историй."), + HumanMessage(content=ending_prompt), + ] + ) + ending = str(ending_response.content).strip() + + return { + "user_choice": user_choice, + "ending": ending, + } + + +def build_graph(): + builder = StateGraph(StoryState) + builder.add_node("generate_scene", generate_scene) + builder.add_node("choice_and_ending", choice_and_ending) + builder.add_edge(START, "generate_scene") + builder.add_edge("generate_scene", "choice_and_ending") + memory = InMemorySaver() + return builder.compile(checkpointer=memory) + + +def _handle_interrupt(interrupts: tuple) -> dict: + first = interrupts[0] + payload = first.value if hasattr(first, "value") else first + if not isinstance(payload, dict): + payload = {"type": "choice", "question": str(payload), "options": []} + + print(f"\n[LLM] {payload.get('question', '')}\n") + + options = payload.get("options") or [] + if options: + choice = questionary.select( + "Выберите действие:", + choices=options, + ).ask() + else: + choice = questionary.text("Ваш выбор:").ask() + + payload = dict(payload) + payload["user_answer"] = choice or "" + return payload + + +def run_story(theme: str = THEME_DEFAULT) -> StoryState: + graph = build_graph() + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + inputs: StoryState | Command = {"theme": theme} + + final_state: StoryState = {"theme": theme} + + while True: + interrupted = False + for chunk in graph.stream(inputs, config=config, stream_mode="updates"): + if "__interrupt__" in chunk: + payload = _handle_interrupt(chunk["__interrupt__"]) + inputs = Command(resume=payload) + interrupted = True + break + + for node_name, update in chunk.items(): + if node_name == "__interrupt__": + continue + if isinstance(update, dict): + final_state.update(update) + if node_name == "generate_scene" and update.get("setup"): + print(f"\nТема: {theme}") + print(f"\n[LLM] {update['setup']}\n") + if node_name == "choice_and_ending" and update.get("ending"): + print(f"\n[LLM] {update['ending']}\n") + + if not interrupted: + break + + print("--- Итог ---") + print(f"Завязка: {final_state.get('setup', '')}") + print(f"Выбор: {final_state.get('user_choice', '')}") + print(f"Концовка: {final_state.get('ending', '')}") + return final_state + + +if __name__ == "__main__": + import sys + + story_theme = sys.argv[1] if len(sys.argv) > 1 else THEME_DEFAULT + run_story(story_theme)