226 lines
6.9 KiB
Python
226 lines
6.9 KiB
Python
"""Интерактивная история: LLM + LangGraph interrupt + questionary."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import uuid
|
|
from typing import TypedDict
|
|
|
|
import questionary
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.constants import START
|
|
from langgraph.graph import StateGraph
|
|
from langgraph.types import Command, interrupt
|
|
|
|
THEME_DEFAULT = "космический кот"
|
|
|
|
LLM_MODEL = os.getenv("OPENROUTER_MODEL", "openai/gpt-oss-20b:free")
|
|
LLM_BASE_URL = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
|
|
|
|
|
|
class StoryState(TypedDict, total=False):
|
|
theme: str
|
|
setup: str
|
|
scene_text: str
|
|
options: list[str]
|
|
user_choice: str
|
|
ending: str
|
|
|
|
|
|
def build_llm() -> ChatOpenAI:
|
|
return ChatOpenAI(
|
|
model=LLM_MODEL,
|
|
base_url=LLM_BASE_URL,
|
|
api_key=os.getenv("OPENAI_API_KEY", "fake"),
|
|
temperature=0.8,
|
|
)
|
|
|
|
|
|
def _parse_scene_response(text: str) -> tuple[str, list[str]]:
|
|
"""Разбирает ответ LLM на завязку и 3 варианта."""
|
|
options: list[str] = []
|
|
setup = text.strip()
|
|
|
|
variants_match = re.search(
|
|
r"ВАРИАНТЫ:\s*(.*)",
|
|
text,
|
|
flags=re.IGNORECASE | re.DOTALL,
|
|
)
|
|
if variants_match:
|
|
setup = text[: variants_match.start()].strip()
|
|
block = variants_match.group(1)
|
|
for line in block.splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
cleaned = re.sub(r"^\d+[\).\]]\s*", "", line).strip()
|
|
if cleaned:
|
|
options.append(cleaned)
|
|
|
|
if len(options) < 3:
|
|
numbered = re.findall(r"^\s*\d+[\).\]]\s*(.+)$", text, flags=re.MULTILINE)
|
|
options = [o.strip() for o in numbered if o.strip()]
|
|
|
|
if len(options) < 3:
|
|
parts = [p.strip() for p in re.split(r"[,;]", text) if p.strip()]
|
|
if len(parts) >= 4:
|
|
setup = parts[0]
|
|
options = parts[1:4]
|
|
|
|
while len(options) < 3:
|
|
options.append(f"Вариант {len(options) + 1}")
|
|
|
|
return setup, options[:3]
|
|
|
|
|
|
def generate_scene(state: StoryState) -> dict:
|
|
theme = state.get("theme", THEME_DEFAULT)
|
|
llm = build_llm()
|
|
prompt = (
|
|
f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) "
|
|
"и ровно 3 варианта поступка героя.\n"
|
|
"Формат ответа строго:\n"
|
|
"ЗАВЯЗКА:\n"
|
|
"<текст завязки>\n"
|
|
"ВАРИАНТЫ:\n"
|
|
"1) <вариант 1>\n"
|
|
"2) <вариант 2>\n"
|
|
"3) <вариант 3>"
|
|
)
|
|
response = llm.invoke(
|
|
[
|
|
SystemMessage(content="Ты автор интерактивных текстовых историй."),
|
|
HumanMessage(content=prompt),
|
|
]
|
|
)
|
|
content = str(response.content)
|
|
setup, options = _parse_scene_response(content)
|
|
return {
|
|
"theme": theme,
|
|
"setup": setup,
|
|
"scene_text": setup,
|
|
"options": options,
|
|
}
|
|
|
|
|
|
def choice_and_ending(state: StoryState) -> dict:
|
|
setup = state.get("setup", "")
|
|
options = state.get("options", [])
|
|
|
|
payload = {
|
|
"type": "choice",
|
|
"question": f"{setup}\n\nЧто делаем?",
|
|
"options": options,
|
|
}
|
|
resumed = interrupt(payload)
|
|
|
|
if isinstance(resumed, dict):
|
|
user_choice = (
|
|
resumed.get("user_answer")
|
|
or resumed.get("answer")
|
|
or resumed.get("choice")
|
|
or ""
|
|
)
|
|
else:
|
|
user_choice = str(resumed)
|
|
|
|
llm = build_llm()
|
|
ending_prompt = (
|
|
f"Завязка: {setup}\n"
|
|
f"Выбор пользователя: {user_choice}\n"
|
|
"Допиши короткую концовку (2–3 предложения). Только текст концовки."
|
|
)
|
|
ending_response = llm.invoke(
|
|
[
|
|
SystemMessage(content="Ты автор интерактивных историй."),
|
|
HumanMessage(content=ending_prompt),
|
|
]
|
|
)
|
|
ending = str(ending_response.content).strip()
|
|
|
|
return {
|
|
"user_choice": user_choice,
|
|
"ending": ending,
|
|
}
|
|
|
|
|
|
def build_graph():
|
|
builder = StateGraph(StoryState)
|
|
builder.add_node("generate_scene", generate_scene)
|
|
builder.add_node("choice_and_ending", choice_and_ending)
|
|
builder.add_edge(START, "generate_scene")
|
|
builder.add_edge("generate_scene", "choice_and_ending")
|
|
memory = InMemorySaver()
|
|
return builder.compile(checkpointer=memory)
|
|
|
|
|
|
def _handle_interrupt(interrupts: tuple) -> dict:
|
|
first = interrupts[0]
|
|
payload = first.value if hasattr(first, "value") else first
|
|
if not isinstance(payload, dict):
|
|
payload = {"type": "choice", "question": str(payload), "options": []}
|
|
|
|
print(f"\n[LLM] {payload.get('question', '')}\n")
|
|
|
|
options = payload.get("options") or []
|
|
if options:
|
|
choice = questionary.select(
|
|
"Выберите действие:",
|
|
choices=options,
|
|
).ask()
|
|
else:
|
|
choice = questionary.text("Ваш выбор:").ask()
|
|
|
|
payload = dict(payload)
|
|
payload["user_answer"] = choice or ""
|
|
return payload
|
|
|
|
|
|
def run_story(theme: str = THEME_DEFAULT) -> StoryState:
|
|
graph = build_graph()
|
|
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
inputs: StoryState | Command = {"theme": theme}
|
|
|
|
final_state: StoryState = {"theme": theme}
|
|
|
|
while True:
|
|
interrupted = False
|
|
for chunk in graph.stream(inputs, config=config, stream_mode="updates"):
|
|
if "__interrupt__" in chunk:
|
|
payload = _handle_interrupt(chunk["__interrupt__"])
|
|
inputs = Command(resume=payload)
|
|
interrupted = True
|
|
break
|
|
|
|
for node_name, update in chunk.items():
|
|
if node_name == "__interrupt__":
|
|
continue
|
|
if isinstance(update, dict):
|
|
final_state.update(update)
|
|
if node_name == "generate_scene" and update.get("setup"):
|
|
print(f"\nТема: {theme}")
|
|
print(f"\n[LLM] {update['setup']}\n")
|
|
if node_name == "choice_and_ending" and update.get("ending"):
|
|
print(f"\n[LLM] {update['ending']}\n")
|
|
|
|
if not interrupted:
|
|
break
|
|
|
|
print("--- Итог ---")
|
|
print(f"Завязка: {final_state.get('setup', '')}")
|
|
print(f"Выбор: {final_state.get('user_choice', '')}")
|
|
print(f"Концовка: {final_state.get('ending', '')}")
|
|
return final_state
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
story_theme = sys.argv[1] if len(sys.argv) > 1 else THEME_DEFAULT
|
|
run_story(story_theme)
|