Обновить solution.py
This commit is contained in:
+165
-95
@@ -1,114 +1,184 @@
|
||||
<|channel|>final solution.py<|message|>import sys
|
||||
from typing import Optional
|
||||
from typing import TypedDict, List
|
||||
|
||||
# Библиотеки, которые должны быть установлены в окружении задачи.
|
||||
# Мы используем langchain для взаимодействия с LLM и rich для красивого вывода.
|
||||
try:
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
except ImportError as exc: # pragma: no cover
|
||||
raise RuntimeError(
|
||||
"Не удалось импортировать langchain. Убедитесь, что пакет установлен."
|
||||
) from exc
|
||||
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.prompt import Prompt
|
||||
from rich.text import Text
|
||||
except ImportError as exc: # pragma: no cover
|
||||
raise RuntimeError(
|
||||
"Не удалось импортировать rich. Убедитесь, что пакет установлен."
|
||||
) from exc
|
||||
import questionary
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain.prompts import ChatPromptTemplate
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.constants import START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.types import interrupt, Command
|
||||
|
||||
|
||||
class TextAdventureGame:
|
||||
# ---------- 1. Состояние графа ----------
|
||||
|
||||
class GraphState(TypedDict):
|
||||
"""Состояние графа с отдельными полями для каждого этапа."""
|
||||
topic: str # Тема (входной параметр)
|
||||
setup: str # Сгенерированная завязка от LLM
|
||||
choices: List[str] # Список вариантов выбора (3 штуки)
|
||||
human_choice: str # Выбранный пользователем вариант
|
||||
ending: str # Концовка от LLM
|
||||
|
||||
|
||||
# ---------- LangChain: промпты и цепочки ----------
|
||||
|
||||
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.7)
|
||||
|
||||
# Цепочка для генерации завязки и вариантов
|
||||
scene_prompt = ChatPromptTemplate.from_template(
|
||||
"Тема: {topic}.\n"
|
||||
"Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
|
||||
"Ответь строго в формате:\n"
|
||||
"ЗАВЯЗКА: <текст завязки>\n"
|
||||
"ВАРИАНТЫ:\n"
|
||||
"1) <первый вариант>\n"
|
||||
"2) <второй вариант>\n"
|
||||
"3) <третий вариант>"
|
||||
)
|
||||
scene_chain = scene_prompt | llm
|
||||
|
||||
# Цепочка для генерации концовки
|
||||
ending_prompt = ChatPromptTemplate.from_template(
|
||||
"Завязка: {setup}\n"
|
||||
"Выбор пользователя: {human_choice}\n\n"
|
||||
"Допиши короткую концовку истории (2–3 предложения)."
|
||||
)
|
||||
ending_chain = ending_prompt | llm
|
||||
|
||||
|
||||
# ---------- 2. Узел: генерация сцены + прерывание + концовка ----------
|
||||
|
||||
def story_node(state: GraphState) -> GraphState:
|
||||
"""
|
||||
Простая текстовая игра на основе LLM.
|
||||
Игрок вводит команды, которые отправляются в модель,
|
||||
а ответ модели выводится пользователю.
|
||||
Шаг 1: генерирует завязку и варианты через LangChain-цепочку.
|
||||
Шаг 2: вызывает interrupt() — граф ставится на паузу.
|
||||
Шаг 3: после resume дописывает концовку через LangChain-цепочку.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "gpt-3.5-turbo", temperature: float = 0.7):
|
||||
"""
|
||||
Инициализация игрового процесса.
|
||||
# --- Шаг 1: генерируем завязку и варианты через LangChain ---
|
||||
scene_response = scene_chain.invoke({"topic": state["topic"]})
|
||||
raw = scene_response.content.strip()
|
||||
|
||||
:param model_name: имя LLM из OpenAI
|
||||
:param temperature: степень случайности генерации текста
|
||||
"""
|
||||
self.console = Console()
|
||||
# Создаём объект модели. Для работы требуется переменная окружения OPENAI_API_KEY.
|
||||
try:
|
||||
self.llm = ChatOpenAI(model=model_name, temperature=temperature)
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError(
|
||||
"Не удалось создать экземпляр LLM. Проверьте ключ OpenAI."
|
||||
) from exc
|
||||
# Парсим ответ: разделяем завязку и варианты
|
||||
setup = ""
|
||||
choices = []
|
||||
|
||||
def _display_welcome(self) -> None:
|
||||
"""Показываем приветственное сообщение."""
|
||||
welcome = Text("Добро пожаловать в текстовую игру!", style="bold green")
|
||||
self.console.print(welcome)
|
||||
lines = raw.splitlines()
|
||||
in_choices = False
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if line.upper().startswith("ЗАВЯЗКА:"):
|
||||
setup = line[len("ЗАВЯЗКА:"):].strip()
|
||||
elif line.upper().startswith("ВАРИАНТЫ:"):
|
||||
in_choices = True
|
||||
elif in_choices and line:
|
||||
# Убираем нумерацию вида "1)", "2)", "3)" или "1.", "2.", "3."
|
||||
for prefix in ("1)", "2)", "3)", "1.", "2.", "3."):
|
||||
if line.startswith(prefix):
|
||||
line = line[len(prefix):].strip()
|
||||
break
|
||||
if line:
|
||||
choices.append(line)
|
||||
|
||||
def _get_player_input(self, prompt: str = "Что вы делаете?") -> Optional[str]:
|
||||
"""
|
||||
Получаем ввод от игрока. Если пользователь введёт 'quit',
|
||||
возвращаем None для завершения игры.
|
||||
"""
|
||||
try:
|
||||
user_input = Prompt.ask(prompt)
|
||||
except KeyboardInterrupt: # pragma: no cover
|
||||
self.console.print("\n[red]Игра прервана пользователем.[/red]")
|
||||
return None
|
||||
# Защита: если парсинг не дал нужных данных
|
||||
if not setup:
|
||||
setup = raw
|
||||
if len(choices) < 3:
|
||||
choices = choices + [f"Вариант {i+1}" for i in range(len(choices), 3)]
|
||||
|
||||
if user_input.strip().lower() in {"quit", "exit"}:
|
||||
return None
|
||||
return user_input
|
||||
state["setup"] = setup
|
||||
state["choices"] = choices
|
||||
|
||||
def _generate_response(self, player_action: str) -> str:
|
||||
"""
|
||||
Отправляем действие игрока в LLM и получаем ответ.
|
||||
# --- Шаг 2: прерывание ---
|
||||
interrupt_payload = {
|
||||
"type": "choice",
|
||||
"question": setup + "\n\nЧто делаем?",
|
||||
"allow_responds": choices,
|
||||
}
|
||||
|
||||
:param player_action: строка с командой игрока
|
||||
:return: текст ответа модели
|
||||
"""
|
||||
# Формируем сообщение для чат‑модели.
|
||||
messages = [
|
||||
{"role": "system", "content": "Вы — мудрый рассказчик, описывающий события в мире приключений."},
|
||||
{"role": "user", "content": player_action},
|
||||
]
|
||||
try:
|
||||
response = self.llm.invoke(messages)
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError("Ошибка при вызове LLM.") from exc
|
||||
# Выполнение останавливается здесь до Command(resume=...)
|
||||
resume_value = interrupt(interrupt_payload)
|
||||
|
||||
return response.content.strip()
|
||||
# После resume в resume_value — словарь с полем "answer"
|
||||
human_choice = resume_value.get("answer", choices[0])
|
||||
state["human_choice"] = human_choice
|
||||
|
||||
def run(self) -> None:
|
||||
"""Основной цикл игры."""
|
||||
self._display_welcome()
|
||||
while True:
|
||||
action = self._get_player_input()
|
||||
if action is None: # пользователь ввёл quit или прервал
|
||||
self.console.print("[bold magenta]Спасибо за игру![/bold magenta]")
|
||||
break
|
||||
# --- Шаг 3: генерируем концовку через LangChain ---
|
||||
ending_response = ending_chain.invoke({
|
||||
"setup": setup,
|
||||
"human_choice": human_choice,
|
||||
})
|
||||
state["ending"] = ending_response.content.strip()
|
||||
|
||||
self.console.print(f"[cyan]Вы сказали:[/cyan] {action}")
|
||||
try:
|
||||
result = self._generate_response(action)
|
||||
except Exception as exc: # pragma: no cover
|
||||
self.console.print(f"[red]Ошибка генерации ответа: {exc}[/red]")
|
||||
continue
|
||||
return state
|
||||
|
||||
self.console.print("\n[green]Ответ модели:[/green]\n")
|
||||
self.console.print(result, style="italic")
|
||||
self.console.print("-" * 40)
|
||||
|
||||
# ---------- 3. Сборка графа ----------
|
||||
|
||||
builder = StateGraph(GraphState)
|
||||
builder.add_node("story", story_node)
|
||||
builder.add_edge(START, "story")
|
||||
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
|
||||
# ---------- 4. Цикл запуска с обработкой прерывания ----------
|
||||
|
||||
def main() -> None:
|
||||
topic = input("Введите тему истории (например, «космический кот»): ").strip()
|
||||
if not topic:
|
||||
topic = "космический кот"
|
||||
|
||||
print(f"\nТема: {topic}\n")
|
||||
|
||||
thread_id = "story_thread_1"
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
|
||||
initial_state: GraphState = {
|
||||
"topic": topic,
|
||||
"setup": "",
|
||||
"choices": [],
|
||||
"human_choice": "",
|
||||
"ending": "",
|
||||
}
|
||||
|
||||
# --- Первый запуск: до прерывания ---
|
||||
for chunk in graph.stream(initial_state, config=config):
|
||||
if "__interrupt__" in chunk:
|
||||
interrupt_obj = chunk["__interrupt__"][0].value # наш payload
|
||||
|
||||
print(f"\n[LLM] {interrupt_obj['question']}\n")
|
||||
|
||||
# Показываем варианты через questionary
|
||||
answer = questionary.select(
|
||||
"Выберите действие:",
|
||||
choices=interrupt_obj["allow_responds"],
|
||||
).ask()
|
||||
|
||||
if answer is None:
|
||||
raise SystemExit("Отмена пользователем.")
|
||||
|
||||
# Добавляем ответ в payload и возобновляем граф
|
||||
interrupt_obj["answer"] = answer
|
||||
|
||||
# --- Возобновление через Command(resume=...) ---
|
||||
for resumed_chunk in graph.stream(
|
||||
Command(resume=interrupt_obj), config=config
|
||||
):
|
||||
# После resume граф дописывает концовку — просто ждём завершения
|
||||
pass
|
||||
|
||||
# Получаем финальное состояние
|
||||
final_state = graph.get_state(config).values
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("=== Итоговая история ===")
|
||||
print("=" * 50)
|
||||
print(f"\nЗавязка:\n{final_state['setup']}")
|
||||
print(f"\nВыбор: {final_state['human_choice']}")
|
||||
print(f"\n[LLM] {final_state['ending']}")
|
||||
print("\n" + "=" * 50)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Запускаем игру. Если переменная окружения OPENAI_API_KEY не задана,
|
||||
# langchain выбросит исключение при создании модели.
|
||||
game = TextAdventureGame()
|
||||
try:
|
||||
game.run()
|
||||
except Exception as exc: # pragma: no cover
|
||||
sys.stderr.write(f"Произошла ошибка: {exc}\n")
|
||||
sys.exit(1)
|
||||
main()
|
||||
Reference in New Issue
Block a user