текстовая игра на основе llm + interrupt: agent.py

2026-05-27 12:24:53 +00:00
parent 4b56939f0f
commit c3c1dac6e6
1 changed files with 167 additions and 0 deletions
@@ -0,0 +1,167 @@
+import os
+from typing import TypedDict, List, Dict, Any
+
+import questionary
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, START
+from langgraph.types import interrupt, Command, InMemorySaver
+from langgraph.checkpoint.memory import MemorySaver
+
+
+# -----------------------------
+# 1. Состояние графа
+# -----------------------------
+class StoryState(TypedDict):
+    theme: str | None
+    hook: str | None          # завязка от LLM
+    options: List[str] | None   # варианты выбора
+    choice: str | None        # выбранный пользователем вариант
+    ending: str | None        # концовка от LLM
+
+
+# -----------------------------
+# 2. Узел генерации сцены и прерывания
+# -----------------------------
+def generate_scene(state: StoryState) -> Dict[str, Any]:
+    """
+    Генерирует завязку и варианты действий.
+    После этого вызывает interrupt для выбора пользователя.
+    """
+    theme = state["theme"]
+    if not theme:
+        raise ValueError("Тема не задана")
+
+    # 2.1 Вызов LLM
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
+    prompt = (
+        f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) и ровно "
+        "три варианта поступка героя. Ответь в формате:\n"
+        "1) Завязка\n"
+        "2) Вариант 1, Вариант 2, Вариант 3"
+    )
+    response = llm.invoke({"input": prompt})
+    text = response.content.strip()
+
+    # 2.2 Парсим ответ
+    try:
+        hook_part, options_part = text.split("\n", 1)
+    except ValueError:
+        raise RuntimeError("LLM не вернул ожидаемый формат")
+
+    hook = hook_part.strip()
+    options_raw = options_part.replace(",", "\n").splitlines()
+    options = [opt.strip() for opt in options_raw if opt.strip()]
+    if len(options) != 3:
+        # fallback: split by commas
+        options = [o.strip() for o in options_part.split(",") if o.strip()]
+
+    state["hook"] = hook
+    state["options"] = options
+
+    # 2.3 Прерывание для выбора пользователя
+    interrupt_payload = {
+        "type": "choice",
+        "question": f"{hook}\n\nЧто делаем?",
+        "options": options,
+    }
+    return interrupt(interrupt_payload)
+
+
+# -----------------------------
+# 3. Узел завершения истории
+# -----------------------------
+def finish_story(state: StoryState) -> Dict[str, Any]:
+    """
+    После получения выбора пользователя генерирует концовку.
+    """
+    hook = state["hook"]
+    choice = state.get("choice")
+    if not hook or not choice:
+        raise RuntimeError("Недостаточно данных для генерации окончания")
+
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
+    prompt = (
+        f"Завязка: {hook}\n"
+        f"Выбор пользователя: {choice}\n"
+        "Допиши короткую концовку (2–3 предложения)."
+    )
+    response = llm.invoke({"input": prompt})
+    ending = response.content.strip()
+    state["ending"] = ending
+    return {"story_state": state}
+
+
+# -----------------------------
+# 4. Сборка графа
+# -----------------------------
+def create_graph() -> StateGraph[StoryState]:
+    graph = StateGraph(StoryState)
+    graph.add_node("generate_scene", generate_scene)
+    graph.add_node("finish_story", finish_story)
+
+    # Переходы: START → generate_scene → finish_story
+    graph.set_entry_point("generate_scene")
+    graph.add_edge("generate_scene", "finish_story")
+
+    # Чекпоинтер для сохранения состояния между прерываниями
+    graph.set_checkpoint_manager(InMemorySaver())
+    return graph
+
+
+# -----------------------------
+# 5. Запуск и обработка прерываний
+# -----------------------------
+def main():
+    theme = questionary.text("Введите тему истории:").ask()
+    if not theme:
+        print("Тема обязательна.")
+        return
+
+    # Инициализируем состояние
+    state: StoryState = {
+        "theme": theme,
+        "hook": None,
+        "options": None,
+        "choice": None,
+        "ending": None,
+    }
+
+    graph = create_graph()
+    config = {"configurable": {"thread_id": "story_thread"}}
+
+    # Запускаем граф
+    stream = graph.stream(state, config=config)
+
+    for chunk in stream:
+        if "__interrupt__" in chunk:
+            interrupt_payload = chunk["__interrupt__"][0].value  # dict с вопросом и вариантами
+            answer = questionary.select(
+                interrupt_payload["question"],
+                choices=interrupt_payload["options"]
+            ).ask()
+            if not answer:
+                print("Выбор не сделан. Завершаем.")
+                return
+
+            # Добавляем ответ в payload и возобновляем граф
+            interrupt_payload["choice"] = answer
+            resume_command = Command(resume=interrupt_payload)
+            stream = graph.stream(resume_command, config=config)
+        else:
+            # Выводим обычный вывод LLM (если есть)
+            if "story_state" in chunk and chunk["story_state"]["ending"]:
+                print("\n[LLM] Концовка:")
+                print(chunk["story_state"]["ending"])
+                break
+
+    # Финальное состояние
+    final_state = stream.final_state()
+    print("\n=== Итоговая история ===")
+    print(f"Тема: {final_state['theme']}")
+    print(f"\n{final_state['hook']}\n")
+    print(f"Выбор: {final_state['choice']}\n")
+    print(final_state["ending"])
+
+
+if __name__ == "__main__":
+    main()