текстовая игра на основе llm + interrupt: client.py
This commit is contained in:
+62
-63
@@ -1,14 +1,16 @@
|
|||||||
import uuid
|
import uuid
|
||||||
from typing import TypedDict, List, Dict, Any
|
from typing import TypedDict, List
|
||||||
|
|
||||||
import questionary
|
import questionary
|
||||||
from langchain_openai import ChatOpenAI
|
from langchain_openai import ChatOpenAI
|
||||||
from langgraph.graph import StateGraph, START, interrupt, Command
|
from langchain.agents import AgentExecutor, create_agent
|
||||||
from langgraph.checkpoint.memory import InMemorySaver
|
from langchain.schema import HumanMessage, SystemMessage
|
||||||
from langgraph.types import GraphState
|
from langchain.tools import BaseTool
|
||||||
|
from langchain.callbacks.human_in_the_loop import HumanInTheLoopMiddleware
|
||||||
|
from langchain.memory import ConversationBufferMemory
|
||||||
|
|
||||||
|
|
||||||
# ---------- 1. Состояние графа ----------
|
# ---------- 1. Состояние (используется в памяти) ----------
|
||||||
class StoryState(TypedDict):
|
class StoryState(TypedDict):
|
||||||
theme: str
|
theme: str
|
||||||
scene_text: str
|
scene_text: str
|
||||||
@@ -17,11 +19,11 @@ class StoryState(TypedDict):
|
|||||||
ending: str
|
ending: str
|
||||||
|
|
||||||
|
|
||||||
# ---------- 2. LLM и узлы ----------
|
# ---------- 2. LLM и вспомогательные функции ----------
|
||||||
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
|
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
|
||||||
|
|
||||||
|
|
||||||
def parse_llm_output(output: str) -> Dict[str, Any]:
|
def parse_llm_output(output: str) -> dict:
|
||||||
"""
|
"""
|
||||||
Ожидаем формат:
|
Ожидаем формат:
|
||||||
<завязка>
|
<завязка>
|
||||||
@@ -32,22 +34,18 @@ def parse_llm_output(output: str) -> Dict[str, Any]:
|
|||||||
"""
|
"""
|
||||||
lines = [line.strip() for line in output.splitlines() if line.strip()]
|
lines = [line.strip() for line in output.splitlines() if line.strip()]
|
||||||
scene_text = lines[0]
|
scene_text = lines[0]
|
||||||
# Остальные строки считаем вариантами
|
|
||||||
raw_choices = "\n".join(lines[1:])
|
raw_choices = "\n".join(lines[1:])
|
||||||
# Разделяем по цифрам или запятой
|
|
||||||
import re
|
import re
|
||||||
|
|
||||||
# Если есть нумерация 1) 2) ...
|
|
||||||
numbered = re.findall(r"\d+\)\s*(.+)", raw_choices)
|
numbered = re.findall(r"\d+\)\s*(.+)", raw_choices)
|
||||||
if numbered:
|
if numbered:
|
||||||
choices = [c.strip() for c in numbered]
|
choices = [c.strip() for c in numbered]
|
||||||
else:
|
else:
|
||||||
# Разделяем по запятой
|
|
||||||
choices = [c.strip() for c in raw_choices.split(",") if c.strip()]
|
choices = [c.strip() for c in raw_choices.split(",") if c.strip()]
|
||||||
return {"scene_text": scene_text, "choices": choices}
|
return {"scene_text": scene_text, "choices": choices}
|
||||||
|
|
||||||
|
|
||||||
def generate_scene(state: StoryState) -> Dict[str, Any]:
|
def generate_scene(state: StoryState) -> dict:
|
||||||
theme = state["theme"]
|
theme = state["theme"]
|
||||||
prompt = (
|
prompt = (
|
||||||
f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя. "
|
f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя. "
|
||||||
@@ -58,21 +56,23 @@ def generate_scene(state: StoryState) -> Dict[str, Any]:
|
|||||||
response = llm.invoke(prompt)
|
response = llm.invoke(prompt)
|
||||||
parsed = parse_llm_output(response.content)
|
parsed = parse_llm_output(response.content)
|
||||||
|
|
||||||
# Сохраняем сцену и варианты
|
|
||||||
state["scene_text"] = parsed["scene_text"]
|
state["scene_text"] = parsed["scene_text"]
|
||||||
state["choices"] = parsed["choices"]
|
state["choices"] = parsed["choices"]
|
||||||
|
|
||||||
# Подготавливаем прерывание
|
# Запрос к пользователю через HumanInTheLoopMiddleware
|
||||||
interrupt_payload = {
|
return {
|
||||||
"type": "choice",
|
"messages": [
|
||||||
"question": f"{parsed['scene_text']}\n\nЧто делаем?",
|
SystemMessage(
|
||||||
"options": parsed["choices"],
|
content=f"{parsed['scene_text']}\n\nЧто делаем?"
|
||||||
|
),
|
||||||
|
*[
|
||||||
|
HumanMessage(content=choice) for choice in parsed["choices"]
|
||||||
|
],
|
||||||
|
]
|
||||||
}
|
}
|
||||||
return interrupt(interrupt_payload)
|
|
||||||
|
|
||||||
|
|
||||||
def add_ending(state: StoryState) -> Dict[str, Any]:
|
def add_ending(state: StoryState) -> dict:
|
||||||
# state уже содержит scene_text и choice_selected
|
|
||||||
prompt = (
|
prompt = (
|
||||||
f"Завязка: {state['scene_text']}\n"
|
f"Завязка: {state['scene_text']}\n"
|
||||||
f"Выбор пользователя: {state['choice_selected']}\n\n"
|
f"Выбор пользователя: {state['choice_selected']}\n\n"
|
||||||
@@ -80,23 +80,25 @@ def add_ending(state: StoryState) -> Dict[str, Any]:
|
|||||||
)
|
)
|
||||||
response = llm.invoke(prompt)
|
response = llm.invoke(prompt)
|
||||||
state["ending"] = response.content.strip()
|
state["ending"] = response.content.strip()
|
||||||
return state
|
return {"messages": [HumanMessage(content=state["ending"])]}
|
||||||
|
|
||||||
|
|
||||||
# ---------- 3. Сборка графа ----------
|
# ---------- 3. Создание агента ----------
|
||||||
def create_graph() -> StateGraph:
|
def create_agent_executor() -> AgentExecutor:
|
||||||
builder = StateGraph(StoryState)
|
# Определяем инструменты (здесь нет внешних, но нужны для агента)
|
||||||
|
tools: List[BaseTool] = []
|
||||||
|
|
||||||
# Узлы
|
agent = create_agent(
|
||||||
builder.add_node("generate_scene", generate_scene)
|
llm=llm,
|
||||||
builder.add_node("add_ending", add_ending)
|
tools=tools,
|
||||||
|
system_message="Ты создаёшь интерактивную историю. После генерации сцены запрашивай выбор у пользователя.",
|
||||||
|
verbose=False,
|
||||||
|
)
|
||||||
|
|
||||||
# Переходы
|
memory = ConversationBufferMemory(return_messages=True)
|
||||||
builder.set_entry_point("generate_scene")
|
|
||||||
builder.add_edge("generate_scene", "add_ending")
|
|
||||||
builder.add_edge("add_ending", END)
|
|
||||||
|
|
||||||
return builder.compile(checkpointer=InMemorySaver())
|
executor = AgentExecutor(agent=agent, tools=tools, memory=memory, verbose=False)
|
||||||
|
return executor
|
||||||
|
|
||||||
|
|
||||||
# ---------- 4. Клиент ----------
|
# ---------- 4. Клиент ----------
|
||||||
@@ -106,8 +108,7 @@ def main():
|
|||||||
print("Тема обязательна.")
|
print("Тема обязательна.")
|
||||||
return
|
return
|
||||||
|
|
||||||
# Инициализируем состояние
|
state: StoryState = {
|
||||||
init_state: StoryState = {
|
|
||||||
"theme": theme,
|
"theme": theme,
|
||||||
"scene_text": "",
|
"scene_text": "",
|
||||||
"choices": [],
|
"choices": [],
|
||||||
@@ -115,38 +116,36 @@ def main():
|
|||||||
"ending": "",
|
"ending": "",
|
||||||
}
|
}
|
||||||
|
|
||||||
graph = create_graph()
|
executor = create_agent_executor()
|
||||||
thread_id = str(uuid.uuid4())
|
|
||||||
config = {"configurable": {"thread_id": thread_id}}
|
|
||||||
|
|
||||||
# Запускаем первый поток
|
# Middleware для прерывания и возобновления
|
||||||
stream = graph.stream(init_state, config)
|
middleware = HumanInTheLoopMiddleware(
|
||||||
|
prompt=lambda x: x["messages"][0].content,
|
||||||
|
choices=lambda x: x["messages"][1:], # список HumanMessage с вариантами
|
||||||
|
)
|
||||||
|
|
||||||
for chunk in stream:
|
# Запускаем генерацию сцены
|
||||||
if "__interrupt__" in chunk:
|
result = executor.invoke({"state": state}, callbacks=[middleware])
|
||||||
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
|
if middleware.interrupted:
|
||||||
resume_command = Command(resume=interrupt_payload)
|
answer = questionary.select(
|
||||||
stream = graph.stream(resume_command, config)
|
middleware.prompt, choices=middleware.choices_texts
|
||||||
|
).ask()
|
||||||
|
if not answer:
|
||||||
|
print("Выбор не сделан. Завершаем.")
|
||||||
|
return
|
||||||
|
|
||||||
elif "__final_state__" in chunk:
|
state["choice_selected"] = answer
|
||||||
final_state: StoryState = chunk["__final_state__"]
|
|
||||||
print("\n--- Итоговая история ---")
|
# Добавляем конец истории
|
||||||
print(f"\n{final_state['scene_text']}\n")
|
result = executor.invoke({"state": state}, callbacks=[middleware])
|
||||||
print(f"Выбор: {final_state['choice_selected']}\n")
|
|
||||||
print(f"{final_state['ending']}")
|
# Вывод финальной истории
|
||||||
else:
|
print("\n--- Итоговая история ---")
|
||||||
# Вывод промежуточных сообщений (если есть)
|
print(f"\n{state['scene_text']}\n")
|
||||||
pass
|
print(f"Выбор: {state['choice_selected']}\n")
|
||||||
|
print(state["ending"])
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Reference in New Issue
Block a user