138 lines
4.8 KiB
Python
138 lines
4.8 KiB
Python
import os
|
||
import asyncio
|
||
import questionary
|
||
from typing import TypedDict, Annotated
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import StateGraph, START, END, interrupt, Command
|
||
from langgraph.graph.message import add_messages
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from langchain_core.messages import HumanMessage, AIMessage
|
||
from deepagents import create_deep_agent
|
||
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
|
||
from deepagents.tools import tool
|
||
|
||
# ---------- LLM ----------
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# ---------- Graph state ----------
|
||
class GameState(TypedDict):
|
||
theme: str
|
||
scene: Annotated[str, add_messages]
|
||
options: list[str]
|
||
choice: str
|
||
ending: Annotated[str, add_messages]
|
||
|
||
# ---------- Graph nodes ----------
|
||
async def generate_scene(state: GameState) -> GameState:
|
||
theme = state["theme"]
|
||
prompt = (
|
||
f"Тема: {theme}.\n"
|
||
"Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
|
||
"Ответь в формате: сначала текст завязки, затем строка с вариантами через запятую."
|
||
)
|
||
response = await llm.ainvoke(HumanMessage(content=prompt))
|
||
text = response.content
|
||
# Разделяем завязку и варианты
|
||
parts = text.split("\n")
|
||
scene_text = parts[0].strip()
|
||
options_line = parts[1] if len(parts) > 1 else ""
|
||
options = [opt.strip() for opt in options_line.split(",") if opt.strip()]
|
||
state["scene"] = [AIMessage(content=scene_text)]
|
||
state["options"] = options
|
||
# Прерываем для выбора
|
||
interrupt_payload = {
|
||
"type": "choice",
|
||
"question": f"{scene_text}\n\nЧто делаем?",
|
||
"options": options,
|
||
}
|
||
interrupt(interrupt_payload)
|
||
return state
|
||
|
||
async def finish_story(state: GameState) -> GameState:
|
||
# state now contains "choice"
|
||
scene_text = state["scene"][0].content
|
||
choice = state["choice"]
|
||
prompt = (
|
||
f"Завязка: {scene_text}\n"
|
||
f"Выбор пользователя: {choice}\n"
|
||
"Допиши короткую концовку (2–3 предложения)."
|
||
)
|
||
response = await llm.ainvoke(HumanMessage(content=prompt))
|
||
ending_text = response.content
|
||
state["ending"] = [AIMessage(content=ending_text)]
|
||
return state
|
||
|
||
# ---------- Build graph ----------
|
||
builder = StateGraph(GameState)
|
||
builder.add_node("generate", generate_scene)
|
||
builder.add_node("finish", finish_story)
|
||
builder.set_entry_point("generate")
|
||
builder.add_edge("generate", "finish")
|
||
builder.add_edge("finish", END)
|
||
|
||
checkpoint = InMemorySaver()
|
||
graph = builder.compile(checkpointer=checkpoint)
|
||
|
||
# ---------- Tool that runs the game ----------
|
||
@tool
|
||
async def play_game(theme: str) -> str:
|
||
"""Play a choose‑your‑adventure game with the given theme."""
|
||
thread_id = f"game-{theme.replace(' ', '-')}-1"
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
# Initial state
|
||
state: GameState = {"theme": theme, "scene": [], "options": [], "choice": "", "ending": []}
|
||
# Start streaming
|
||
stream = graph.stream(state, config)
|
||
async for chunk in stream:
|
||
if "__interrupt__" in chunk:
|
||
# Extract payload
|
||
payload = chunk["__interrupt__"][0].value
|
||
# Show question and options
|
||
answer = questionary.select(
|
||
payload["question"],
|
||
choices=payload["options"],
|
||
).ask()
|
||
# Resume with answer
|
||
payload["choice"] = answer
|
||
resume_cmd = Command(resume=payload)
|
||
stream = graph.stream(resume_cmd, config)
|
||
continue
|
||
# Collect messages
|
||
if "scene" in chunk:
|
||
state["scene"] = chunk["scene"]
|
||
if "ending" in chunk:
|
||
state["ending"] = chunk["ending"]
|
||
# Build final story
|
||
story = "\n\n".join([m.content for m in state["scene"] + state["ending"]])
|
||
return story
|
||
|
||
# ---------- DeepAgent setup ----------
|
||
backend = CompositeBackend([
|
||
LocalShellBackend(workspace_dir="./workspace"),
|
||
FilesystemBackend(),
|
||
])
|
||
|
||
agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[play_game],
|
||
backend=backend,
|
||
system_prompt="You are a helpful assistant that can play choose‑your‑adventure games.",
|
||
)
|
||
|
||
async def main():
|
||
theme = questionary.text("Enter a theme for the adventure:").ask()
|
||
result = await agent.ainvoke(
|
||
{"messages": [HumanMessage(content=f"play_game {theme}")]},
|
||
{"configurable": {"thread_id": "session-1"}},
|
||
)
|
||
print("\n\n--- Final Story ---\n")
|
||
print(result["messages"][-1].content)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|