173 lines
6.5 KiB
Python
173 lines
6.5 KiB
Python
import os
|
||
from typing import TypedDict, List, Optional
|
||
|
||
import questionary
|
||
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 interrupt, Command
|
||
|
||
|
||
# ---------- 1. Состояние графа ----------
|
||
|
||
class GraphState(TypedDict):
|
||
"""Состояние графа с отдельными полями для каждого этапа."""
|
||
topic: str # Тема / завязка (входной параметр)
|
||
setup: str # Сгенерированная завязка от LLM
|
||
choices: List[str] # Список вариантов выбора (3 штуки)
|
||
human_choice: str # Выбранный пользователем вариант
|
||
ending: str # Концовка от LLM
|
||
|
||
|
||
# ---------- 2. Узел: генерация сцены + прерывание + концовка ----------
|
||
|
||
def story_node(state: GraphState) -> GraphState:
|
||
"""
|
||
Шаг 1: генерирует завязку и варианты через LLM.
|
||
Шаг 2: вызывает interrupt() — граф ставится на паузу.
|
||
Шаг 3: после resume дописывает концовку через LLM.
|
||
"""
|
||
llm = ChatOpenAI(temperature=0.7)
|
||
|
||
# --- Шаг 1: генерируем завязку и варианты ---
|
||
topic = state["topic"]
|
||
scene_prompt = (
|
||
f"Тема: {topic}.\n"
|
||
"Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
|
||
"Ответь строго в формате:\n"
|
||
"ЗАВЯЗКА: <текст завязки>\n"
|
||
"ВАРИАНТЫ:\n"
|
||
"1) <первый вариант>\n"
|
||
"2) <второй вариант>\n"
|
||
"3) <третий вариант>"
|
||
)
|
||
scene_response = llm.invoke(scene_prompt)
|
||
raw = scene_response.content.strip()
|
||
|
||
# Парсим ответ: разделяем завязку и варианты
|
||
setup = ""
|
||
choices = []
|
||
|
||
lines = raw.splitlines()
|
||
in_choices = False
|
||
for line in lines:
|
||
line = line.strip()
|
||
if line.upper().startswith("ЗАВЯЗКА:"):
|
||
setup = line[len("ЗАВЯЗКА:"):].strip()
|
||
elif line.upper().startswith("ВАРИАНТЫ:"):
|
||
in_choices = True
|
||
elif in_choices and line:
|
||
# Убираем нумерацию вида "1)", "2)", "3)" или "1.", "2.", "3."
|
||
for prefix in ("1)", "2)", "3)", "1.", "2.", "3."):
|
||
if line.startswith(prefix):
|
||
line = line[len(prefix):].strip()
|
||
break
|
||
if line:
|
||
choices.append(line)
|
||
|
||
# Защита: если парсинг не дал нужных данных — берём весь текст как завязку
|
||
if not setup:
|
||
setup = raw
|
||
if len(choices) < 3:
|
||
choices = choices + [f"Вариант {i+1}" for i in range(len(choices), 3)]
|
||
|
||
# Сохраняем в состояние
|
||
state["setup"] = setup
|
||
state["choices"] = choices
|
||
|
||
# --- Шаг 2: прерывание ---
|
||
interrupt_payload = {
|
||
"type": "choice",
|
||
"question": setup + "\n\nЧто делаем?",
|
||
"allow_responds": choices,
|
||
}
|
||
|
||
# Выполнение останавливается здесь до Command(resume=...)
|
||
resume_value = interrupt(interrupt_payload)
|
||
|
||
# После resume в resume_value — словарь с полем "answer"
|
||
human_choice = resume_value.get("answer", choices[0])
|
||
state["human_choice"] = human_choice
|
||
|
||
# --- Шаг 3: генерируем концовку ---
|
||
ending_prompt = (
|
||
f"Завязка: {setup}\n"
|
||
f"Выбор пользователя: {human_choice}\n\n"
|
||
"Допиши короткую концовку истории (2–3 предложения)."
|
||
)
|
||
ending_response = llm.invoke(ending_prompt)
|
||
state["ending"] = ending_response.content.strip()
|
||
|
||
return state
|
||
|
||
|
||
# ---------- 3. Сборка графа ----------
|
||
|
||
builder = StateGraph(GraphState)
|
||
builder.add_node("story", story_node)
|
||
builder.add_edge(START, "story")
|
||
|
||
graph = builder.compile(checkpointer=InMemorySaver())
|
||
|
||
|
||
# ---------- 4. Цикл запуска с обработкой прерывания ----------
|
||
|
||
def main() -> None:
|
||
topic = input("Введите тему истории (например, «космический кот»): ").strip()
|
||
if not topic:
|
||
topic = "космический кот"
|
||
|
||
print(f"\nТема: {topic}\n")
|
||
|
||
thread_id = "story_thread_1"
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
|
||
initial_state: GraphState = {
|
||
"topic": topic,
|
||
"setup": "",
|
||
"choices": [],
|
||
"human_choice": "",
|
||
"ending": "",
|
||
}
|
||
|
||
# --- Первый запуск: до прерывания ---
|
||
for chunk in graph.stream(initial_state, config=config):
|
||
if "__interrupt__" in chunk:
|
||
interrupt_obj = chunk["__interrupt__"][0].value # наш payload
|
||
|
||
print(f"\n[LLM] {interrupt_obj['question']}\n")
|
||
|
||
# Показываем варианты через questionary
|
||
answer = questionary.select(
|
||
"Выберите действие:",
|
||
choices=interrupt_obj["allow_responds"],
|
||
).ask()
|
||
|
||
if answer is None:
|
||
raise SystemExit("Отмена пользователем.")
|
||
|
||
# Добавляем ответ в payload и возобновляем граф
|
||
interrupt_obj["answer"] = answer
|
||
|
||
# --- Возобновление через Command(resume=...) ---
|
||
for resumed_chunk in graph.stream(
|
||
Command(resume=interrupt_obj), config=config
|
||
):
|
||
# После resume граф дописывает концовку — просто ждём завершения
|
||
pass
|
||
|
||
# Получаем финальное состояние
|
||
final_state = graph.get_state(config).values
|
||
|
||
print("\n" + "=" * 50)
|
||
print("=== Итоговая история ===")
|
||
print("=" * 50)
|
||
print(f"\nЗавязка:\n{final_state['setup']}")
|
||
print(f"\nВыбор: {final_state['human_choice']}")
|
||
print(f"\n[LLM] {final_state['ending']}")
|
||
print("\n" + "=" * 50)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |