123 lines
4.5 KiB
Python
123 lines
4.5 KiB
Python
import os
|
|
import asyncio
|
|
from typing import TypedDict, List, Dict, Any
|
|
|
|
# LangGraph imports
|
|
from langgraph.graph import StateGraph, START
|
|
from langgraph.types import interrupt, Command
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
# LangChain LLM
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
# Console interaction
|
|
import questionary
|
|
|
|
# ---------- 1. Define state ---------------------------------
|
|
class GameState(TypedDict):
|
|
theme: str
|
|
intro: str
|
|
options: List[str]
|
|
choice: str | None
|
|
ending: str
|
|
|
|
# ---------- 2. LLM client -----------------------------------
|
|
llm = ChatOpenAI(
|
|
model="openai/gpt-oss-20b:free",
|
|
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
|
api_key=os.getenv("JOURNAL_MCP_PAT"),
|
|
temperature=0.7,
|
|
)
|
|
|
|
# ---------- 3. Node to generate intro and options ---------
|
|
async def generate_intro(state: GameState) -> Dict:
|
|
theme = state["theme"]
|
|
prompt = (
|
|
f"You are a creative storyteller.\n"
|
|
f"Theme: {theme}.\n"
|
|
"Generate a short opening paragraph (2-3 sentences).\n"
|
|
"Then provide exactly three distinct actions the protagonist can take, one per line, prefixed with numbers 1) 2) 3).\n"
|
|
"Return only the intro followed by the numbered list."
|
|
)
|
|
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
|
text = response.content
|
|
# Split into intro and options
|
|
parts = text.split("\n")
|
|
intro_lines = []
|
|
options: List[str] = []
|
|
for line in parts:
|
|
if line.strip().startswith(tuple(str(i)+")" for i in range(1,4)):
|
|
options.append(line.strip())
|
|
else:
|
|
intro_lines.append(line)
|
|
state["intro"] = " ".join(intro_lines).strip()
|
|
state["options"] = options
|
|
# Trigger interrupt to ask user
|
|
payload = {
|
|
"type": "choice",
|
|
"question": f"{state['intro']}\n\nWhat do you do?",
|
|
"choices": options,
|
|
}
|
|
return interrupt(payload)
|
|
|
|
# ---------- 4. Node to process choice and generate ending -----
|
|
async def generate_ending(state: GameState) -> Dict:
|
|
# state now contains 'choice'
|
|
theme = state["theme"]
|
|
intro = state["intro"]
|
|
choice = state.get("choice", "")
|
|
prompt = (
|
|
f"You are a creative storyteller.\n"
|
|
f"Theme: {theme}.\n"
|
|
f"Intro: {intro}.\n"
|
|
f"The player chose: {choice}.\n"
|
|
"Write a short ending (2-3 sentences) that follows from this choice."
|
|
)
|
|
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
|
state["ending"] = response.content.strip()
|
|
return state
|
|
|
|
# ---------- 5. Build graph ---------------------------------
|
|
builder = StateGraph(GameState)
|
|
builder.add_node("intro", generate_intro)
|
|
builder.add_node("ending", generate_ending)
|
|
builder.set_entry_point("intro")
|
|
builder.add_edge(START, "intro")
|
|
builder.add_edge("intro", "ending")
|
|
# No edge from ending; graph ends after ending node
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
# ---------- 6. Run loop with interrupt handling ---------
|
|
async def run_game(theme: str):
|
|
thread_id = f"thread-{theme.replace(' ', '_')}"
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
# Start stream
|
|
async for event in graph.astream({"theme": theme}, config, stream_mode="messages"):
|
|
if isinstance(event, tuple):
|
|
msg, _meta = event
|
|
print(msg.content)
|
|
elif "__interrupt__" in event:
|
|
interrupt_payload = event["__interrupt__"][0].value
|
|
# Show question and choices using questionary
|
|
answer = questionary.select(
|
|
interrupt_payload["question"],
|
|
choices=[c.split(') ',1)[1] if ') ' in c else c for c in interrupt_payload["choices"]]
|
|
).ask()
|
|
# Map back to full choice string
|
|
full_choice = next(c for c in interrupt_payload["choices"] if c.endswith(answer))
|
|
# Resume with user answer
|
|
resume_payload = {**interrupt_payload, "choice": full_choice}
|
|
async for sub_event in graph.astream(Command(resume=resume_payload), config, stream_mode="messages"):
|
|
if isinstance(sub_event, tuple):
|
|
msg, _meta = sub_event
|
|
print(msg.content)
|
|
# After completion, print final state
|
|
final_state = await graph.aget_state(config)
|
|
print("\n--- Final Story ---")
|
|
print(final_state["intro"])
|
|
print(f"Choice: {final_state['choice']}")
|
|
print(final_state["ending"])
|
|
|
|
if __name__ == "__main__":
|
|
theme = questionary.text("Enter a theme for the adventure:").ask()
|
|
asyncio.run(run_game(theme)) |