74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
from typing import TypedDict, Optional
|
||
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from langgraph.types import Command, interrupt
|
||
from langchain_openai import ChatOpenAI
|
||
from pydantic import SecretStr
|
||
|
||
|
||
class GraphState(TypedDict):
|
||
human_value: Optional[str] = None
|
||
|
||
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b",
|
||
base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
|
||
api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
|
||
temperature=0.7,
|
||
)
|
||
|
||
|
||
def ask(state: GraphState) -> Command:
|
||
# Trigger an interrupt with a question and options
|
||
return interrupt(
|
||
type="question",
|
||
question="Choose your favorite color:",
|
||
options=["red", "green", "blue"],
|
||
)
|
||
|
||
|
||
# Build the graph
|
||
builder = StateGraph(GraphState)
|
||
builder.add_node("ask", ask)
|
||
|
||
# Set entry point correctly (only one argument – node name)
|
||
builder.set_entry_point("ask")
|
||
|
||
# Define edges
|
||
builder.add_edge(START, "ask")
|
||
builder.add_edge("ask", END)
|
||
|
||
# Add a checkpoint
|
||
checkpoint = InMemorySaver()
|
||
graph = builder.compile(checkpointer=checkpoint)
|
||
|
||
|
||
def main():
|
||
state: GraphState = {"human_value": None}
|
||
while True:
|
||
result = graph.invoke(state)
|
||
# If the graph finished without interruption
|
||
if "messages" in result and not any(
|
||
getattr(msg, "type", "") == "__interrupt__" for msg in result["messages"]
|
||
):
|
||
print("Graph finished. Result:", state)
|
||
break
|
||
|
||
# Handle interrupt
|
||
for msg in result["messages"]:
|
||
if getattr(msg, "type", "") == "__interrupt__":
|
||
question = msg.content.get("question")
|
||
options = msg.content.get("options")
|
||
print(question)
|
||
for i, opt in enumerate(options, 1):
|
||
print(f"{i}. {opt}")
|
||
choice_idx = int(input("Select option number: ")) - 1
|
||
state["human_value"] = options[choice_idx]
|
||
# Resume the graph with updated state
|
||
# The resume is implicit by invoking again after setting the value
|
||
break
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |