add main.py

main.py

@@ -0,0 +1,119 @@
+import os
+import questionary
+from typing import TypedDict, List, Optional
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.types import interrupt, Command
+from langgraph.checkpoint.memory import InMemorySaver
+
+from langchain_openai import ChatOpenAI
+
+# LLM setup – OpenRouter
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.7,
+)
+
+# State definition
+class GameState(TypedDict):
+    theme: str
+    scene: Optional[str]
+    options: Optional[List[str]]
+    choice: Optional[str]
+    ending: Optional[str]
+
+# Node: generate scene and options
+async def generate_scene(state: GameState) -> GameState:
+    prompt = (
+        f"Тема: {state['theme']}\n"
+        "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
+        "Ответь в формате: сначала текст завязки, затем каждая строка с вариантом,\n"
+        "по одному на строку, без нумерации."
+    )
+    response = await llm.ainvoke([{"role": "user", "content": prompt}])
+    text = response.content.strip()
+    parts = text.split("\n")
+    # first non-empty line is scene
+    scene = None
+    options: List[str] = []
+    for line in parts:
+        line = line.strip()
+        if not line:
+            continue
+        if scene is None:
+            scene = line
+        else:
+            options.append(line)
+    state["scene"] = scene
+    state["options"] = options
+    # interrupt for choice
+    payload = {
+        "type": "choice",
+        "question": f"{scene}\n\nЧто делаем?",
+        "options": options,
+    }
+    return interrupt(payload)
+
+# Node: finish story after choice
+async def finish_story(state: GameState) -> GameState:
+    # state now contains 'choice'
+    prompt = (
+        f"Завязка: {state['scene']}\n"
+        f"Выбор пользователя: {state['choice']}\n"
+        "Допиши короткую концовку (2–3 предложения)."
+    )
+    response = await llm.ainvoke([{"role": "user", "content": prompt}])
+    ending = response.content.strip()
+    state["ending"] = ending
+    return state
+
+# Build graph
+builder = StateGraph(GameState)
+builder.add_node("generate", generate_scene)
+builder.add_node("finish", finish_story)
+builder.add_edge(START, "generate")
+builder.add_edge("generate", "finish")
+builder.add_edge("finish", END)
+
+graph = builder.compile(checkpointer=InMemorySaver())
+
+# Runner
+async def run_game(theme: str):
+    state: GameState = {
+        "theme": theme,
+        "scene": None,
+        "options": None,
+        "choice": None,
+        "ending": None,
+    }
+    config = {"configurable": {"thread_id": "game1"}}
+    # start stream
+    stream = graph.stream(state, config)
+    async for chunk in stream:
+        if "__interrupt__" in chunk:
+            # handle interrupt
+            interrupt_payload = chunk["__interrupt__"][0].value
+            # show question and options
+            answer = questionary.select(
+                interrupt_payload["question"],
+                choices=interrupt_payload["options"],
+            ).ask()
+            # resume with answer
+            resume_payload = {**interrupt_payload, "choice": answer}
+            stream = graph.stream(Command(resume=resume_payload), config)
+            continue
+        # print normal messages
+        if "messages" in chunk:
+            for msg in chunk["messages"]:
+                if msg.role == "assistant":
+                    print(msg.content)
+    # after completion, print final state
+    print("\nИтоговое состояние:")
+    print(state)
+
+if __name__ == "__main__":
+    import asyncio
+    theme = questionary.text("Введите тему игры:").ask()
+    asyncio.run(run_game(theme))