Add main.py

main.py

@@ -0,0 +1,67 @@
+import os
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, START
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.prebuilt import interrupt_and_resume
+from typing import TypedDict, List
+import questionary
+
+class GameState(TypedDict):
+    intro: str | None
+    options: List[str] | None
+    choice: str | None
+    ending: str | None
+
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
+
+async def generate_intro_and_options(state: GameState):
+    prompt = (
+        "You are a creative storyteller. Generate an exciting intro for a choose-your-own-adventure game and 3 distinct options the player can choose from. Return JSON with keys 'intro' and 'options'."
+    )
+    res = await llm.ainvoke(prompt)
+    import json
+    data = json.loads(res.content)
+    state["intro"] = data.get("intro")
+    state["options"] = data.get("options")
+    return interrupt_and_resume(state, {
+        "type": "question",
+        "text": data.get("intro"),
+        "choices": data.get("options"),
+    })
+
+async def generate_ending(state: GameState):
+    prompt = (
+        f"Player chose: {state['choice']}. Write a short ending that follows the intro and incorporates the choice."
+    )
+    res = await llm.ainvoke(prompt)
+    state["ending"] = res.content
+    return state
+
+builder = StateGraph(GameState)
+builder.add_node("intro", generate_intro_and_options)
+builder.add_node("ending", generate_ending)
+builder.set_entry_point("intro")
+builder.add_edge(START, "intro")
+builder.add_conditional_edges(
+    "intro",
+    lambda x: "choice" if isinstance(x.get("__interrupt__"), dict) else None,
+)
+builder.add_edge("choice", "ending")
+graph = builder.compile(checkpointer=InMemorySaver())
+
+async def main():
+    state = await graph.astart()
+    while True:
+        if "__interrupt__" in state:
+            intr = state["__interrupt__"]
+            answer = questionary.select(intr["text"], choices=intr["choices"]).ask()
+            state["choice"] = answer
+            state = await graph.arun(state)
+        else:
+            break
+    print("\n--- Ending ---")
+    print(state.get("ending"))
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())