Add agent.py
This commit is contained in:
@@ -0,0 +1,56 @@
|
||||
import questionary
|
||||
from typing import TypedDict, Optional
|
||||
|
||||
from langgraph.graph import StateGraph, START, END, interrupt, Command
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
# Define state schema
|
||||
class State(TypedDict, total=False):
|
||||
human_value: Optional[str]
|
||||
foo: Optional[str]
|
||||
|
||||
# Node that triggers interrupt
|
||||
async def interrupt_node(state: State) -> State:
|
||||
payload = {
|
||||
"type": "choice",
|
||||
"question": "Уверены, что хотите продолжить?",
|
||||
"options": ["approve", "reject"],
|
||||
}
|
||||
# Trigger interrupt
|
||||
return interrupt(payload)
|
||||
|
||||
# Build graph
|
||||
graph_builder = StateGraph(State)
|
||||
graph_builder.add_node("interrupt", interrupt_node)
|
||||
graph_builder.set_entry_point("interrupt")
|
||||
graph_builder.add_edge("interrupt", END)
|
||||
|
||||
# Compile graph with in-memory checkpointing
|
||||
checkpoint = InMemorySaver()
|
||||
graph = graph_builder.compile(checkpointer=checkpoint)
|
||||
|
||||
if __name__ == "__main__":
|
||||
initial_state: State = {"human_value": None, "foo": None}
|
||||
thread_id = "main-thread"
|
||||
|
||||
# Run graph to detect interrupt
|
||||
for chunk in graph.stream(initial_state, thread_id=thread_id):
|
||||
if "__interrupt__" in chunk:
|
||||
payload = chunk["__interrupt__"]
|
||||
try:
|
||||
answer = questionary.select(
|
||||
payload["question"], choices=payload["options"]
|
||||
).ask()
|
||||
except Exception as e:
|
||||
print(f"Error during user input: {e}")
|
||||
answer = None
|
||||
# Store answer in state
|
||||
if answer is not None:
|
||||
initial_state["human_value"] = answer
|
||||
# Resume graph
|
||||
final_state = graph.invoke(Command(resume=payload), thread_id=thread_id)
|
||||
print("Final state:", final_state)
|
||||
break
|
||||
else:
|
||||
# No interrupt, graph finished normally
|
||||
print("Graph finished without interrupt.")
|
||||
Reference in New Issue
Block a user