from typing import Annotated, TypedDict

from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages


class State(TypedDict):
    messages: Annotated[list, add_messages]


model = ChatOpenAI(
    model="deepseek-v4-flash",
    temperature=0,
    base_url="https://api.deepseek.com/v1"
)


def chatbot(state: State):
    return {
        "messages": [model.invoke(state["messages"])]
    }


builder = StateGraph(State)
builder.add_node("chatbot", chatbot)
builder.add_edge(START, "chatbot")
builder.add_edge("chatbot", END)

checkpointer = MemorySaver()
graph = builder.compile(checkpointer=checkpointer)

config = {
    "configurable": {
        "thread_id": "user-001",
    }
}

first = graph.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "我叫小明，请记住。",
            }
        ]
    },
    config=config,
)

second = graph.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "我叫什么名字？",
            }
        ]
    },
    config=config,
)

print(second["messages"][-1].content)