from typing import Literal, TypedDict

from langgraph.graph import StateGraph, START, END


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


def evaluate(state: State):
    return {}


def success(state: State):
    return {"result": "通过"}


def retry(state: State):
    return {"result": "需要重试"}


def route_after_evaluate(
    state: State,
) -> Literal["success", "retry"]:
    if state["score"] >= 60:
        return "success"
    return "retry"


builder = StateGraph(State)

builder.add_node("evaluate", evaluate)
builder.add_node("success", success)
builder.add_node("retry", retry)

builder.add_edge(START, "evaluate")

builder.add_conditional_edges(
    "evaluate",
    route_after_evaluate,
    {
        "success": "success",
        "retry": "retry",
    },
)

builder.add_edge("success", END)
builder.add_edge("retry", END)

graph = builder.compile()

print(graph.invoke({"score": 80, "result": ""}))
print(graph.invoke({"score": 40, "result": ""}))