add: main.py

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())