текстовая игра на основе llm + interrupt: client.py
This commit is contained in:
+62
-63
@@ -1,14 +1,16 @@
|
||||
import uuid
|
||||
from typing import TypedDict, List, Dict, Any
|
||||
from typing import TypedDict, List
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
|
||||
# ---------- 1. Состояние графа ----------
|
||||
# ---------- 1. Состояние (используется в памяти) ----------
|
||||
class StoryState(TypedDict):
|
||||
theme: str
|
||||
scene_text: str
|
||||
@@ -17,11 +19,11 @@ class StoryState(TypedDict):
|
||||
ending: str
|
||||
|
||||
|
||||
# ---------- 2. LLM и узлы ----------
|
||||
# ---------- 2. LLM и вспомогательные функции ----------
|
||||
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()]
|
||||
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]:
|
||||
def generate_scene(state: StoryState) -> dict:
|
||||
theme = state["theme"]
|
||||
prompt = (
|
||||
f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя. "
|
||||
@@ -58,21 +56,23 @@ def generate_scene(state: StoryState) -> Dict[str, Any]:
|
||||
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"],
|
||||
# Запрос к пользователю через HumanInTheLoopMiddleware
|
||||
return {
|
||||
"messages": [
|
||||
SystemMessage(
|
||||
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]:
|
||||
# state уже содержит scene_text и choice_selected
|
||||
def add_ending(state: StoryState) -> dict:
|
||||
prompt = (
|
||||
f"Завязка: {state['scene_text']}\n"
|
||||
f"Выбор пользователя: {state['choice_selected']}\n\n"
|
||||
@@ -80,23 +80,25 @@ def add_ending(state: StoryState) -> Dict[str, Any]:
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
state["ending"] = response.content.strip()
|
||||
return state
|
||||
return {"messages": [HumanMessage(content=state["ending"])]}
|
||||
|
||||
|
||||
# ---------- 3. Сборка графа ----------
|
||||
def create_graph() -> StateGraph:
|
||||
builder = StateGraph(StoryState)
|
||||
# ---------- 3. Создание агента ----------
|
||||
def create_agent_executor() -> AgentExecutor:
|
||||
# Определяем инструменты (здесь нет внешних, но нужны для агента)
|
||||
tools: List[BaseTool] = []
|
||||
|
||||
# Узлы
|
||||
builder.add_node("generate_scene", generate_scene)
|
||||
builder.add_node("add_ending", add_ending)
|
||||
agent = create_agent(
|
||||
llm=llm,
|
||||
tools=tools,
|
||||
system_message="Ты создаёшь интерактивную историю. После генерации сцены запрашивай выбор у пользователя.",
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
# Переходы
|
||||
builder.set_entry_point("generate_scene")
|
||||
builder.add_edge("generate_scene", "add_ending")
|
||||
builder.add_edge("add_ending", END)
|
||||
memory = ConversationBufferMemory(return_messages=True)
|
||||
|
||||
return builder.compile(checkpointer=InMemorySaver())
|
||||
executor = AgentExecutor(agent=agent, tools=tools, memory=memory, verbose=False)
|
||||
return executor
|
||||
|
||||
|
||||
# ---------- 4. Клиент ----------
|
||||
@@ -106,8 +108,7 @@ def main():
|
||||
print("Тема обязательна.")
|
||||
return
|
||||
|
||||
# Инициализируем состояние
|
||||
init_state: StoryState = {
|
||||
state: StoryState = {
|
||||
"theme": theme,
|
||||
"scene_text": "",
|
||||
"choices": [],
|
||||
@@ -115,38 +116,36 @@ def main():
|
||||
"ending": "",
|
||||
}
|
||||
|
||||
graph = create_graph()
|
||||
thread_id = str(uuid.uuid4())
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
executor = create_agent_executor()
|
||||
|
||||
# Запускаем первый поток
|
||||
stream = graph.stream(init_state, config)
|
||||
# Middleware для прерывания и возобновления
|
||||
middleware = HumanInTheLoopMiddleware(
|
||||
prompt=lambda x: x["messages"][0].content,
|
||||
choices=lambda x: x["messages"][1:], # список HumanMessage с вариантами
|
||||
)
|
||||
|
||||
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
|
||||
# Запускаем генерацию сцены
|
||||
result = executor.invoke({"state": state}, callbacks=[middleware])
|
||||
|
||||
# Добавляем ответ в payload и возобновляем граф
|
||||
interrupt_payload["answer"] = answer
|
||||
resume_command = Command(resume=interrupt_payload)
|
||||
stream = graph.stream(resume_command, config)
|
||||
# После прерывания пользователь выберет вариант
|
||||
if middleware.interrupted:
|
||||
answer = questionary.select(
|
||||
middleware.prompt, choices=middleware.choices_texts
|
||||
).ask()
|
||||
if not answer:
|
||||
print("Выбор не сделан. Завершаем.")
|
||||
return
|
||||
|
||||
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
|
||||
state["choice_selected"] = answer
|
||||
|
||||
# Добавляем конец истории
|
||||
result = executor.invoke({"state": state}, callbacks=[middleware])
|
||||
|
||||
# Вывод финальной истории
|
||||
print("\n--- Итоговая история ---")
|
||||
print(f"\n{state['scene_text']}\n")
|
||||
print(f"Выбор: {state['choice_selected']}\n")
|
||||
print(state["ending"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user