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