67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TypedDict
|
|
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.types import Command, interrupt
|
|
|
|
|
|
class GameState(TypedDict):
|
|
scene: str
|
|
inventory: list[str]
|
|
decision: str
|
|
|
|
|
|
def intro(state: GameState) -> GameState:
|
|
return {
|
|
"scene": "Вы стоите у двери лаборатории AI-агента.",
|
|
"inventory": state.get("inventory", []),
|
|
"decision": "",
|
|
}
|
|
|
|
|
|
def ask_human(state: GameState) -> GameState:
|
|
decision = interrupt(
|
|
{
|
|
"scene": state["scene"],
|
|
"question": "Открыть дверь или осмотреть коридор?",
|
|
"options": ["open", "look"],
|
|
}
|
|
)
|
|
return {**state, "decision": str(decision)}
|
|
|
|
|
|
def resolve(state: GameState) -> GameState:
|
|
inventory = list(state.get("inventory", []))
|
|
if state.get("decision") == "open":
|
|
inventory.append("access_card")
|
|
scene = "Дверь открыта, агент получил access_card."
|
|
else:
|
|
scene = "Вы нашли подсказку на стене и вернулись к двери."
|
|
return {**state, "scene": scene, "inventory": inventory}
|
|
|
|
|
|
def build_graph():
|
|
graph = StateGraph(GameState)
|
|
graph.add_node("intro", intro)
|
|
graph.add_node("ask_human", ask_human)
|
|
graph.add_node("resolve", resolve)
|
|
graph.add_edge(START, "intro")
|
|
graph.add_edge("intro", "ask_human")
|
|
graph.add_edge("ask_human", "resolve")
|
|
graph.add_edge("resolve", END)
|
|
return graph.compile(checkpointer=InMemorySaver())
|
|
|
|
|
|
def main() -> None:
|
|
graph = build_graph()
|
|
config = {"configurable": {"thread_id": "demo-game"}}
|
|
print(graph.invoke({"scene": "", "inventory": [], "decision": ""}, config=config))
|
|
print("Resume example:")
|
|
print(graph.invoke(Command(resume="open"), config=config)) # type: ignore[name-defined]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|