diff --git a/main.py b/main.py new file mode 100644 index 0000000..24fb66d --- /dev/null +++ b/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())