agent-memory

为 Agent 添加记忆

注意： 此模板默认不包含记忆功能。使用此技能可添加记忆能力。如需预配置的记忆模板，请参阅：

• agent-langgraph-advanced - 支持短期和长期记忆，并包含长时间运行的后台任务

记忆类型

类型	使用场景	存储	标识符
短期	会话内的对话历史	AsyncCheckpointSaver	thread_id
长期	可跨会话保留的用户事实	AsyncDatabricksStore	user_id

前置条件

1. 将记忆依赖项添加到 pyproject.toml：

   dependencies = [
       "databricks-langchain[memory]",
   ]
   

   然后运行 uv sync

2. 配置 Lakebase - 参阅 lakebase-setup 技能，了解如何：

   • 创建/配置 Lakebase 实例
   • 初始化表（首次设置时的关键步骤）

---

快速设置摘要

添加记忆需要修改 4 个文件：

文件	要添加的内容
pyproject.toml	记忆依赖项
.env	Lakebase 环境变量（用于本地开发）
databricks.yml	Lakebase 数据库资源，以及配置块中的环境变量
agent_server/agent.py	记忆工具和 AsyncDatabricksStore

---

关键原则

在实现记忆功能之前，请先了解生产实现中的这些模式。

1. 工厂函数模式

记忆工具应由工厂函数返回，而不是定义为独立函数：

def memory_tools():
    @tool
    async def get_user_memory(query: str, config: RunnableConfig) -> str:
        ...
    @tool
    async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str:
        ...
    @tool
    async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str:
        ...
    return [get_user_memory, save_user_memory, delete_user_memory]

2. 提取用户 ID

从请求中提取 user_id，优先检查 custom_inputs。返回 None（而非默认值），由调用方决定如何处理：

def get_user_id(request: ResponsesAgentRequest) -> Optional[str]:
    custom_inputs = dict(request.custom_inputs or {})
    if "user_id" in custom_inputs:
        return custom_inputs["user_id"]
    if request.context and getattr(request.context, "user_id", None):
        return request.context.user_id
    return None

3. 分开处理错误

分别检查 user_id 和 store，并提供不同的错误消息：

user_id = config.get("configurable", {}).get("user_id")
if not user_id:
    return "Memory not available - no user_id provided."

store: Optional[BaseStore] = config.get("configurable", {}).get("store")
if not store:
    return "Memory not available - store not configured."

4. 保存时验证 JSON

存储前验证 JSON 输入——LLM 可能会传入无效的 JSON：

try:
    memory_data = json.loads(memory_data_json)
    if not isinstance(memory_data, dict):
        return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}"
    await store.aput(namespace, memory_key, memory_data)
except json.JSONDecodeError as e:
    return f"Failed to save memory: Invalid JSON - {e}"

5. 通过 RunnableConfig 传递 Store

通过配置传递 Store，而不是将其作为函数参数：

config = {"configurable": {"user_id": user_id, "store": store}}
# Tools access via: config.get("configurable", {}).get("store")

---

完整示例

此技能的 examples 文件夹中提供了完整实现：

# Copy to your project
cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/

参阅 examples/memory_tools.py，了解包含所有辅助函数、可用于生产环境的代码。

生产环境参考

预置模板中的实现：

文件	说明
agent-langgraph-advanced/agent_server/utils_memory.py	记忆工具工厂、辅助函数、错误处理
agent-langgraph-advanced/agent_server/agent.py	与 Agent 集成、Store 初始化

关键函数：

• memory_tools() - 返回 get/save/delete 工具的工厂函数
• get_user_id() - 从请求中提取 user_id
• init_lakebase_config() - 将自动扩缩 Lakebase 环境变量读入 LakebaseConfig
• get_lakebase_access_error_message() - 提供有用的错误消息

---

配置文件

步骤 1：databricks.yml（Lakebase 资源）

将 Lakebase 数据库资源添加到应用中：

resources:
  apps:
    agent_langgraph:
      name: "your-app-name"
      source_code_path: ./

      resources:
        # ... other resources (experiment, UC functions, etc.) ...

        # Autoscaling Lakebase instance for long-term memory
        - name: 'postgres'
          postgres:
            branch: "projects/<project-name>/branches/<branch-name>"
            database: "projects/<project-name>/branches/<branch-name>/databases/<database-id>"
            permission: 'CAN_CONNECT_AND_CREATE'

重要提示： name: 'postgres' 必须与 databricks.yml 的 config.env 块中的 value_from 引用相匹配。

步骤 2：databricks.yml 配置块（环境变量）

将 Lakebase 环境变量添加到 databricks.yml 中应用的 config.env：

      config:
        command: ["uv", "run", "start-app"]
        env:
          # ... other env vars ...

          # Autoscaling Lakebase endpoint (resolved from postgres resource)
          - name: LAKEBASE_AUTOSCALING_ENDPOINT
            value_from: "postgres"

          # Embedding configuration
          - name: EMBEDDING_ENDPOINT
            value: "databricks-gte-large-en"
          - name: EMBEDDING_DIMS
            value: "1024"

重要提示： LAKEBASE_AUTOSCALING_ENDPOINT 使用 value_from: "postgres"，以便在部署时从 postgres 资源解析。

步骤 3：.env（本地开发）

# Lakebase configuration for long-term memory
LAKEBASE_AUTOSCALING_ENDPOINT=<your-endpoint>
EMBEDDING_ENDPOINT=databricks-gte-large-en
EMBEDDING_DIMS=1024

---

集成示例

展示如何将记忆功能集成到流式函数中的最简示例：

from agent_server.utils_memory import memory_tools, get_user_id

@stream()
async def streaming(request: ResponsesAgentRequest):
    user_id = get_user_id(request)

    async with AsyncDatabricksStore(
        autoscaling_endpoint=LAKEBASE_AUTOSCALING_ENDPOINT,
        embedding_endpoint=EMBEDDING_ENDPOINT,
        embedding_dims=EMBEDDING_DIMS,
    ) as store:
        await store.setup()  # Creates tables if needed

        tools = await mcp_client.get_tools() + memory_tools()
        config = {"configurable": {"user_id": user_id, "store": store}}

        agent = create_react_agent(model=model, tools=tools)
        async for event in agent.astream(messages, config):
            yield event

---

初始化表并部署

初始化 Lakebase 表（仅首次需要）

部署前，请在本地初始化表：

uv run python -c "$(cat <<'EOF'
import asyncio
from databricks_langchain import AsyncDatabricksStore

async def setup():
    async with AsyncDatabricksStore(
        autoscaling_endpoint="<your-endpoint>",
        embedding_endpoint="databricks-gte-large-en",
        embedding_dims=1024,
    ) as store:
        await store.setup()
        print("Tables created!")

asyncio.run(setup())
EOF
)"

部署

初始化表后，部署 Agent。完整说明请参阅 deploy 技能。

---

短期记忆

如需保存会话内的对话历史，请使用 AsyncCheckpointSaver：

from databricks_langchain import AsyncCheckpointSaver

async with AsyncCheckpointSaver(autoscaling_endpoint=LAKEBASE_AUTOSCALING_ENDPOINT) as checkpointer:
    agent = create_react_agent(
        model=model,
        tools=tools,
        checkpointer=checkpointer,
    )

    config = {"configurable": {"thread_id": thread_id}}
    async for event in agent.astream(messages, config):
        yield event

如需查看完整实现，请参阅 agent-langgraph-advanced 模板。

---

测试记忆

本地测试

# Start the server
uv run start-app

# Save a memory
curl -X POST http://localhost:8000/invocations \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "Remember that I am on the shipping team"}],
      "custom_inputs": {"user_id": "[email protected]"}
  }'

# Recall the memory
curl -X POST http://localhost:8000/invocations \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "What team am I on?"}],
      "custom_inputs": {"user_id": "[email protected]"}
  }'

# Delete a memory
curl -X POST http://localhost:8000/invocations \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "Forget what team I am on"}],
      "custom_inputs": {"user_id": "[email protected]"}
  }'

测试已部署的应用

# Get OAuth token (PATs don't work for apps)
TOKEN=$(databricks auth token --host <workspace-url> | jq -r '.access_token')

# Test memory save
curl -X POST https://<app-url>/invocations \
  -H "Authorization: Bearer $TOKEN" \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}],
      "custom_inputs": {"user_id": "[email protected]"}
  }'

---

首次设置检查清单

• 已将 databricks-langchain[memory] 添加到 pyproject.toml
• 已运行 uv sync 安装依赖项
• 已创建或找到 Lakebase 实例
• 已将 Lakebase 环境变量添加到 .env（用于本地开发）
• 已将 postgres 资源添加到 databricks.yml
• 已将 LAKEBASE_AUTOSCALING_ENDPOINT 添加到 databricks.yml 的 config.env
• 已在本地初始化表，运行 await store.setup()
• 已通过运行 databricks bundle deploy && databricks bundle run 完成部署

---

故障排除

问题	原因	解决方案
"embedding_dims is required"	缺少参数	将 embedding_dims=1024 添加到 AsyncDatabricksStore
"relation 'store' does not exist"	尚未创建表	先在本地运行 await store.setup()
"Unable to resolve Lakebase endpoint 'None'"	缺少环境变量	检查 databricks.yml config.env 中的 LAKEBASE_AUTOSCALING_ENDPOINT
"permission denied for table store"	缺少授权	将 postgres 资源添加到 databricks.yml
"Memory not available - no user_id"	缺少 user_id	在请求中传入 custom_inputs.user_id
记忆未持久化	user_id 不同	在各个请求中使用一致的 user_id
部署后应用未更新	忘记运行 bundle	部署后运行 databricks bundle run agent_langgraph

---

预构建记忆模板

如需无需手动设置即可使用的完整配置实现：

模板	记忆类型	主要功能
agent-langgraph-advanced	短期 + 长期	AsyncCheckpointSaver、AsyncDatabricksStore、记忆工具

---

后续步骤

• 配置 Lakebase：请参阅 lakebase-setup 技能
• 本地测试：请参阅 run-locally 技能
• 部署：请参阅 deploy 技能