1
from typing import NotRequired
2

3
from langchain.agents import AgentState, create_agent
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

19
    Returns:
20
        str: 保存状态
21
    """
22
    runtime.store.put(("users",), runtime.state["user_id"], {"name": name})
23
    return "saved"
24

25
@tool(parse_docstring=True)
26
def get_user_info(runtime: ToolRuntime) -> str:
27
    """从长期记忆中读取用户信息
28

29
    Returns:
30
        str: 用户信息
31
    """
32
    item = runtime.store.get(("users",), runtime.state["user_id"])
33
    return str(item.value) if item else "unknown"
34

35
agent = create_agent(
36
    model=model,
37
    tools=[save_user_info, get_user_info],
38
    store=store,                              # ← 把 store 交给 Agent
39
    state_schema=CustomState,                 # ← 用自定义状态把 user_id 带进来
40
    system_prompt="用户提及个人信息时及时记录，用户询问个人信息时尝试用工具检索",
41
)
42

43
agent.invoke({"messages": ["我叫韩立"], "user_id": "user_1"})

1
from dataclasses import dataclass
2

3
from langchain.agents import create_agent
4

5
@dataclass
6
class UserContext:
7
    username: str
8

9
agent = create_agent(
10
    model=model,
11
    store=store,
12
    context_schema=UserContext,     # ← 声明用哪个 schema
13
)
14

15
agent.invoke(
16
    {"messages": ["你好啊，你知道 Ada Lovelace 的贡献吗？"]},
17
    context=UserContext(username="Ada Lovelace"),   # ← 每次调用传入具体值
18
)

1
from dataclasses import dataclass
2
from typing import Any
3

4
from langchain.agents import create_agent
5
from langchain.agents.middleware import AgentMiddleware, AgentState, hook_config
6
from langchain_core.messages import AIMessage
7
from langgraph.runtime import Runtime
8
from langgraph.store.memory import InMemoryStore
9

10
store = InMemoryStore()
11
namespace = ("users", "credits")
12
store.put(namespace, "Ada Lovelace", {"tokens_credit_left": 5000, "user_level": 5})
13
store.put(namespace, "Blackwell",    {"tokens_credit_left": 2999, "user_level": 5})
14
store.put(namespace, "Ampere",       {"tokens_credit_left": 1000, "user_level": 5})
15

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

20
class CheckCredit(AgentMiddleware):
21
    @hook_config(can_jump_to=["end"])          # ← 类写法要跳转，必须加这个装饰器
22
    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
23
        store = runtime.store
24
        context = runtime.context               # ← 静态上下文
25

26
        username = context.username
27
        credit = store.get(namespace, username)
28

29
        if not credit:                          # 查无此人
30
            return {"jump_to": "end", "messages": AIMessage("您尚未注册~")}
31

32
        if credit.value["tokens_credit_left"] < 3000:     # 额度不足
33
            return {"jump_to": "end", "messages": AIMessage(f"{username}额度不足，请充值")}
34

35
        return None
36

37
    def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
38
        store = runtime.store
39
        context = runtime.context
40

41
        username = context.username
42
        credit = store.get(namespace, username)
43

44
        # 用上一条 AIMessage 的 usage_metadata 算本次消耗
45
        usage_metadata = state["messages"][-1].usage_metadata
46
        token_usage = usage_metadata["input_tokens"] + usage_metadata["output_tokens"] * 6
47
        credit.value["tokens_credit_left"] -= token_usage
48

49
        # 更新长期记忆
50
        store.put(namespace, username, credit.value)
51
        return None
52

53
    def after_agent(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
54
        store = runtime.store
55
        context = runtime.context
56

57
        username = context.username
58
        item = store.get(namespace, username)
59
        print(f"{username} 当前剩余额度：{item.value['tokens_credit_left']}")
60
        return None
61

62
agent = create_agent(
63
    model=model,
64
    middleware=[CheckCredit()],
65
    store=store,
66
    context_schema=UserContext,
67
)

1
from dataclasses import dataclass
2

3
from langchain.agents import create_agent
4
from langchain.agents.middleware import dynamic_prompt, ModelRequest
5
from langgraph.store.memory import InMemoryStore
6

7
store = InMemoryStore()
8
namespace = ("users", "preferences")
9
store.put(namespace, "Ada Lovelace", {"chat_preferences": ["不喜欢啰嗦", "尽可能用最精简的文字解释清楚"]})
10
store.put(namespace, "Blackwell",    {"chat_preferences": ["喜欢长篇大论、引经据典"]})
11

12
@dataclass
13
class UserContext:
14
    username: str
15

16
@dynamic_prompt
17
def personalized_prompt(request: ModelRequest) -> str:
18
    username = request.runtime.context.username
19
    preferences = request.runtime.store.get(namespace, username).value["chat_preferences"]
20
    custom_prompt = "# 用户偏好\n" + "\n".join(preferences)
21
    print(f"{username} 自定义系统提示词\n{custom_prompt}")
22
    return custom_prompt            # ← 返回的就是这次的 system prompt
23

24
agent = create_agent(
25
    model=model,
26
    middleware=[personalized_prompt],
27
    store=store,
28
    context_schema=UserContext,
29
)

1
from typing import List
2

3
from pydantic import BaseModel, Field
4
from langchain.tools import tool
5
from langgraph.prebuilt.tool_node import ToolRuntime
6
from langchain_core.utils.function_calling import convert_to_openai_tool
7

8
class UserInfo(BaseModel):
9
    username: str = Field(description="用户名")
10
    age: int = Field(description="年龄")
11
    hobbies: List[str] = Field(description="兴趣爱好")
12

13
@tool(parse_docstring=True)
14
def save_user_info(user_info: UserInfo, runtime: ToolRuntime) -> str:
15
    """保存用户信息
16

17
    Args:
18
        user_info: 用户信息对象
19
    """
20
    return "ok"
21

22
print(convert_to_openai_tool(save_user_info))

1
{
2
  "type": "function",
3
  "function": {
4
    "name": "save_user_info",
5
    "description": "保存用户信息",
6
    "parameters": {
7
      "properties": {
8
        "user_info": {
9
          "properties": {
10
            "username": { "description": "用户名", "type": "string" },
11
            "age":      { "description": "年龄",   "type": "integer" },
12
            "hobbies":  { "description": "兴趣爱好", "items": {"type": "string"}, "type": "array" }
13
          },
14
          "required": ["username", "age", "hobbies"],
15
          "type": "object",
16
          "description": "用户信息对象"
17
        }
18
      },
19
      "required": ["user_info"],
20
      "type": "object"
21
    }
22
  }
23
}

1
from dataclasses import dataclass
2
from typing import Any, List
3

4
from pydantic import BaseModel, Field
5

6
from langchain.agents import create_agent
7
from langchain.tools import tool
8
from langchain_core.messages import HumanMessage, SystemMessage
9
from langgraph.checkpoint.memory import InMemorySaver
10
from langgraph.prebuilt.tool_node import ToolRuntime
11
from langgraph.store.memory import InMemoryStore
12

13
@dataclass
14
class UserContext:
15
    user_id: str
16

17
store = InMemoryStore()
18
namespace = ("users", "user_info")
19

20
class UserInfo(BaseModel):
21
    username: str = Field(description="用户名", default="unknown")
22
    age: int = Field(description="年龄", default=0)
23
    # 这里相当于告诉模型，每次创建新实例时，调用 list() 创建一个新的 list 传进去
24
    hobbies: List[str] = Field(description="兴趣爱好", default_factory=list)
25
    # 这种写法是用可变对象做默认值，会有多个对象共享一个列表对象的风险
26
    # hobbies: List[str] = Field(description="兴趣爱好", default=[])
27

28
@tool(parse_docstring=True)
29
def read_user_info(runtime: ToolRuntime[UserContext, Any]) -> UserInfo | str:
30
    """读取用户信息
31

32
    Returns:
33
        UserInfo | str: 如果找到用户信息，返回 UserInfo 对象；否则返回空字符串。
34
    """
35
    user_id = runtime.context.user_id            # ← 从 context 拿用户 id
36
    item = runtime.store.get(namespace, user_id)
37
    return UserInfo(**item.value) if item else ''
38

39
@tool(parse_docstring=True)
40
def write_user_info(user_info: UserInfo, runtime: ToolRuntime[UserContext, Any]) -> bool:
41
    """将用户信息写入长期记忆
42

43
    Args:
44
        user_info：用户信息对象
45

46
    Returns:
47
        bool: 如果成功写入则返回 True，否则返回 False
48
    """
49
    user_id = runtime.context.user_id
50
    runtime.store.put(namespace, user_id, user_info.model_dump())
51
    return True
52

53
agent = create_agent(
54
    model=model,
55
    tools=[read_user_info, write_user_info],
56
    store=store,
57
    checkpointer=InMemorySaver(),
58
    context_schema=UserContext,
59
)