add main.py

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)