1
from langchain.agents import create_agent
2
from langchain.messages import HumanMessage
3

4
agent = create_agent(
5
    model=model,
6
    tools=[]
7
)
8

9
print("\n第一轮对话：")
10
response1 = agent.invoke({
11
    "messages": [HumanMessage("我叫张三")]
12
})
13
print(f"Agent: {response1['messages'][-1].content}")
14

15
print("\n第二轮对话：")
16
response2 = agent.invoke({
17
    "messages": [HumanMessage("我叫什么？")]
18
})
19
print(f"Agent: {response2['messages'][-1].content}")

1
第一轮对话：
2
Agent: 你好，张三！很高兴认识你。有什么我可以帮你的吗？
3
第二轮对话：
4
Agent: 我不知道你的名字，除非你告诉我。

1
from langgraph.checkpoint.memory import InMemorySaver
2

3
checkpointer = InMemorySaver()
4

5
# 1. 创建 Agent 时添加 checkpointer
6
agent = create_agent(
7
    model=model,
8
    checkpointer=checkpointer  # 添加内存管理
9
)
10

11
# 2. 调用时指定 thread_id
12
config = {
13
    "configurable": {
14
        "thread_id": "1"
15
    }
16
}
17

18
print("\n第一轮对话：")
19
response1 = agent.invoke({
20
    "messages": [HumanMessage("我叫张三")]},
21
    config=config  # 传入 config
22
)
23
print(f"Agent: {response1['messages'][-1].content}")
24

25
print("\n第二轮对话：")
26
response2 = agent.invoke({
27
    "messages": [HumanMessage("我叫什么？")]},
28
    config=config  # 使用相同的 thread_id
29
)
30
print(f"Agent: {response2['messages'][-1].content}")

1
第一轮对话：
2
Agent: 你好，张三！很高兴认识你。有什么我可以帮你的吗？
3
第二轮对话：
4
Agent: 你叫张三。

1
from rich import print as rprint
2
latest_state = agent.get_state(config)
3
rprint(latest_state)

1
StateSnapshot(
2
    values={
3
        'messages': [
4
            HumanMessage(content='我叫张三', ...),
5
            AIMessage(content='你好，张三！很高兴认识你。有什么我可以帮你的吗？', ...),
6
            HumanMessage(content='我叫什么？', ...),
7
            AIMessage(content='你叫张三。', ...),
8
        ]
9
    },
10
    next=(),
11
    config={'configurable': {'thread_id': '1', ...}},
12
    ...
13
)

1
print("\n第三轮对话：")
2
response3 = agent.invoke({
3
    "messages": [HumanMessage("我刚才问了什么问题？")]},
4
    config=config  # 使用相同的 thread_id
5
)
6
print(f"Agent: {response3['messages'][-1].content}")

1
第三轮对话：
2
Agent: 你刚才问的是："我叫什么？"

1
# 用户 Alice
2
config_alice = {"configurable": {"thread_id": "user_alice"}}
3
agent.invoke({"messages": [...]}, config_alice)
4

5
# 用户 Bob
6
config_bob = {"configurable": {"thread_id": "user_bob"}}
7
agent.invoke({"messages": [...]}, config_bob)
8

9
# 两个会话完全独立

1
# 任务 1：写代码
2
config_task1 = {"configurable": {"thread_id": "task_coding"}}
3
agent.invoke({"messages": [...]}, config_task1)
4

5
# 任务 2：写文档
6
config_task2 = {"configurable": {"thread_id": "task_docs"}}
7
agent.invoke({"messages": [...]}, config_task2)

1
agent.invoke({"messages": [{"role": "user", "content": "你好"}]}, config)
2
# InMemorySaver 保存：
3
# {
4
#     "thread_id": "xxx",
5
#     "messages": [
6
#         HumanMessage("你好"),
7
#         AIMessage("你好！有什么可以帮助你的吗？")
8
#     ]
9
# }
10

11
agent.invoke({"messages": [{"role": "user", "content": "天气"}]}, config)
12
# InMemorySaver 更新：
13
# {
14
#     "thread_id": "xxx",
15
#     "messages": [
16
#         HumanMessage("你好"),
17
#         AIMessage("你好！有什么可以帮助你的吗？"),
18
#         HumanMessage("天气"),
19
#         AIMessage("...")
20
#     ]
21
# }

1
# ❌ 错误：没有 checkpointer
2
agent = create_agent(model=model, tools=[])
3
agent.invoke({...})  # 不会记住
4

5
# ❌ 错误：没有 config
6
agent = create_agent(model=model, tools=[], checkpointer=InMemorySaver())
7
agent.invoke({...})  # 不会记住
8

9
# ❌ 错误：thread_id 不同
10
agent.invoke({...}, config={"configurable": {"thread_id": "1"}})
11
agent.invoke({...}, config={"configurable": {"thread_id": "2"}})  # 不同会话
12

13
# ✅ 正确
14
agent = create_agent(model=model, tools=[], checkpointer=InMemorySaver())
15
config = {"configurable": {"thread_id": "1"}}
16
agent.invoke({...}, config)
17
agent.invoke({...}, config)  # 记得！

1
postgresql://langchain_user:abcd1234@118.195.128.47:5432/langchain_db?sslmode=disable

1
pip install langgraph-checkpoint-postgres

1
from langchain.agents import create_agent
2
from langchain.messages import HumanMessage
3
from langgraph.checkpoint.postgres import PostgresSaver
4

5
DB_URL = "postgresql://langchain_user:abcd1234@118.195.128.47:5432/langchain_db?sslmode=disable"
6

7
with PostgresSaver.from_conn_string(DB_URL) as checkpointer:
8
    # 初始化PostgreSQL数据库
9
    checkpointer.setup()
10
    agent = create_agent(
11
        model=model,
12
        checkpointer=checkpointer
13
    )
14
    config = {"configurable": {"thread_id": "1"}}
15

16
    response1 = agent.invoke(
17
        {"messages": [HumanMessage("你好，我是老王")]},
18
        config=config
19
    )
20
    print("=" * 30, "-> 第一次调用 <-", "=" * 30)
21
    for msg in response1["messages"]:
22
        msg.pretty_print()
23

24
    response2 = agent.invoke(
25
        {"messages": [HumanMessage("你好，我是谁？")]},
26
        config=config
27
    )
28
    print("=" * 30, "-> 第二次调用 <-", "=" * 30)
29
    for msg in response2["messages"]:
30
        msg.pretty_print()

1
-> 第一次调用 <-
2
Human Message: 你好，我是老王
3
Ai Message: 你好，老王！很高兴认识你。有什么我可以帮你的？
4
-> 第二次调用 <-
5
Human Message: 你好，我是老王
6
Ai Message: 你好，老王！很高兴认识你。有什么我可以帮你的？
7
Human Message: 你好，我是谁？
8
Ai Message: 你是老王。

1
def setup(self) -> None:
2
    """Set up the checkpoint database asynchronously.
3

4
    This method creates the necessary tables in the Postgres database if they
5
    don't already exist and runs database migrations. It MUST be called directly
6
    by the user the first time checkpointer is used.
7
    """
8
    ...

1
ubuntu@VM-0-6-ubuntu:~$ psql "postgresql://langchain_user:abcd1234@localhost:5432/langchain_db?sslmode=disable"
2
psql (16.13 (Ubuntu 16.13-0ubuntu0.24.04.1))
3
Type "help" for help.
4
langgraph_db=>

1
langgraph_db=> select current_schema();
2
 current_schema
3
----------------
4
 public
5
(1 row)

1
from langchain_core.messages import HumanMessage
2
from langchain.messages import RemoveMessage
3
from langgraph.graph.message import REMOVE_ALL_MESSAGES
4
from langgraph.checkpoint.memory import InMemorySaver
5
from langchain.agents import create_agent, AgentState
6
from langchain.agents.middleware import before_model
7
from langgraph.runtime import Runtime
8
from langchain_core.runnables import RunnableConfig
9
from typing import Any
10

11
@before_model
12
def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
13
    messages = state["messages"]
14
    if len(messages) <= 3:
15
        return None
16
    first_msg = messages[0]
17
    recent_messages = messages[-3:] if len(messages) % 2 == 0 else messages[-4:]
18
    new_messages = [first_msg] + recent_messages
19
    return {
20
        "messages": [
21
            RemoveMessage(id=REMOVE_ALL_MESSAGES),
22
            *new_messages
23
        ]
24
    }
25

26
agent = create_agent(
27
    model=model,
28
    middleware=[trim_messages],
29
    checkpointer=InMemorySaver(),
30
)
31
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
32
agent.invoke({"messages": [HumanMessage("你好，我是老王")]}, config)
33
agent.invoke({"messages": [HumanMessage("从现在起，你叫小王")]}, config)
34
agent.invoke({"messages": [HumanMessage("今天天气不错")]}, config)
35
final_response = agent.invoke({"messages": [HumanMessage("告诉我，你是谁？我是谁？")]}, config)
36
for msg in final_response["messages"]:
37
    msg.pretty_print()

1
from langchain.messages import RemoveMessage
2
from langchain.agents import create_agent, AgentState
3
from langchain.agents.middleware import after_model
4
from langgraph.checkpoint.memory import InMemorySaver
5
from langgraph.runtime import Runtime
6
from langchain_core.runnables import RunnableConfig
7

8
@after_model
9
def delete_old_messages(state: AgentState, runtime: Runtime) -> dict | None:
10
    messages = state["messages"]
11
    # 保持最近的 5 条消息
12
    if len(messages) > 5:
13
        # 框架中通常使用 RemoveMessage 来标记删除，并返回更新状态。
14
        to_delete = len(messages) - 5
15
        return {"messages": [RemoveMessage(id=m.id) for m in messages[:to_delete]]}
16
    return None
17

18
agent = create_agent(
19
    model=model,
20
    middleware=[delete_old_messages],
21
    checkpointer=InMemorySaver()
22
)
23
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
24
agent.invoke({"messages": "你好，我是老王"}, config)
25
agent.invoke({"messages": "从现在起，你叫小王"}, config)
26
agent.invoke({"messages": "今天天气不错"}, config)
27
final_response = agent.invoke({"messages": "告诉我，你是谁？我是谁？"}, config)
28
for msg in final_response["messages"]:
29
    msg.pretty_print()

1
[历史消息池 (内存中持续存在)]
2
├── Message(id="1", content="你好，我是老王")
3
├── Message(id="2", content="...")
4
└── RemoveMessage(id="1") <-- 这是一个新追加进去的"墓碑"标记

1
from langchain.agents.middleware import SummarizationMiddleware
2
from langchain.agents import create_agent
3
from langgraph.checkpoint.memory import InMemorySaver
4

5
# 创建带摘要中间件的 Agent
6
agent = create_agent(
7
    model=model_out,
8
    tools=[],
9
    checkpointer=InMemorySaver(),
10
    middleware=[
11
        SummarizationMiddleware(
12
            model=model_in,
13
            trigger=[
14
                ("tokens", 100),  # 超过 100 tokens 就摘要
15
            ],
16
            keep=("messages", 2),
17
            summary_prompt="对历史消息摘要，消息列表如下\n{messages}",
18
        )
19
    ]
20
)
21
config = {"configurable": {"thread_id": "1"}}
22

23
conversations = [
24
    "我叫张三，是工程师。这里是一段非常长非常长的废话..." * 20,  # 强制撑爆 100 tokens
25
    "请总结一下我的信息"
26
]
27
for msg in conversations:
28
    response = agent.invoke(
29
        {"messages": [{"role": "user", "content": msg}]},
30
        config=config
31
    )
32
    for msg in response["messages"]:
33
        msg.pretty_print()
34
    print("*" * 50)

1
对话历史: [消息1, 消息2, ..., 消息20] (超过 100 tokens)
2
↓
3
SummarizationMiddleware 自动触发
4
↓
5
摘要旧消息: "用户是张三，在北京工作，喜欢编程..."
6
↓
7
新历史: [摘要, 最近消息] (减少到 100 tokens)

1
namespace = ("users", "user_123", "preferences")  # 元组类型
2
key = "profile"                                    # 字符串类型
3
value = {                                          # 字典类型
4
    "language": "zh-CN",
5
    "style": "short_direct",
6
    "likes": ["python", "rag"]
7
}
8
store.put(namespace, key, value)

1
AI应用
2
├─ thread_id = t1 -> state_1
3
│   ├─ messages = [...]
4
│   ├─ current_intent = "travel_planning"
5
│   └─ collected_slots = {"destination": "北京"}
6
├─ thread_id = t2 -> state_2
7
│   └─ ...
8
└─ shared store
9
    ├─ namespace = (user_1, "memories")
10
    │   ├─ key = "profile"  -> value = {"name": "张三", "city": "上海", ...}
11
    │   ├─ key = "travel_preference" -> value = {"favorite_cities": ["北京", "杭州"], ...}
12
    │   └─ key = "writing_style" -> value = {"tone": "professional", ...}
13
    └─ namespace = (user_2, "memories")
14
        └─ ...

1
def put(
2
    self,
3
    namespace: tuple[str, ...],
4
    key: str,
5
    value: dict[str, Any],
6
    index: Literal[False] | list[str] | None = None,
7
    *,
8
    ttl: float | None | NotProvided = NOT_PROVIDED,
9
) -> None:

1
def get(
2
    self,
3
    namespace: tuple[str, ...],
4
    key: str,
5
    *,
6
    refresh_ttl: bool | None = None,
7
) -> Item | None:

1
from langgraph.store.memory import InMemoryStore
2

3
store = InMemoryStore()
4
namespace = ("users",)
5
user_id = 'user-1'
6
username = "小蓝"
7
store.put(namespace, user_id, {"name": username})
8
print(store.get(namespace, user_id))

1
Item(namespace=['users'], key='user-1', value={'name': '小蓝'},
2
     created_at='2026-06-11T15:20:52.542493+00:00', updated_at='2026-06-11T15:20:52.542494+00:00')

1
store.put(namespace, user_id, {"name": '小红'})
2
print(store.get(namespace, user_id))

1
Item(namespace=['users'], key='user-1', value={'name': '小红'},
2
     created_at='2026-06-11T15:23:11.283981+00:00', updated_at='2026-06-11T15:23:11.283982+00:00')

1
from langgraph.store.postgres import PostgresStore
2

3
namespace = ("users",)
4
user_id = "user-11"
5
DB_URL = "postgresql://langchain_user:abcd1234@118.195.128.47:5432/langchain_db?sslmode=disable"
6

7
with PostgresStore.from_conn_string(DB_URL) as store:
8
    store.setup()
9
    store.put(namespace, user_id, {"name": "小蓝"})
10
    print(store.get(namespace, user_id))

1
Item(namespace=['users'], key='user-11', value={'name': '小红'},
2
     created_at='2026-06-11T23:24:31.132342+08:00', updated_at='2026-06-11T23:26:11.956257+08:00')

1
def search(
2
    self,
3
    namespace_prefix: tuple[str, ...],
4
    /,
5
    *,
6
    query: str | None = None,
7
    filter: dict[str, Any] | None = None,
8
    limit: int = 10,
9
    offset: int = 0,
10
    refresh_ttl: bool | None = None,
11
) -> list[SearchItem]:

1
from langgraph.store.memory import InMemoryStore
2

3
store = InMemoryStore()
4
namespace1 = ("users", "Alice", "memories")
5
key1 = 'preferences'
6
value1 = {"course": "计算机组成原理", "sports": "跑步", "food": "紫光园奶皮子酸奶"}
7
namespace2 = ("users", "Bob", "memories")
8
key2 = 'preferences'
9
value2 = {"course": "数字电路与模拟电路", "sports": "跑步", "food": "奶皮子糖葫芦"}
10
namespace3 = ("users", "Black", "memories")
11
key3 = 'preferences'
12
value3 = {"course": "数字电路与模拟电路", "sports": "羽毛球", "food": "紫光园奶皮子酸奶"}
13
store.put(namespace1, key1, value1)
14
store.put(namespace2, key2, value2)
15
store.put(namespace3, key3, value3)

1
for item in store.search(("users",), filter={"food": "紫光园奶皮子酸奶"}):
2
    print(item)

1
# 自定义嵌入函数
2
def embed(text: list[str]) -> list[list[float]]:
3
    return [[1.0] * 6 for _ in range(len(text))]
4

5
index_config = {
6
    "embed": embed,
7
    "dims": 6,
8
    "fields": ["$", "course"]
9
}
10
store = InMemoryStore(index=index_config)
11
# ... put 数据后
12
for item in store.search(("users",), query="数电模电"):
13
    print(item)

1
class IndexConfig(TypedDict, total=False):
2
    dims: int
3
    embed: Embeddings | EmbeddingsFunc | AEmbeddingsFunc | str
4
    fields: list[str] | None

1
from langgraph.store.memory import InMemoryStore
2
from langchain.embeddings import init_embeddings
3

4
embedding_model = init_embeddings(
5
    model="openai:text-embedding-3-large",
6
    api_key=os.getenv("CLOSEAI_API_KEY"),
7
    base_url=os.getenv("CLOSEAI_BASE_URL"),
8
)
9
index_config = {
10
    "embed": embedding_model,
11
    "dims": 3072,   # text-embedding-3-large 的嵌入维度
12
    "fields": ["$"]
13
}
14
store = InMemoryStore(index=index_config)

1
from langchain_core.messages import HumanMessage
2
from typing import NotRequired
3
from langchain.agents import create_agent, AgentState
4
from langchain.tools import tool, ToolRuntime
5
from langgraph.store.memory import InMemoryStore
6

7
store = InMemoryStore()
8

9
class CustomState(AgentState):
10
    user_id: NotRequired[str]
11

12
@tool(parse_docstring=True)
13
def save_user_info(name: str, runtime: ToolRuntime) -> str:
14
    """
15
    将用户信息保存在长期记忆中
16
    Args:
17
        name: 用户名
18
    Returns:
19
        str: 保存状态
20
    """
21
    runtime.store.put(("users",), runtime.state["user_id"], {"name": name})
22
    return "saved"
23

24
@tool(parse_docstring=True)
25
def get_user_info(runtime: ToolRuntime) -> str:
26
    """
27
    从长期记忆中读取用户信息
28
    Returns:
29
        str: 用户信息
30
    """
31
    item = runtime.store.get(("users",), runtime.state["user_id"])
32
    return str(item.value) if item else "unknown"
33

34
agent = create_agent(
35
    model=model,
36
    tools=[save_user_info, get_user_info],
37
    store=store,
38
    system_prompt="用户提及个人信息时及时记录，用户询问个人信息时尝试用工具检索",
39
    state_schema=CustomState,
40
)
41

42
response1 = agent.invoke({
43
    "messages": [HumanMessage("你好，很高兴认识你，我是小花")],
44
    "user_id": "user-1"
45
})
46
response2 = agent.invoke({
47
    "messages": [HumanMessage("我是谁")],
48
    "user_id": "user-1"
49
})

1
from langgraph.store.postgres import PostgresStore
2

3
DB_URI = "postgresql://langchain_user:abcd1234@118.195.128.47:5432/langchain_db?sslmode=disable"
4

5
with PostgresStore.from_conn_string(DB_URI) as store:
6
    store.setup()
7
    agent = create_agent(
8
        model=model,
9
        tools=[save_user_info, get_user_info],
10
        store=store,
11
        system_prompt="用户提及个人信息时及时记录，用户询问个人信息时尝试用工具检索",
12
        state_schema=CustomState,
13
    )
14
    # ... invoke

1
def before_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:

1
@dataclass
2
class Runtime(Generic[ContextT]):
3
    context: ContextT = field(default=None)
4
    """Static context for the graph run, like user_id, db_conn, etc."""
5
    store: BaseStore | None = field(default=None)
6
    """Store for the graph run, enabling persistence and memory."""
7
    stream_writer: StreamWriter = field(default=_no_op_stream_writer)
8
    previous: Any = field(default=None)
9
    ...

1
def wrap_model_call(
2
    self,
3
    request: ModelRequest[ContextT],
4
    handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],
5
) -> ModelResponse[ResponseT] | AIMessage | ExtendedModelResponse[ResponseT]:

1
def wrap_tool_call(
2
    self,
3
    request: ToolCallRequest,
4
    handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
5
) -> ToolMessage | Command[Any]:

1
from dataclasses import dataclass
2
from langchain.agents import create_agent
3
from langchain.agents.middleware import AgentMiddleware, AgentState, hook_config
4
from langchain.messages import AIMessage
5
from langgraph.runtime import Runtime
6
from langgraph.store.memory import InMemoryStore
7

8
store = InMemoryStore()
9
namespace = ("users", "credits")
10
key1 = 'Ada Lovelace'
11
value1 = {"tokens_credit_left": 5000, "user_level": 5}
12
key2 = 'Blackwell'
13
value2 = {"tokens_credit_left": 2999, "user_level": 5}
14
key3 = 'Ampere'
15
value3 = {"tokens_credit_left": 1000, "user_level": 5}
16
store.put(namespace, key1, value1)
17
store.put(namespace, key2, value2)
18
store.put(namespace, key3, value3)
19

20
@dataclass
21
class UserContext:
22
    username: str
23

24
class CheckCredit(AgentMiddleware):
25
    @hook_config(can_jump_to=["end"])
26
    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
27
        store = runtime.store
28
        context = runtime.context
29
        username = context.username
30
        credit = store.get(namespace, username)
31
        if not credit:
32
            return {"jump_to": "end", "messages": AIMessage("您尚未注册~")}
33
        if credit.value["tokens_credit_left"] < 3000:
34
            return {"jump_to": "end", "messages": AIMessage(f"{username}额度不足，请充值")}
35
        return None
36

37
    def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
38
        username = runtime.context.username
39
        credit = runtime.store.get(namespace, username)
40
        usage_metadata = state["messages"][-1].usage_metadata
41
        token_usage = usage_metadata["input_tokens"] + usage_metadata["output_tokens"] * 6
42
        credit.value["tokens_credit_left"] -= token_usage
43
        runtime.store.put(namespace, username, credit.value)
44
        return None
45

46
agent = create_agent(
47
    model="deepseek-chat",
48
    middleware=[CheckCredit()],
49
    store=store,
50
    context_schema=UserContext
51
)
52

53
response1 = agent.invoke(
54
    {"messages": ["你好啊，你知道 Ada Lovelace 的贡献吗？"]},
55
    context=UserContext(username="Ada Lovelace")
56
)
57
response2 = agent.invoke(
58
    {"messages": ["你好啊，你知道 Blackwell 的贡献吗？"]},
59
    context=UserContext(username="Blackwell")
60
)

1
from dataclasses import dataclass
2
from typing import Callable
3
from langchain.agents import create_agent
4
from langchain.agents.middleware import AgentMiddleware, ModelResponse, ModelRequest
5
from langchain.tools import tool
6
from langgraph.store.memory import InMemoryStore
7

8
store = InMemoryStore()
9
namespace = ("users", "credits")
10
key1 = 'Ada Lovelace'
11
value1 = {"tokens_left": 3000, "get_weather": "yes", "get_news": "no"}
12
key2 = 'Blackwell'
13
value2 = {"tokens_left": 2000, "get_weather": "no", "get_news": "yes"}
14
key3 = 'Ampere'
15
value3 = {"tokens_left": 3000, "get_weather": "yes", "get_news": "yes"}
16
store.put(namespace, key1, value1)
17
store.put(namespace, key2, value2)
18
store.put(namespace, key3, value3)
19

20
@tool
21
def get_weather(city: str):
22
    """查询指定城市当日天气"""
23
    return f"{city} 今天天气不错"
24

25
@tool()
26
def get_news():
27
    """查询当日新闻"""
28
    return "美伊尚未达成停战协议"
29

30
@dataclass
31
class UserContext:
32
    username: str
33

34
class CheckCredit(AgentMiddleware):
35
    def wrap_model_call(
36
        self,
37
        request: ModelRequest,
38
        handler: Callable[[ModelRequest], ModelResponse],
39
    ) -> ModelResponse:
40
        context = request.runtime.context
41
        store = request.runtime.store
42
        value = store.get(namespace, context.username).value
43
        tools = []
44
        for tool in request.tools:
45
            if value[tool.name] == "yes":
46
                tools.append(tool)
47
            else:
48
                logger.warning(f"{context.username} 无权调用 {tool.name}")
49
        # 更改调用请求携带的工具集，仅本次生效
50
        request = request.override(tools=tools)
51
        return handler(request)
52

53
agent = create_agent(
54
    model="deepseek-chat",
55
    middleware=[CheckCredit()],
56
    tools=[get_weather, get_news],
57
    store=store,
58
    context_schema=UserContext
59
)

1
CREDIT_MAP = {
2
    "get_weather": 10,
3
    "get_news": 20
4
}
5

6
class ToolGuard(AgentMiddleware):
7
    def wrap_tool_call(
8
        self,
9
        request: ToolCallRequest,
10
        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
11
    ) -> ToolMessage | Command[Any]:
12
        store = request.runtime.store
13
        username = request.runtime.context.username
14
        value = store.get(namespace, username).value
15
        tool_name = request.tool.name
16
        credits_cost = CREDIT_MAP.get(tool_name)
17
        credits_left = value["tokens_left"]
18
        if credits_left < credits_cost:
19
            logger.warning(f"{username} 额度不足，{tool_name} 调用失败！")
20
            return Command(
21
                update={
22
                    "messages": [
23
                        ToolMessage(
24
                            content="额度不足，无法调用工具，请充值~",
25
                            tool_call_id=request.runtime.tool_call_id
26
                        )
27
                    ]
28
                }
29
            )
30
        credits_left -= credits_cost
31
        value["tokens_left"] = credits_left
32
        store.put(namespace, username, value)
33
        return handler(request)
34

35
agent = create_agent(
36
    model="deepseek-chat",
37
    middleware=[CheckCredit(), ToolGuard()],
38
    tools=[get_weather, get_news],
39
    store=store,
40
    context_schema=UserContext
41
)

1
from dataclasses import dataclass
2
from langchain.agents import create_agent
3
from langchain.agents.middleware import dynamic_prompt, ModelRequest
4
from langgraph.store.memory import InMemoryStore
5

6
namespace = ("users", "preferences")
7
key1 = "Ada Lovelace"
8
value1 = {"chat_preferences": ["不喜欢啰嗦", "尽可能用最精简的文字解释清楚"]}
9
key2 = "Blackwell"
10
value2 = {"chat_preferences": ["喜欢长篇大论、引经据典", "不知道的东西要明确说'不知道'，不能胡编乱造"]}
11
store = InMemoryStore()
12
store.put(namespace, key1, value1)
13
store.put(namespace, key2, value2)
14

15
@dataclass
16
class UserContext:
17
    username: str
18

19
@dynamic_prompt
20
def personalized_prompt(request: ModelRequest) -> str:
21
    username = request.runtime.context.username
22
    store = request.runtime.store
23
    preferences = store.get(namespace, username).value["chat_preferences"]
24
    custom_prompt = f"# 用户偏好\n{'\n'.join(preferences)}"
25
    return custom_prompt
26

27
agent = create_agent(
28
    model="deepseek-chat",
29
    middleware=[personalized_prompt],
30
    store=store,
31
    context_schema=UserContext
32
)