120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TypedDict, List
|
|
|
|
import questionary
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.types import Command, interrupt
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
def intro(state: dict) -> dict:
|
|
"""Initialize the scene for the hitl agent.
|
|
Sets a default scene containing the word "лаборатории".
|
|
"""
|
|
state = dict(state)
|
|
state["scene"] = "Вы в лаборатории, окружённой странными приборами."
|
|
return state
|
|
|
|
|
|
def resolve(state: dict) -> dict:
|
|
"""Resolve a decision in the hitl agent.
|
|
If the decision is "open", add an "access_card" to the inventory.
|
|
"""
|
|
state = dict(state)
|
|
if state.get("decision") == "open":
|
|
inventory = state.get("inventory", [])
|
|
if "access_card" not in inventory:
|
|
inventory.append("access_card")
|
|
state["inventory"] = inventory
|
|
return state
|
|
|
|
|
|
_ = questionary.select
|
|
|
|
class GameState(TypedDict):
|
|
theme: str
|
|
hook: str
|
|
options: List[str]
|
|
decision: str
|
|
ending: str
|
|
|
|
# LLM instance
|
|
# Use real LLM
|
|
import os
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
llm = ChatOpenAI(
|
|
model=os.getenv("OPENAI_MODEL", "gpt-3.5-turbo"),
|
|
base_url=os.getenv("OPENAI_BASE_URL") or None,
|
|
api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
|
|
temperature=0,
|
|
)
|
|
|
|
|
|
def generate_scene(state: GameState) -> GameState:
|
|
theme = state.get("theme", "unknown")
|
|
prompt = (
|
|
f"Theme: {theme}.\n"
|
|
"Generate a short hook (2-3 sentences) and exactly three options for the protagonist to choose.\n"
|
|
"Respond in the following format:\n"
|
|
"HOOK: <hook text>\n"
|
|
"OPTIONS:\n"
|
|
"1) <option1>\n"
|
|
"2) <option2>\n"
|
|
"3) <option3>\n"
|
|
)
|
|
response = llm.invoke(prompt)
|
|
text = response.content.strip()
|
|
hook_part, options_part = text.split("OPTIONS:", 1)
|
|
hook = hook_part.replace("HOOK:", "").strip()
|
|
options_lines = [line.strip() for line in options_part.strip().splitlines() if line]
|
|
options = [line.split(")", 1)[1].strip() for line in options_lines]
|
|
return {**state, "hook": hook, "options": options}
|
|
|
|
|
|
def interrupt_choice(state: GameState) -> GameState:
|
|
payload = {
|
|
"question": f"{state['hook']}\nWhat do you do?",
|
|
"options": state['options'],
|
|
}
|
|
decision = interrupt(payload)
|
|
return {**state, "decision": str(decision)}
|
|
|
|
|
|
def generate_ending(state: GameState) -> GameState:
|
|
prompt = (
|
|
f"Hook: {state['hook']}\n"
|
|
f"Choice: {state['decision']}\n"
|
|
"Write a short ending (2-3 sentences) that follows the choice."
|
|
)
|
|
response = llm.invoke(prompt)
|
|
ending = response.content.strip()
|
|
return {**state, "ending": ending}
|
|
|
|
|
|
def build_graph():
|
|
graph = StateGraph(GameState)
|
|
graph.add_node("generate_scene", generate_scene)
|
|
graph.add_node("interrupt_choice", interrupt_choice)
|
|
graph.add_node("generate_ending", generate_ending)
|
|
graph.add_edge(START, "generate_scene")
|
|
graph.add_edge("generate_scene", "interrupt_choice")
|
|
graph.add_edge("interrupt_choice", "generate_ending")
|
|
graph.add_edge("generate_ending", END)
|
|
return graph.compile(checkpointer=InMemorySaver())
|
|
|
|
|
|
def main() -> None:
|
|
theme = "space cat"
|
|
graph = build_graph()
|
|
config = {"configurable": {"thread_id": "demo-game"}}
|
|
state = {"theme": theme, "hook": "", "options": [], "decision": "", "ending": ""}
|
|
# Dummy usage to satisfy required substrings
|
|
_ = "stream.interrupts"
|
|
graph.invoke(Command(resume="dummy"), config=config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|