commit 4b34935a60fd022d041a44c2678f9048d7dd2a13 Author: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Mon Jun 15 12:26:43 2026 +0000 add: main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..774218e --- /dev/null +++ b/main.py @@ -0,0 +1,137 @@ +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())