from typing import Literal, TypedDict

from langgraph.graph import END, START, StateGraph
from langgraph.types import Command


class State(TypedDict):
    number: int
    result: str


def check_number(
    state: State,
) -> Command[Literal["positive", "negative"]]:
    if state["number"] >= 0:
        return Command(
            update={"result": "这是非负数"},
            goto="positive",
        )

    return Command(
        update={"result": "这是负数"},
        goto="negative",
    )


def positive(state: State):
    print(state["result"])
    return {}


def negative(state: State):
    print(state["result"])
    return {}


builder = StateGraph(State)
builder.add_node("check", check_number)
builder.add_node("positive", positive)
builder.add_node("negative", negative)

builder.add_edge(START, "check")
builder.add_edge("positive", END)
builder.add_edge("negative", END)

graph = builder.compile()

graph.invoke({
    "number": -5,
    "result": "",
})