feat: solution for 'текстовая игра на основе llm + interrupt'
This commit is contained in:
+48
-172
@@ -1,180 +1,56 @@
|
||||
```python
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Interactive choose-your-own-adventure using LangGraph, OpenAI LLM, and console interrupts.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import TypedDict, List
|
||||
|
||||
import questionary
|
||||
from langgraph.graph import StateGraph
|
||||
import uuid
|
||||
from dotenv import load_dotenv
|
||||
from graph import build_graph
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.types import interrupt
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain.prompts import PromptTemplate
|
||||
from langchain.schema import HumanMessage
|
||||
import questionary
|
||||
|
||||
# Ensure OpenAI key is set
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("Please set the OPENAI_API_KEY environment variable.")
|
||||
|
||||
# -----------------------------
|
||||
# State definition
|
||||
# -----------------------------
|
||||
class AdventureState(TypedDict):
|
||||
topic: str
|
||||
opening: str
|
||||
choices: List[str]
|
||||
choice: str
|
||||
ending: str
|
||||
|
||||
# -----------------------------
|
||||
# LLM setup
|
||||
# -----------------------------
|
||||
llm = ChatOpenAI(
|
||||
model="gpt-4o-mini",
|
||||
temperature=0.7,
|
||||
api_key=OPENAI_API_KEY,
|
||||
)
|
||||
|
||||
# -----------------------------
|
||||
# Node: Generate opening and choices
|
||||
# -----------------------------
|
||||
def generate_scene_and_choices(state: AdventureState) -> AdventureState:
|
||||
topic = state["topic"]
|
||||
prompt = PromptTemplate(
|
||||
input_variables=["topic"],
|
||||
template=(
|
||||
"You are a creative storyteller. "
|
||||
"Topic: {topic}. "
|
||||
"Write a short opening (2–3 sentences) and exactly three distinct choices for the hero. "
|
||||
"Respond with the opening first, then a numbered list of the choices. "
|
||||
"Example format:\n"
|
||||
"Opening: ...\n"
|
||||
"1. ...\n"
|
||||
"2. ...\n"
|
||||
"3. ..."
|
||||
),
|
||||
)
|
||||
messages = [HumanMessage(content=prompt.format(topic=topic))]
|
||||
response = llm.invoke(messages)
|
||||
text = response.content.strip()
|
||||
|
||||
# Parse opening and choices
|
||||
lines = [line.strip() for line in text.splitlines() if line.strip()]
|
||||
opening_line = lines[0]
|
||||
if opening_line.lower().startswith("opening:"):
|
||||
opening = opening_line[len("opening:"):].strip()
|
||||
else:
|
||||
opening = opening_line
|
||||
|
||||
choices = []
|
||||
for line in lines[1:]:
|
||||
if len(line) >= 2 and line[0].isdigit() and line[1] in {".", ":"}:
|
||||
choice_text = line.split(" ", 1)[1].strip()
|
||||
choices.append(choice_text)
|
||||
else:
|
||||
choices.append(line)
|
||||
|
||||
state["opening"] = opening
|
||||
state["choices"] = choices
|
||||
return state
|
||||
|
||||
# -----------------------------
|
||||
# Node: Interrupt for user choice
|
||||
# -----------------------------
|
||||
def interrupt_choice(state: AdventureState) -> AdventureState:
|
||||
question = f"{state['opening']}\n\nWhat do you do?"
|
||||
payload = {
|
||||
"type": "choice",
|
||||
"question": question,
|
||||
"options": state["choices"],
|
||||
}
|
||||
# The interrupt will pause the graph and return a dict with "__interrupt__" key
|
||||
return interrupt(payload)
|
||||
|
||||
# -----------------------------
|
||||
# Node: Generate ending
|
||||
# -----------------------------
|
||||
def generate_ending(state: AdventureState) -> AdventureState:
|
||||
prompt = PromptTemplate(
|
||||
input_variables=["opening", "choice"],
|
||||
template=(
|
||||
"You are a creative storyteller. "
|
||||
"Opening: {opening}\n"
|
||||
"The hero chose: {choice}\n"
|
||||
"Write a short ending (2–3 sentences) that concludes the story."
|
||||
),
|
||||
)
|
||||
messages = [HumanMessage(content=prompt.format(opening=state["opening"], choice=state["choice"]))]
|
||||
response = llm.invoke(messages)
|
||||
state["ending"] = response.content.strip()
|
||||
return state
|
||||
|
||||
# -----------------------------
|
||||
# Build the graph
|
||||
# -----------------------------
|
||||
def build_graph() -> StateGraph[AdventureState]:
|
||||
graph = StateGraph(AdventureState)
|
||||
|
||||
graph.add_node("generate_scene_and_choices", generate_scene_and_choices)
|
||||
graph.add_node("interrupt_choice", interrupt_choice)
|
||||
graph.add_node("generate_ending", generate_ending)
|
||||
|
||||
graph.set_entry_point("generate_scene_and_choices")
|
||||
graph.add_edge("generate_scene_and_choices", "interrupt_choice")
|
||||
graph.add_edge("interrupt_choice", "generate_ending")
|
||||
graph.add_edge("generate_ending", "__end__")
|
||||
|
||||
return graph
|
||||
|
||||
# -----------------------------
|
||||
# Main execution
|
||||
# -----------------------------
|
||||
def main() -> None:
|
||||
topic = questionary.text("Enter a topic for your adventure:").ask()
|
||||
if not topic:
|
||||
print("No topic provided. Exiting.")
|
||||
return
|
||||
|
||||
# Initial state
|
||||
state: AdventureState = {
|
||||
"topic": topic,
|
||||
"opening": "",
|
||||
"choices": [],
|
||||
"choice": "",
|
||||
"ending": "",
|
||||
}
|
||||
|
||||
graph = build_graph()
|
||||
# Use an in-memory checkpoint to allow resuming after interrupt
|
||||
def main():
|
||||
load_dotenv()
|
||||
checkpoint = InMemorySaver()
|
||||
compiled = graph.compile(checkpointer=checkpoint)
|
||||
|
||||
# Run the graph, handling interrupts manually
|
||||
config = {"configurable": {"thread_id": "adventure_thread"}}
|
||||
while True:
|
||||
result = compiled.run(state, config=config)
|
||||
# If an interrupt occurs, handle it
|
||||
if "__interrupt__" in result:
|
||||
interrupt_data = result["__interrupt__"]
|
||||
# Prompt user for choice
|
||||
choice = questionary.select(
|
||||
interrupt_data["question"],
|
||||
choices=interrupt_data["options"]
|
||||
).ask()
|
||||
if not choice:
|
||||
print("No choice selected. Exiting.")
|
||||
graph = build_graph(checkpoint)
|
||||
thread_id = str(uuid.uuid4())
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
print("=== Начало истории ===")
|
||||
# Start the graph and handle interrupt
|
||||
stream = graph.stream(config=config)
|
||||
for chunk in stream:
|
||||
if "__interrupt__" in chunk:
|
||||
payload = chunk["__interrupt__"]
|
||||
if payload.get("type") == "choice":
|
||||
question = payload.get("question", "")
|
||||
options = payload.get("options", [])
|
||||
if not options:
|
||||
print("Нет вариантов выбора.")
|
||||
return
|
||||
choice = questionary.select(question, choices=options).ask()
|
||||
if choice is None:
|
||||
print("Выход из игры.")
|
||||
return
|
||||
# Resume graph with the chosen option
|
||||
resume_config = {"configurable": {"thread_id": thread_id, "choice": choice}}
|
||||
for resume_chunk in graph.stream(resume_config):
|
||||
if "__interrupt__" in resume_chunk:
|
||||
print("Unexpected interrupt during ending.")
|
||||
return
|
||||
text = resume_chunk.get("text", "")
|
||||
if text:
|
||||
print(text, end="")
|
||||
break
|
||||
state["choice"] = choice
|
||||
continue
|
||||
# If the graph has finished, print the ending
|
||||
if "ending" in result:
|
||||
print("\n" + result["ending"])
|
||||
break
|
||||
else:
|
||||
text = chunk.get("text", "")
|
||||
if text:
|
||||
print(text, end="")
|
||||
# After finishing, print the ending
|
||||
final_state = checkpoint.get_state(thread_id)
|
||||
ending = final_state.get("ending", "")
|
||||
if ending:
|
||||
print("\n\n=== Концовка ===")
|
||||
print(ending)
|
||||
else:
|
||||
print("\nКонцовка не найдена.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
main()
|
||||
Reference in New Issue
Block a user