add main.py

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()