161 lines
5.1 KiB
Python
161 lines
5.1 KiB
Python
import os
|
|
import asyncio
|
|
from typing import TypedDict, List, Annotated
|
|
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.checkpoint import InMemorySaver
|
|
from langgraph.types import interrupt, Command
|
|
from langgraph.graph.message import add_messages
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import HumanMessage
|
|
from langchain.tools import tool
|
|
from deepagents import create_deep_agent
|
|
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
|
|
|
|
import questionary
|
|
|
|
# ---------- 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,
|
|
)
|
|
|
|
# ---------- Backend & Tools (required by deepagents) ----------
|
|
backend = CompositeBackend(
|
|
[
|
|
LocalShellBackend(workspace_dir="./workspace"),
|
|
FilesystemBackend(),
|
|
]
|
|
)
|
|
|
|
@tool
|
|
def dummy_tool(query: str) -> str:
|
|
"""A placeholder tool required by the deep agent."""
|
|
return f"tool result for {query}"
|
|
|
|
agent = create_deep_agent(
|
|
model=llm,
|
|
tools=[dummy_tool],
|
|
backend=backend,
|
|
system_prompt="You are a helpful storytelling assistant.",
|
|
)
|
|
|
|
# ---------- State ----------
|
|
class StoryState(TypedDict):
|
|
messages: Annotated[List, add_messages]
|
|
topic: str
|
|
intro: str
|
|
options: List[str]
|
|
choice: str
|
|
ending: str
|
|
|
|
# ---------- Node ----------
|
|
async def story_node(state: StoryState):
|
|
# Phase 1: generate intro and options
|
|
if "intro" not in state or not state["intro"]:
|
|
prompt = (
|
|
f"Topic: {state['topic']}. "
|
|
"Create a short story beginning (2-3 sentences) and exactly three possible actions for the hero. "
|
|
"Respond in the following format:\n"
|
|
"Intro: <your text>\n"
|
|
"Options:\n"
|
|
"1) <first action>\n"
|
|
"2) <second action>\n"
|
|
"3) <third action>"
|
|
)
|
|
response = await agent.ainvoke(
|
|
{"messages": [HumanMessage(content=prompt)]},
|
|
{"configurable": {"thread_id": "agent-1"}},
|
|
)
|
|
text = response["messages"][-1].content
|
|
|
|
# Simple parsing
|
|
intro_part, options_part = text.split("Options:", 1)
|
|
intro = intro_part.replace("Intro:", "").strip()
|
|
raw_options = options_part.strip().splitlines()
|
|
options = [line.split(")", 1)[1].strip() for line in raw_options if ")" in line]
|
|
|
|
# Store and interrupt
|
|
state["intro"] = intro
|
|
state["options"] = options
|
|
payload = {
|
|
"type": "choice",
|
|
"question": f"{intro}\n\nWhat does the hero do?",
|
|
"options": options,
|
|
}
|
|
return interrupt(payload)
|
|
|
|
# Phase 2: after user choice, generate ending
|
|
if "choice" in state and state["choice"]:
|
|
prompt = (
|
|
f"Intro: {state['intro']}\n"
|
|
f"User choice: {state['choice']}\n"
|
|
"Write a short conclusion (2-3 sentences) that follows from this choice."
|
|
)
|
|
response = await agent.ainvoke(
|
|
{"messages": [HumanMessage(content=prompt)]},
|
|
{"configurable": {"thread_id": "agent-2"}},
|
|
)
|
|
ending = response["messages"][-1].content.strip()
|
|
state["ending"] = ending
|
|
return state
|
|
|
|
# Should not reach here
|
|
return state
|
|
|
|
# ---------- Graph ----------
|
|
graph = StateGraph(StoryState)
|
|
graph.add_node("story", story_node)
|
|
graph.add_edge(START, "story")
|
|
graph.add_edge("story", END)
|
|
graph.set_entry_point("story")
|
|
checkpointer = InMemorySaver()
|
|
graph = graph.compile(checkpointer=checkpointer)
|
|
|
|
# ---------- Main loop ----------
|
|
async def main():
|
|
topic = questionary.text("Enter a story theme (e.g., 'space cat'):", default="space cat").ask()
|
|
thread_id = "session-1"
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
|
|
# Initial empty state
|
|
initial_state: StoryState = {
|
|
"messages": [],
|
|
"topic": topic,
|
|
"intro": "",
|
|
"options": [],
|
|
"choice": "",
|
|
"ending": "",
|
|
}
|
|
|
|
# First run - will pause at interrupt
|
|
async for event in graph.stream(initial_state, config):
|
|
if "__interrupt__" in event:
|
|
payload = event["__interrupt__"][0].value
|
|
question = payload["question"]
|
|
choices = payload["options"]
|
|
answer = questionary.select(question, choices=choices).ask()
|
|
payload["choice"] = answer
|
|
|
|
# Resume graph with the updated payload
|
|
resume_cmd = Command(resume=payload)
|
|
async for resume_event in graph.stream(resume_cmd, config):
|
|
if "__interrupt__" in resume_event:
|
|
# No further interrupts expected
|
|
continue
|
|
# Continue until END
|
|
break
|
|
|
|
# Retrieve final state
|
|
final_state = await graph.aget_state(thread_id)
|
|
print("\n--- Final Story ---")
|
|
print(f"Theme: {final_state.state['topic']}\n")
|
|
print(final_state.state["intro"])
|
|
print(f"\nYou chose: {final_state.state['choice']}")
|
|
print(f"\n{final_state.state['ending']}\n")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |