103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
from langchain_openai import ChatOpenAI
|
|
from pydantic import SecretStr
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.types import interrupt, Command
|
|
import questionary
|
|
|
|
# LLM configuration with placeholders
|
|
llm = ChatOpenAI(
|
|
model="openai/gpt-oss-20b",
|
|
base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
|
|
api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
|
|
temperature=0.7,
|
|
)
|
|
|
|
# State definition
|
|
class GameState(dict):
|
|
topic: str | None = None
|
|
intro: str | None = None
|
|
options: list[str] | None = None
|
|
choice: str | None = None
|
|
ending: str | None = None
|
|
|
|
# Node to generate intro and options
|
|
def generate_intro(state: GameState) -> GameState:
|
|
topic = state.get("topic", "неизвестная тема")
|
|
prompt = f"""Тема: {topic}. Придумай короткую завязку (2–3 предложения) и ровно 3 варианта действий героя. Ответь в формате:
|
|
Завязка:
|
|
<текст>
|
|
Варианты:
|
|
1) <вариант1>
|
|
2) <вариант2>
|
|
3) <вариант3>"""
|
|
response = llm.invoke(prompt).content
|
|
# Parse response
|
|
parts = response.split("Варианты:")
|
|
intro_part = parts[0].replace("Завязка:", "").strip()
|
|
options_part = parts[1] if len(parts) > 1 else ""
|
|
options = [opt.strip() for opt in options_part.splitlines() if opt.strip()]
|
|
state["intro"] = intro_part
|
|
state["options"] = options
|
|
# Interrupt to ask user
|
|
return interrupt(
|
|
{
|
|
"type": "choice",
|
|
"question": f"{intro_part}\n\nЧто делаем?",
|
|
"options": options,
|
|
}
|
|
)
|
|
|
|
# Node to handle resume after choice and generate ending
|
|
def generate_ending(state: GameState) -> GameState:
|
|
intro = state.get("intro", "")
|
|
choice = state.get("choice", "")
|
|
prompt = f"""Завязка: {intro}
|
|
Выбор пользователя: {choice}
|
|
Допиши короткую концовку (2–3 предложения)."""
|
|
ending = llm.invoke(prompt).content.strip()
|
|
state["ending"] = ending
|
|
return state
|
|
|
|
# Build graph
|
|
graph_builder = StateGraph(GameState)
|
|
graph_builder.add_node("generate_intro", generate_intro)
|
|
graph_builder.add_node("generate_ending", generate_ending)
|
|
graph_builder.set_entry_point("generate_intro")
|
|
graph_builder.add_edge(START, "generate_intro")
|
|
graph_builder.add_edge("generate_intro", "generate_ending") # continue after interrupt
|
|
graph_builder.add_edge("generate_ending", END)
|
|
|
|
# Compile with checkpoint
|
|
checkpoint = InMemorySaver()
|
|
app = graph_builder.compile(checkpointer=checkpoint)
|
|
|
|
def run_game(topic: str):
|
|
state = GameState({"topic": topic})
|
|
config = {"configurable": {"thread_id": "game_thread"}}
|
|
# First stream until interrupt
|
|
for chunk in app.stream(state, config):
|
|
if "__interrupt__" in chunk:
|
|
payload = chunk["__interrupt__"][0].value # dict from interrupt()
|
|
answer = questionary.select(
|
|
payload["question"],
|
|
choices=payload["options"]
|
|
).ask()
|
|
# Resume with user's choice
|
|
resume_payload = {"choice": answer}
|
|
for res_chunk in app.stream(Command(resume=resume_payload), config):
|
|
if "ending" in res_chunk:
|
|
print("\n[LLM] Концовка:", res_chunk["ending"])
|
|
# Final state
|
|
final_state = checkpoint.get("game_thread")
|
|
print("\nИтоговое состояние:")
|
|
print("Тема:", final_state.get("topic"))
|
|
print("Завязка:", final_state.get("intro"))
|
|
print("Выбор:", final_state.get("choice"))
|
|
print("Концовка:", final_state.get("ending"))
|
|
|
|
if __name__ == "__main__":
|
|
topic_input = input("Введите тему истории: ").strip()
|
|
if not topic_input:
|
|
topic_input = "космический кот"
|
|
run_game(topic_input) |