120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
import os
|
||
import questionary
|
||
from typing import TypedDict, List, Optional
|
||
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.types import interrupt, Command
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
# LLM setup – 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.7,
|
||
)
|
||
|
||
# State definition
|
||
class GameState(TypedDict):
|
||
theme: str
|
||
scene: Optional[str]
|
||
options: Optional[List[str]]
|
||
choice: Optional[str]
|
||
ending: Optional[str]
|
||
|
||
# Node: generate scene and options
|
||
async def generate_scene(state: GameState) -> GameState:
|
||
prompt = (
|
||
f"Тема: {state['theme']}\n"
|
||
"Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
|
||
"Ответь в формате: сначала текст завязки, затем каждая строка с вариантом,\n"
|
||
"по одному на строку, без нумерации."
|
||
)
|
||
response = await llm.ainvoke([{"role": "user", "content": prompt}])
|
||
text = response.content.strip()
|
||
parts = text.split("\n")
|
||
# first non-empty line is scene
|
||
scene = None
|
||
options: List[str] = []
|
||
for line in parts:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
if scene is None:
|
||
scene = line
|
||
else:
|
||
options.append(line)
|
||
state["scene"] = scene
|
||
state["options"] = options
|
||
# interrupt for choice
|
||
payload = {
|
||
"type": "choice",
|
||
"question": f"{scene}\n\nЧто делаем?",
|
||
"options": options,
|
||
}
|
||
return interrupt(payload)
|
||
|
||
# Node: finish story after choice
|
||
async def finish_story(state: GameState) -> GameState:
|
||
# state now contains 'choice'
|
||
prompt = (
|
||
f"Завязка: {state['scene']}\n"
|
||
f"Выбор пользователя: {state['choice']}\n"
|
||
"Допиши короткую концовку (2–3 предложения)."
|
||
)
|
||
response = await llm.ainvoke([{"role": "user", "content": prompt}])
|
||
ending = response.content.strip()
|
||
state["ending"] = ending
|
||
return state
|
||
|
||
# Build graph
|
||
builder = StateGraph(GameState)
|
||
builder.add_node("generate", generate_scene)
|
||
builder.add_node("finish", finish_story)
|
||
builder.add_edge(START, "generate")
|
||
builder.add_edge("generate", "finish")
|
||
builder.add_edge("finish", END)
|
||
|
||
graph = builder.compile(checkpointer=InMemorySaver())
|
||
|
||
# Runner
|
||
async def run_game(theme: str):
|
||
state: GameState = {
|
||
"theme": theme,
|
||
"scene": None,
|
||
"options": None,
|
||
"choice": None,
|
||
"ending": None,
|
||
}
|
||
config = {"configurable": {"thread_id": "game1"}}
|
||
# start stream
|
||
stream = graph.stream(state, config)
|
||
async for chunk in stream:
|
||
if "__interrupt__" in chunk:
|
||
# handle interrupt
|
||
interrupt_payload = chunk["__interrupt__"][0].value
|
||
# show question and options
|
||
answer = questionary.select(
|
||
interrupt_payload["question"],
|
||
choices=interrupt_payload["options"],
|
||
).ask()
|
||
# resume with answer
|
||
resume_payload = {**interrupt_payload, "choice": answer}
|
||
stream = graph.stream(Command(resume=resume_payload), config)
|
||
continue
|
||
# print normal messages
|
||
if "messages" in chunk:
|
||
for msg in chunk["messages"]:
|
||
if msg.role == "assistant":
|
||
print(msg.content)
|
||
# after completion, print final state
|
||
print("\nИтоговое состояние:")
|
||
print(state)
|
||
|
||
if __name__ == "__main__":
|
||
import asyncio
|
||
theme = questionary.text("Введите тему игры:").ask()
|
||
asyncio.run(run_game(theme))
|