68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
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())
|