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

solutions/69b1a07c67bbf488a1177da4_текстовая_игра_на_основе_llm___interrupt/client.py

@@ -3,14 +3,12 @@ from typing import TypedDict, List
 
 import questionary
 from langchain_openai import ChatOpenAI
-from langchain.agents import AgentExecutor, create_agent
-from langchain.schema import HumanMessage, SystemMessage
-from langchain.tools import BaseTool
-from langchain.callbacks.human_in_the_loop import HumanInTheLoopMiddleware
-from langchain.memory import ConversationBufferMemory
+from langgraph.graph import StateGraph, START
+from langgraph.types import interrupt, Command
+from langgraph.checkpoint.memory import InMemorySaver
 
 
-# ---------- 1. Состояние (используется в памяти) ----------
+# ---------- 1. Состояние графа ----------
 class StoryState(TypedDict):
     theme: str
     scene_text: str
@@ -59,17 +57,14 @@ def generate_scene(state: StoryState) -> dict:
     state["scene_text"] = parsed["scene_text"]
     state["choices"] = parsed["choices"]
 
-    # Запрос к пользователю через HumanInTheLoopMiddleware
-    return {
-        "messages": [
-            SystemMessage(
-                content=f"{parsed['scene_text']}\n\nЧто делаем?"
-            ),
-            *[
-                HumanMessage(content=choice) for choice in parsed["choices"]
-            ],
-        ]
-    }
+    # Создаём прерывание с вопросом и вариантами
+    return interrupt(
+        {
+            "type": "choice",
+            "question": f"{parsed['scene_text']}\n\nЧто делаем?",
+            "options": parsed["choices"],
+        }
+    )
 
 
 def add_ending(state: StoryState) -> dict:
@@ -80,34 +75,35 @@ def add_ending(state: StoryState) -> dict:
     )
     response = llm.invoke(prompt)
     state["ending"] = response.content.strip()
-    return {"messages": [HumanMessage(content=state["ending"])]}
+    return {"messages": [response]}
 
 
-# ---------- 3. Создание агента ----------
-def create_agent_executor() -> AgentExecutor:
-    # Определяем инструменты (здесь нет внешних, но нужны для агента)
-    tools: List[BaseTool] = []
+# ---------- 3. Создание графа ----------
+def create_graph() -> StateGraph:
+    graph = StateGraph(StoryState)
 
-    agent = create_agent(
-        llm=llm,
-        tools=tools,
-        system_message="Ты создаёшь интерактивную историю. После генерации сцены запрашивай выбор у пользователя.",
-        verbose=False,
-    )
+    # Узел генерации сцены
+    def scene_node(state: StoryState) -> dict:
+        return generate_scene(state)
 
-    memory = ConversationBufferMemory(return_messages=True)
+    # Узел добавления концовки
+    def ending_node(state: StoryState) -> dict:
+        return add_ending(state)
 
-    executor = AgentExecutor(agent=agent, tools=tools, memory=memory, verbose=False)
-    return executor
+    graph.add_node("scene", scene_node)
+    graph.add_node("ending", ending_node)
+
+    # Переходы
+    graph.set_entry_point("scene")
+    graph.add_edge("scene", "ending")
+
+    # Чекпоинтер для возобновления после прерывания
+    graph.compile(checkpointer=InMemorySaver())
+    return graph
 
 
-# ---------- 4. Клиент ----------
-def main():
-    theme = questionary.text("Введите тему истории:").ask()
-    if not theme:
-        print("Тема обязательна.")
-        return
-
+# ---------- 4. Запуск и обработка прерываний ----------
+def run_story(theme: str):
     state: StoryState = {
         "theme": theme,
         "scene_text": "",
@@ -116,37 +112,51 @@ def main():
         "ending": "",
     }
 
-    executor = create_agent_executor()
+    graph = create_graph()
+    config = {"thread_id": str(uuid.uuid4())}
 
-    # Middleware для прерывания и возобновления
-    middleware = HumanInTheLoopMiddleware(
-        prompt=lambda x: x["messages"][0].content,
-        choices=lambda x: x["messages"][1:],  # список HumanMessage с вариантами
-    )
+    # Запускаем граф в режиме stream
+    stream = graph.stream(state, config=config)
 
-    # Запускаем генерацию сцены
-    result = executor.invoke({"state": state}, callbacks=[middleware])
+    for chunk in stream:
+        if "__interrupt__" in chunk:
+            interrupt_payload = chunk["__interrupt__"][0]["value"]
+            question = interrupt_payload["question"]
+            options = interrupt_payload["options"]
 
-    # После прерывания пользователь выберет вариант
-    if middleware.interrupted:
-        answer = questionary.select(
-            middleware.prompt, choices=middleware.choices_texts
-        ).ask()
-        if not answer:
-            print("Выбор не сделан. Завершаем.")
-            return
+            answer = questionary.select(question, choices=options).ask()
+            if not answer:
+                print("Выбор не сделан. Завершаем.")
+                return
 
-        state["choice_selected"] = answer
+            # Возобновляем граф с ответом пользователя
+            resume_payload = {
+                "type": "choice",
+                "question": question,
+                "options": options,
+                "answer": answer,
+            }
+            stream = graph.stream(Command(resume=resume_payload), config=config)
+            continue
 
-        # Добавляем конец истории
-        result = executor.invoke({"state": state}, callbacks=[middleware])
+        if "messages" in chunk:
+            for msg in chunk["messages"]:
+                print(msg.content)
 
-    # Вывод финальной истории
+    # После завершения выводим итоговое состояние
     print("\n--- Итоговая история ---")
     print(f"\n{state['scene_text']}\n")
     print(f"Выбор: {state['choice_selected']}\n")
     print(state["ending"])
 
 
+def main():
+    theme = questionary.text("Введите тему истории:").ask()
+    if not theme:
+        print("Тема обязательна.")
+        return
+    run_story(theme)
+
+
 if __name__ == "__main__":
     main()