137 lines
4.5 KiB
Python
137 lines
4.5 KiB
Python
import os
|
||
import uuid
|
||
import asyncio
|
||
import questionary
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage, AIMessage
|
||
from langchain.tools import tool
|
||
from langgraph.graph import StateGraph, START, END, Command
|
||
from langgraph.graph.message import add_messages
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from deepagents import create_deep_agent
|
||
|
||
# LLM configuration – OpenRouter
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# ---------- Graph State ----------
|
||
class GameState(TypedDict):
|
||
messages: Annotated[list, add_messages]
|
||
theme: str
|
||
intro: str
|
||
options: list[str]
|
||
choice: str
|
||
ending: str
|
||
|
||
# ---------- Graph Nodes ----------
|
||
async def generate_scene(state: GameState) -> GameState:
|
||
theme = state["theme"]
|
||
prompt = (
|
||
f"Theme: {theme}\n"
|
||
"Generate a short introduction (2-3 sentences) followed by exactly three numbered options for the hero. "
|
||
"Output format: first line is the intro, next three lines are options numbered 1) 2) 3)."
|
||
)
|
||
response = await llm.ainvoke(HumanMessage(content=prompt))
|
||
text = response.content.strip()
|
||
lines = text.splitlines()
|
||
intro = lines[0].strip()
|
||
options = [line.strip() for line in lines[1:4]]
|
||
return {**state, "intro": intro, "options": options}
|
||
|
||
async def interrupt_choice(state: GameState) -> GameState:
|
||
payload = {
|
||
"type": "choice",
|
||
"question": f"{state['intro']}\nWhat do you do?",
|
||
"options": state["options"],
|
||
}
|
||
# Pause execution until user responds
|
||
await interrupt(payload)
|
||
# After resume, the payload will contain 'choice'
|
||
return {**state, "choice": payload["choice"]}
|
||
|
||
async def generate_ending(state: GameState) -> GameState:
|
||
prompt = (
|
||
f"Intro: {state['intro']}\n"
|
||
f"Choice: {state['choice']}\n"
|
||
"Write a short ending (2-3 sentences) for this story."
|
||
)
|
||
response = await llm.ainvoke(HumanMessage(content=prompt))
|
||
ending = response.content.strip()
|
||
return {**state, "ending": ending}
|
||
|
||
# ---------- Build Graph ----------
|
||
graph = StateGraph(GameState)
|
||
graph.add_node("scene", generate_scene)
|
||
graph.add_node("choice", interrupt_choice)
|
||
graph.add_node("ending", generate_ending)
|
||
|
||
graph.set_entry_point("scene")
|
||
graph.add_edge("scene", "choice")
|
||
graph.add_edge("choice", "ending")
|
||
graph.add_edge("ending", END)
|
||
|
||
checkpoint = InMemorySaver()
|
||
graph.compile(checkpointer=checkpoint)
|
||
|
||
# ---------- Game Runner ----------
|
||
async def run_game(theme: str) -> dict:
|
||
thread_id = str(uuid.uuid4())
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
state: GameState = {"messages": [], "theme": theme, "intro": "", "options": [], "choice": "", "ending": ""}
|
||
stream = graph.stream(state, config)
|
||
last_state = state
|
||
async for chunk in stream:
|
||
# Handle interrupt
|
||
if "__interrupt__" in chunk:
|
||
interrupt_payload = chunk["__interrupt__"][0].value
|
||
# Show options to user
|
||
answer = questionary.select(
|
||
interrupt_payload["question"],
|
||
choices=interrupt_payload["options"],
|
||
).ask()
|
||
# Resume with user's choice
|
||
resume_payload = {**interrupt_payload, "choice": answer}
|
||
stream = graph.stream(Command(resume=resume_payload), config)
|
||
continue
|
||
# Capture state when available
|
||
if "state" in chunk:
|
||
last_state = chunk["state"]
|
||
return last_state
|
||
|
||
# ---------- DeepAgent Tool ----------
|
||
@tool
|
||
def play_game(theme: str) -> str:
|
||
"""Play a choose-your-own-adventure game with the given theme."""
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
final_state = loop.run_until_complete(run_game(theme))
|
||
loop.close()
|
||
intro = final_state["intro"]
|
||
ending = final_state["ending"]
|
||
return f"\n{intro}\n\n{ending}\n"
|
||
|
||
# ---------- Create DeepAgent ----------
|
||
agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[play_game],
|
||
backend=None,
|
||
system_prompt="You are a game master that can play choose-your-own-adventure games.",
|
||
)
|
||
|
||
# ---------- CLI ----------
|
||
async def main():
|
||
theme = input("Enter a theme for the adventure: ")
|
||
result = await agent.ainvoke(
|
||
{"messages": [HumanMessage(content=f"Play a game with theme: {theme}")]},
|
||
{"configurable": {"thread_id": "cli-session"}},
|
||
)
|
||
# The tool returns the full story
|
||
print(result["messages"][-1].content)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|