From 43b257360efd49f13a440f1f18497450d587fa22 Mon Sep 17 00:00:00 2001 From: huangmeng Date: Thu, 23 Jul 2026 13:32:19 +0800 Subject: [PATCH] feat(memory): add TOS ContextBucket long-term memory backend Introduce a `tos_context` long-term memory backend backed by Volcengine TOS ContextBucket, wired through the runtime direct-connect path (LongTermMemory(backend="tos_context")), configuration, tests, and docs. Backend & wiring - Add TosContextBucketLTMBackend that maps VeADK's memory abstraction onto ContextBucket coordinates (app_name/index -> context bucket, user_id -> context set), with lazy bucket ensure and the standard save/search hooks. - Register the backend under backend="tos_context" and strip the injected session_id / app_name kwargs the SDK does not accept. - Add TOSContextBucketConfig so the backend can be configured from config.yaml. Credentials - Resolve credentials via _get_ak_sk_sts(): use explicit AK/SK (with the STS token from VOLCENGINE_SESSION_TOKEN) when both are present, otherwise fall back to the VeFaaS IAM credential file, mirroring VikingDBLTMBackend so keyless cloud (VeFaaS / AgentKit) deployments authenticate correctly. - Use the global VOLCENGINE_SESSION_TOKEN for the STS token instead of a standalone DATABASE_TOS_CONTEXT_SECURITY_TOKEN, matching TOS object storage and other components. Fail-closed dependency handling - ContextBucket APIs require tos>=2.9.4b1 (beta). Keep the core dependency at tos>=2.8.4 and probe TosClientV2 for the required methods at init, raising a clear RuntimeError with the upgrade command when the installed SDK is too old. Tests & docs - Add backend tests covering the explicit / IAM-fallback / VOLCENGINE_SESSION_TOKEN-default credential paths. - Add zh/en wiki pages, update the long-term memory index and sidebar, and note in the backend comparison table that temporary credentials need VOLCENGINE_SESSION_TOKEN (with an anchor to the configuration section). - Extend config.yaml.full with the tos_context section; the STS token uses the global VOLCENGINE_SESSION_TOKEN documented on the backend page rather than a dedicated placeholder. Note: the backend targets the runtime direct-connect path only; the CLI harness template and cloud harness env mapping are intentionally not wired for tos_context. --- config.yaml.full | 8 + .../framework/memory/long-term/index.en.mdx | 5 +- .../docs/framework/memory/long-term/index.mdx | 5 +- .../framework/memory/long-term/meta.en.json | 2 +- .../docs/framework/memory/long-term/meta.json | 2 +- .../memory/long-term/tos-context.en.mdx | 91 +++++ .../memory/long-term/tos-context.mdx | 91 +++++ pyproject.toml | 2 +- .../test_tos_context_bucket_backend.py | 310 +++++++++++++++++ veadk/configs/database_configs.py | 12 + veadk/memory/long_term_memory.py | 9 +- .../tos_context_bucket_backend.py | 324 ++++++++++++++++++ 12 files changed, 853 insertions(+), 8 deletions(-) create mode 100644 docs/content/docs/framework/memory/long-term/tos-context.en.mdx create mode 100644 docs/content/docs/framework/memory/long-term/tos-context.mdx create mode 100644 tests/memory/long_term_memory_backends/test_tos_context_bucket_backend.py create mode 100644 veadk/memory/long_term_memory_backends/tos_context_bucket_backend.py diff --git a/config.yaml.full b/config.yaml.full index b13d0b3d..013a5061 100644 --- a/config.yaml.full +++ b/config.yaml.full @@ -198,6 +198,14 @@ database: region: cn-beijing # default Volcengine TOS Vector region bucket: account_id: + # [optional] for long-term memory backend="tos_context" (TOS ContextBucket). Requires tos>=2.9.4b1. + # Note: distinct from `tos` (object storage) and `tos_vector` (vector store); this is the ContextBucket memory service. + tos_context: + account_id: # [required] account ID for ContextBucket control-plane + control_endpoint: # [required] ContextBucket controller (control-plane) endpoint + bucket_name: # ContextBucket name; falls back to LongTermMemory `index`/`app_name` + endpoint: tos-cn-beijing.volces.com # default Volcengine TOS data-plane endpoint + region: cn-beijing # default Volcengine TOS region # Dynamic config nacos: diff --git a/docs/content/docs/framework/memory/long-term/index.en.mdx b/docs/content/docs/framework/memory/long-term/index.en.mdx index 2495fb5c..984ff70d 100644 --- a/docs/content/docs/framework/memory/long-term/index.en.mdx +++ b/docs/content/docs/framework/memory/long-term/index.en.mdx @@ -25,7 +25,7 @@ ltm = LongTermMemory(backend="viking", app_name="ltm_demo") | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `backend` | `"local" \| "opensearch" \| "redis" \| "viking" \| "mem0"` or a backend instance | `"opensearch"` | Selects the backend, or pass a `BaseLongTermMemoryBackend` instance directly. `viking_mem` is deprecated and maps to `viking`. | +| `backend` | `"local" \| "opensearch" \| "redis" \| "viking" \| "mem0" \| "tos_context"` or a backend instance | `"opensearch"` | Selects the backend, or pass a `BaseLongTermMemoryBackend` instance directly. `viking_mem` is deprecated and maps to `viking`. | | `backend_config` | `dict` | `{}` | Config passed to the backend constructor; if it lacks `index`, it's filled from `index` or `app_name`. | | `top_k` | `int` | `5` | Number of most-similar chunks to retrieve. | | `index` | `str` | `""` | Index/collection name for storing memories. Falls back to `app_name`, then `default_app`. | @@ -33,7 +33,7 @@ ltm = LongTermMemory(backend="viking", app_name="ltm_demo") | `user_id` | `str` | `""` | **Deprecated**, kept only for backward compatibility. | -Vector backends (`local`, `opensearch`, `redis`) embed memories, which requires `pip install "veadk-python[extensions]"` and an embedding model (env prefix `MODEL_EMBEDDING_`, falling back to `MODEL_AGENT_API_KEY`). `viking` and `mem0` are managed services and need no local embedding. +Vector backends (`local`, `opensearch`, `redis`) embed memories, which requires `pip install "veadk-python[extensions]"` and an embedding model (env prefix `MODEL_EMBEDDING_`, falling back to `MODEL_AGENT_API_KEY`). `viking`, `mem0`, and `tos_context` are managed services and need no local embedding. ## Backend contract: `BaseLongTermMemoryBackend` @@ -70,6 +70,7 @@ Use `local` for debugging; prefer `viking` or `mem0` in production. | `mem0` | Mem0 memory (managed) | Mem0 API key | Recommended for production | [Mem0](/en/docs/framework/memory/long-term/mem0) | | `opensearch` | OpenSearch vector store | OpenSearch + embedding | Self-hosted vector search | [OpenSearch](/en/docs/framework/memory/long-term/opensearch) | | `redis` | Redis vector store | Redis (RediSearch) + embedding | Self-hosted vector search | [Redis](/en/docs/framework/memory/long-term/redis) | +| `tos_context` | TOS ContextBucket (managed) | Volcengine account (AK/SK; temporary credentials also need `VOLCENGINE_SESSION_TOKEN`) + `tos>=2.9.4b1` | Server-side inference/retrieval, per-user isolation | [TOS ContextBucket](/en/docs/framework/memory/long-term/tos-context#configuration) | ## Binding to an Agent diff --git a/docs/content/docs/framework/memory/long-term/index.mdx b/docs/content/docs/framework/memory/long-term/index.mdx index d78b23f2..8460e96d 100644 --- a/docs/content/docs/framework/memory/long-term/index.mdx +++ b/docs/content/docs/framework/memory/long-term/index.mdx @@ -25,7 +25,7 @@ ltm = LongTermMemory(backend="viking", app_name="ltm_demo") | 参数 | 类型 | 默认值 | 说明 | | :--- | :--- | :--- | :--- | -| `backend` | `"local" \| "opensearch" \| "redis" \| "viking" \| "mem0"` 或后端实例 | `"opensearch"` | 选择后端实现，也可直接传入一个 `BaseLongTermMemoryBackend` 实例。`viking_mem` 已废弃，自动转为 `viking`。 | +| `backend` | `"local" \| "opensearch" \| "redis" \| "viking" \| "mem0" \| "tos_context"` 或后端实例 | `"opensearch"` | 选择后端实现，也可直接传入一个 `BaseLongTermMemoryBackend` 实例。`viking_mem` 已废弃，自动转为 `viking`。 | | `backend_config` | `dict` | `{}` | 传给后端构造函数的配置；若不含 `index`，会用 `index` 或 `app_name` 补齐。 | | `top_k` | `int` | `5` | 检索时返回最相似的片段数量。 | | `index` | `str` | `""` | 存储记忆所用的索引/集合名。为空时回退到 `app_name`，再为空则用 `default_app`。 | @@ -33,7 +33,7 @@ ltm = LongTermMemory(backend="viking", app_name="ltm_demo") | `user_id` | `str` | `""` | **已废弃**，仅为向后兼容保留。 | -向量类后端（`local`、`opensearch`、`redis`）会对记忆做向量化（embedding），需要安装扩展依赖：`pip install "veadk-python[extensions]"`，并配置 embedding 模型（环境变量前缀 `MODEL_EMBEDDING_`，缺省时复用 `MODEL_AGENT_API_KEY`）。`viking` 与 `mem0` 是托管服务，无需本地 embedding。 +向量类后端（`local`、`opensearch`、`redis`）会对记忆做向量化（embedding），需要安装扩展依赖：`pip install "veadk-python[extensions]"`，并配置 embedding 模型（环境变量前缀 `MODEL_EMBEDDING_`，缺省时复用 `MODEL_AGENT_API_KEY`）。`viking`、`mem0` 与 `tos_context` 是托管服务，无需本地 embedding。 ## 后端契约：`BaseLongTermMemoryBackend` @@ -70,6 +70,7 @@ class BaseLongTermMemoryBackend(ABC, BaseModel): | `mem0` | Mem0 记忆库（托管） | Mem0 API Key | 生产推荐 | [Mem0](/cn/docs/framework/memory/long-term/mem0) | | `opensearch` | OpenSearch 向量库 | OpenSearch + embedding | 自建向量检索 | [OpenSearch](/cn/docs/framework/memory/long-term/opensearch) | | `redis` | Redis 向量库 | Redis(RediSearch) + embedding | 自建向量检索 | [Redis](/cn/docs/framework/memory/long-term/redis) | +| `tos_context` | TOS ContextBucket（托管） | 火山引擎账号（AK/SK，临时凭证另需 `VOLCENGINE_SESSION_TOKEN`）+ `tos>=2.9.4b1` | 服务端推理与检索、按用户隔离 | [TOS ContextBucket](/cn/docs/framework/memory/long-term/tos-context#配置) | ## 绑定到 Agent diff --git a/docs/content/docs/framework/memory/long-term/meta.en.json b/docs/content/docs/framework/memory/long-term/meta.en.json index 924656a6..6268ac5e 100644 --- a/docs/content/docs/framework/memory/long-term/meta.en.json +++ b/docs/content/docs/framework/memory/long-term/meta.en.json @@ -1,4 +1,4 @@ { "title": "Long-term memory", - "pages": ["index", "local", "vikingdb", "mem0", "opensearch", "redis"] + "pages": ["index", "local", "vikingdb", "mem0", "opensearch", "redis", "tos-context"] } diff --git a/docs/content/docs/framework/memory/long-term/meta.json b/docs/content/docs/framework/memory/long-term/meta.json index fdb152e7..6698dbae 100644 --- a/docs/content/docs/framework/memory/long-term/meta.json +++ b/docs/content/docs/framework/memory/long-term/meta.json @@ -1,4 +1,4 @@ { "title": "长期记忆", - "pages": ["index", "local", "vikingdb", "mem0", "opensearch", "redis"] + "pages": ["index", "local", "vikingdb", "mem0", "opensearch", "redis", "tos-context"] } diff --git a/docs/content/docs/framework/memory/long-term/tos-context.en.mdx b/docs/content/docs/framework/memory/long-term/tos-context.en.mdx new file mode 100644 index 00000000..c2ecbae9 --- /dev/null +++ b/docs/content/docs/framework/memory/long-term/tos-context.en.mdx @@ -0,0 +1,91 @@ +--- +title: "TOS ContextBucket (tos_context)" +--- + +The `tos_context` backend (`TosContextBucketLTMBackend`) uses the ContextBucket memory capability of Volcengine [TOS](https://www.volcengine.com/product/TOS) as long-term storage. It is a managed service: the server performs memory inference and retrieval, so no self-hosted vector store or local embedding is required. + +The backend maps `index` (or `app_name`) to a single **ContextBucket**, and maps the runtime `user_id` to a **ContextSet** under that bucket, so memories are naturally isolated per user. The corresponding ContextBucket and ContextSet are created automatically on first use (lazy ensure). + +## When to use + +- Production with persistence and managed operations; +- When you want the ContextBucket memory capability of Volcengine TOS (server-side inference and retrieval); +- When you want strong isolation by `user_id` (one ContextSet per user). + +## Requirements and installation + +The ContextBucket API is currently only available in a **pre-release (beta)** TOS SDK: `tos>=2.9.4b1`. VeADK's core dependency is the stable `tos>=2.8.4` (used by TOS object storage and Viking DB), so it does **not** install the beta for you. Before using `backend="tos_context"`, upgrade the TOS SDK explicitly: + +```bash +pip install --upgrade "tos>=2.9.4b1" --pre +``` + + +pip / uv **ignore pre-releases by default**, so you must pass `--pre` (uv uses `--prerelease=allow`). Otherwise `2.9.4b1` won't be installed and you'll still hit the errors below at runtime. + + +## Common runtime errors and how to handle them + +`tos_context` is fail-closed: when a dependency or config requirement is not met, it **raises at initialization** rather than degrading silently or writing bad data. Common errors: + +| Error | Trigger | How to fix | +| :--- | :--- | :--- | +| `ImportError` | The `tos` SDK is not installed. | Install it: `pip install --upgrade "tos>=2.9.4b1" --pre`. | +| `RuntimeError` | `tos` is installed but too old (e.g. `2.8.4`, or stable `2.9.2`); `TosClientV2` lacks the ContextBucket methods. | Upgrade to the pre-release: `pip install --upgrade "tos>=2.9.4b1" --pre`. The error message includes the installed version and this command. | +| `ValueError` | `account_id` or `control_endpoint` is not configured. | Set the env vars `DATABASE_TOS_CONTEXT_ACCOUNT_ID` and `DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT` (see the configuration table below). | + + +**Why a `RuntimeError` instead of an `ImportError` when `tos` is installed?** ContextBucket support was added as new methods on the existing `TosClientV2` class, so an older release still imports fine — it just lacks those methods. At initialization VeADK probes for the required methods (`hasattr`) to decide whether the installed SDK actually supports ContextBucket, turning a silent "imports but fails on call" problem into a clear, actionable error at startup. + + +## Configuration + +Authenticated via Volcengine AK/SK. Credentials are resolved in this order: explicit AK/SK **first** (env `VOLCENGINE_ACCESS_KEY` / `VOLCENGINE_SECRET_KEY`, with the STS token taken from `VOLCENGINE_SESSION_TOKEN`); **if both are not provided**, it falls back to the VeFaaS IAM credential file (`/var/run/secrets/iam/credential`) for the temporary AK/SK and STS token — the path used for keyless cloud deployments (VeFaaS / AgentKit). Other connection settings use the env prefix `DATABASE_TOS_CONTEXT_`; locally they can be placed under `database.tos_context.*` in `config.yaml` (flattened into the matching env vars). + +| Field | Env var | Default | Description | +| :--- | :--- | :--- | :--- | +| `volcengine_access_key` | `VOLCENGINE_ACCESS_KEY` | — | Volcengine Access Key (long-term or STS temporary AK). | +| `volcengine_secret_key` | `VOLCENGINE_SECRET_KEY` | — | Volcengine Secret Key (long-term or STS temporary SK). | +| `session_token` | `VOLCENGINE_SESSION_TOKEN` | — | STS security token for temporary credentials (optional; only needed with a temporary AK/SK). | +| `account_id` | `DATABASE_TOS_CONTEXT_ACCOUNT_ID` | — | **Required**. Account ID for the ContextBucket control-plane APIs. | +| `control_endpoint` | `DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT` | — | **Required**. ContextBucket controller (control-plane) endpoint. | +| `context_bucket_name` | `DATABASE_TOS_CONTEXT_BUCKET_NAME` | falls back to `index` | ContextBucket name. If unset, uses the `LongTermMemory` `index` (or `app_name`). | +| `endpoint` | `DATABASE_TOS_CONTEXT_ENDPOINT` | `tos-cn-beijing.volces.com` | TOS data-plane endpoint. | +| `region` | `DATABASE_TOS_CONTEXT_REGION` | `cn-beijing` | Region. | + + +`account_id` and `control_endpoint` are required; a missing value raises `ValueError` at initialization (fail-closed — no silent degradation). + + + +The STS token reuses the global `VOLCENGINE_SESSION_TOKEN` convention (consistent with TOS object storage and other components); the standalone `DATABASE_TOS_CONTEXT_SECURITY_TOKEN` is no longer used. It is not needed when using long-term AK/SK. If `VOLCENGINE_ACCESS_KEY` / `VOLCENGINE_SECRET_KEY` are not both provided, the VeFaaS IAM credential file is read instead; a missing file raises `FileNotFoundError`. + + +## Usage + +```python +from veadk import Agent +from veadk.memory.long_term_memory import LongTermMemory + +# index maps to the ContextBucket name and must satisfy the bucket naming rules (see Callout below) +ltm = LongTermMemory( + backend="tos_context", + index="demo-agent-memory", + top_k=5, +) +agent = Agent( + name="demo", + long_term_memory=ltm, + auto_save_session=True, # persist to long-term memory at session end +) +``` + +At init, if the ContextBucket doesn't exist, VeADK creates it; the first write/search for each `user_id` also creates the corresponding ContextSet if missing (enabling the `memory` scene). + + +`index` (the ContextBucket name) must follow TOS bucket naming rules: **length 3–63**, containing only **lowercase letters, digits, and hyphens (`-`)**, and not starting or ending with a hyphen. Use `DATABASE_TOS_CONTEXT_BUCKET_NAME` to override the bucket name derived from `app_name`/`index`. + + + +Memories are isolated by `user_id`: different `user_id`s write to different ContextSets, and retrieval only matches memories of the same `user_id`. Use a consistent `user_id` for cross-session recall. + diff --git a/docs/content/docs/framework/memory/long-term/tos-context.mdx b/docs/content/docs/framework/memory/long-term/tos-context.mdx new file mode 100644 index 00000000..06efbfdb --- /dev/null +++ b/docs/content/docs/framework/memory/long-term/tos-context.mdx @@ -0,0 +1,91 @@ +--- +title: "TOS ContextBucket（tos_context）" +--- + +`tos_context` 后端（`TosContextBucketLTMBackend`）使用火山引擎 [TOS](https://www.volcengine.com/product/TOS) 的 ContextBucket 记忆能力作为长期记忆存储。它是托管服务，由服务端完成记忆推理（inference）与检索，无需自建向量库或本地 embedding。 + +后端把 `index`（或 `app_name`）映射为一个 **ContextBucket**，把运行时的 `user_id` 映射为该桶下的一个 **ContextSet**，从而按用户天然隔离记忆。首次使用时会自动创建对应的 ContextBucket 与 ContextSet（Lazy Ensure）。 + +## 何时使用 + +- 生产环境、需要持久化与托管运维； +- 希望使用火山引擎 TOS 的 ContextBucket 记忆能力（服务端推理与检索）； +- 希望按 `user_id` 做强隔离（每个用户对应一个 ContextSet）。 + +## 依赖要求与安装 + +ContextBucket API 目前只在**预发布(beta)版本**的 TOS SDK 中提供:`tos>=2.9.4b1`。VeADK 的核心依赖为稳定版 `tos>=2.8.4`(供 TOS 对象存储、Viking DB 使用),**不会**自动为你安装 beta 版。因此在使用 `backend="tos_context"` 前，需要显式升级 TOS SDK： + +```bash +pip install --upgrade "tos>=2.9.4b1" --pre +``` + + +pip / uv 默认会**忽略预发布版本**，必须显式加 `--pre`（uv 使用 `--prerelease=allow`），否则不会安装上 `2.9.4b1`，运行时仍会命中下方的报错。 + + +## 常见运行时报错与处理 + +`tos_context` 采用 fail-closed 设计：依赖或配置不满足时会**在初始化阶段直接抛错**，不会静默降级或写入错误数据。常见错误及处理如下： + +| 报错类型 | 触发条件 | 处理方式 | +| :--- | :--- | :--- | +| `ImportError` | 环境中未安装 `tos` SDK。 | 安装 SDK：`pip install --upgrade "tos>=2.9.4b1" --pre`。 | +| `RuntimeError` | 已安装 `tos`，但版本过旧（如 `2.8.4`、`2.9.2` 稳定版），`TosClientV2` 缺少 ContextBucket 相关方法。 | 升级到预发布版：`pip install --upgrade "tos>=2.9.4b1" --pre`。报错信息会附带当前已安装版本号与该命令。 | +| `ValueError` | `account_id` 或 `control_endpoint` 未配置。 | 补齐环境变量 `DATABASE_TOS_CONTEXT_ACCOUNT_ID`、`DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT`（见下方配置表）。 | + + +**为什么装了 `tos` 还会报 `RuntimeError` 而不是 `ImportError`？** 因为 ContextBucket 能力是在既有的 `TosClientV2` 类上**新增方法**实现的，旧版本仍可 `import tos` 成功、只是缺少这些方法。VeADK 在初始化时会通过方法探测（`hasattr`）判断当前 SDK 是否真正支持 ContextBucket，从而把"能 import 但调用时静默失败"的隐性问题，转化为初始化阶段清晰、可操作的报错。 + + +## 配置 + +通过火山引擎 AK/SK 鉴权。凭据解析顺序为：**优先**使用显式的 AK/SK（环境变量 `VOLCENGINE_ACCESS_KEY` / `VOLCENGINE_SECRET_KEY`，此时 STS Token 取自 `VOLCENGINE_SESSION_TOKEN`）；**若二者未同时提供**，则回退到 VeFaaS IAM 凭证文件（`/var/run/secrets/iam/credential`），从中读取临时 AK/SK 与 STS Token——这是云上部署（VeFaaS / AgentKit）免密运行的路径。其余连接项使用环境变量前缀 `DATABASE_TOS_CONTEXT_`，本地可写在 `config.yaml` 的 `database.tos_context.*` 下（会被拍平为对应环境变量）。 + +| 配置项 | 环境变量 | 默认值 | 说明 | +| :--- | :--- | :--- | :--- | +| `volcengine_access_key` | `VOLCENGINE_ACCESS_KEY` | — | 火山引擎 Access Key（可为长期或 STS 临时 AK）。 | +| `volcengine_secret_key` | `VOLCENGINE_SECRET_KEY` | — | 火山引擎 Secret Key（可为长期或 STS 临时 SK）。 | +| `session_token` | `VOLCENGINE_SESSION_TOKEN` | — | STS 临时凭证的 Security Token（可选，仅在使用临时 AK/SK 时需要）。 | +| `account_id` | `DATABASE_TOS_CONTEXT_ACCOUNT_ID` | — | **必填**。账号 ID，用于 ContextBucket 控制面接口。 | +| `control_endpoint` | `DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT` | — | **必填**。ContextBucket 控制面（Controller）域名。 | +| `context_bucket_name` | `DATABASE_TOS_CONTEXT_BUCKET_NAME` | 回退到 `index` | ContextBucket 名。未设置时使用 `LongTermMemory` 的 `index`（或 `app_name`）。 | +| `endpoint` | `DATABASE_TOS_CONTEXT_ENDPOINT` | `tos-cn-beijing.volces.com` | TOS 数据面域名。 | +| `region` | `DATABASE_TOS_CONTEXT_REGION` | `cn-beijing` | 区域。 | + + +`account_id` 与 `control_endpoint` 为必填项，缺失会在初始化时直接抛出 `ValueError`（fail-closed，不会静默降级）。 + + + +STS Token 复用全局约定 `VOLCENGINE_SESSION_TOKEN`（与 TOS 对象存储等其他组件一致），不再使用独立的 `DATABASE_TOS_CONTEXT_SECURITY_TOKEN`。使用长期 AK/SK 时无需设置。若 `VOLCENGINE_ACCESS_KEY` / `VOLCENGINE_SECRET_KEY` 未同时提供，会尝试读取 VeFaaS IAM 凭证文件；文件不存在时抛出 `FileNotFoundError`。 + + +## 用法 + +```python +from veadk import Agent +from veadk.memory.long_term_memory import LongTermMemory + +# index 会映射为 ContextBucket 名，必须满足桶命名规则（见下方 Callout） +ltm = LongTermMemory( + backend="tos_context", + index="demo-agent-memory", + top_k=5, +) +agent = Agent( + name="demo", + long_term_memory=ltm, + auto_save_session=True, # 会话结束自动写入长期记忆 +) +``` + +初始化时若 ContextBucket 不存在，VeADK 会自动创建；每个 `user_id` 首次写入/检索时，若对应的 ContextSet 不存在也会自动创建（并启用 `memory` 场景）。 + + +`index`（即 ContextBucket 名）需满足 TOS 桶命名规则：**长度 3–63**，仅包含**小写字母、数字和连字符（`-`）**，且不能以连字符开头或结尾。可用 `DATABASE_TOS_CONTEXT_BUCKET_NAME` 覆盖由 `app_name`/`index` 推导出的桶名。 + + + +记忆按 `user_id` 隔离：不同 `user_id` 写入的是不同的 ContextSet，检索时也只会命中同一 `user_id` 的记忆。跨会话召回时请确保使用一致的 `user_id`。 + diff --git a/pyproject.toml b/pyproject.toml index 5053e6c8..f41bb9b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ dependencies = [ "websockets>=15,<16", # Stream temporary AgentKit Sandbox conversations "openviking-sdk>=0.1.3", "python-frontmatter==1.1.0", - "tos>=2.8.4", # For TOS storage and Viking DB + "tos>=2.8.4", # For TOS storage and Viking DB ] [project.scripts] diff --git a/tests/memory/long_term_memory_backends/test_tos_context_bucket_backend.py b/tests/memory/long_term_memory_backends/test_tos_context_bucket_backend.py new file mode 100644 index 00000000..6fafa003 --- /dev/null +++ b/tests/memory/long_term_memory_backends/test_tos_context_bucket_backend.py @@ -0,0 +1,310 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from types import SimpleNamespace + +import pytest + +from veadk.configs.database_configs import TOSContextBucketConfig +from veadk.memory.long_term_memory import LongTermMemory +from veadk.memory.long_term_memory_backends import tos_context_bucket_backend +from veadk.memory.long_term_memory_backends.tos_context_bucket_backend import ( + TosContextBucketLTMBackend, +) + + +class _ServerError(Exception): + def __init__(self, status_code: int): + self.status_code = status_code + + +class _Client: + def __init__(self): + self.buckets: set[str] = set() + self.context_sets: dict[str, SimpleNamespace] = {} + self.created_memories: list[dict] = [] + self.search_output = SimpleNamespace(results=[]) + self.get_context_set_calls = 0 + + def get_context_bucket(self, *, context_bucket_name, **_kwargs): + if context_bucket_name not in self.buckets: + raise _ServerError(404) + + def create_context_bucket(self, *, context_bucket_name, **_kwargs): + self.buckets.add(context_bucket_name) + + def get_context_set(self, *, context_set_name, **_kwargs): + self.get_context_set_calls += 1 + if context_set_name not in self.context_sets: + raise _ServerError(404) + return self.context_sets[context_set_name] + + def create_context_set(self, *, context_set_name, enable, scenes, **_kwargs): + self.context_sets[context_set_name] = SimpleNamespace( + enable=enable, scenes=scenes + ) + + def create_context_bucket_memory(self, **kwargs): + self.created_memories.append(kwargs) + + def search_context_bucket_memory(self, **_kwargs): + return self.search_output + + +@pytest.fixture +def client(monkeypatch): + fake_client = _Client() + monkeypatch.setattr(tos_context_bucket_backend, "TosServerError", _ServerError) + # Give every test a clean credential baseline. Other test modules (e.g. + # tests/test_cloud.py) set VOLCENGINE_ACCESS_KEY / VOLCENGINE_SECRET_KEY at + # import time via os.environ (never restored), which leaks into this process + # under pytest-xdist and would otherwise flip the IAM fallback tests into the + # explicit-credentials branch. Tests that need credentials set them + # explicitly, so clearing here keeps this file self-contained. + for _var in ( + "VOLCENGINE_ACCESS_KEY", + "VOLCENGINE_SECRET_KEY", + "VOLCENGINE_SESSION_TOKEN", + ): + monkeypatch.delenv(_var, raising=False) + + def _make_client(*args, **kwargs): + fake_client.init_args = args + fake_client.init_kwargs = kwargs + return fake_client + + monkeypatch.setattr(tos_context_bucket_backend, "TosClientV2", _make_client) + # These tests exercise the backend logic against a fake client, so they must + # not depend on the real `tos` SDK version. Bypass the ContextBucket support + # probe here; the probe itself is covered by + # `test_asserts_when_context_bucket_unsupported`. + monkeypatch.setattr( + tos_context_bucket_backend, "_assert_context_bucket_supported", lambda: None + ) + # Default: pretend the VeFaaS IAM file supplies credentials so tests that + # don't set explicit AK/SK don't touch the filesystem. + monkeypatch.setattr( + tos_context_bucket_backend, + "get_credential_from_vefaas_iam", + lambda: SimpleNamespace( + access_key_id="iam-ak", + secret_access_key="iam-sk", + session_token="iam-sts", + ), + ) + return fake_client + + +def _backend(**kwargs) -> TosContextBucketLTMBackend: + defaults = { + "index": "agent-memory", + "account_id": "2100000000", + "tos_context_config": TOSContextBucketConfig( + endpoint="tos.example.com", + region="cn-beijing", + control_endpoint="tosapi-controller.example.com", + ), + } + defaults.update(kwargs) + return TosContextBucketLTMBackend(**defaults) + + +def test_initialization_creates_missing_context_bucket(client): + backend = _backend() + + assert backend.context_bucket_name == "agent-memory" + assert client.buckets == {"agent-memory"} + + +def test_backend_requires_controller_configuration(client): + with pytest.raises(ValueError, match="CONTROL_ENDPOINT"): + TosContextBucketLTMBackend(index="agent-memory", account_id="2100000000") + + +def test_backend_rejects_invalid_context_bucket_name(client): + with pytest.raises(ValueError, match="lowercase letters"): + _backend(index="Agent_Memory") + + +def test_save_lazily_creates_and_caches_context_set(client): + backend = _backend() + event = json.dumps({"role": "user", "parts": [{"text": "I like Go."}]}) + + assert backend.save_memory("user-1", [event]) is True + assert backend.save_memory("user-1", [event]) is True + + assert client.context_sets["user-1"].scenes == ["memory"] + assert client.get_context_set_calls == 1 + assert [memory["content"] for memory in client.created_memories] == [ + "I like Go.", + "I like Go.", + ] + assert all(memory["infer"] is True for memory in client.created_memories) + + +def test_save_returns_false_when_tos_request_fails(client): + backend = _backend() + + def fail(**_kwargs): + raise RuntimeError("service unavailable") + + client.create_context_bucket_memory = fail + + assert backend.save_memory("user-1", ["plain text"]) is False + + +def test_save_rejects_existing_context_set_without_memory_scene(client): + backend = _backend() + client.context_sets["user-1"] = SimpleNamespace(enable=True, scenes=[]) + + assert backend.save_memory("user-1", ["plain text"]) is False + assert client.created_memories == [] + + +def test_search_returns_extracted_memories_and_is_user_scoped(client): + backend = _backend() + client.search_output = SimpleNamespace( + results=[ + SimpleNamespace(memory="User likes basketball."), + SimpleNamespace(memory=None), + ] + ) + captured = {} + + def search(**kwargs): + captured.update(kwargs) + return client.search_output + + client.search_context_bucket_memory = search + + assert backend.search_memory("user-2", "hobbies", 3) == ["User likes basketball."] + assert captured["context_set_name"] == "user-2" + assert captured["limit"] == 3 + + +def test_search_returns_empty_list_when_tos_request_fails(client): + backend = _backend() + + def fail(**_kwargs): + raise RuntimeError("service unavailable") + + client.search_context_bucket_memory = fail + + assert backend.search_memory("user-1", "hobbies", 3) == [] + + +def test_long_term_memory_registers_tos_backend(client, monkeypatch): + monkeypatch.setenv("DATABASE_TOS_CONTEXT_ACCOUNT_ID", "2100000000") + monkeypatch.setenv( + "DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT", "tosapi-controller.example.com" + ) + + memory = LongTermMemory(backend="tos_context", app_name="agent-memory") + + assert isinstance(memory._backend, TosContextBucketLTMBackend) + assert memory._backend.context_bucket_name == "agent-memory" + + +@pytest.mark.asyncio +async def test_long_term_memory_forwards_tos_save_kwargs(client, monkeypatch): + monkeypatch.setenv("DATABASE_TOS_CONTEXT_ACCOUNT_ID", "2100000000") + monkeypatch.setenv( + "DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT", "tosapi-controller.example.com" + ) + memory = LongTermMemory(backend="tos_context", app_name="agent-memory") + monkeypatch.setattr( + memory, "_filter_and_convert_events", lambda _events, **_kwargs: ["text"] + ) + + await memory.add_session_to_memory( + SimpleNamespace(id="session-1", user_id="user-1", events=[]), + infer=False, + agent_id="agent-1", + ) + + # The backend strips the framework-injected `session_id` / `app_name` + # kwargs that the TOS SDK does not accept, so they must not appear here. + assert client.created_memories == [ + { + "account_id": "2100000000", + "context_bucket_name": "agent-memory", + "context_set_name": "user-1", + "content": "text", + "infer": False, + "agent_id": "agent-1", + } + ] + + +def test_uses_explicit_credentials_and_session_token(client): + _backend( + volcengine_access_key="env-ak", + volcengine_secret_key="env-sk", + session_token="env-sts", + ) + + assert client.init_args[:2] == ("env-ak", "env-sk") + assert client.init_kwargs["security_token"] == "env-sts" + + +def test_falls_back_to_vefaas_iam_when_credentials_missing(client): + # `client` fixture stubs get_credential_from_vefaas_iam -> iam-ak/sk/sts. + _backend() + + assert client.init_args[:2] == ("iam-ak", "iam-sk") + assert client.init_kwargs["security_token"] == "iam-sts" + + +def test_session_token_defaults_to_volcengine_env(client, monkeypatch): + monkeypatch.setenv("VOLCENGINE_SESSION_TOKEN", "global-sts") + _backend( + volcengine_access_key="env-ak", + volcengine_secret_key="env-sk", + ) + + assert client.init_kwargs["security_token"] == "global-sts" + + +def test_asserts_when_context_bucket_unsupported(monkeypatch): + # Simulate an older `tos` SDK: `TosClientV2` exists but lacks the + # ContextBucket methods. The probe must fail closed with an upgrade hint. + # Note: this test intentionally does NOT use the `client` fixture, which + # bypasses the probe. + class _OldTosClientV2: + pass + + monkeypatch.setattr(tos_context_bucket_backend, "TosClientV2", _OldTosClientV2) + + with pytest.raises(RuntimeError) as excinfo: + tos_context_bucket_backend._assert_context_bucket_supported() + + message = str(excinfo.value) + assert "create_context_bucket_memory" in message + assert "search_context_bucket_memory" in message + assert "tos>=2.9.4b1" in message + + +def test_probe_passes_when_context_bucket_methods_present(monkeypatch): + class _NewTosClientV2: + def create_context_bucket_memory(self): # pragma: no cover - stub + ... + + def search_context_bucket_memory(self): # pragma: no cover - stub + ... + + monkeypatch.setattr(tos_context_bucket_backend, "TosClientV2", _NewTosClientV2) + + # Should not raise. + tos_context_bucket_backend._assert_context_bucket_supported() diff --git a/veadk/configs/database_configs.py b/veadk/configs/database_configs.py index 94111949..d914cc86 100644 --- a/veadk/configs/database_configs.py +++ b/veadk/configs/database_configs.py @@ -224,6 +224,18 @@ class TOSVectorConfig(BaseSettings): user_agent_customized_key_values: dict[str, str] | None = None +class TOSContextBucketConfig(BaseSettings): + """Configuration for the TOS ContextBucket controller APIs.""" + + model_config = SettingsConfigDict(env_prefix="DATABASE_TOS_CONTEXT_") + + endpoint: str = "tos-cn-beijing.volces.com" + + region: str = "cn-beijing" + + control_endpoint: str | None = None + + class MSENacosConfig(BaseSettings): model_config = SettingsConfigDict(env_prefix="NACOS_") diff --git a/veadk/memory/long_term_memory.py b/veadk/memory/long_term_memory.py index cbb7e313..3b885758 100644 --- a/veadk/memory/long_term_memory.py +++ b/veadk/memory/long_term_memory.py @@ -78,6 +78,12 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]: ) return OpenVikingLTMBackend + case "tos_context": + from veadk.memory.long_term_memory_backends.tos_context_bucket_backend import ( + TosContextBucketLTMBackend, + ) + + return TosContextBucketLTMBackend case _: raise ValueError(f"Unsupported long term memory backend: {backend}") except ImportError as e: @@ -97,7 +103,7 @@ class LongTermMemory(BaseMemoryService, BaseModel): It supports configuration of the backend service and retrieval behavior. Attributes: - backend (Union[Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0", "openviking"], BaseLongTermMemoryBackend]): + backend (Union[Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0", "openviking", "tos_context"], BaseLongTermMemoryBackend]): The type or instance of the long-term memory backend. Defaults to "opensearch". backend_config (dict): @@ -128,6 +134,7 @@ class LongTermMemory(BaseMemoryService, BaseModel): "viking_mem", "mem0", "openviking", + "tos_context", ], BaseLongTermMemoryBackend, ] = "opensearch" diff --git a/veadk/memory/long_term_memory_backends/tos_context_bucket_backend.py b/veadk/memory/long_term_memory_backends/tos_context_bucket_backend.py new file mode 100644 index 00000000..3d6337ad --- /dev/null +++ b/veadk/memory/long_term_memory_backends/tos_context_bucket_backend.py @@ -0,0 +1,324 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TOS ContextBucket implementation of VeADK long-term memory.""" + +import json +import os +import threading +from typing import Any + +from pydantic import Field, PrivateAttr + +import veadk.config # noqa: F401 # Load .env and config.yaml before settings. +from veadk.auth.veauth.utils import get_credential_from_vefaas_iam +from veadk.configs.database_configs import TOSContextBucketConfig +from veadk.memory.long_term_memory_backends.base_backend import ( + BaseLongTermMemoryBackend, +) +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + +try: + from tos import TosClientV2 + from tos.exceptions import TosServerError +except ImportError as exc: + raise ImportError( + "TOS ContextBucket long-term memory requires the pre-release " + "`tos>=2.9.4b1`, which is not installed. Install it with:\n" + ' pip install --upgrade "tos>=2.9.4b1" --pre' + ) from exc + + +# ContextBucket support was added to `TosClientV2` as new methods on an existing +# class, so importing `tos` succeeds even on older releases (e.g. 2.8.4) that lack +# the feature. Probe for the required methods explicitly to fail closed with a +# clear message instead of surfacing a swallowed AttributeError at call time. +_REQUIRED_CONTEXT_BUCKET_METHODS = ( + "create_context_bucket_memory", + "search_context_bucket_memory", +) + + +def _resolve_tos_version() -> str: + try: + import tos + + version = getattr(tos, "__version__", None) + # `tos.__version__` may itself be a module exposing `__version__`. + return getattr(version, "__version__", None) or str(version) or "unknown" + except Exception: # pragma: no cover - best-effort version reporting + return "unknown" + + +def _assert_context_bucket_supported() -> None: + missing = [ + method + for method in _REQUIRED_CONTEXT_BUCKET_METHODS + if not hasattr(TosClientV2, method) + ] + if missing: + raise RuntimeError( + f"The installed `tos` SDK ({_resolve_tos_version()}) does not support " + "TOS ContextBucket long-term memory " + f"(missing methods: {', '.join(missing)}). " + "This feature requires the pre-release `tos>=2.9.4b1`. Upgrade with:\n" + ' pip install --upgrade "tos>=2.9.4b1" --pre' + ) + + +class TosContextBucketLTMBackend(BaseLongTermMemoryBackend): + """Persist VeADK long-term memory through TOS ContextBucketMemory. + + The backend maps its index to one ContextBucket and each runtime user ID to + a ContextSet. The service then performs memory inference and retrieval. + """ + + volcengine_access_key: str | None = Field( + default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY") + ) + volcengine_secret_key: str | None = Field( + default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY") + ) + session_token: str = Field( + default_factory=lambda: os.getenv("VOLCENGINE_SESSION_TOKEN", "") + ) + context_bucket_name: str | None = Field( + default_factory=lambda: os.getenv("DATABASE_TOS_CONTEXT_BUCKET_NAME") + ) + account_id: str | None = Field( + default_factory=lambda: os.getenv("DATABASE_TOS_CONTEXT_ACCOUNT_ID") + ) + tos_context_config: TOSContextBucketConfig = Field( + default_factory=TOSContextBucketConfig + ) + + _client: Any = PrivateAttr() + _ensured_sets: set[str] = PrivateAttr(default_factory=set) + _ensure_lock: threading.Lock = PrivateAttr(default_factory=threading.Lock) + + def model_post_init(self, __context: Any) -> None: + _assert_context_bucket_supported() + + if not self.account_id: + raise ValueError( + "DATABASE_TOS_CONTEXT_ACCOUNT_ID is required for the TOS " + "ContextBucket long-term memory backend." + ) + if not self.tos_context_config.control_endpoint: + raise ValueError( + "DATABASE_TOS_CONTEXT_CONTROL_ENDPOINT is required for the TOS " + "ContextBucket long-term memory backend." + ) + + self.context_bucket_name = self.context_bucket_name or self.index + self.precheck_index_naming() + + ak, sk, sts_token = self._get_ak_sk_sts() + self._client = TosClientV2( + ak, + sk, + self.tos_context_config.endpoint, + self.tos_context_config.region, + security_token=sts_token or None, + control_endpoint=self.tos_context_config.control_endpoint, + ) + self._ensure_context_bucket() + + def _get_ak_sk_sts(self) -> tuple[str, str, str]: + """Resolve (access_key, secret_key, sts_token) for the TOS client. + + Mirrors the credential strategy of the VikingDB backend: use the + explicitly provided AK/SK (with the STS token from + ``VOLCENGINE_SESSION_TOKEN``) when both are present, otherwise fall back + to the VeFaaS IAM credential file so cloud deployments authenticate + with the instance role without static secrets. + """ + if self.volcengine_access_key and self.volcengine_secret_key: + logger.debug( + "Using Volcengine credentials from environment for " + "TosContextBucketLTMBackend." + ) + return ( + self.volcengine_access_key, + self.volcengine_secret_key, + self.session_token, + ) + + cred = get_credential_from_vefaas_iam() + logger.debug( + "Using Volcengine credentials from VeFaaS IAM file for " + "TosContextBucketLTMBackend." + ) + return cred.access_key_id, cred.secret_access_key, cred.session_token + + def precheck_index_naming(self) -> None: + name = self.context_bucket_name + if not isinstance(name, str) or not 3 <= len(name) <= 63: + raise ValueError( + "Invalid TOS ContextBucket name: it must be 3-63 characters long. " + "Set DATABASE_TOS_CONTEXT_BUCKET_NAME to override the app name." + ) + if ( + name.startswith("-") + or name.endswith("-") + or any(char not in "abcdefghijklmnopqrstuvwxyz0123456789-" for char in name) + ): + raise ValueError( + "Invalid TOS ContextBucket name: use lowercase letters, digits, " + "and hyphens, without a leading or trailing hyphen. Set " + "DATABASE_TOS_CONTEXT_BUCKET_NAME to override the app name." + ) + + @staticmethod + def _status_code(error: Exception) -> int | None: + return getattr(error, "status_code", None) + + def _ensure_context_bucket(self) -> None: + try: + self._client.get_context_bucket( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + ) + except TosServerError as error: + if self._status_code(error) != 404: + raise + try: + self._client.create_context_bucket( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + ) + except TosServerError as create_error: + if self._status_code(create_error) != 409: + raise + self._client.get_context_bucket( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + ) + + def _context_set_name(self, user_id: str) -> str: + """Return the ContextSet name for a user. + + Keeping this mapping in one place permits a future deterministic encoding + if the service imposes stricter user-ID naming constraints. + """ + return user_id + + @staticmethod + def _validate_memory_context_set(context_set: Any, context_set_name: str) -> None: + if not getattr(context_set, "enable", False) or "memory" not in ( + getattr(context_set, "scenes", None) or [] + ): + raise ValueError( + f"ContextSet {context_set_name!r} is not enabled for memory." + ) + + def _ensure_context_set(self, context_set_name: str) -> None: + if context_set_name in self._ensured_sets: + return + + with self._ensure_lock: + if context_set_name in self._ensured_sets: + return + try: + context_set = self._client.get_context_set( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + context_set_name=context_set_name, + ) + self._validate_memory_context_set(context_set, context_set_name) + except TosServerError as error: + if self._status_code(error) != 404: + raise + try: + self._client.create_context_set( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + context_set_name=context_set_name, + enable=True, + scenes=["memory"], + ) + except TosServerError as create_error: + if self._status_code(create_error) != 409: + raise + context_set = self._client.get_context_set( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + context_set_name=context_set_name, + ) + self._validate_memory_context_set(context_set, context_set_name) + self._ensured_sets.add(context_set_name) + + @staticmethod + def _to_content(event_string: str) -> str: + try: + payload = json.loads(event_string) + parts = payload.get("parts") or [] + text = parts[0].get("text") if parts else None + return text or event_string + except (AttributeError, IndexError, TypeError, json.JSONDecodeError): + return event_string + + def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: + if not event_strings: + return True + + try: + context_set_name = self._context_set_name(user_id) + self._ensure_context_set(context_set_name) + infer = kwargs.pop("infer", True) + # Framework-level bookkeeping kwargs are not accepted by the TOS + # ContextBucket SDK; drop them so they are not forwarded downstream. + kwargs.pop("session_id", None) + kwargs.pop("app_name", None) + for event_string in event_strings: + self._client.create_context_bucket_memory( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + context_set_name=context_set_name, + content=self._to_content(event_string), + infer=infer, + **kwargs, + ) + return True + except Exception as error: + logger.error(f"Failed to save memory to TOS ContextBucket: {error}") + return False + + def search_memory( + self, user_id: str, query: str, top_k: int, **kwargs + ) -> list[str]: + try: + context_set_name = self._context_set_name(user_id) + self._ensure_context_set(context_set_name) + # Framework-level bookkeeping kwargs are not accepted by the TOS + # ContextBucket SDK; drop them so they are not forwarded downstream. + kwargs.pop("app_name", None) + output = self._client.search_context_bucket_memory( + account_id=self.account_id, + context_bucket_name=self.context_bucket_name, + context_set_name=context_set_name, + query=query, + limit=top_k, + **kwargs, + ) + return [ + result.memory + for result in getattr(output, "results", []) + if getattr(result, "memory", None) + ] + except Exception as error: + logger.error(f"Failed to search memory from TOS ContextBucket: {error}") + return []