用户输入 → [节点 1] → 💾 保存状态 → [节点 2] → 💾 保存状态 → 输出
↓  ↓
Checkpoint 1   Checkpoint 2

from langchain_openai import ChatOpenAI
from langchain.chat_models import init_chat_model
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent
 
# 初始化检查点保存器
checkpointer = InMemorySaver()

BASE_URL="" 
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

agent = create_react_agent(
    model=model,
    tools=[],
    # 传入检查点，是将持久化能力“注入”图的关键步骤。编译后的graph对象现在具备了状态管理的能力。
    checkpointer=checkpointer
)

config = {"configurable": {"thread_id": "1"}}  # 激活记忆机制的核心。如果没有提供thread_id，每次invoke调用都将是无状态的，只要使用相同的thread_id，LangGraph就会在多次调用之间维持对话状态

response = agent.invoke(
    {"messages": [{"role": "user", "content": "你好，我叫ada！"}]},
    config
)

print(f"thread1_bot_answer：{response['messages'][-1].content}")

response = agent.invoke(
    {"messages": [{"role": "user", "content": "你好，请问你还记得我叫什么名字么？"}]},
    config
)

print('------------线程1------------------')
print(f"thread1_bot_answer：{response['messages'][-1].content}")

new_config = {"configurable": {"thread_id": "2"}}
response = agent.invoke(
    {"messages": [{"role": "user", "content": "你好，请问你还记得我叫什么名字么？"}]},
    new_config
)
print('------------线程2------------------')
print(f"thread2_bot_answer：{response['messages'][-1].content}")

thread1_bot_answer：你好，Ada！很高兴认识你！😊
 
这是一个很美的名字呢！有什么我可以帮助你的吗？无论是想聊聊天，还是有任何问题需要解答，我都很乐意为你提供帮助。
------------线程1------------------
thread1_bot_answer：当然记得！你刚才告诉我你叫 Ada～很高兴再次和你打招呼！😊
------------线程2------------------
thread2_bot_answer：你好！很抱歉，我无法记住之前对话中的个人信息，比如你的名字。这是为了保护你的隐私，所以我不会保留这类数据。你可以告诉我你的名字，或者任何你想让我称呼你的方式，我会很乐意在这次的对话中使用它！😊

pip install -U "psycopg[binary,pool]" langgraph-checkpoint-postgres

from langchain.chat_models import init_chat_model
from langgraph.graph import StateGraph, MessagesState, START

from langgraph.checkpoint.postgres import PostgresSaver

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

DB_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"

with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    checkpointer.setup()  # 第一次调用时必须要setup()
    
    def call_model(state: MessagesState):
   response = model.invoke(state["messages"])
   return {"messages": response}
    
    builder = StateGraph(MessagesState)
    builder.add_node(call_model)
    builder.add_edge(START, "call_model")
    
    graph = builder.compile(checkpointer=checkpointer)
    
    config = {
   "configurable": {
"thread_id": "1"
   }
    }

    response = graph.invoke(
   {"messages": [{"role": "user", "content": "你好，我叫ada！"}]},
   config
    )

    print(response['messages'][-1].content)

    response = graph.invoke(
   {"messages": [{"role": "user", "content": "你好，请问你还记得我叫什么名字么？"}]},
   config
    )

    print(response['messages'][-1].content)

bot_answer_1：你好，Ada！很高兴再次见到你！😊
 
你的名字真动听！今天有什么我可以帮你解答或者想聊的话题吗？
bot_answer_2：当然记得！你告诉我你叫 **Ada**。很高兴再次和你打招呼！😊

from langgraph.graph import START, StateGraph
from langgraph.checkpoint.memory import InMemorySaver
from typing import TypedDict

class State(TypedDict):
    foo: str

# 子图
def subgraph_node_1(state: State):
    return {"foo": state["foo"] + "bar"}

subgraph_builder = StateGraph(State)
subgraph_builder.add_node(subgraph_node_1)
subgraph_builder.add_edge(START, "subgraph_node_1")
subgraph = subgraph_builder.compile()  

# 父图
builder = StateGraph(State)
builder.add_node("node_1", subgraph)  
builder.add_edge(START, "node_1")

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

subgraph_builder = StateGraph(...)
subgraph = subgraph_builder.compile(checkpointer=True)

from typing import Annotated
from langchain.chat_models import init_chat_model
from langgraph.prebuilt import InjectedState, create_react_agent
from langgraph.prebuilt.chat_agent_executor import AgentState

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

class CustomState(AgentState):
    user_id: str

def get_user_info(
    state: Annotated[CustomState, InjectedState]
) -> str:
    """查询用户信息"""
    user_id = state["user_id"]
    return"ada"if state["user_id"] == "user_ada"else"Unknown"

agent = create_react_agent(
    model=model,
    tools=[get_user_info],
    state_schema=CustomState,
)

agent.invoke({
    "messages": "查询用户信息",
    "user_id": "user_ada"
})

根据查询结果，您的用户信息显示用户名为：**ada**\n\n这是您的用户信息查询结果。如果您需要了解其他信息或有其他需求，请告诉我。

from typing import Annotated
from langchain_core.tools import InjectedToolCallId
from langchain_core.runnables import RunnableConfig
from langchain_core.messages import ToolMessage
from langgraph.prebuilt import InjectedState, create_react_agent
from langgraph.prebuilt.chat_agent_executor import AgentState
from langgraph.types import Command

class CustomState(AgentState):
    user_name: str

def update_user_info(
    tool_call_id: Annotated[str, InjectedToolCallId],
    config: RunnableConfig
) -> Command:
    """查询并更新用户信息"""
    user_id = config["configurable"].get("user_id")
    name = "ada"if user_id == "user_123"else"Unknown user"
    return Command(update={
   "user_name": name,
   # 更新消息历史记录
   "messages": [
ToolMessage(
    "成功查询到用户信息",
    tool_call_id=tool_call_id
)
   ]
    })

def greet(
    state: Annotated[CustomState, InjectedState]
) -> str:
    """找到用户信息后，使用此方式向用户问好。"""
    user_name = state["user_name"]
    return f"你好 {user_name}！"

agent = create_react_agent(
    model=model,
    tools=[update_user_info, greet],
    state_schema=CustomState
)

agent.invoke(
    {"messages": [{"role": "user", "content": "向用户打招呼"}]},
    config={"configurable": {"user_id": "user_123"}}
)

from langchain.chat_models import init_chat_model
from langchain_core.runnables import RunnableConfig
from langgraph.store.memory import InMemoryStore
from langgraph.config import get_store
from langgraph.prebuilt import create_react_agent

store = InMemoryStore()

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

store.put(
    ("users",),  # 命名空间：元组类型，类比文件系统中的文件夹，支持分层组织结构
    "user_123",  # 键: 字符串，是命名空间内的唯一标识符，一般推荐使用uuid库生成唯一标识符
    {
   "name": "ada",
   "language": "中文",
    }  # 值：Python字典类型，比如保存公共角色资料时可以是包含姓名、偏好等键值对的字典
)

def get_user_info(config: RunnableConfig) -> str:
    """查找用户信息的函数，可以查看长期记忆中储存的用户信息"""
    store = get_store()  # 获取上下文中可用的store实例
    user_id = config["configurable"].get("user_id")
    user_info = store.get(("users",), user_id)  # 输入命名空间和键进行精确查询
    return str(user_info.value) if user_info else"Unknown user"

agent = create_react_agent(
    model=model,
    tools=[get_user_info],
    # 传入store
    store=store
)

response = agent.invoke(
    {"messages": [{"role": "user", "content": "帮我查找长期记忆中储存的用户信息"}]},
    config={"configurable": {"user_id": "user_123"}}
)

print(response['messages'])

# ------------------第一次对话----------------------
你好ada！很高兴认识你～我已经记住你的名字啦！✨
 
有什么我可以帮你的吗？
# ------------------第二次对话----------------------
你的名字是ada

from typing_extensions import TypedDict
 
from langchain.chat_models import init_chat_model
from langgraph.config import get_store
from langchain_core.runnables import RunnableConfig
from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore

store = InMemoryStore() 

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

class UserInfo(TypedDict): 
    name: str

def save_user_info(user_info: UserInfo, config: RunnableConfig) -> str: 
    """将用户信息保存到store"""
    store = get_store() 
    user_id = config["configurable"].get("user_id")
    store.put(("users",), user_id, user_info) 
    return"成功保存了用户信息"

agent = create_react_agent(
    model=model,
    tools=[save_user_info],
    store=store
)

agent.invoke(
    {"messages": [{"role": "user", "content": "我叫ada！请你记住我的名字"}]},
    config={"configurable": {"user_id": "user_123"}} 
)

store.get(("users",), "user_123").value
 
# 输出：{'name': 'ada'}

import requests
from typing import List, Optional, Dict
from langchain.embeddings.base import Embeddings

class SelfAPIEmbeddings(Embeddings):
    """
    一个自定义的 Embedding 类，用于调用自部署的 embedding API 服务。
    """
    def __init__(self):
   """
   初始化函数。
   """
   self.token = ""
   self.url = ""
   self.model_id = ""

    def _call_api(self, texts: List[str]) -> List[List[float]]:
   """
   内部方法，用于调用 API。
   *** 您需要根据您自己服务的实际 API 格式来修改这部分 ***
   """
   try:
payload = {
    'model': self.model_id,
    'input': texts
}

headers = {
    'Content-Type': 'application/json',
    'Authorization': f'Bearer {self.token}'
}

response = requests.post(self.url, headers=headers, data=json.dumps(payload))

# 判断是否异常
if response.status_code != 200:
    print(response.json())
    exit()

res = response.json()
return res['data']
   except requests.exceptions.RequestException as e:
print(f"Error calling embedding API: {e}")
# 可以选择返回空列表或重新抛出异常
raise e

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
   """
   为文档列表生成 embeddings。
   """
   if not texts:
return []
   # 为了避免请求体过大，可以分批处理
   # 这里为了简单起见，一次性发送所有文本
   new_texts = []
   for i in texts:
i = eval(i)
new_texts.append(i['text'])
   
   res = self._call_api(new_texts)
   new_res = []
   for i in res:
new_res.append(i['embedding'])
   return new_res

    def embed_query(self, text: str) -> List[float]:
   """
   为单个查询文本生成 embedding。
   """
   if not text:
return []
   result = self._call_api([text])
   return result[0]['embedding']

from langchain.embeddings import init_embeddings
from langgraph.store.memory import InMemoryStore
 
custom_embeddings = SelfAPIEmbeddings()
store = InMemoryStore(
    index={
   "embed": custom_embeddings,
   "dims": 2560,
    }
)

store.put(("user_123", "memories"), "1", {"text": "我喜欢吃披萨"})
store.put(("user_123", "memories"), "2", {"text": "我是一名程序员"})
 
items = store.search(
    ("user_123", "memories"), query="我肚子饿了", limit=1
)

[Item(namespace=['user_123', 'memories'], key='1', value={'text': '我喜欢吃披萨'}, created_at='2025-11-12T09:59:55.097931+00:00', updated_at='2025-11-12T09:59:55.097937+00:00', score=0.6804530799409887)]

from langchain_core.messages.utils import (
    trim_messages,
    count_tokens_approximately
)
from langchain.chat_models import init_chat_model
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, START, MessagesState

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

summarization_model = model.bind(max_tokens=128)

def call_model(state: MessagesState):
    # 保留最近消息，总 token ≤ 128
    messages = trim_messages(
   state["messages"],
   strategy="last",  # 保留最后的消息
   token_counter=count_tokens_approximately,
   max_tokens=128,
   start_on="human",  # 确保第一条消息（不包括系统消息）是从human消息开始保留
   end_on=("human", "tool"),  # 保留到human或tool消息为止
   allow_partial=False,  # 不允许分割消息内容
   include_system=True  # 保留system prompt
    )
    # --- 在这里打印传入模型的内容 ---
    print("-" * 20)
    print(f"Messages being sent to the model (trimmed to <= 128 tokens): {len(messages)}")
    for msg in messages:
   print(f"  [{msg.type.upper()}]: {msg.content}")
    print("-" * 20)
    response = model.invoke(messages)
    return {"messages": [response]}

checkpointer = InMemorySaver()
builder = StateGraph(MessagesState)
builder.add_node(call_model)
builder.add_edge(START, "call_model")
graph = builder.compile(checkpointer=checkpointer)

# 如果需要在agent中修剪，那么需要将pre_model_hook和trim_messages结合使用
"""
def call_model(state)
    messages = trim_messages(...)
 
agent = create_react_agent(
 model,
 tools,
 pre_model_hook=call_model,
 checkpointer=checkpointer,
)
"""

config = {"configurable": {"thread_id": "1"}}
 
graph.invoke({"messages": "你好，我叫ada"}, config)
graph.invoke({"messages": "请写一首诗，关于小狗的"}, config)
graph.invoke({"messages": "再写一首关于小猫的"}, config)
final_response = graph.invoke({"messages": "我叫什么名字呢？"}, config)

--------------------
Messages being sent to the model (trimmed to <= 128 tokens): 1
  [HUMAN]: 你好，我叫ada
--------------------

--------------------
Messages being sent to the model (trimmed to <= 128 tokens): 3
  [HUMAN]: 你好，我叫ada
  [AI]: 你好，Ada！很高兴认识你！😊

这是一个很美的名字呢！有什么我可以帮助你的吗？无论是想聊天、有问题需要解答，还是需要任何形式的帮助，我都很乐意为你服务。
  [HUMAN]: 请写一首诗，关于小狗的
--------------------

--------------------
Messages being sent to the model (trimmed to <= 128 tokens): 5
  [HUMAN]: 你好，我叫ada
  [AI]: 你好，Ada！很高兴认识你！😊

这是一个很美的名字呢！有什么我可以帮助你的吗？无论是想聊天、有问题需要解答，还是需要任何形式的帮助，我都很乐意为你服务。
  [HUMAN]: 请写一首诗，关于小狗的
  [AI]: 好的，Ada！为你写一首关于小狗的可爱小诗，希望你喜欢：

**《小狗的约定》**
.....  （此处省略诗的内容）
这首诗捕捉了小狗的活泼、忠诚和它们带给我们的温暖。你觉得怎么样？😊
  [HUMAN]: 再写一首关于小猫的
--------------------

--------------------
Messages being sent to the model (trimmed to <= 128 tokens): 3
  [HUMAN]: 再写一首关于小猫的
  [AI]: 好的，Ada！这首关于小猫的诗，希望同样能带给你一丝轻盈与温柔：

**《小猫的遐想》**
.....  （此处省略诗的内容）
希望你喜欢这首小诗！🐾
  [HUMAN]: 我叫什么名字呢？
--------------------

print("\n" + "="*30)
print("  查看 thread_id='1' 的完整对话历史")
print("="*30)
 
current_state = graph.get_state(config)
conversation_history = current_state.values["messages"]
 
for message in conversation_history:
    print(f"[{message.type.upper()}]: {message.content}")

from langchain_core.messages import RemoveMessage
from langchain.chat_models import init_chat_model
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, START, MessagesState

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

def delete_messages(state):
    messages = state["messages"]
    if len(messages) > 2:
   # 删除最早的两条消息
   return {"messages": [RemoveMessage(id=m.id) for m in messages[:2]]}
    
def call_model(state: MessagesState):
    response = model.invoke(state["messages"])
    return {"messages": response}

builder = StateGraph(MessagesState)
builder.add_sequence([call_model, delete_messages])
builder.add_edge(START, "call_model")

checkpointer = InMemorySaver()
app = builder.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}}

for event in app.stream(
    {"messages": [{"role": "user", "content": "你好，我是ada"}]},
    config,
    stream_mode="values"
):
    print([(message.type, message.content) for message in event["messages"]])

for event in app.stream(
    {"messages": [{"role": "user", "content": "我叫什么名字"}]},
    config,
    stream_mode="values"
):
    print([(message.type, message.content) for message in event["messages"]])

from langgraph.graph.message import REMOVE_ALL_MESSAGES
 
def clear_messages(state):
    return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}

builder = StateGraph(MessagesState)
builder.add_sequence([call_model, clear_messages])
builder.add_edge(START, "call_model")

checkpointer = InMemorySaver()
app = builder.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}}

for event in app.stream(
    {"messages": [{"role": "user", "content": "你好，我是ada"}]},
    config,
    stream_mode="values"
):
    print([(message.type, message.content) for message in event["messages"]])

for event in app.stream(
    {"messages": [{"role": "user", "content": "我叫什么名字"}]},
    config,
    stream_mode="values"
):
    print([(message.type, message.content) for message in event["messages"]])

pip install -U langmem

import tiktoken
from typing import Any, TypedDict

from langchain.chat_models import init_chat_model
from langchain_core.messages import AnyMessage, BaseMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.graph import StateGraph, START, MessagesState
from langgraph.checkpoint.memory import InMemorySaver
from langmem.short_term import SummarizationNode, RunningSummary

summary_prompt = ChatPromptTemplate.from_messages(
    [
   MessagesPlaceholder(variable_name="messages"),
   # 使用 HumanMessage 模拟用户在最后发出总结指令
   ("human", "请根据以上对话，生成一段简洁、连贯的中文摘要，注意不要丢失细节"),
    ]
)

update_summary_prompt = ChatPromptTemplate.from_messages(
    [
   MessagesPlaceholder(variable_name="messages"),
   # 使用 HumanMessage 模拟用户在最后发出总结指令
   ("human", "以下是目前为止的对话摘要：{existing_summary}\n\n请根据以上新消息扩展此摘要："),
    ]
)

BASE_URL=""
TOKEN=""
MODEL_NAME=""

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)
summarization_model = model.bind(max_tokens=128)

# count_tokens_approximately更适合英文分词，中文这里使用tiktoken来计算token数量
encoding = tiktoken.get_encoding("cl100k_base")

def count_tokens_accurately(messages: list[BaseMessage]) -> int:
    """使用 tiktoken 精确计算消息列表的 token 总数"""
    # 注意：langmem 的 token_counter 期望的输入是消息列表
    text_content = " ".join([msg.content for msg in messages if isinstance(msg.content, str)])
    return len(encoding.encode(text_content))

class State(MessagesState):
    context: dict[str, RunningSummary]  

class LLMInputState(TypedDict):  
    summarized_messages: list[AnyMessage]
    context: dict[str, RunningSummary]

summarization_node = SummarizationNode(  
#token_counter=count_tokens_approximately,
    token_counter=count_tokens_accurately,  # 更换为自定义的token计算工具
    model=summarization_model,
    max_tokens=256,
    max_tokens_before_summary=256,
    max_summary_tokens=128,
    initial_summary_prompt=summary_prompt,  # 使用自定义prompt，默认为英文
    existing_summary_prompt=update_summary_prompt
)

def call_model(state: LLMInputState):
    response = model.invoke(state["summarized_messages"])
    return {"messages": [response]}

checkpointer = InMemorySaver()
builder = StateGraph(State)
builder.add_node(call_model)
builder.add_node("summarize", summarization_node)  
builder.add_edge(START, "summarize")
builder.add_edge("summarize", "call_model")
graph = builder.compile(checkpointer=checkpointer)

# Invoke the graph
config = {"configurable": {"thread_id": "1"}}
graph.invoke({"messages": "你好，我叫ada"}, config)
graph.invoke({"messages": "请写一首诗，关于小狗的"}, config)
graph.invoke({"messages": "再写一首关于小猫的"}, config)
final_response = graph.invoke({"messages": "我叫什么名字？"}, config)

final_response["messages"][-1].pretty_print()
# 检查摘要是否生成
if"running_summary"in final_response["context"]:
    print("\n生成的摘要:", final_response["context"]["running_summary"].summary)
else:
    print("\n对话较短，尚未生成摘要。")

graph.get_state(config=config)

graph.get_state_history(config=config)

checkpointer.delete_thread(thread_id)

pip install langchain-mcp-adapters

BASE_URL="" 
TOKEN=""
MODEL_NAME=""

from langchain_openai import ChatOpenAI
from langchain.chat_models import init_chat_model

model = init_chat_model(
    model=MODEL_NAME,
    model_provider="openai", 
    base_url=BASE_URL,
    api_key=TOKEN,
    temperature=0,
)

from langgraph.store.memory import InMemoryStore
from langgraph.checkpoint.memory import InMemorySaver
 
store = InMemoryStore()
checkpointer = InMemorySaver()

from typing import List
from typing_extensions import TypedDict

from langchain_core.messages import BaseMessage
from langchain_core.runnables import RunnableConfig
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langgraph.config import get_store

from langgraph.graph import StateGraph, END
from langgraph.types import interrupt
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage

# 搜索功能
url = ''
TOKEN = ''
search_client = MultiServerMCPClient(
    {
   "other_search": {
"url": url,
"headers": {
    "Authorization": f"Bearer {TOKEN}"
},
"transport": "sse"
   }
    }
)
search_tools = await search_client.get_tools()
search_agent = create_react_agent(
    model,
    search_tools,
    name="search_assistant",
    prompt="你是一个能搜索各种信息的助手。"
)

class UserInfo(TypedDict): 
    user_id: str
    hotel_name: str
    date: str
    num_guests: int

def book_hotel(user_info: UserInfo, config: RunnableConfig):
    """处理酒店预订并更新长期记忆"""
    user_id = config["configurable"].get("user_id")
    print(user_info)
    
    hotel_name = user_info.get("hotel_name")
    date = user_info.get("date")
    num_guests = user_info.get("num_guests")
    
    # 存储用户个人预订历史
    namespace = ("user_bookings",)
    
    user_bookings = store.get(namespace, user_id) or []
    user_bookings.append(user_info)
    store.put(namespace, user_id, user_bookings)
    
    # 更新总预订计数
    namespace = ("total_hotel_bookings",)
    total_bookings = store.get(namespace, 'total_bookings_num') or 0
    store.put(namespace, 'total_bookings_num', total_bookings + 1)
    
    return f"成功为用户 {user_id} 预订了 {hotel_name}，入住日期：{date}，入住人数：{num_guests}"

book_hotel_agent = create_react_agent(
    model=model,
    tools=[book_hotel],
    store=store,
    name="hotel_assistant",
    prompt="你是一个酒店预定助手。不需要用户ID、身份证号码、姓名和联系方式，就可以预定，请直接预定！"
)

def query_booking_from_store(config: RunnableConfig) -> str:
    """根据用户ID从存储中查询酒店预订信息。"""
    store = get_store()
    user_id = config["configurable"].get("user_id")
    booking_info = store.get(("user_bookings",), user_id)
 
    if booking_info and booking_info.value:
   return f"已找到预订信息：{str(booking_info.value)}"
    else:
   return "未找到该用户的预订信息"

class SubgraphState(TypedDict):
    messages: List[BaseMessage]

def authentication_and_query_node(state: SubgraphState, config: RunnableConfig):
    """
    这个节点首先中断图的执行以请求管理员ID，
    然后在恢复后验证ID，并调用工具查询信息。
    """
    # 核心：调用 interrupt() 来暂停图的执行
    admin_input = interrupt("请输入管理员id，如需退出查询，请输入exit")
    
    # 当图被恢复时，admin_input 将会获得传入的值
    if admin_input == "exit":
   result = "用户已退出查询。"
    elif admin_input == "admin_123":
   # 验证成功，调用真正的查询工具
   result = query_booking_from_store(config)
    else:
   # 验证失败
   result = f"没有权限查询：admin_id 不匹配 (输入为: '{admin_input}')"

    return {"messages": [AIMessage(content=result)]}

query_workflow = StateGraph(SubgraphState)
query_workflow.add_node("auth_and_query", authentication_and_query_node)
query_workflow.set_entry_point("auth_and_query")
query_workflow.add_edge("auth_and_query", END)
 
booking_query_subgraph = query_workflow.compile(checkpointer=checkpointer, store=store)
 
# 为子图命名，以便Supervisor可以调用它
booking_query_subgraph.name = "booking_info_assistant"

from langgraph_supervisor import create_supervisor
 
workflow = create_supervisor(
    [search_agent, book_hotel_agent, booking_query_subgraph],
    model=model,
    prompt=(
   "您是团队主管，负责管理信息搜索助手、酒店预订助手、以及用户信息查询助手。"
   "如需搜索各种信息，请交由 search_assistant 处理。"
   "如需预定酒店，请交由 hotel_assistant 处理。"
   "如需查询用户的酒店预定信息，请交由 booking_info_assistant 处理。"
   "**注意**，你每次只能调用一个助理agent！"
    ),
)

supervisor = workflow.compile(checkpointer=checkpointer, store=store)

config = {"configurable": {"thread_id": "1", "user_id": "user_123"}}
 
print("--- 第一次交互 ：查询北京火锅店---")
async for chunk in supervisor.astream(
    {"messages": [("user", "北京最出名的老北京火锅是哪家？")]},
    config
):
    # 打印每个数据块的内容
    for key, value in chunk.items():
   print(f"Node: '{key}'")
   if value:
print("  value:")
print(value)
    print("----")

北京最出名的老北京火锅有很多家，其中比较有名的包括：\n\n1. **东来顺**：东来顺是北京最著名的老北京涮羊肉火锅店之一，历史悠久，以其独特的涮羊肉和秘制的调料而闻名。\n\n2. **南门涮肉**：南门涮肉也是一家老字号的火锅店，以其传统的涮羊肉和地道的北京风味而受到欢迎。\n\n3. **老北京涮肉馆**：这是一家专注于传统老北京涮肉的火锅店，以其正宗的口味和优质的服务而受到食客的喜爱。\n\n4. **聚宝源**：聚宝源是一家以涮羊肉为主的火锅店，以其新鲜的食材和独特的调料而受到欢迎。\n\n这些火锅店都有各自的特色和忠实的顾客群体，您可以根据自己的口味和需求选择合适的火锅店。

print("--- 第二次交互 ：根据上一步推荐的火锅店查询酒店---")
async for chunk in supervisor.astream(
    {"messages": [("user", "那第一个推荐的火锅店附近有哪些酒店呀")]},
    config
):
    # 打印每个数据块的内容
    for key, value in chunk.items():
   print(f"Node: '{key}'")
   if value:
print("  value:")
print(value)
    print("----")

东来顺火锅店位于北京市东城区东华门大街。以下是东来顺火锅店附近的一些酒店推荐：\n\n1. **北京王府井希尔顿酒店**：这是一家豪华酒店，距离东来顺火锅店步行约10分钟，提供高品质的住宿和服务。\n\n2. **北京东方君悦大酒店**：这家五星级酒店位于王府井大街，距离东来顺火锅店步行约15分钟，设施齐全，服务优质。\n\n3. **北京华尔道夫酒店**：这是一家高端酒店，距离东来顺火锅店步行约10分钟，提供豪华的住宿体验和优质的服务。\n\n4. **北京诺富特和平宾馆**：这家四星级酒店位于王府井大街，距离东来顺火锅店步行约15分钟，性价比较高，适合商务和休闲旅行。\n\n5. **北京天伦王朝酒店**：这家四星级酒店距离东来顺火锅店步行约10分钟，提供舒适的住宿环境和便捷的交通。\n\n这些酒店都位于东来顺火锅店附近，您可以根据自己的需求和预算选择合适的酒店。

print("--- 第三次交互 ：预定酒店---")
 
async for chunk in supervisor.astream(
    {"messages": [("user", "帮我预定北京王府井希尔顿酒店酒店，预定日期：2025-11-13到2025-11-14，入住人数1")]},
    config
):
    for key, value in chunk.items():
   print(f"Node: '{key}'")
   if value:
print("  value:")
print(value)
    print("----")

已成功为您预订了北京王府井希尔顿酒店，入住日期为2025年11月13日至2025年11月14日，入住人数为1人。祝您旅途愉快！

print("--- 第四次交互 ：管理员查询预定信息---")
config = {"configurable": {"thread_id": "2", "user_id": "user_123"}}  # 更换管理员操作线程

interrupt_data = None
interrupt_input = None
print("--- 第一次运行，将会触发中断 ---")
async for chunk in supervisor.astream(
    {"messages": [("user", "查询用户预定酒店信息")]},
    config,
):
    for key, value in chunk.items():
   print(f"Node: '{key}'")
   if value:
print("  value:")
print(value)

   if key == "__interrupt__":
print("\n======= 图已成功中断！=======")
interrupt_data = value[0] 
print(f"中断信息: {interrupt_data.value}")
break

    if interrupt_data:
   break
    print("----")

if interrupt_data:
    # 模拟管理员输入正确的密码
    interrupt_input = Command(resume="admin_123")
    # 如果想测试错误的密码，可以使用下面这行
#interrupt_input = Command(resume="wrong_password")
    
    print(f"\n--- 接收到中断输入 '{interrupt_input}'，继续执行图 ---")
    # 恢复图的执行
    async for chunk in supervisor.astream(
   interrupt_input,
   config,
    ):
   for key, value in chunk.items():
print(f"Node: '{key}'")
if value:
    print("  value:")
    print(value)
   print("----")

请输入管理员id，如需退出查询，请输入exit

已找到预订信息：[{'user_id': 'user_123', 'hotel_name': '北京王府井希尔顿酒店', 'date': '2025-11-13到2025-11-14', 'num_guests': 1}]

用户预定的酒店信息如下：\n- 用户ID: user_123\n- 酒店名称: 北京王府井希尔顿酒店\n- 预定日期: 2025-11-13到2025-11-14\n- 客人数: 1