add main.py
This commit is contained in:
@@ -0,0 +1,225 @@
|
||||
"""Интерактивная история: 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)
|
||||
Reference in New Issue
Block a user