add: main.py

main.py

@@ -0,0 +1,136 @@
+import os
+import uuid
+import asyncio
+import questionary
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain.tools import tool
+from langgraph.graph import StateGraph, START, END, Command
+from langgraph.graph.message import add_messages
+from langgraph.checkpoint.memory import InMemorySaver
+from deepagents import create_deep_agent
+
+# LLM configuration – OpenRouter
+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):
+    messages: Annotated[list, add_messages]
+    theme: str
+    intro: str
+    options: list[str]
+    choice: str
+    ending: str
+
+# ---------- Graph Nodes ----------
+async def generate_scene(state: GameState) -> GameState:
+    theme = state["theme"]
+    prompt = (
+        f"Theme: {theme}\n"
+        "Generate a short introduction (2-3 sentences) followed by exactly three numbered options for the hero. "
+        "Output format: first line is the intro, next three lines are options numbered 1) 2) 3)."
+    )
+    response = await llm.ainvoke(HumanMessage(content=prompt))
+    text = response.content.strip()
+    lines = text.splitlines()
+    intro = lines[0].strip()
+    options = [line.strip() for line in lines[1:4]]
+    return {**state, "intro": intro, "options": options}
+
+async def interrupt_choice(state: GameState) -> GameState:
+    payload = {
+        "type": "choice",
+        "question": f"{state['intro']}\nWhat do you do?",
+        "options": state["options"],
+    }
+    # Pause execution until user responds
+    await interrupt(payload)
+    # After resume, the payload will contain 'choice'
+    return {**state, "choice": payload["choice"]}
+
+async def generate_ending(state: GameState) -> GameState:
+    prompt = (
+        f"Intro: {state['intro']}\n"
+        f"Choice: {state['choice']}\n"
+        "Write a short ending (2-3 sentences) for this story."
+    )
+    response = await llm.ainvoke(HumanMessage(content=prompt))
+    ending = response.content.strip()
+    return {**state, "ending": ending}
+
+# ---------- Build Graph ----------
+graph = StateGraph(GameState)
+graph.add_node("scene", generate_scene)
+graph.add_node("choice", interrupt_choice)
+graph.add_node("ending", generate_ending)
+
+graph.set_entry_point("scene")
+graph.add_edge("scene", "choice")
+graph.add_edge("choice", "ending")
+graph.add_edge("ending", END)
+
+checkpoint = InMemorySaver()
+graph.compile(checkpointer=checkpoint)
+
+# ---------- Game Runner ----------
+async def run_game(theme: str) -> dict:
+    thread_id = str(uuid.uuid4())
+    config = {"configurable": {"thread_id": thread_id}}
+    state: GameState = {"messages": [], "theme": theme, "intro": "", "options": [], "choice": "", "ending": ""}
+    stream = graph.stream(state, config)
+    last_state = state
+    async for chunk in stream:
+        # Handle interrupt
+        if "__interrupt__" in chunk:
+            interrupt_payload = chunk["__interrupt__"][0].value
+            # Show options to user
+            answer = questionary.select(
+                interrupt_payload["question"],
+                choices=interrupt_payload["options"],
+            ).ask()
+            # Resume with user's choice
+            resume_payload = {**interrupt_payload, "choice": answer}
+            stream = graph.stream(Command(resume=resume_payload), config)
+            continue
+        # Capture state when available
+        if "state" in chunk:
+            last_state = chunk["state"]
+    return last_state
+
+# ---------- DeepAgent Tool ----------
+@tool
+def play_game(theme: str) -> str:
+    """Play a choose-your-own-adventure game with the given theme."""
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    final_state = loop.run_until_complete(run_game(theme))
+    loop.close()
+    intro = final_state["intro"]
+    ending = final_state["ending"]
+    return f"\n{intro}\n\n{ending}\n"
+
+# ---------- Create DeepAgent ----------
+agent = create_deep_agent(
+    model=llm,
+    tools=[play_game],
+    backend=None,
+    system_prompt="You are a game master that can play choose-your-own-adventure games.",
+)
+
+# ---------- CLI ----------
+async def main():
+    theme = input("Enter a theme for the adventure: ")
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=f"Play a game with theme: {theme}")]},
+        {"configurable": {"thread_id": "cli-session"}},
+    )
+    # The tool returns the full story
+    print(result["messages"][-1].content)
+
+if __name__ == "__main__":
+    asyncio.run(main())