Memory Agent：task_mAIstro

0Overall architecture overview

task_mAIstro is a ToDo management assistant. The biggest difference from the previous passive memory Chatbot is:It decides when and what kind of memories to save, instead of automatically saving everything after each round of dialogue.

1Imports and dependencies

import uuid
from datetime import datetime

from pydantic           import BaseModel, Field
from trustcall           import create_extractor
from typing              import Literal, Optional, TypedDict

from langchain_core.runnables import RunnableConfig
from langchain_core.messages  import merge_message_runs
from langchain_core.messages  import SystemMessage, HumanMessage

from langchain_deepseek  import ChatDeepSeek

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph          import StateGraph, MessagesState, START, END
from langgraph.store.base     import BaseStore
from langgraph.store.memory   import InMemoryStore

import configuration

2Spy tool class: monitor the internal calls of Trustcall

Trustcall internally operates memory via utility calls (PatchDocUpdate, schema name is newly created).SpyIs a LangChain Runnable listener, passedwith_listeners(on_end=spy)Mount on the extractor and capture all internal tool call records after the call ends.

class Spy:
    def __init__(self):
        self.called_tools = []

    def __call__(self, run):
        q = [run]
        while q:
            r = q.pop()
            if r.child_runs:
                q.extend(r.child_runs)
            if r.run_type == "chat_model":
                self.called_tools.append(
                    r.outputs["generations"][0][0]["message"]["kwargs"]["tool_calls"]
                )

extract_tool_info: Convert Spy data into readable strings

This function parses the tool calls captured by Spy, distinguishes "update existing memory (PatchDoc)" and "new memory (schema name)", formats the result into a string, and finally asToolMessageThe content is returned to the task_mAIstro node.

def extract_tool_info(tool_calls, schema_name="Memory"):
    changes = []

    for call_group in tool_calls:
        for call in call_group:
            if call['name'] == 'PatchDoc':
                changes.append({
                    'type': 'update',
                    'doc_id': call['args']['json_doc_id'],
                    'planned_edits': call['args']['planned_edits'],
                    'value': call['args']['patches'][0]['value']
                })
            elif call['name'] == schema_name:
                changes.append({'type': 'new', 'value': call['args']})

    result_parts = []
    for change in changes:
        if change['type'] == 'update':
            result_parts.append(
                f"Document {change['doc_id']} updated:\n"
                f"Plan: {change['planned_edits']}\n"
                f"Added content: {change['value']}"
            )
        else:
            result_parts.append(
                f"New {schema_name} created:\n"
                f"Content: {change['value']}"
            )

    return "\n\n".join(result_parts)

3Three types of memory Schema: Profile, ToDo, UpdateMemory

User portrait: Profile

class Profile(BaseModel):
    """This is the profile of the user you are chatting with"""
    name: Optional[str]  = Field(description="The user's name",     default=None)
    location: Optional[str] = Field(description="The user's location", default=None)
    job: Optional[str]   = Field(description="The user's job",       default=None)
    connections: list[str] = Field(
        description="Personal connections: family, friends, coworkers",
        default_factory=list
    )
    interests: list[str] = Field(
        description="Interests that the user has",
        default_factory=list
    )

All fields are optional (Optionalor bringdefault_factory) to avoid errors due to incomplete information when extracting for the first time. Trustcall's JSON Patch only updates the changed fields and retains other known information.

To-do items: ToDo

class ToDo(BaseModel):
    task: str = Field(description="The task to be completed.")
    time_to_complete: Optional[int] = Field(
        description="Estimated time to complete (minutes)."
    )
    deadline: Optional[datetime] = Field(
        description="When the task needs to be completed by", default=None
    )
    solutions: list[str] = Field(
        description="Specific, actionable solutions (ideas / providers / options)",
        min_items=1, default_factory=list
    )
    status: Literal["not started", "in progress", "done", "archived"] = Field(
        description="Current status of the task", default="not started"
    )

solutionsField requires at least one specific course of action (min_items=1). Combined with the following Instructions (users can request "include specific local businesses"), this field will become more and more abundant.

Routing tool: UpdateMemory

class UpdateMemory(TypedDict):
    """Decision on what memory type to update"""
    update_type: Literal['user', 'todo', 'instructions']

Initialize model with profile_extractor

model = ChatDeepSeek(model="deepseek-v4-pro", temperature=0)

# Profile 提取器：全局复用，无 Spy，不需要 enable_inserts（每用户只有一条记录）
profile_extractor = create_extractor(
    model,
    tools=[Profile],
    tool_choice="Profile",
)

profile_extractorIt is a global singleton: because Profile updates do not require Spy monitoring, the same instance can be safely reused. In contrast, ToDo's extractor needs to bind a new one each timeSpy()instance, so created inside the node function.

4Three-segment system prompt word

The behavior of the entire agent is controlled by three prompt words, which are used in three different scenarios.

① MODEL_SYSTEM_MESSAGE — Main agent prompt word

MODEL_SYSTEM_MESSAGE = """You are a helpful chatbot.
You are designed to be a companion to a user, helping them keep track of their ToDo list.

You have a long term memory which keeps track of three things:
1. The user's profile (general information about them)
2. The user's ToDo list
3. General instructions for updating the ToDo list

Here is the current User Profile (may be empty):
<user_profile>
{user_profile}
</user_profile>

Here is the current ToDo List (may be empty):
<todo>
{todo}
</todo>

Here are the current user-specified preferences (may be empty):
<instructions>
{instructions}
</instructions>

Instructions for reasoning:
1. Reason carefully about the user's messages.
2. Decide whether any long-term memory should be updated:
   - Personal info → UpdateMemory(update_type='user')
   - Tasks mentioned → UpdateMemory(update_type='todo')
   - Preferences for ToDo updates → UpdateMemory(update_type='instructions')
3. Do not tell the user when you update the profile or instructions.
   Tell the user when you update the todo list.
4. Err on the side of updating the todo list.
5. Respond naturally after tool calls."""

② TRUSTCALL_INSTRUCTION — extractor instruction

TRUSTCALL_INSTRUCTION = """Reflect on following interaction.
Use the provided tools to retain any necessary memories about the user.
Use parallel tool calling to handle updates and insertions simultaneously.
System Time: {time}"""

This prompt word serves asSystemMessageInserted at the top of the Trustcall's message list to instruct the extractor to focus on the memory value of the conversation.{time}Inject the current time to provide a contextual baseline for the timestamped ToDo deadline.

③ CREATE_INSTRUCTIONS — Instruction update prompt word

CREATE_INSTRUCTIONS = """Reflect on the following interaction.
Based on this interaction, update your instructions for how to update ToDo list items.
Use any feedback from the user to update how they like to have items added, etc.

Your current instructions are:
<current_instructions>
{current_instructions}
</current_instructions>"""

only inupdate_instructionsused in nodes. Instead of using Trustcall, LLM directly generates new natural language instructions and overwrites the old version in the Store.

5Node ①: task_mAIstro — reading memory and routing decisions

def task_mAIstro(state: MessagesState, config: RunnableConfig, store: BaseStore):
    # 从 config 解析 user_id（由 configuration.Configuration 定义）
    configurable = configuration.Configuration.from_runnable_config(config)
    user_id = configurable.user_id

    # 1. 读取用户画像
    namespace = ("profile", user_id)
    memories = store.search(namespace)
    user_profile = memories[0].value if memories else None

    # 2. 读取所有 ToDo，拼接为多行字符串
    namespace = ("todo", user_id)
    memories = store.search(namespace)
    todo = "\n".join(f"{mem.value}" for mem in memories)

    # 3. 读取操作偏好指令
    namespace = ("instructions", user_id)
    memories = store.search(namespace)
    instructions = memories[0].value if memories else ""

    # 4. 构建系统提示，注入三类记忆
    system_msg = MODEL_SYSTEM_MESSAGE.format(
        user_profile=user_profile, todo=todo, instructions=instructions
    )

    # 5. 绑定 UpdateMemory 工具，串行调用（确保路由逻辑简单）
    response = model.bind_tools(
        [UpdateMemory], parallel_tool_calls=False
    ).invoke([SystemMessage(content=system_msg)] + state["messages"])

    return {"messages": [response]}

6Node ②: update_profile — Trustcall incremental update profile

def update_profile(state: MessagesState, config: RunnableConfig, store: BaseStore):
    configurable = configuration.Configuration.from_runnable_config(config)
    user_id = configurable.user_id
    namespace = ("profile", user_id)

    # 从 Store 读取已有 Profile，格式化为 Trustcall 三元组
    existing_items = store.search(namespace)
    tool_name = "Profile"
    existing_memories = (
        [(item.key, tool_name, item.value) for item in existing_items]
        if existing_items else None
    )

    # 注入当前时间，去掉最后一条消息（task_mAIstro 的工具调用消息）
    instr = TRUSTCALL_INSTRUCTION.format(time=datetime.now().isoformat())
    updated_messages = list(merge_message_runs(
        messages=[SystemMessage(content=instr)] + state["messages"][:-1]
    ))

    # Trustcall 提取：有 existing 时做 JSON Patch 增量更新
    result = profile_extractor.invoke({
        "messages": updated_messages,
        "existing": existing_memories
    })

    # 保存到 Store（json_doc_id 存在则更新同一条记录，否则新建）
    for r, rmeta in zip(result["responses"], result["response_metadata"]):
        store.put(namespace, rmeta.get("json_doc_id", str(uuid.uuid4())),
                  r.model_dump(mode="json"))

    # 必须返回 ToolMessage：响应 task_mAIstro 的 UpdateMemory 工具调用
    tool_calls = state['messages'][-1].tool_calls
    return {"messages": [{"role": "tool", "content": "updated profile",
                           "tool_call_id": tool_calls[0]['id']}]}

7Node ③: update_todos — ToDo update with Spy

The biggest difference between ToDo update and Profile update: Requiredenable_inserts=TrueSupports new tasks and requires Spy to capture change details and put them into ToolMessage.

def update_todos(state: MessagesState, config: RunnableConfig, store: BaseStore):
    configurable = configuration.Configuration.from_runnable_config(config)
    user_id = configurable.user_id
    namespace = ("todo", user_id)

    existing_items = store.search(namespace)
    tool_name = "ToDo"
    existing_memories = (
        [(item.key, tool_name, item.value) for item in existing_items]
        if existing_items else None
    )

    instr = TRUSTCALL_INSTRUCTION.format(time=datetime.now().isoformat())
    updated_messages = list(merge_message_runs(
        messages=[SystemMessage(content=instr)] + state["messages"][:-1]
    ))

    # 每次调用创建新 Spy（防止跨调用数据混用）
    spy = Spy()

    # ToDo 提取器：enable_inserts=True 允许新建任务，挂载 Spy 监听器
    todo_extractor = create_extractor(
        model, tools=[ToDo], tool_choice=tool_name, enable_inserts=True
    ).with_listeners(on_end=spy)

    result = todo_extractor.invoke({
        "messages": updated_messages, "existing": existing_memories
    })

    for r, rmeta in zip(result["responses"], result["response_metadata"]):
        store.put(namespace, rmeta.get("json_doc_id", str(uuid.uuid4())),
                  r.model_dump(mode="json"))

    # 把 Spy 捕获的变更摘要作为 ToolMessage 内容，让 task_mAIstro 能告知用户
    tool_calls = state['messages'][-1].tool_calls
    todo_update_msg = extract_tool_info(spy.called_tools, tool_name)
    return {"messages": [{"role": "tool", "content": todo_update_msg,
                           "tool_call_id": tool_calls[0]['id']}]}

8Node ④: update_instructions — plain text instruction update

Instructions are the user's natural language preference description of "how to manage ToDo". There is no fixed Schema, no Trustcall, and LLM can directly generate the text and then overwrite it.

def update_instructions(state: MessagesState, config: RunnableConfig, store: BaseStore):
    configurable = configuration.Configuration.from_runnable_config(config)
    user_id = configurable.user_id
    namespace = ("instructions", user_id)

    # 用 get（精确 key）而非 search（全量），因为 Instructions 始终只有一条
    existing_memory = store.get(namespace, "user_instructions")

    system_msg = CREATE_INSTRUCTIONS.format(
        current_instructions=existing_memory.value if existing_memory else None
    )

    # 直接调用 LLM 生成新指令文本（不用 Trustcall）
    new_memory = model.invoke(
        [SystemMessage(content=system_msg)] +
        state['messages'][:-1] +
        [HumanMessage(content="Please update the instructions based on the conversation")]
    )

    # 固定 key 覆写：Instructions 只有一条，无需追踪 json_doc_id
    store.put(namespace, "user_instructions", {"memory": new_memory.content})

    tool_calls = state['messages'][-1].tool_calls
    return {"messages": [{"role": "tool", "content": "updated instructions",
                           "tool_call_id": tool_calls[0]['id']}]}

9Conditional edge: route_message

def route_message(
    state: MessagesState, config: RunnableConfig, store: BaseStore
) -> Literal[END, "update_todos", "update_instructions", "update_profile"]:

    message = state['messages'][-1]          # task_mAIstro 的输出

    if len(message.tool_calls) == 0:           # 无工具调用 → 直接结束
        return END

    tool_call = message.tool_calls[0]          # 只取第一个（已强制串行）
    update_type = tool_call['args']['update_type']

    if   update_type == "user":         return "update_profile"
    elif update_type == "todo":         return "update_todos"
    elif update_type == "instructions": return "update_instructions"
    else:
        raise ValueError

10Graph assembly and compilation

# config_schema 将 Configuration（含 user_id）注册为可注入的运行时配置
builder = StateGraph(MessagesState, config_schema=configuration.Configuration)

# add_node(fn) 自动用函数名作为节点名
builder.add_node(task_mAIstro)
builder.add_node(update_todos)
builder.add_node(update_profile)
builder.add_node(update_instructions)

builder.add_edge(START, "task_mAIstro")
builder.add_conditional_edges("task_mAIstro", route_message)

# 更新节点执行完后，返回 task_mAIstro（携带 ToolMessage）
builder.add_edge("update_todos",        "task_mAIstro")
builder.add_edge("update_profile",     "task_mAIstro")
builder.add_edge("update_instructions", "task_mAIstro")

# 双存储模式：线程内短期记忆 + 跨线程长期记忆
across_thread_memory = InMemoryStore()
within_thread_memory = MemorySaver()
graph = builder.compile(
    checkpointer=within_thread_memory,
    store=across_thread_memory
)

11Full run demo

config = {"configurable": {"thread_id": "1", "user_id": "Lance"}}