192 lines
6.5 KiB
Python
192 lines
6.5 KiB
Python
"""
|
||
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:
|
||
# <story>
|
||
# 1) <option1>
|
||
# 2) <option2>
|
||
# 3) <option3>
|
||
# 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"
|
||
"<story>\n"
|
||
"1) <option1>\n"
|
||
"2) <option2>\n"
|
||
"3) <option3>\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) |