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

2026-05-27 12:24:53 +00:00
parent c3c1dac6e6
commit c0ae33275b
1 changed files with 153 additions and 0 deletions
@@ -0,0 +1,153 @@
+import uuid
+from typing import TypedDict, List, Dict, Any
+
+import questionary
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, START, interrupt, Command
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.types import GraphState
+
+
+# ---------- 1. Состояние графа ----------
+class StoryState(TypedDict):
+    theme: str
+    scene_text: str
+    choices: List[str]
+    choice_selected: str
+    ending: str
+
+
+# ---------- 2. LLM и узлы ----------
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
+
+
+def parse_llm_output(output: str) -> Dict[str, Any]:
+    """
+    Ожидаем формат:
+        <завязка>
+        1) вариант
+        2) вариант
+        3) вариант
+    Или варианты через запятую.
+    """
+    lines = [line.strip() for line in output.splitlines() if line.strip()]
+    scene_text = lines[0]
+    # Остальные строки считаем вариантами
+    raw_choices = "\n".join(lines[1:])
+    # Разделяем по цифрам или запятой
+    import re
+
+    # Если есть нумерация 1) 2) ...
+    numbered = re.findall(r"\d+\)\s*(.+)", raw_choices)
+    if numbered:
+        choices = [c.strip() for c in numbered]
+    else:
+        # Разделяем по запятой
+        choices = [c.strip() for c in raw_choices.split(",") if c.strip()]
+    return {"scene_text": scene_text, "choices": choices}
+
+
+def generate_scene(state: StoryState) -> Dict[str, Any]:
+    theme = state["theme"]
+    prompt = (
+        f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя. "
+        "Ответь в формате:\n"
+        "1) Завязка\n"
+        "2) Вариант 1, вариант 2, вариант 3"
+    )
+    response = llm.invoke(prompt)
+    parsed = parse_llm_output(response.content)
+
+    # Сохраняем сцену и варианты
+    state["scene_text"] = parsed["scene_text"]
+    state["choices"] = parsed["choices"]
+
+    # Подготавливаем прерывание
+    interrupt_payload = {
+        "type": "choice",
+        "question": f"{parsed['scene_text']}\n\nЧто делаем?",
+        "options": parsed["choices"],
+    }
+    return interrupt(interrupt_payload)
+
+
+def add_ending(state: StoryState) -> Dict[str, Any]:
+    # state уже содержит scene_text и choice_selected
+    prompt = (
+        f"Завязка: {state['scene_text']}\n"
+        f"Выбор пользователя: {state['choice_selected']}\n\n"
+        "Допиши короткую концовку (2–3 предложения)."
+    )
+    response = llm.invoke(prompt)
+    state["ending"] = response.content.strip()
+    return state
+
+
+# ---------- 3. Сборка графа ----------
+def create_graph() -> StateGraph:
+    builder = StateGraph(StoryState)
+
+    # Узлы
+    builder.add_node("generate_scene", generate_scene)
+    builder.add_node("add_ending", add_ending)
+
+    # Переходы
+    builder.set_entry_point("generate_scene")
+    builder.add_edge("generate_scene", "add_ending")
+    builder.add_edge("add_ending", END)
+
+    return builder.compile(checkpointer=InMemorySaver())
+
+
+# ---------- 4. Клиент ----------
+def main():
+    theme = questionary.text("Введите тему истории:").ask()
+    if not theme:
+        print("Тема обязательна.")
+        return
+
+    # Инициализируем состояние
+    init_state: StoryState = {
+        "theme": theme,
+        "scene_text": "",
+        "choices": [],
+        "choice_selected": "",
+        "ending": "",
+    }
+
+    graph = create_graph()
+    thread_id = str(uuid.uuid4())
+    config = {"configurable": {"thread_id": thread_id}}
+
+    # Запускаем первый поток
+    stream = graph.stream(init_state, config)
+
+    for chunk in stream:
+        if "__interrupt__" in chunk:
+            interrupt_payload = chunk["__interrupt__"][0].value  # dict with question and options
+            answer = questionary.select(
+                interrupt_payload["question"],
+                choices=interrupt_payload["options"],
+            ).ask()
+            if not answer:
+                print("Выбор не сделан. Завершаем.")
+                return
+
+            # Добавляем ответ в payload и возобновляем граф
+            interrupt_payload["answer"] = answer
+            resume_command = Command(resume=interrupt_payload)
+            stream = graph.stream(resume_command, config)
+
+        elif "__final_state__" in chunk:
+            final_state: StoryState = chunk["__final_state__"]
+            print("\n--- Итоговая история ---")
+            print(f"\n{final_state['scene_text']}\n")
+            print(f"Выбор: {final_state['choice_selected']}\n")
+            print(f"{final_state['ending']}")
+        else:
+            # Вывод промежуточных сообщений (если есть)
+            pass
+
+
+if __name__ == "__main__":
+    main()