diff --git a/game.py b/game.py new file mode 100644 index 0000000..7e37f40 --- /dev/null +++ b/game.py @@ -0,0 +1,192 @@ +""" +Interactive choose-your-own-adventure game using LangGraph and OpenAI. + +Run with: + python game.py + +Make sure you have an OpenAI API key set in the environment variable OPENAI_API_KEY. +""" + +import os +import re +from typing import TypedDict, List, Dict, Any + +from langgraph.graph import StateGraph, START +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.types import interrupt, Command +from langchain_openai import ChatOpenAI +import questionary + +# --------------------------------------------------------------------------- +# 1. Define the graph state +# --------------------------------------------------------------------------- +class GameState(TypedDict): + theme: str + story: str + options: List[str] + choice: str + ending: str + +# --------------------------------------------------------------------------- +# 2. Helper functions +# --------------------------------------------------------------------------- + +# Simple parser that expects the LLM to return: +# +# 1) +# 2) +# 3) +# The story may span multiple lines. + +def parse_story_and_options(text: str) -> tuple[str, List[str]]: + """Parse the LLM output into a story and a list of options. + + The function is tolerant to small formatting variations. + """ + # Split into lines + lines = text.strip().splitlines() + # Find the first line that starts with a digit and a closing parenthesis or dot + option_pattern = re.compile(r"^\s*\d+[\).]\s*(.*)") + story_lines = [] + options: List[str] = [] + for line in lines: + m = option_pattern.match(line) + if m: + options.append(m.group(1).strip()) + else: + story_lines.append(line) + story = "\n".join(story_lines).strip() + return story, options + +# --------------------------------------------------------------------------- +# 3. Graph nodes +# --------------------------------------------------------------------------- + +# LLM model +llm = ChatOpenAI(temperature=0.7, model="gpt-4o-mini") + +def generate_scene_and_interrupt(state: GameState) -> GameState: + """Generate the opening scene and options, then interrupt for user choice. + + The node returns a new state with story and options populated, and then + calls interrupt(). After the user responds, the graph resumes in this + same node with the updated state containing the choice. + """ + theme = state["theme"] + + # If we have already received a choice, generate the ending and finish. + if state.get("choice"): + # Generate ending + prompt = ( + f"Theme: {theme}\n" + f"Story: {state['story']}\n" + f"User choice: {state['choice']}\n" + "Write a short ending (2–3 sentences) that follows from the choice." + ) + ending = llm.invoke(prompt).content + state["ending"] = ending.strip() + return state + + # No choice yet: generate scene and options + prompt = ( + f"Theme: {theme}\n" + "Create a short opening scene (2–3 sentences) and exactly three numbered options for the hero to choose.\n" + "Respond in the following format:\n" + "\n" + "1) \n" + "2) \n" + "3) \n" + ) + raw = llm.invoke(prompt).content + story, options = parse_story_and_options(raw) + state["story"] = story + state["options"] = options + + # Prepare interrupt payload + interrupt_payload = { + "type": "choice", + "question": story, + "options": options, + } + # Interrupt the graph; the returned state will contain the user response + interrupt(interrupt_payload) + # After interrupt, the graph will resume in this same node with the updated state. + return state + +# --------------------------------------------------------------------------- +# 4. Build the graph +# --------------------------------------------------------------------------- + +graph_builder = StateGraph(GameState) +graph_builder.add_node("scene", generate_scene_and_interrupt) +# All paths go through the same node; we finish when ending is set. +graph_builder.set_entry_point("scene") +# The node will keep looping until an ending is produced. +# We use a simple condition: if ending exists, we finish. + +def is_finished(state: GameState) -> bool: + return bool(state.get("ending")) + +graph_builder.add_conditional_edges("scene", lambda s: "finished" if is_finished(s) else "scene") +graph_builder.add_edge("finished", "finished") # terminal + +# Compile the graph with a checkpoint +checkpoint = InMemorySaver() +graph = graph_builder.compile(checkpointer=checkpoint) + +# --------------------------------------------------------------------------- +# 5. Main loop handling interrupts +# --------------------------------------------------------------------------- + +def run_game(theme: str): + # Initial state + state: GameState = { + "theme": theme, + "story": "", + "options": [], + "choice": "", + "ending": "", + } + + thread_id = os.urandom(8).hex() + config = {"configurable": {"thread_id": thread_id}} + + # Start streaming + stream = graph.stream(Command(state), config) + try: + for chunk in stream: + # The chunk may contain an interrupt + if "__interrupt__" in chunk: + interrupt_payload = chunk["__interrupt__"][0].value + # Show question and options to user + answer = questionary.select( + interrupt_payload["question"], + choices=interrupt_payload["options"], + ).ask() + # Add the answer to the payload + interrupt_payload["choice"] = answer + # Resume the graph with the updated payload + stream = graph.stream(Command(resume=interrupt_payload), config) + continue + # Normal output: print to console + if "story" in chunk: + print(chunk["story"], end="\n\n") + if "ending" in chunk: + print("\n[LLM] " + chunk["ending"], end="\n\n") + except KeyboardInterrupt: + print("\nGame interrupted by user.") + + # Final state + final_state = graph.get_state(config) + print("--- Final state ---") + print("Theme:", final_state["theme"]) + print("Story:", final_state["story"]) + print("Choice:", final_state["choice"]) + print("Ending:", final_state["ending"]) + +if __name__ == "__main__": + print("Welcome to the LLM choose‑your‑own‑adventure game!") + theme = questionary.text("Enter a theme for the story:").ask() + if not theme: + theme = "A mysterious space cat" + run_game(theme) \ No newline at end of file