update main.py

main.py

@@ -0,0 +1,123 @@
+import os
+import asyncio
+from typing import TypedDict, List, Dict, Any
+
+# LangGraph imports
+from langgraph.graph import StateGraph, START
+from langgraph.types import interrupt, Command
+from langgraph.checkpoint.memory import InMemorySaver
+
+# LangChain LLM
+from langchain_openai import ChatOpenAI
+
+# Console interaction
+import questionary
+
+# ---------- 1. Define state ---------------------------------
+class GameState(TypedDict):
+    theme: str
+    intro: str
+    options: List[str]
+    choice: str | None
+    ending: str
+
+# ---------- 2. LLM client -----------------------------------
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
+    api_key=os.getenv("JOURNAL_MCP_PAT"),
+    temperature=0.7,
+)
+
+# ---------- 3. Node to generate intro and options ---------
+async def generate_intro(state: GameState) -> Dict:
+    theme = state["theme"]
+    prompt = (
+        f"You are a creative storyteller.\n"
+        f"Theme: {theme}.\n"
+        "Generate a short opening paragraph (2-3 sentences).\n"
+        "Then provide exactly three distinct actions the protagonist can take, one per line, prefixed with numbers 1) 2) 3).\n"
+        "Return only the intro followed by the numbered list."
+    )
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    text = response.content
+    # Split into intro and options
+    parts = text.split("\n")
+    intro_lines = []
+    options: List[str] = []
+    for line in parts:
+        if line.strip().startswith(tuple(str(i)+")" for i in range(1,4)):
+            options.append(line.strip())
+        else:
+            intro_lines.append(line)
+    state["intro"] = " ".join(intro_lines).strip()
+    state["options"] = options
+    # Trigger interrupt to ask user
+    payload = {
+        "type": "choice",
+        "question": f"{state['intro']}\n\nWhat do you do?",
+        "choices": options,
+    }
+    return interrupt(payload)
+
+# ---------- 4. Node to process choice and generate ending -----
+async def generate_ending(state: GameState) -> Dict:
+    # state now contains 'choice'
+    theme = state["theme"]
+    intro = state["intro"]
+    choice = state.get("choice", "")
+    prompt = (
+        f"You are a creative storyteller.\n"
+        f"Theme: {theme}.\n"
+        f"Intro: {intro}.\n"
+        f"The player chose: {choice}.\n"
+        "Write a short ending (2-3 sentences) that follows from this choice."
+    )
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    state["ending"] = response.content.strip()
+    return state
+
+# ---------- 5. Build graph ---------------------------------
+builder = StateGraph(GameState)
+builder.add_node("intro", generate_intro)
+builder.add_node("ending", generate_ending)
+builder.set_entry_point("intro")
+builder.add_edge(START, "intro")
+builder.add_edge("intro", "ending")
+# No edge from ending; graph ends after ending node
+graph = builder.compile(checkpointer=InMemorySaver())
+
+# ---------- 6. Run loop with interrupt handling ---------
+async def run_game(theme: str):
+    thread_id = f"thread-{theme.replace(' ', '_')}"
+    config = {"configurable": {"thread_id": thread_id}}
+    # Start stream
+    async for event in graph.astream({"theme": theme}, config, stream_mode="messages"):
+        if isinstance(event, tuple):
+            msg, _meta = event
+            print(msg.content)
+        elif "__interrupt__" in event:
+            interrupt_payload = event["__interrupt__"][0].value
+            # Show question and choices using questionary
+            answer = questionary.select(
+                interrupt_payload["question"],
+                choices=[c.split(') ',1)[1] if ') ' in c else c for c in interrupt_payload["choices"]]
+            ).ask()
+            # Map back to full choice string
+            full_choice = next(c for c in interrupt_payload["choices"] if c.endswith(answer))
+            # Resume with user answer
+            resume_payload = {**interrupt_payload, "choice": full_choice}
+            async for sub_event in graph.astream(Command(resume=resume_payload), config, stream_mode="messages"):
+                if isinstance(sub_event, tuple):
+                    msg, _meta = sub_event
+                    print(msg.content)
+    # After completion, print final state
+    final_state = await graph.aget_state(config)
+    print("\n--- Final Story ---")
+    print(final_state["intro"])
+    print(f"Choice: {final_state['choice']}")
+    print(final_state["ending"])
+
+if __name__ == "__main__":
+    theme = questionary.text("Enter a theme for the adventure:").ask()
+    asyncio.run(run_game(theme))