Добавлен solution.py

solution.py

@@ -1,131 +1,114 @@
-# solution.py
+from typing import TypedDict
+import json
 
-"""
+import questionary
-Мини‑игра «выбери свою историю» с прерыванием.
-Использует:
-- LangGraph для построения графа и управления прерываниями
-- LangChain + langchain-openai для вызова LLM (ChatOpenAI)
-- questionary для интерактивного выбора в консоли
-"""
-
-# ────────────────────── Импорты ──────────────────────
-
-from typing import TypedDict, List, Dict, Any
-
-import questionary  # pip install questionary
-
-# LangGraph
 from langgraph.graph import StateGraph
 from langgraph.constants import START
 from langgraph.types import interrupt, Command
 from langgraph.checkpoint.memory import InMemorySaver
-
-# LangChain + OpenAI LLM
 from langchain_openai import ChatOpenAI
+from rich.console import Console
+
 
 # ────────────────────── Состояние графа ──────────────────────
 
 class State(TypedDict):
     """Состояние графа."""
-    # Текст истории, генерируемый LLM
+    story: str                # Текст истории, генерируемый LLM
-    story: str
+    options: list[str]        # Варианты действий
-    # Ответ пользователя после прерывания
+    choice: str | None        # Выбор пользователя
-    human_value: str | None
+    ending: str | None        # Концовка
 
 
-# ────────────────────── Узел с прерыванием ──────────────────────
+# ────────────────────── Инструмент LLM ──────────────────────
 
-def interrupt_node(state: State) -> State:
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
-    """
+console = Console()
-    Генерирует часть истории и ставит граф на паузу,
+
-    ожидая выбора игрока.
+
-    """
+def generate_story(state: State) -> State:
-    # Запрос к LLM для создания начала истории
+    """Генерирует начало истории и варианты действий."""
-    llm = ChatOpenAI(model="gpt-4o-mini")  # можно заменить на любой доступный LLM
     prompt = (
-        "Создай короткую интригующую историю про героя, "
+        "Generate a short adventure intro and three choices in JSON format.\n"
-        "который стоит перед выбором. В конце предложи три варианта действий."
+        "The JSON should have keys 'intro' (string) and 'choices' (list of strings)."
     )
-    response = llm.invoke(prompt)
+    response = llm.invoke({"messages": [{"role": "system", "content": prompt}]}).content
-    state["story"] = response.content
+    data = json.loads(response)
 
-    # Прерывание: пользователь выбирает вариант
+    state["story"] = data["intro"]
-    interrupt_payload = {
+    state["options"] = data["choices"]
-        "type": "alert",
+
-        "question": "Выберите действие героя:",
+    # Пауза для выбора игрока
-        "allow_responds": ["Выйти", "Бороться", "Спрятаться"],
+    return interrupt(state, {
-    }
+        "type": "choice",
-    # Возврат прерывания; выполнение не продолжится до возобновления
+        "question": f"{state['story']}\nВыберите действие:",
-    return interrupt(interrupt_payload)
+        "options": state["options"],
+    })
+
+
+def finish_story(state: State) -> State:
+    """Генерирует окончание истории на основе выбранного варианта."""
+    prompt = (
+        f"Write an ending for the story based on the choice '{state['choice']}'. "
+        "Keep it concise and satisfying."
+    )
+    response = llm.invoke({"messages": [{"role": "system", "content": prompt}]}).content
+    state["ending"] = response
+    return state
 
 
 # ────────────────────── Граф ──────────────────────
 
-graph_builder = StateGraph(State)
+builder = StateGraph(State)
+builder.add_node("generate", generate_story)
+builder.add_node("finish", finish_story)
 
-# Добавляем узел и переход из START в него
+builder.set_entry_point("generate")
-graph_builder.add_node("interrupt", interrupt_node)
+builder.add_edge(START, "generate")
-graph_builder.set_entry_point(START)
+builder.add_conditional_edges(
-graph_builder.add_edge(START, "interrupt")
+    "generate",
-# После прерывания возвращаемся к тому же узлу для завершения
+    lambda x: True if x.get("choice") else None,
-graph_builder.add_edge("interrupt", "interrupt")
+    {True: "finish"},
+)
+builder.add_edge("finish", START)  # цикл можно завершить здесь
 
-# Чекпоинтер (InMemorySaver) сохраняет состояние между паузами
+graph = builder.compile(checkpointer=InMemorySaver())
-checkpoint = InMemorySaver()
-graph = graph_builder.compile(checkpointer=checkpoint)
 
 
 # ────────────────────── Запуск и обработка прерываний ──────────────────────
 
 def main() -> None:
-    """
-    Основной цикл запуска графа с обработкой интерактивных прерываний.
-    """
-    # Уникальный идентификатор потока (можно использовать любой строковый ID)
     thread_id = "interactive_story"
 
-    # Инициализируем состояние
+    init_state: State = {"story": "", "options": [], "choice": None, "ending": None}
-    init_state: State = {"story": "", "human_value": None}
+    config = {"configurable": {"thread_id": thread_id}}
 
     # Запускаем поток генерации
-    stream = graph.stream(
+    stream = graph.stream(Command(resume=init_state), config=config)
-        Command(resume=init_state),
-        configurable={"thread_id": thread_id},
-    )
 
     for chunk in stream:
-        # Если в чанке есть прерывание, обрабатываем его
         if "__interrupt__" in chunk:
-            interrupt_info = chunk["__interrupt__"][0].value  # dict с payload
+            payload = chunk["__interrupt__"][0].value
-            rprint(f"\n=== Прервано: {interrupt_info['type']} ===")
-            rprint(interrupt_info["question"])
-            # Выводим варианты и получаем выбор пользователя
             answer = questionary.select(
-                interrupt_info["question"],
+                payload["question"],
-                choices=interrupt_info["allow_responds"],
+                choices=payload["options"],
             ).ask()
             if answer is None:
                 raise RuntimeError("Пользователь отменил ввод.")
-            rprint(f"\nВы выбрали: {answer}")
+            # Возобновляем граф с выбранным вариантом
-
+            stream = graph.stream(Command(resume={"choice": answer}), config=config)
-            # Добавляем ответ в payload и возобновляем граф
-            interrupt_info["human_value"] = answer
-            stream = graph.stream(
-                Command(resume=interrupt_info),
-                configurable={"thread_id": thread_id},
-            )
             continue
 
-        # Выводим обычные сообщения от LLM (если есть)
+        # Выводим сообщения от LLM
-        if "story" in chunk:
+        if "story" in chunk and chunk["story"]:
-            rprint(f"\nИстория: {chunk['story']}")
+            console.print(f"[bold cyan]История:[/]\n{chunk['story']}")
-        if "human_value" in chunk and chunk["human_value"]:
+        if "ending" in chunk and chunk["ending"]:
-            rprint(f"\nОтвет игрока сохранён в состоянии: {chunk['human_value']}")
+            console.print(f"\n[bold green]Концовка:[/]\n{chunk['ending']}")
 
-    # После завершения выводим итоговое состояние
+    # После завершения выводим итоговую историю
-    final_state = graph.get_state(configurable={"thread_id": thread_id})
+    final_state = graph.get_state(config=config)
-    rprint("\n=== Итоговое состояние ===")
+    full_story = f"{final_state['story']}\n\n{final_state['ending']}"
-    rprint(final_state)
+    console.print("\n[bold magenta]Итоговая история:[/]\n" + full_story)
 
 
 if __name__ == "__main__":