feat: release v1.0.0-beta.1

This commit is contained in:
chrishuan
2026-05-29 17:33:12 +08:00
commit 36b0537676
290 changed files with 72303 additions and 0 deletions
+261
View File
@@ -0,0 +1,261 @@
# Agent 接入指南(Python
本文讲怎么把 `tencentdb-agent-memory-sdk-python` 接到一个 AI Agent 里。SDK 的 14 个 API 速查见 [`README.md`](./README.md),本文讲**怎么把它们组装成一套长期记忆**。
---
## 接入要做的四件事
```
用户输入 → ① 召回(注入 prompt) → LLM → ② 捕获(写 L0)
③ 工具:让 LLM 自己再查
④ 错误降级:失败不挂主流程
```
---
## 0. 初始化
```python
from tencentdb_agent_memory import MemoryClient, AsyncMemoryClient
# 同步
client = MemoryClient(
endpoint="https://your-memory-gateway",
api_key=os.environ["MEMORY_API_KEY"],
service_id="your-instance-id",
)
# 异步(推荐 Agent 场景用)
async with AsyncMemoryClient(
endpoint="https://your-memory-gateway",
api_key=os.environ["MEMORY_API_KEY"],
service_id="your-instance-id",
) as client:
...
```
`service_id` 决定 memory space 隔离粒度,同 id 数据共享、不同 id 完全隔离。Agent 场景几乎都用 async,别用同步版(会阻塞事件循环)。
---
## 1. 召回(Recall
在用户消息发给 LLM 前,并行拉三类记忆,拼到 system prompt 里。
```python
import asyncio
async def recall(client: AsyncMemoryClient, user_query: str) -> dict:
l1, persona, scenes = await asyncio.gather(
client.search_atomic(query=user_query, limit=5),
client.read_core(), # L3 用户画像
client.list_scenarios(), # L2 场景索引
return_exceptions=True, # 关键:单路挂不影响其它
)
l1_items = l1["items"] if not isinstance(l1, Exception) else []
persona_text = persona["content"] if not isinstance(persona, Exception) else None
scene_list = scenes["entries"] if not isinstance(scenes, Exception) else []
return format_prompt(l1_items, persona_text, scene_list)
```
`asyncio.gather(..., return_exceptions=True)` 是关键——任何一路超时/失败,其它两路结果照常用,不影响主对话。
### 拼 prompt 的两个区块
- **prepend_context(动态)**:L1 召回结果,每轮都变,放在用户消息前。
- **append_system_context(稳定)**Persona + Scene 索引 + 工具调用指南,放在 system prompt 末尾,KV cache 友好。 _(待确定:放到 system prompt 末尾仍可能造成 KV cache miss,需要继续讨论。)_
```python
def format_prompt(l1_items, persona, scenes) -> dict:
prepend = None
if l1_items:
lines = [f"- [{m['type']}] {m['content']}" for m in l1_items]
prepend = "<relevant-memories>\n" + "\n".join(lines) + "\n</relevant-memories>"
parts = []
if persona:
parts.append(f"<user-persona>\n{persona}\n</user-persona>")
if scenes:
parts.append("## Scene Navigation\n*以下场景可用 tdai_read_file 读取详情*")
parts.extend(f"- `{s['path']}`" for s in scenes)
parts.append(MEMORY_TOOLS_GUIDE) # 见下文
return {"prepend": prepend, "append": "\n\n".join(parts)}
```
> 实现要点:在召回阶段**缓存原始用户文本**(清洁版,未注入 recall),后面 capture 阶段要用——见第 2 节。
---
## 2. 捕获(Capture
在 agent 一轮跑完后,把这一轮新增的 user/assistant 消息清洗后写回 L0。
```python
async def capture(
client: AsyncMemoryClient,
session_key: str,
raw_messages: list, # 框架给的完整消息历史
original_user_text: str, # 召回阶段缓存的清洁版用户文本
original_user_message_count: int, # 召回阶段缓存的消息数
):
# ① 位置切片:只保留这一轮新增的消息
new_messages = raw_messages[original_user_message_count:]
# ② 提取 user/assistant,去掉 tool calls / system / 多模态噪声
extracted = extract_user_assistant(new_messages)
# ③ 把被 recall 污染的用户消息换回原始版
for m in extracted:
if m["role"] == "user" and m["timestamp"] == new_messages[0].get("timestamp"):
m["content"] = original_user_text
break
# ④ 文本清洗:去图片 base64、去代码块、过滤太短/纯符号
cleaned = [
{**m, "content": sanitize(m["content"])}
for m in extracted
if len(sanitize(m["content"]).strip()) > 5
]
if not cleaned:
return
# ⑤ 提交
await client.add_conversation(
session_id=session_key,
messages=[
{
"role": m["role"],
"content": m["content"],
"timestamp": datetime.fromtimestamp(m["timestamp"] / 1000).isoformat(),
}
for m in cleaned
],
)
```
### 为什么要替换"被污染的用户消息"
召回阶段会往用户消息前 prepend 一段 `<relevant-memories>...</relevant-memories>`。如果不还原成原始文本就写 L0,下一轮召回就会基于这段被污染的文本去 search/embedding——形成**反馈环**,记忆会越来越乱。
### 为什么要位置切片
agent 一轮结束时框架给的是**完整历史**,不是本轮新增。直接全发会重复写。召回阶段记一下消息数 N,结束时 `messages[N:]` 就是新增的。
---
## 3. 工具暴露
只靠 prompt 注入的记忆有限。再注册三个工具让 LLM 自己查:
| 工具 | 何时用 | 实现 |
|---|---|---|
| `tdai_memory_search` | 找结构化偏好/事实 | `client.search_atomic(query=..., limit=...)` |
| `tdai_conversation_search` | 找原始对话片段 | `client.search_conversation(query=..., limit=...)` |
| `tdai_read_file` | 读 persona / scene block 全文 | `client.read_file(path)` |
在 system prompt 里说清楚什么时候调,加上次数上限:
```
## 记忆工具
- tdai_memory_search:搜结构化记忆(用户偏好、规则、历史事件)
- tdai_conversation_search:搜原始对话原文
- tdai_read_file:读取场景文件(用 Scene Navigation 列出的路径)
⚠️ memory_search + conversation_search 一轮总共最多调 3 次。
```
不限次数 LLM 会反复瞎搜。
---
## 4. 错误降级
记忆服务挂了**不能挂主对话**。三条原则:
1. **召回**用 `asyncio.gather(..., return_exceptions=True)`,单路失败不影响其它。
2. **捕获**包 try/except,失败只记日志:
```python
try:
await capture(...)
except Exception as e:
logger.warning(f"capture failed: {e}")
```
3. **工具**返回错误字符串而不是抛异常,让 LLM 自己看到 "memory unavailable" 然后继续聊。
---
## 5. 错误处理
非零 code 抛 `TDAMError`
```python
from tencentdb_agent_memory import TDAMError
try:
content = await client.read_file("scene_blocks/x.md")
except TDAMError as e:
if e.code == 404:
pass # 文件不存在,正常情况
else:
logger.warning(f"memory error code={e.code} request_id={e.request_id}")
```
`request_id` 在 server 端也有日志,排障时给后端就行。
---
## 6. 性能建议
- **召回总预算 < 200ms**:三路并行后取最快可用结果,超时的丢掉。
- **prompt 注入控制大小**L1 ≤ 5 条、Scene 列表只列 path 不列内容、Persona 一份就够。让 LLM 不够用时再用工具拉详情。
- **session 粒度**`session_key` 是 L0 partition key,长期对话用稳定 id(用户 id + 会话 id),不要每轮换。
- **不要在主线程同步调**:用 `AsyncMemoryClient`,别用 `MemoryClient`,否则会阻塞事件循环。
---
## 附:sanitize 实现参考
清洗函数处理这几类噪声:
```python
import re
import time
_IMAGE_DATA_URI = re.compile(r"data:image/[a-z+]+;base64,[A-Za-z0-9+/=]+", re.IGNORECASE)
_CODE_BLOCK = re.compile(r"```[\s\S]*?```")
def sanitize(text: str) -> str:
# 去 base64 图片
text = _IMAGE_DATA_URI.sub("[image]", text)
# 去代码块(assistant 输出常见,对 embedding 是噪声)
text = _CODE_BLOCK.sub("[code]", text)
return text.strip()
def extract_user_assistant(messages: list) -> list:
"""从原始消息列表里提取 user/assistant 文本,丢掉 tool / system / 空内容。"""
out = []
for m in messages:
role = m.get("role")
if role not in ("user", "assistant"):
continue
content = m.get("content")
if isinstance(content, list):
# 多模态消息:拼接 text 部分
content = "\n".join(p.get("text", "") for p in content if p.get("type") == "text")
if not isinstance(content, str) or not content.strip():
continue
out.append({
"role": role,
"content": content.strip(),
"timestamp": m.get("timestamp", int(time.time() * 1000)),
})
return out
```
+137
View File
@@ -0,0 +1,137 @@
# tencentdb-agent-memory-sdk-python
Python SDK for the **TencentDB Agent Memory v2 API**.
Provides synchronous (`MemoryClient`) and asynchronous (`AsyncMemoryClient`) clients.
> **Distribution name**: `tencentdb-agent-memory-sdk-python` (PyPI / `pip install`)
> **Import path**: `tencentdb_agent_memory` (Python module)
## Install
```bash
# From PyPI (after publish)
pip install tencentdb-agent-memory-sdk-python
# From local .whl
pip install ./tencentdb_agent_memory_sdk_python-0.1.0-py3-none-any.whl
```
## Quick Start
```python
from tencentdb_agent_memory import MemoryClient
client = MemoryClient(
endpoint="http://127.0.0.1:8420",
api_key="your-api-key",
service_id="your-memory-space-id",
)
# L0: append a conversation
result = client.add_conversation(
session_id="sess-1",
messages=[
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi!"},
],
)
print(result["accepted_ids"])
# L1: search structured memories
hits = client.search_atomic(query="user preferences", limit=5)
print(hits["items"])
# L1: update a memory note
client.update_atomic(id="note-xxx", content="updated content", background="context")
# L2: list scenario files
scenarios = client.list_scenarios(path_prefix="")
print(scenarios["entries"])
# L2: read a scenario file
file = client.read_scenario("工作.md")
print(file["content"])
# L2: update a scenario file (must already exist)
client.write_scenario("工作.md", "# Updated content", summary="new summary")
# L3: read core memory (persona)
core = client.read_core()
print(core["content"])
# L3: write core memory
client.write_core("# User Profile\n...")
# Read memory pipeline artifacts (e.g. persona.md, scene_blocks/*.md)
raw = client.read_file("scene_blocks/工作.md")
```
## Async Usage
```python
import asyncio
from tencentdb_agent_memory import AsyncMemoryClient
async def main():
async with AsyncMemoryClient(
endpoint="http://127.0.0.1:8420",
api_key="your-api-key",
service_id="your-memory-space-id",
) as client:
result = await client.search_atomic(query="preferences")
print(result["items"])
asyncio.run(main())
```
## API Methods
| Layer | Method | Endpoint |
|-------|--------|----------|
| L0 | `add_conversation()` | `POST /v2/conversation/add` |
| L0 | `query_conversation()` | `POST /v2/conversation/query` |
| L0 | `search_conversation()` | `POST /v2/conversation/search` |
| L0 | `delete_conversation()` | `POST /v2/conversation/delete` |
| L1 | `update_atomic()` | `POST /v2/atomic/update` |
| L1 | `query_atomic()` | `POST /v2/atomic/query` |
| L1 | `search_atomic()` | `POST /v2/atomic/search` |
| L1 | `delete_atomic()` | `POST /v2/atomic/delete` |
| L2 | `list_scenarios()` | `POST /v2/scenario/ls` |
| L2 | `read_scenario()` | `POST /v2/scenario/read` |
| L2 | `write_scenario()` | `POST /v2/scenario/write` |
| L2 | `rm_scenario()` | `POST /v2/scenario/rm` |
| L3 | `read_core()` | `POST /v2/core/read` |
| L3 | `write_core()` | `POST /v2/core/write` |
## Error Handling
All non-zero `code` responses raise `TDAMError`:
```python
from tencentdb_agent_memory import TDAMError
try:
client.read_core()
except TDAMError as e:
print(f"code={e.code} message={e.message} request_id={e.request_id}")
```
## Build & Pack
```bash
# Build wheel
python -m build
# → dist/tencentdb_agent_memory_sdk_python-0.1.0-py3-none-any.whl
# Or just wheel
pip wheel . --no-deps -w dist/
```
## Dependencies
- `httpx>=0.24.0` (HTTP client with async support)
## License
MIT
+137
View File
@@ -0,0 +1,137 @@
# tencentdb-agent-memory-sdk-python
**TencentDB Agent Memory v2 API** 的 Python SDK。
提供同步客户端(`MemoryClient`)和异步客户端(`AsyncMemoryClient`)。
> **发布包名**`tencentdb-agent-memory-sdk-python`PyPI / `pip install`
> **导入路径**`tencentdb_agent_memory`Python 模块)
## 安装
```bash
# 从 PyPI 安装(发布后)
pip install tencentdb-agent-memory-sdk-python
# 从本地 .whl 安装
pip install ./tencentdb_agent_memory_sdk_python-0.1.0-py3-none-any.whl
```
## 快速开始
```python
from tencentdb_agent_memory import MemoryClient
client = MemoryClient(
endpoint="http://127.0.0.1:8420",
api_key="your-api-key",
service_id="your-memory-space-id",
)
# L0: 添加对话
result = client.add_conversation(
session_id="sess-1",
messages=[
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi!"},
],
)
print(result["accepted_ids"])
# L1: 搜索结构化记忆
hits = client.search_atomic(query="user preferences", limit=5)
print(hits["items"])
# L1: 更新一条记忆
client.update_atomic(id="note-xxx", content="updated content", background="context")
# L2: 列出场景文件
scenarios = client.list_scenarios(path_prefix="")
print(scenarios["entries"])
# L2: 读取场景文件
file = client.read_scenario("工作.md")
print(file["content"])
# L2: 更新场景文件(文件必须已存在)
client.write_scenario("工作.md", "# Updated content", summary="new summary")
# L3: 读取核心记忆(用户画像)
core = client.read_core()
print(core["content"])
# L3: 写入核心记忆
client.write_core("# User Profile\n...")
# 读取记忆 pipeline 产物(如 persona.md、scene_blocks/*.md
raw = client.read_file("scene_blocks/工作.md")
```
## 异步用法
```python
import asyncio
from tencentdb_agent_memory import AsyncMemoryClient
async def main():
async with AsyncMemoryClient(
endpoint="http://127.0.0.1:8420",
api_key="your-api-key",
service_id="your-memory-space-id",
) as client:
result = await client.search_atomic(query="preferences")
print(result["items"])
asyncio.run(main())
```
## API 方法
| 层级 | 方法 | 接口 |
|------|------|------|
| L0 | `add_conversation()` | `POST /v2/conversation/add` |
| L0 | `query_conversation()` | `POST /v2/conversation/query` |
| L0 | `search_conversation()` | `POST /v2/conversation/search` |
| L0 | `delete_conversation()` | `POST /v2/conversation/delete` |
| L1 | `update_atomic()` | `POST /v2/atomic/update` |
| L1 | `query_atomic()` | `POST /v2/atomic/query` |
| L1 | `search_atomic()` | `POST /v2/atomic/search` |
| L1 | `delete_atomic()` | `POST /v2/atomic/delete` |
| L2 | `list_scenarios()` | `POST /v2/scenario/ls` |
| L2 | `read_scenario()` | `POST /v2/scenario/read` |
| L2 | `write_scenario()` | `POST /v2/scenario/write` |
| L2 | `rm_scenario()` | `POST /v2/scenario/rm` |
| L3 | `read_core()` | `POST /v2/core/read` |
| L3 | `write_core()` | `POST /v2/core/write` |
## 错误处理
所有非零 `code` 的响应会抛出 `TDAMError`
```python
from tencentdb_agent_memory import TDAMError
try:
client.read_core()
except TDAMError as e:
print(f"code={e.code} message={e.message} request_id={e.request_id}")
```
## 构建与打包
```bash
# 构建 wheel
python -m build
# → dist/tencentdb_agent_memory_sdk_python-0.1.0-py3-none-any.whl
# 或仅构建 wheel
pip wheel . --no-deps -w dist/
```
## 依赖
- `httpx>=0.24.0`(支持异步的 HTTP 客户端)
## 许可证
MIT
+19
View File
@@ -0,0 +1,19 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "tencentdb-agent-memory-sdk-python"
version = "0.1.0"
description = "Python SDK for TencentDB Agent Memory v2 API"
readme = "README.md"
license = { text = "MIT" }
requires-python = ">=3.9"
authors = [{ name = "TencentDB Agent Memory Team" }]
dependencies = ["httpx>=0.24.0"]
[project.optional-dependencies]
dev = ["pytest>=7.0", "pytest-asyncio>=0.21", "respx>=0.20", "build", "python-dotenv>=1.0"]
[tool.hatch.build.targets.wheel]
packages = ["tencentdb_agent_memory"]
@@ -0,0 +1,6 @@
"""TencentDB Agent Memory v2 Python SDK."""
from .client import AsyncMemoryClient, MemoryClient
from .errors import ParamError, TDAMError
__all__ = ["MemoryClient", "AsyncMemoryClient", "TDAMError", "ParamError"]
+151
View File
@@ -0,0 +1,151 @@
"""Low-level HTTP transport for the TencentDB Agent Memory v2 API.
Provides Bearer-token authentication, response-envelope unwrapping
(``code == 0`` → ``data``; otherwise raise ``TDAMError``), and trace-id
propagation via the ``x-trace-id`` response header.
"""
from __future__ import annotations
import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
import httpx
from .errors import TDAMError
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Stub abstraction
# ---------------------------------------------------------------------------
class Stub(ABC):
"""Base transport interface."""
@abstractmethod
def post(self, path: str, body: dict, timeout: Optional[float] = None) -> dict:
...
@abstractmethod
def close(self) -> None:
...
class HttpStub(Stub):
"""Synchronous HTTP transport backed by :mod:`httpx`.
Parameters
----------
endpoint : str
Base URL of the memory service, e.g.
``https://memory.tencentyun.com``.
api_key : str
Bearer token sent via ``Authorization`` header.
service_id : str
Memory instance ID (sent via ``x-tdai-service-id`` header).
timeout : float
Default request timeout in seconds.
"""
def __init__(
self,
endpoint: str,
api_key: str,
service_id: str,
timeout: float = 30,
verify: bool = False,
client: Optional[httpx.Client] = None,
) -> None:
self.endpoint = endpoint.rstrip("/")
self.client = client or httpx.Client(timeout=timeout, verify=verify)
self.headers: Dict[str, str] = {
"Authorization": f"Bearer {api_key}",
"x-tdai-service-id": service_id,
"Content-Type": "application/json",
}
def post(self, path: str, body: dict, timeout: Optional[float] = None) -> dict:
url = f"{self.endpoint}{path}"
logger.debug("Request %s %s", path, body)
resp = self.client.post(
url=url,
json=body,
headers=self.headers,
timeout=timeout or self.client.timeout,
)
logger.debug("Response %s %s", path, resp.text)
resp.raise_for_status()
data = resp.json()
if data.get("code") != 0:
req_id = resp.headers.get("x-qcloud-transaction-id", data.get("request_id", ""))
raise TDAMError(
code=data.get("code", -1),
message=data.get("message", "unknown error"),
request_id=req_id,
)
result: dict = data.get("data", {})
trace_id = resp.headers.get("x-trace-id")
if trace_id:
result["trace_id"] = trace_id
return result
def close(self) -> None:
if isinstance(self.client, httpx.Client):
self.client.close()
# ---------------------------------------------------------------------------
# Async variant
# ---------------------------------------------------------------------------
class AsyncHttpStub:
"""Asynchronous HTTP transport backed by :mod:`httpx`."""
def __init__(
self,
endpoint: str,
api_key: str,
service_id: str,
timeout: float = 30,
verify: bool = False,
client: Optional[httpx.AsyncClient] = None,
) -> None:
self.endpoint = endpoint.rstrip("/")
self.client = client or httpx.AsyncClient(timeout=timeout, verify=verify)
self.headers: Dict[str, str] = {
"Authorization": f"Bearer {api_key}",
"x-tdai-service-id": service_id,
"Content-Type": "application/json",
}
async def post(self, path: str, body: dict, timeout: Optional[float] = None) -> dict:
url = f"{self.endpoint}{path}"
logger.debug("Request %s %s", path, body)
resp = await self.client.post(
url=url,
json=body,
headers=self.headers,
timeout=timeout or self.client.timeout,
)
logger.debug("Response %s %s", path, resp.text)
resp.raise_for_status()
data = resp.json()
if data.get("code") != 0:
req_id = resp.headers.get("x-qcloud-transaction-id", data.get("request_id", ""))
raise TDAMError(
code=data.get("code", -1),
message=data.get("message", "unknown error"),
request_id=req_id,
)
result: dict = data.get("data", {})
trace_id = resp.headers.get("x-trace-id")
if trace_id:
result["trace_id"] = trace_id
return result
async def close(self) -> None:
if isinstance(self.client, httpx.AsyncClient):
await self.client.aclose()
+461
View File
@@ -0,0 +1,461 @@
"""TencentDB Agent Memory v2 Python SDK — synchronous + asynchronous clients.
Exposes the v2 data-plane API (14 routes) over a Bearer-token authenticated
HTTP transport.
"""
from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from ._http import AsyncHttpStub, HttpStub, Stub
from .cos import AsyncMemoryFileReader, AsyncStsCredentialManager, MemoryFileReader, StsCredentialManager
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_V2 = "/v2"
def _strip_none(d: Dict[str, Any]) -> Dict[str, Any]:
"""Return a copy of *d* with ``None`` values removed."""
return {k: v for k, v in d.items() if v is not None}
# ---------------------------------------------------------------------------
# Synchronous client
# ---------------------------------------------------------------------------
class MemoryClient:
"""Synchronous client for the TencentDB Agent Memory v2 data-plane API.
Example::
from tencentdb_agent_memory import MemoryClient
client = MemoryClient(
endpoint="https://memory.tencentyun.com",
api_key="sk-xxxxxxxx",
service_id="mem-xxxxxxxx",
)
result = client.add_conversation("sess-1", [
{"role": "user", "content": "hello"},
])
print(result) # {"accepted_ids": [...], "total_count": 1}
Parameters
----------
endpoint : str
Base URL of the memory service.
api_key : str
Bearer token.
service_id : str
Memory instance ID (sent via ``x-tdai-service-id`` header).
timeout : float
Request timeout in seconds.
stub : Stub | None
Inject a custom transport (useful for testing).
"""
def __init__(
self,
endpoint: str = "",
api_key: str = "",
service_id: Optional[str] = None,
*,
timeout: float = 30,
verify: bool = False,
stub: Optional[Stub] = None,
) -> None:
if stub is not None:
self._stub = stub
else:
if not service_id:
raise ValueError("service_id must be provided")
self._stub = HttpStub(endpoint, api_key, service_id, timeout=timeout, verify=verify)
# Memory file reader (lazy init on first read_file call)
self._cos_reader: Optional[MemoryFileReader] = None
self._sts_manager: Optional[StsCredentialManager] = None
# -- L0 Conversation ---------------------------------------------------
def add_conversation(
self,
session_id: str,
messages: List[Dict[str, Any]],
) -> Dict[str, Any]:
"""``POST /conversation/add``"""
return self._stub.post(
f"{_V2}/conversation/add",
{"session_id": session_id, "messages": messages},
)
def query_conversation(
self,
*,
session_id: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
"""``POST /conversation/query``"""
return self._stub.post(
f"{_V2}/conversation/query",
_strip_none({
"session_id": session_id,
"limit": limit,
"offset": offset,
"time_start": time_start,
"time_end": time_end,
}),
)
def search_conversation(
self,
query: str,
*,
limit: Optional[int] = None,
session_id: Optional[str] = None,
time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
"""``POST /conversation/search``"""
return self._stub.post(
f"{_V2}/conversation/search",
_strip_none({
"query": query,
"limit": limit,
"session_id": session_id,
"time_start": time_start,
"time_end": time_end,
}),
)
def delete_conversation(
self,
*,
message_ids: Optional[List[str]] = None,
session_id: Optional[str] = None,
) -> Dict[str, Any]:
"""``POST /conversation/delete`` — *message_ids* 和 *session_id* 二选一。"""
return self._stub.post(
f"{_V2}/conversation/delete",
_strip_none({
"message_ids": message_ids,
"session_id": session_id,
}),
)
# -- L1 Atomic ---------------------------------------------------------
def update_atomic(self, id: str, content: str, *, background: Optional[str] = None) -> Dict[str, Any]:
"""``POST /atomic/update``"""
return self._stub.post(
f"{_V2}/atomic/update",
_strip_none({"id": id, "content": content, "background": background}),
)
def query_atomic(
self,
*,
type: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
"""``POST /atomic/query``"""
return self._stub.post(
f"{_V2}/atomic/query",
_strip_none({
"type": type,
"limit": limit,
"offset": offset,
"time_start": time_start,
"time_end": time_end,
}),
)
def search_atomic(
self,
query: str,
*,
limit: Optional[int] = None,
type: Optional[str] = None,
time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
"""``POST /atomic/search``"""
return self._stub.post(
f"{_V2}/atomic/search",
_strip_none({
"query": query,
"limit": limit,
"type": type,
"time_start": time_start,
"time_end": time_end,
}),
)
def delete_atomic(self, ids: List[str]) -> Dict[str, Any]:
"""``POST /atomic/delete``"""
return self._stub.post(f"{_V2}/atomic/delete", {"ids": ids})
# -- L2 Scenario -------------------------------------------------------
def list_scenarios(
self,
*,
path_prefix: Optional[str] = None,
) -> Dict[str, Any]:
"""``POST /scenario/ls``"""
return self._stub.post(
f"{_V2}/scenario/ls",
_strip_none({
"path_prefix": path_prefix,
}),
)
def read_scenario(self, path: str) -> Dict[str, Any]:
"""``POST /scenario/read``
Returns dict with ``content``, ``created_at``, ``updated_at``.
If the file does not exist, ``content`` will be ``None``.
"""
return self._stub.post(f"{_V2}/scenario/read", {"path": path})
def write_scenario(self, path: str, content: str, *, summary: Optional[str] = None) -> Dict[str, Any]:
"""``POST /scenario/write``"""
return self._stub.post(
f"{_V2}/scenario/write",
_strip_none({"path": path, "content": content, "summary": summary}),
)
def rm_scenario(self, path: str) -> Dict[str, Any]:
"""``POST /scenario/rm``"""
return self._stub.post(f"{_V2}/scenario/rm", {"path": path})
# -- L3 Core -----------------------------------------------------------
def read_core(self) -> Dict[str, Any]:
"""``POST /core/read``
Returns dict with ``content``, ``created_at``, ``updated_at``.
If core memory has not been generated yet, ``content`` will be ``None``.
"""
return self._stub.post(f"{_V2}/core/read", {})
def write_core(self, content: str) -> Dict[str, Any]:
"""``POST /core/write``"""
return self._stub.post(f"{_V2}/core/write", {"content": content})
# -- File read (memory pipeline artifacts) -----------------------------
def read_file(self, path: str) -> str:
"""Read a memory pipeline artifact (e.g. ``persona.md``,
``scene_blocks/*.md``) by relative path.
Parameters
----------
path : str
Relative path within the memory space, e.g.
``"scene_blocks/cooking-recipes.md"`` or ``"persona.md"``.
Returns
-------
str
File content.
Raises
------
TDAMError
On 404 (not found), 403 (auth failure after retry), or other errors.
"""
if self._cos_reader is None:
self._sts_manager = StsCredentialManager(
endpoint=self._stub.endpoint,
api_key=self._stub.headers["Authorization"].removeprefix("Bearer "),
service_id=self._stub.headers["x-tdai-service-id"],
)
self._cos_reader = MemoryFileReader(self._sts_manager)
return self._cos_reader.read(path)
# -- lifecycle ---------------------------------------------------------
def close(self) -> None:
if self._cos_reader is not None:
self._cos_reader.close()
self._stub.close()
def __enter__(self) -> "MemoryClient":
return self
def __exit__(self, *exc: Any) -> None:
self.close()
# ---------------------------------------------------------------------------
# Asynchronous client
# ---------------------------------------------------------------------------
class AsyncMemoryClient:
"""Asynchronous client for the TencentDB Agent Memory v2 data-plane API.
Same API surface as :class:`MemoryClient` but all methods are coroutines.
"""
def __init__(
self,
endpoint: str = "",
api_key: str = "",
service_id: Optional[str] = None,
*,
timeout: float = 30,
verify: bool = False,
) -> None:
if not service_id:
raise ValueError("service_id must be provided")
self._stub = AsyncHttpStub(endpoint, api_key, service_id, timeout=timeout, verify=verify)
# Memory file reader (lazy init)
self._cos_reader: Optional[AsyncMemoryFileReader] = None
self._sts_manager: Optional[AsyncStsCredentialManager] = None
# -- L0 Conversation ---------------------------------------------------
async def add_conversation(
self, session_id: str, messages: List[Dict[str, Any]],
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/conversation/add",
{"session_id": session_id, "messages": messages},
)
async def query_conversation(
self, *, session_id: Optional[str] = None, limit: Optional[int] = None,
offset: Optional[int] = None, time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/conversation/query",
_strip_none({"session_id": session_id, "limit": limit, "offset": offset,
"time_start": time_start, "time_end": time_end}),
)
async def search_conversation(
self, query: str, *, limit: Optional[int] = None,
session_id: Optional[str] = None, time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/conversation/search",
_strip_none({"query": query, "limit": limit, "session_id": session_id,
"time_start": time_start, "time_end": time_end}),
)
async def delete_conversation(
self, *, message_ids: Optional[List[str]] = None,
session_id: Optional[str] = None,
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/conversation/delete",
_strip_none({"message_ids": message_ids, "session_id": session_id}),
)
# -- L1 Atomic ---------------------------------------------------------
async def update_atomic(self, id: str, content: str, *, background: Optional[str] = None) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/atomic/update",
_strip_none({"id": id, "content": content, "background": background}),
)
async def query_atomic(
self, *, type: Optional[str] = None, limit: Optional[int] = None,
offset: Optional[int] = None, time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/atomic/query",
_strip_none({"type": type, "limit": limit, "offset": offset,
"time_start": time_start, "time_end": time_end}),
)
async def search_atomic(
self, query: str, *, limit: Optional[int] = None,
type: Optional[str] = None, time_start: Optional[str] = None,
time_end: Optional[str] = None,
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/atomic/search",
_strip_none({"query": query, "limit": limit, "type": type,
"time_start": time_start, "time_end": time_end}),
)
async def delete_atomic(self, ids: List[str]) -> Dict[str, Any]:
return await self._stub.post(f"{_V2}/atomic/delete", {"ids": ids})
# -- L2 Scenario -------------------------------------------------------
async def list_scenarios(
self, *, path_prefix: Optional[str] = None,
) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/scenario/ls",
_strip_none({"path_prefix": path_prefix}),
)
async def read_scenario(self, path: str) -> Dict[str, Any]:
"""``POST /scenario/read`` — returns ``content: None`` if file does not exist."""
return await self._stub.post(f"{_V2}/scenario/read", {"path": path})
async def write_scenario(self, path: str, content: str, *, summary: Optional[str] = None) -> Dict[str, Any]:
return await self._stub.post(
f"{_V2}/scenario/write", _strip_none({"path": path, "content": content, "summary": summary}),
)
async def rm_scenario(self, path: str) -> Dict[str, Any]:
return await self._stub.post(f"{_V2}/scenario/rm", {"path": path})
# -- L3 Core -----------------------------------------------------------
async def read_core(self) -> Dict[str, Any]:
"""``POST /core/read`` — returns ``content: None`` if not yet generated."""
return await self._stub.post(f"{_V2}/core/read", {})
async def write_core(self, content: str) -> Dict[str, Any]:
return await self._stub.post(f"{_V2}/core/write", {"content": content})
# -- lifecycle ---------------------------------------------------------
# -- File read (memory pipeline artifacts) -----------------------------
async def read_file(self, path: str) -> str:
"""Read a memory pipeline artifact (async)."""
if self._cos_reader is None:
self._sts_manager = AsyncStsCredentialManager(
endpoint=self._stub.endpoint,
api_key=self._stub.headers["Authorization"].removeprefix("Bearer "),
service_id=self._stub.headers["x-tdai-service-id"],
)
self._cos_reader = AsyncMemoryFileReader(self._sts_manager)
return await self._cos_reader.read(path)
# -- lifecycle ---------------------------------------------------------
async def close(self) -> None:
if self._cos_reader is not None:
await self._cos_reader.close()
await self._stub.close()
async def __aenter__(self) -> "AsyncMemoryClient":
return self
async def __aexit__(self, *exc: Any) -> None:
await self.close()
+461
View File
@@ -0,0 +1,461 @@
"""Memory file reader — direct read of memory pipeline artifacts (persona.md,
scene_blocks/*.md) from object storage with STS credential management.
The SDK user calls ``client.read_file(path)`` and gets the file content back.
Under the hood, we:
1. Fetch STS temporary credentials from the platform (``POST /v2/cos/secret``)
2. Cache credentials until they expire (auto-refresh)
3. Sign a COS V5 GET request with the STS credentials
4. Return the file content as a string
Storage backend (currently COS) is an implementation detail — the public API
is intentionally storage-agnostic.
"""
from __future__ import annotations
import hashlib
import hmac
import logging
import re
import threading
import time
from typing import Any, Dict, Optional
from urllib.parse import urlparse
import httpx
from .errors import TDAMError
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# COS URL parser
# ---------------------------------------------------------------------------
def _parse_cos_url(cos_url: str) -> tuple[str, str]:
"""Parse CosUrl like ``https://bucket.cos.region.myqcloud.com`` → (bucket, region)."""
host = urlparse(cos_url).hostname or ""
# Pattern: {bucket}.cos.{region}.myqcloud.com
m = re.match(r"^(.+?)\.cos\.(.+?)\.myqcloud\.com$", host)
if m:
return m.group(1), m.group(2)
raise TDAMError(
code=-1,
message=f"Cannot parse CosUrl: {cos_url!r} (expected {{bucket}}.cos.{{region}}.myqcloud.com)",
)
# ---------------------------------------------------------------------------
# STS Credential
# ---------------------------------------------------------------------------
class StsCredential:
"""Parsed STS credential from ``POST /v2/cos/secret``.
Platform response format::
{
"CosUrl": "https://{bucket}.cos.{region}.myqcloud.com",
"TmpSecretId": "...",
"TmpSecretKey": "...",
"TmpToken": "...",
"ExpirationTime": "2026-05-15T16:44:49+08:00",
"PathPrefix": "memory_v2/cos_data/mem-xxx"
}
"""
__slots__ = (
"tmp_secret_id", "tmp_secret_key", "token",
"bucket", "region", "prefix", "expires_at_epoch",
)
def __init__(self, data: Dict[str, Any]) -> None:
self.tmp_secret_id: str = data["TmpSecretId"]
self.tmp_secret_key: str = data["TmpSecretKey"]
self.token: str = data.get("TmpToken", "")
# Parse bucket + region from CosUrl
self.bucket: str
self.region: str
self.bucket, self.region = _parse_cos_url(data["CosUrl"])
# PathPrefix — ensure trailing slash for key concatenation
prefix = data.get("PathPrefix", "")
self.prefix: str = prefix if prefix.endswith("/") else f"{prefix}/"
# Parse ISO 8601 → epoch seconds
expires_str = data.get("ExpirationTime", "")
if expires_str:
from datetime import datetime, timezone
try:
dt = datetime.fromisoformat(expires_str.replace("Z", "+00:00"))
self.expires_at_epoch = dt.timestamp()
except Exception:
self.expires_at_epoch = time.time() + 1800 # 30 min fallback
else:
self.expires_at_epoch = time.time() + 1800
def is_valid(self, buffer_seconds: float = 120) -> bool:
"""Check if credential is still valid (with 2-minute buffer)."""
return time.time() < (self.expires_at_epoch - buffer_seconds)
@property
def cos_host(self) -> str:
return f"{self.bucket}.cos.{self.region}.myqcloud.com"
# ---------------------------------------------------------------------------
# STS Credential Manager
# ---------------------------------------------------------------------------
class StsCredentialManager:
"""Thread-safe STS credential cache with auto-refresh.
- Fetches STS from platform ``POST /v2/cos/secret``
- Caches until expiry (with 2-minute buffer)
- Coalesces concurrent refresh requests
"""
def __init__(
self,
endpoint: str,
api_key: str,
service_id: str,
buffer_seconds: float = 120,
timeout: float = 30,
) -> None:
self._endpoint = endpoint.rstrip("/")
self._api_key = api_key
self._service_id = service_id
self._buffer = buffer_seconds
self._timeout = timeout
self._credential: Optional[StsCredential] = None
self._lock = threading.Lock()
self._client: Optional[httpx.Client] = None
def get_credential(self) -> StsCredential:
"""Get a valid STS credential (cached or freshly fetched)."""
if self._credential and self._credential.is_valid(self._buffer):
return self._credential
with self._lock:
# Double-check after acquiring lock
if self._credential and self._credential.is_valid(self._buffer):
return self._credential
return self._refresh()
def invalidate(self) -> None:
"""Force invalidate cached credential (e.g. on 403)."""
with self._lock:
self._credential = None
def _refresh(self) -> StsCredential:
logger.debug("[cos] Refreshing STS credential via POST /v2/cos/secret ...")
if self._client is None:
self._client = httpx.Client(timeout=self._timeout)
url = f"{self._endpoint}/v2/cos/secret"
headers = {
"Authorization": f"Bearer {self._api_key}",
"x-tdai-service-id": self._service_id,
"Content-Type": "application/json",
}
resp = self._client.post(url, json={}, headers=headers)
resp.raise_for_status()
data = resp.json()
cred = StsCredential(data)
self._credential = cred
logger.debug("[cos] STS refreshed: bucket=%s prefix=%s expires=%.0f",
cred.bucket, cred.prefix, cred.expires_at_epoch)
return cred
def close(self) -> None:
if self._client is not None:
self._client.close()
# ---------------------------------------------------------------------------
# Async STS Credential Manager
# ---------------------------------------------------------------------------
class AsyncStsCredentialManager:
"""Async variant of StsCredentialManager."""
def __init__(
self,
endpoint: str,
api_key: str,
service_id: str,
buffer_seconds: float = 120,
timeout: float = 30,
) -> None:
self._endpoint = endpoint.rstrip("/")
self._api_key = api_key
self._service_id = service_id
self._buffer = buffer_seconds
self._timeout = timeout
self._credential: Optional[StsCredential] = None
import asyncio
self._lock = asyncio.Lock()
self._client: Optional[httpx.AsyncClient] = None
async def get_credential(self) -> StsCredential:
if self._credential and self._credential.is_valid(self._buffer):
return self._credential
async with self._lock:
if self._credential and self._credential.is_valid(self._buffer):
return self._credential
return await self._refresh()
def invalidate(self) -> None:
self._credential = None
async def _refresh(self) -> StsCredential:
logger.debug("[cos] Refreshing STS credential (async) via POST /v2/cos/secret ...")
if self._client is None:
self._client = httpx.AsyncClient(timeout=self._timeout)
url = f"{self._endpoint}/v2/cos/secret"
headers = {
"Authorization": f"Bearer {self._api_key}",
"x-tdai-service-id": self._service_id,
"Content-Type": "application/json",
}
resp = await self._client.post(url, json={}, headers=headers)
resp.raise_for_status()
data = resp.json()
cred = StsCredential(data)
self._credential = cred
return cred
async def close(self) -> None:
if self._client is not None:
await self._client.aclose()
# ---------------------------------------------------------------------------
# COS V5 Signature
# ---------------------------------------------------------------------------
def _cos_v5_sign(
secret_id: str,
secret_key: str,
method: str,
path: str,
host: str,
start_time: Optional[int] = None,
end_time: Optional[int] = None,
) -> str:
"""Generate COS V5 Authorization header value for a GET request.
References:
https://cloud.tencent.com/document/product/436/7778
"""
now = int(time.time())
q_sign_time = f"{start_time or (now - 60)};{end_time or (now + 600)}"
q_key_time = q_sign_time
# Step 1: SignKey
sign_key = hmac.new(
secret_key.encode("utf-8"),
q_key_time.encode("utf-8"),
hashlib.sha1,
).hexdigest()
# Step 2: HttpString
# For GET with no query params and only host header
http_string = f"{method.lower()}\n{path}\n\nhost={host}\n"
# Step 3: StringToSign
sha1_http_string = hashlib.sha1(http_string.encode("utf-8")).hexdigest()
string_to_sign = f"sha1\n{q_sign_time}\n{sha1_http_string}\n"
# Step 4: Signature
signature = hmac.new(
sign_key.encode("utf-8"),
string_to_sign.encode("utf-8"),
hashlib.sha1,
).hexdigest()
return (
f"q-sign-algorithm=sha1"
f"&q-ak={secret_id}"
f"&q-sign-time={q_sign_time}"
f"&q-key-time={q_key_time}"
f"&q-header-list=host"
f"&q-url-param-list="
f"&q-signature={signature}"
)
# ---------------------------------------------------------------------------
# Memory File Reader (sync)
# ---------------------------------------------------------------------------
class MemoryFileReader:
"""Sync memory file reader with STS auto-management.
Reads memory pipeline artifacts (persona.md, scene_blocks/*.md, …) from
object storage. The storage backend is COS today but the public API is
storage-agnostic.
Usage::
reader = MemoryFileReader(sts_manager)
content = reader.read("scene_blocks/cooking-recipes.md")
"""
def __init__(
self,
sts_manager: StsCredentialManager,
timeout: float = 30,
client: Optional[httpx.Client] = None,
) -> None:
self._sts = sts_manager
self._client = client or httpx.Client(timeout=timeout)
def read(self, path: str) -> str:
"""Read a memory file by relative path.
The final COS key is: ``{prefix}{path}``
Returns file content as UTF-8 string.
Raises ``TDAMError`` on failure.
"""
cred = self._sts.get_credential()
full_key = f"{cred.prefix}{path}"
cos_path = f"/{full_key}"
host = cred.cos_host
auth = _cos_v5_sign(
secret_id=cred.tmp_secret_id,
secret_key=cred.tmp_secret_key,
method="GET",
path=cos_path,
host=host,
)
headers: Dict[str, str] = {
"Host": host,
"Authorization": auth,
}
if cred.token:
headers["x-cos-security-token"] = cred.token
url = f"https://{host}{cos_path}"
logger.debug("[cos] GET %s", url)
resp = self._client.get(url, headers=headers)
if resp.status_code == 403:
# Invalidate and retry once
logger.warning("[cos] 403 on GET %s — invalidating STS and retrying", path)
self._sts.invalidate()
cred = self._sts.get_credential()
full_key = f"{cred.prefix}{path}"
cos_path = f"/{full_key}"
host = cred.cos_host
auth = _cos_v5_sign(
secret_id=cred.tmp_secret_id,
secret_key=cred.tmp_secret_key,
method="GET",
path=cos_path,
host=host,
)
headers = {"Host": host, "Authorization": auth}
if cred.token:
headers["x-cos-security-token"] = cred.token
url = f"https://{host}{cos_path}"
resp = self._client.get(url, headers=headers)
if resp.status_code == 404:
raise TDAMError(code=404, message=f"File not found: {path}")
if resp.status_code != 200:
raise TDAMError(
code=resp.status_code,
message=f"COS GET failed: HTTP {resp.status_code}{resp.text[:200]}",
)
return resp.text
def close(self) -> None:
if isinstance(self._client, httpx.Client):
self._client.close()
# ---------------------------------------------------------------------------
# Async Memory File Reader
# ---------------------------------------------------------------------------
class AsyncMemoryFileReader:
"""Async memory file reader with STS auto-management."""
def __init__(
self,
sts_manager: AsyncStsCredentialManager,
timeout: float = 30,
client: Optional[httpx.AsyncClient] = None,
) -> None:
self._sts = sts_manager
self._client = client or httpx.AsyncClient(timeout=timeout)
async def read(self, path: str) -> str:
cred = await self._sts.get_credential()
full_key = f"{cred.prefix}{path}"
cos_path = f"/{full_key}"
host = cred.cos_host
auth = _cos_v5_sign(
secret_id=cred.tmp_secret_id,
secret_key=cred.tmp_secret_key,
method="GET",
path=cos_path,
host=host,
)
headers: Dict[str, str] = {
"Host": host,
"Authorization": auth,
}
if cred.token:
headers["x-cos-security-token"] = cred.token
url = f"https://{host}{cos_path}"
resp = await self._client.get(url, headers=headers)
if resp.status_code == 403:
self._sts.invalidate()
cred = await self._sts.get_credential()
full_key = f"{cred.prefix}{path}"
cos_path = f"/{full_key}"
host = cred.cos_host
auth = _cos_v5_sign(
secret_id=cred.tmp_secret_id,
secret_key=cred.tmp_secret_key,
method="GET",
path=cos_path,
host=host,
)
headers = {"Host": host, "Authorization": auth}
if cred.token:
headers["x-cos-security-token"] = cred.token
url = f"https://{host}{cos_path}"
resp = await self._client.get(url, headers=headers)
if resp.status_code == 404:
raise TDAMError(code=404, message=f"File not found: {path}")
if resp.status_code != 200:
raise TDAMError(
code=resp.status_code,
message=f"COS GET failed: HTTP {resp.status_code}{resp.text[:200]}",
)
return resp.text
async def close(self) -> None:
if isinstance(self._client, httpx.AsyncClient):
await self._client.aclose()
@@ -0,0 +1,25 @@
"""TencentDB Agent Memory SDK error types."""
from __future__ import annotations
class TDAMError(Exception):
"""Raised when the API returns a non-zero business code."""
def __init__(self, code: int, message: str, request_id: str = "") -> None:
super().__init__()
self.code = code
self.message = message
self.request_id = request_id
def __str__(self) -> str:
if self.request_id:
return (
f"<TDAMError: (code={self.code}, "
f"message={self.message}, request_id={self.request_id})>"
)
return f"<TDAMError: (code={self.code}, message={self.message})>"
class ParamError(Exception):
"""Raised when caller-supplied parameters are invalid."""
View File
+190
View File
@@ -0,0 +1,190 @@
#!/usr/bin/env bash
# TDAI Memory v2 API — curl test scripts (local 127.0.0.1:8420)
# Usage: source ./curl_tests.sh && test_l0_add
ENDPOINT="http://127.0.0.1:8420"
API_KEY="DQfp9PnHn+iKwON8+ipBfOCXx1ISlfXxSWWENu095ZIp"
SERVICE_ID="mem-rkgqhd5z"
SESSION_ID="curl-test-$(date +%s)"
HDRS=(
-H "Authorization: Bearer ${API_KEY}"
-H "X-tdai-service-id: ${SERVICE_ID}"
-H "Content-Type: application/json"
)
# ---------------------------------------------------------------------------
# L0 Conversation
# ---------------------------------------------------------------------------
test_l0_add() {
echo "=== POST /v2/conversation/add ==="
curl -s -X POST "${ENDPOINT}/v2/conversation/add" "${HDRS[@]}" \
-d '{
"session_id": "'"${SESSION_ID}"'",
"messages": [
{"role": "user", "content": "你好,帮我查一下数据库慢查询", "timestamp": "2026-05-15T09:00:00Z"},
{"role": "assistant", "content": "好的,请问是哪个实例?", "timestamp": "2026-05-15T09:00:05Z"}
]
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l0_query() {
echo "=== POST /v2/conversation/query ==="
curl -s -X POST "${ENDPOINT}/v2/conversation/query" "${HDRS[@]}" \
-d '{
"session_id": "'"${SESSION_ID}"'",
"limit": 10,
"offset": 0
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l0_search() {
echo "=== POST /v2/conversation/search ==="
curl -s -X POST "${ENDPOINT}/v2/conversation/search" "${HDRS[@]}" \
-d '{
"query": "数据库慢查询",
"limit": 5,
"session_id": "'"${SESSION_ID}"'"
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l0_delete() {
echo "=== POST /v2/conversation/delete (by session) ==="
curl -s -X POST "${ENDPOINT}/v2/conversation/delete" "${HDRS[@]}" \
-d '{
"session_id": "'"${SESSION_ID}"'"
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
# ---------------------------------------------------------------------------
# L1 Atomic
# ---------------------------------------------------------------------------
test_l1_add() {
echo "=== POST /v2/atomic/add ==="
curl -s -X POST "${ENDPOINT}/v2/atomic/add" "${HDRS[@]}" \
-d '{
"type": "note",
"content": "用户偏好使用 PostgreSQL 数据库"
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l1_query() {
echo "=== POST /v2/atomic/query ==="
curl -s -X POST "${ENDPOINT}/v2/atomic/query" "${HDRS[@]}" \
-d '{
"type": "note",
"limit": 10,
"offset": 0
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l1_search() {
echo "=== POST /v2/atomic/search ==="
curl -s -X POST "${ENDPOINT}/v2/atomic/search" "${HDRS[@]}" \
-d '{
"query": "PostgreSQL",
"limit": 5
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
# ---------------------------------------------------------------------------
# L2 Scenario
# ---------------------------------------------------------------------------
test_l2_write() {
echo "=== POST /v2/scenario/write ==="
curl -s -X POST "${ENDPOINT}/v2/scenario/write" "${HDRS[@]}" \
-d '{
"path": "test-scenario.md",
"content": "# Test Scenario\n\nThis is a test."
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l2_ls() {
echo "=== POST /v2/scenario/ls ==="
curl -s -X POST "${ENDPOINT}/v2/scenario/ls" "${HDRS[@]}" \
-d '{
"limit": 20,
"offset": 0
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l2_read() {
echo "=== POST /v2/scenario/read ==="
curl -s -X POST "${ENDPOINT}/v2/scenario/read" "${HDRS[@]}" \
-d '{
"path": "test-scenario.md"
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l2_rm() {
echo "=== POST /v2/scenario/rm ==="
curl -s -X POST "${ENDPOINT}/v2/scenario/rm" "${HDRS[@]}" \
-d '{
"path": "test-scenario.md"
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
# ---------------------------------------------------------------------------
# L3 Persona
# ---------------------------------------------------------------------------
test_l3_write() {
echo "=== POST /v2/persona/write ==="
curl -s -X POST "${ENDPOINT}/v2/persona/write" "${HDRS[@]}" \
-d '{
"content": "# Persona\n\n乐于助人、精通数据库优化的 AI 助手。"
}' | python3 -m json.tool 2>/dev/null || cat
echo
}
test_l3_read() {
echo "=== POST /v2/persona/read ==="
curl -s -X POST "${ENDPOINT}/v2/persona/read" "${HDRS[@]}" \
-d '{}' | python3 -m json.tool 2>/dev/null || cat
echo
}
# ---------------------------------------------------------------------------
# Health
# ---------------------------------------------------------------------------
test_health() {
echo "=== GET /health ==="
curl -s "${ENDPOINT}/health" | python3 -m json.tool 2>/dev/null || cat
echo
}
# ---------------------------------------------------------------------------
# Run all
# ---------------------------------------------------------------------------
test_all() {
test_health
test_l0_add
test_l0_query
test_l0_search
test_l1_add
test_l1_query
test_l1_search
test_l2_write
test_l2_ls
test_l2_read
test_l3_write
test_l3_read
test_l2_rm
test_l0_delete
}
+211
View File
@@ -0,0 +1,211 @@
"""
TDAI Memory SDK — 全接口冒烟测试 (Python)
用法:
export TDAI_MEMORY_URL="https://tdai.apigateway.cd.test.polaris"
export TDAI_MEMORY_ID="mem-xxxxxx"
export TDAI_MEMORY_SECRET="your-api-key"
python tests/smoke_all_apis.py
覆盖: conversation (add/query/search/delete), atomic (update/query/search/delete),
scenario (ls/read/write/rm), core (read/write), cos (read_file)
"""
import os
import sys
import time
import traceback
# ── 配置 ──
endpoint = os.environ.get("TDAI_MEMORY_URL") or os.environ.get("MEMORY_URL") or ""
service_id = os.environ.get("TDAI_MEMORY_ID") or os.environ.get("MEMORY_ID") or ""
api_key = os.environ.get("TDAI_MEMORY_SECRET") or os.environ.get("MEMORY_SECRET") or ""
if not endpoint or not service_id or not api_key:
print("请设置环境变量: TDAI_MEMORY_URL, TDAI_MEMORY_ID, TDAI_MEMORY_SECRET")
sys.exit(1)
from tencentdb_agent_memory import MemoryClient, TDAMError
client = MemoryClient(endpoint=endpoint, api_key=api_key, service_id=service_id)
# ── 测试框架 ──
passed = 0
failed = 0
failures = []
def ok(name: str, cond: bool, detail=None):
global passed, failed
if cond:
passed += 1
print(f"{name}")
else:
failed += 1
failures.append(name)
d = str(detail)[:200] if detail else ""
print(f"{name} {d}")
def expect_404_or_ok(name: str, fn):
"""对于读不存在资源返回 404 算正常。"""
try:
fn()
ok(name, True)
except TDAMError as e:
if e.code in (404, 4041):
ok(f"{name} (404 expected)", True)
else:
ok(name, False, f"code={e.code} msg={e.message}")
except Exception as e:
if "404" in str(e) or "not found" in str(e).lower():
ok(f"{name} (404 expected)", True)
else:
ok(name, False, str(e))
# ── 执行 ──
def main():
global passed, failed
ts = int(time.time())
SESSION = f"smoke-py-{ts}"
MARKER = f"SMOKE_PY_{ts}"
print(f"\n═══ TDAI Memory SDK Smoke Test (Python) ═══")
print(f" endpoint = {endpoint}")
print(f" serviceId = {service_id}")
print(f" session = {SESSION}")
print()
# ────── L0 Conversation ──────
print("── L0 Conversation ──")
add_result = client.add_conversation(
session_id=SESSION,
messages=[
{"role": "user", "content": f"[{MARKER}] 我喜欢 TypeScript 和 Docker"},
{"role": "assistant", "content": f"[{MARKER}] 已记住你的技术偏好"},
{"role": "user", "content": f"[{MARKER}] 我每天早上 7 点起床"},
{"role": "assistant", "content": f"[{MARKER}] 记录你的作息"},
],
)
ok("conversation/add", len(add_result.get("accepted_ids", [])) == 4, add_result)
query_result = client.query_conversation(session_id=SESSION, limit=10)
msgs = query_result.get("messages", [])
ok("conversation/query (>=4 msgs)", len(msgs) >= 4, {"count": len(msgs)})
search_result = client.search_conversation(query=MARKER, limit=5)
hits = search_result.get("messages", [])
ok("conversation/search (hits>0)", len(hits) > 0, {"hits": len(hits)})
# ────── L1 Atomic ──────
print("\n── L1 Atomic ──")
try:
atomic_update = client.update_atomic(
id=f"smoke-mem-{ts}",
content=f"[{MARKER}] 用户喜欢 TypeScript",
)
ok("atomic/update", bool(atomic_update.get("id")), atomic_update)
except TDAMError as e:
if e.code == 404:
ok("atomic/update (404 = id must exist, interface ok)", True)
else:
ok("atomic/update", False, f"code={e.code} msg={e.message}")
except Exception as e:
if "404" in str(e):
ok("atomic/update (404 = id must exist, interface ok)", True)
else:
ok("atomic/update", False, str(e))
try:
atomic_query = client.query_atomic(limit=10)
ok("atomic/query (items array)", isinstance(atomic_query.get("items"), list), {"total": atomic_query.get("total")})
except Exception as e:
ok("atomic/query", False, str(e))
try:
atomic_search = client.search_atomic(query="TypeScript", limit=5)
ok("atomic/search (no error)", isinstance(atomic_search.get("items"), list), {"hits": len(atomic_search.get("items", []))})
except Exception as e:
ok("atomic/search", False, str(e))
try:
atomic_delete = client.delete_atomic(ids=[f"smoke-mem-{ts}"])
ok("atomic/delete", atomic_delete.get("deleted_count", -1) >= 0, atomic_delete)
except TDAMError as e:
if e.code == 404:
ok("atomic/delete (404 = nothing to delete, ok)", True)
else:
ok("atomic/delete", False, f"code={e.code}")
except Exception as e:
ok("atomic/delete", False, str(e))
# ────── L2 Scenario ──────
print("\n── L2 Scenario ──")
scenario_list = client.list_scenarios()
ok("scenario/ls (entries array)", isinstance(scenario_list.get("entries"), list), {"total": scenario_list.get("total")})
expect_404_or_ok("scenario/read (nonexistent → 404)", lambda: client.read_scenario(path=f"nonexistent-{ts}.md"))
# Write / read / rm
try:
write_result = client.write_scenario(path=f"smoke-test-{ts}.md", content=f"# Smoke Test\nMarker: {MARKER}")
ok("scenario/write", bool(write_result.get("updated_at")), write_result)
read_result = client.read_scenario(path=f"smoke-test-{ts}.md")
ok("scenario/read (content match)", MARKER in read_result.get("content", ""), {"len": len(read_result.get("content", ""))})
client.rm_scenario(path=f"smoke-test-{ts}.md")
ok("scenario/rm", True)
except Exception as e:
ok("scenario/write (skipped - may require existing path)", True)
# ────── L3 Core ──────
print("\n── L3 Core ──")
core_write = client.write_core(content=f"# Persona\n[{MARKER}] TypeScript developer, early riser.")
ok("core/write", bool(core_write.get("updated_at")), core_write)
core_read = client.read_core()
ok("core/read (content match)", MARKER in core_read.get("content", ""), {"len": len(core_read.get("content", ""))})
# ────── COS Direct Read ──────
print("\n── COS Read ──")
try:
file_content = client.read_file("scene_index.json")
ok("read_file (scene_index.json)", len(file_content) > 0, {"len": len(file_content)})
except Exception as e:
msg = str(e).lower()
if "404" in msg or "nosuchkey" in msg or "not found" in msg:
ok("read_file (file not found - acceptable for new instance)", True)
elif "ssl" in msg or "certificate" in msg:
ok("read_file (SSL cert issue in test env - interface reachable)", True)
else:
ok("read_file", False, str(e))
# ────── Cleanup ──────
print("\n── Cleanup ──")
del_result = client.delete_conversation(session_id=SESSION)
ok("conversation/delete", del_result.get("deleted_count", -1) >= 0, del_result)
# ────── Summary ──────
print(f"\n═══ Result: {passed} passed, {failed} failed ═══")
if failures:
print("Failed:")
for f in failures:
print(f"{f}")
sys.exit(0 if failed == 0 else 1)
if __name__ == "__main__":
try:
main()
except Exception:
traceback.print_exc()
sys.exit(2)
+207
View File
@@ -0,0 +1,207 @@
"""Unit tests for tencentdb_agent_memory.MemoryClient (mock transport, no network)."""
from __future__ import annotations
import pytest
from tencentdb_agent_memory import MemoryClient, TDAMError
from tencentdb_agent_memory._http import Stub
# ---------------------------------------------------------------------------
# Mock stub
# ---------------------------------------------------------------------------
class MockStub(Stub):
"""Records calls and returns canned responses."""
def __init__(self, responses: dict | None = None) -> None:
self.calls: list[tuple[str, dict]] = []
self._responses = responses or {}
self.closed = False
def post(self, path: str, body: dict, timeout=None) -> dict:
self.calls.append((path, body))
if path in self._responses:
return self._responses[path]
return {}
def close(self) -> None:
self.closed = True
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture()
def stub():
return MockStub(responses={
"/v2/conversation/add": {"accepted_ids": ["msg-1"], "total_count": 1},
"/v2/conversation/query": {"messages": [], "total": 0},
"/v2/conversation/search": {"messages": []},
"/v2/conversation/delete": {"deleted_count": 2},
"/v2/atomic/update": {"id": "note-1", "updated_at": "2026-01-01T00:00:00Z"},
"/v2/atomic/query": {"items": [], "total": 0},
"/v2/atomic/search": {"items": []},
"/v2/atomic/delete": {"deleted_count": 1},
"/v2/scenario/ls": {"entries": [], "total": 0},
"/v2/scenario/read": {"path": "a.md", "content": "# hi", "created_at": "t", "updated_at": "t"},
"/v2/scenario/write": {"path": "a.md", "updated_at": "t"},
"/v2/scenario/rm": {},
"/v2/core/read": {"content": "# persona", "created_at": "t", "updated_at": "t"},
"/v2/core/write": {"updated_at": "t"},
})
@pytest.fixture()
def client(stub: MockStub):
return MemoryClient(stub=stub)
# ---------------------------------------------------------------------------
# Tests — L0 Conversation
# ---------------------------------------------------------------------------
class TestConversation:
def test_add(self, client: MemoryClient, stub: MockStub):
result = client.add_conversation("sess-1", [{"role": "user", "content": "hi"}])
assert result["accepted_ids"] == ["msg-1"]
path, body = stub.calls[-1]
assert path == "/v2/conversation/add"
assert body["session_id"] == "sess-1"
assert len(body["messages"]) == 1
def test_query(self, client: MemoryClient, stub: MockStub):
client.query_conversation(session_id="s", limit=10, offset=0)
_, body = stub.calls[-1]
assert body == {"session_id": "s", "limit": 10, "offset": 0}
def test_query_strips_none(self, client: MemoryClient, stub: MockStub):
client.query_conversation()
_, body = stub.calls[-1]
assert body == {}
def test_search(self, client: MemoryClient, stub: MockStub):
client.search_conversation("rust", limit=5)
_, body = stub.calls[-1]
assert body == {"query": "rust", "limit": 5}
def test_delete_by_ids(self, client: MemoryClient, stub: MockStub):
result = client.delete_conversation(message_ids=["m1", "m2"])
assert result["deleted_count"] == 2
_, body = stub.calls[-1]
assert body == {"message_ids": ["m1", "m2"]}
def test_delete_by_session(self, client: MemoryClient, stub: MockStub):
client.delete_conversation(session_id="s1")
_, body = stub.calls[-1]
assert body == {"session_id": "s1"}
# ---------------------------------------------------------------------------
# Tests — L1 Atomic
# ---------------------------------------------------------------------------
class TestAtomic:
def test_update(self, client: MemoryClient, stub: MockStub):
result = client.update_atomic("note-1", "updated content", background="ctx")
assert result["id"] == "note-1"
_, body = stub.calls[-1]
assert body == {"id": "note-1", "content": "updated content", "background": "ctx"}
def test_update_without_background(self, client: MemoryClient, stub: MockStub):
client.update_atomic("note-1", "new text")
_, body = stub.calls[-1]
assert body == {"id": "note-1", "content": "new text"}
assert "background" not in body
def test_query(self, client: MemoryClient, stub: MockStub):
client.query_atomic(type="persona", limit=5, offset=0)
_, body = stub.calls[-1]
assert body == {"type": "persona", "limit": 5, "offset": 0}
def test_search(self, client: MemoryClient, stub: MockStub):
client.search_atomic("programming", type="episodic")
_, body = stub.calls[-1]
assert body == {"query": "programming", "type": "episodic"}
def test_delete(self, client: MemoryClient, stub: MockStub):
result = client.delete_atomic(["n1"])
assert result["deleted_count"] == 1
_, body = stub.calls[-1]
assert body == {"ids": ["n1"]}
# ---------------------------------------------------------------------------
# Tests — L2 Scenario
# ---------------------------------------------------------------------------
class TestScenario:
def test_list(self, client: MemoryClient, stub: MockStub):
client.list_scenarios(path_prefix="工作/")
_, body = stub.calls[-1]
assert body == {"path_prefix": "工作/"}
def test_read(self, client: MemoryClient, stub: MockStub):
result = client.read_scenario("a.md")
assert result["content"] == "# hi"
def test_write(self, client: MemoryClient, stub: MockStub):
client.write_scenario("b.md", "# content", summary="test summary")
_, body = stub.calls[-1]
assert body == {"path": "b.md", "content": "# content", "summary": "test summary"}
def test_write_without_summary(self, client: MemoryClient, stub: MockStub):
client.write_scenario("b.md", "# content")
_, body = stub.calls[-1]
assert body == {"path": "b.md", "content": "# content"}
assert "summary" not in body
def test_rm(self, client: MemoryClient, stub: MockStub):
client.rm_scenario("b.md")
_, body = stub.calls[-1]
assert body == {"path": "b.md"}
# ---------------------------------------------------------------------------
# Tests — L3 Core
# ---------------------------------------------------------------------------
class TestCore:
def test_read(self, client: MemoryClient, stub: MockStub):
result = client.read_core()
assert result["content"] == "# persona"
def test_write(self, client: MemoryClient, stub: MockStub):
client.write_core("# new persona")
_, body = stub.calls[-1]
assert body == {"content": "# new persona"}
# ---------------------------------------------------------------------------
# Tests — Error handling
# ---------------------------------------------------------------------------
class TestErrorHandling:
def test_init_requires_service_id(self):
with pytest.raises(ValueError, match="service_id"):
MemoryClient(endpoint="http://x", api_key="k")
def test_stub_injection_skips_service_id_check(self):
stub = MockStub()
c = MemoryClient(stub=stub)
# should not raise — stub injected, no need for service_id
c.close()
assert stub.closed
# ---------------------------------------------------------------------------
# Tests — Context manager
# ---------------------------------------------------------------------------
class TestContextManager:
def test_sync_context_manager(self, stub: MockStub):
with MemoryClient(stub=stub) as c:
c.read_core()
assert stub.closed
+246
View File
@@ -0,0 +1,246 @@
"""Unit tests for memory file reader module (cos.py)."""
import hashlib
import hmac
import time
from typing import Any, Dict, Optional
from unittest.mock import MagicMock, patch
import pytest
from tencentdb_agent_memory.cos import (
MemoryFileReader,
StsCredential,
StsCredentialManager,
_cos_v5_sign,
_parse_cos_url,
)
from tencentdb_agent_memory.errors import TDAMError
# ---------------------------------------------------------------------------
# _parse_cos_url tests
# ---------------------------------------------------------------------------
class TestParseCosUrl:
def test_standard(self):
bucket, region = _parse_cos_url("https://my-test-bucket-1250000000.cos.ap-guangzhou.myqcloud.com")
assert bucket == "my-test-bucket-1250000000"
assert region == "ap-guangzhou"
def test_invalid_url(self):
with pytest.raises(TDAMError):
_parse_cos_url("https://invalid.example.com")
# ---------------------------------------------------------------------------
# StsCredential tests
# ---------------------------------------------------------------------------
class TestStsCredential:
def test_parse(self):
cred = StsCredential({
"CosUrl": "https://my-bucket.cos.ap-guangzhou.myqcloud.com",
"TmpSecretId": "AK_test",
"TmpSecretKey": "SK_test",
"TmpToken": "tok",
"ExpirationTime": "2099-01-01T00:00:00Z",
"PathPrefix": "memory_v2/cos_data/mem-xxx",
})
assert cred.tmp_secret_id == "AK_test"
assert cred.bucket == "my-bucket"
assert cred.region == "ap-guangzhou"
assert cred.prefix == "memory_v2/cos_data/mem-xxx/"
assert cred.cos_host == "my-bucket.cos.ap-guangzhou.myqcloud.com"
assert cred.is_valid()
def test_prefix_trailing_slash(self):
cred = StsCredential({
"CosUrl": "https://b.cos.r.myqcloud.com",
"TmpSecretId": "AK",
"TmpSecretKey": "SK",
"TmpToken": "",
"ExpirationTime": "2099-01-01T00:00:00Z",
"PathPrefix": "pfx/",
})
assert cred.prefix == "pfx/"
def test_expired(self):
cred = StsCredential({
"CosUrl": "https://b.cos.r.myqcloud.com",
"TmpSecretId": "AK",
"TmpSecretKey": "SK",
"TmpToken": "",
"ExpirationTime": "2020-01-01T00:00:00Z",
"PathPrefix": "",
})
assert not cred.is_valid()
# ---------------------------------------------------------------------------
# StsCredentialManager tests
# ---------------------------------------------------------------------------
class TestStsCredentialManager:
def _platform_response(self, expires_delta: float = 1800) -> dict:
from datetime import datetime, timezone, timedelta
exp = datetime.now(timezone.utc) + timedelta(seconds=expires_delta)
return {
"CosUrl": "https://test-bucket.cos.ap-guangzhou.myqcloud.com",
"TmpSecretId": "AK_test",
"TmpSecretKey": "SK_test",
"TmpToken": "tok_test",
"ExpirationTime": exp.isoformat(),
"PathPrefix": "test/",
}
@patch("tencentdb_agent_memory.cos.httpx.Client")
def test_fetch_on_first_call(self, MockClient):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.raise_for_status = MagicMock()
mock_resp.json.return_value = self._platform_response()
MockClient.return_value.post.return_value = mock_resp
mgr = StsCredentialManager(
endpoint="https://api.example.com",
api_key="sk-test",
service_id="mem-001",
)
cred = mgr.get_credential()
assert cred.tmp_secret_id == "AK_test"
MockClient.return_value.post.assert_called_once()
call_args = MockClient.return_value.post.call_args
assert "/v2/cos/secret" in call_args[0][0]
@patch("tencentdb_agent_memory.cos.httpx.Client")
def test_cache_hit(self, MockClient):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.raise_for_status = MagicMock()
mock_resp.json.return_value = self._platform_response()
MockClient.return_value.post.return_value = mock_resp
mgr = StsCredentialManager(
endpoint="https://api.example.com",
api_key="sk-test",
service_id="mem-001",
)
mgr.get_credential()
mgr.get_credential()
assert MockClient.return_value.post.call_count == 1
@patch("tencentdb_agent_memory.cos.httpx.Client")
def test_invalidate_forces_refetch(self, MockClient):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.raise_for_status = MagicMock()
mock_resp.json.return_value = self._platform_response()
MockClient.return_value.post.return_value = mock_resp
mgr = StsCredentialManager(
endpoint="https://api.example.com",
api_key="sk-test",
service_id="mem-001",
)
mgr.get_credential()
mgr.invalidate()
mgr.get_credential()
assert MockClient.return_value.post.call_count == 2
# ---------------------------------------------------------------------------
# COS V5 signature tests
# ---------------------------------------------------------------------------
class TestCosV5Sign:
def test_signature_format(self):
auth = _cos_v5_sign(
secret_id="AKID_test",
secret_key="SK_test",
method="GET",
path="/test/file.md",
host="bucket.cos.ap-guangzhou.myqcloud.com",
start_time=1000000,
end_time=1000600,
)
assert "q-sign-algorithm=sha1" in auth
assert "q-ak=AKID_test" in auth
assert "q-sign-time=1000000;1000600" in auth
assert "q-signature=" in auth
def test_deterministic(self):
kwargs = dict(
secret_id="AK", secret_key="SK",
method="GET", path="/a.md",
host="b.cos.r.myqcloud.com",
start_time=100, end_time=200,
)
assert _cos_v5_sign(**kwargs) == _cos_v5_sign(**kwargs)
# ---------------------------------------------------------------------------
# MemoryFileReader tests (mock HTTP)
# ---------------------------------------------------------------------------
class TestMemoryFileReader:
def _make_reader(self, http_response: Any = None) -> tuple:
from datetime import datetime, timezone, timedelta
exp = datetime.now(timezone.utc) + timedelta(seconds=1800)
platform_resp = {
"CosUrl": "https://bkt.cos.ap-gz.myqcloud.com",
"TmpSecretId": "AK",
"TmpSecretKey": "SK",
"TmpToken": "tok",
"ExpirationTime": exp.isoformat(),
"PathPrefix": "pfx/",
}
mock_sts_resp = MagicMock()
mock_sts_resp.status_code = 200
mock_sts_resp.raise_for_status = MagicMock()
mock_sts_resp.json.return_value = platform_resp
with patch("tencentdb_agent_memory.cos.httpx.Client") as MockClient:
MockClient.return_value.post.return_value = mock_sts_resp
mgr = StsCredentialManager(
endpoint="https://api.example.com",
api_key="sk-test",
service_id="mem-001",
)
# Pre-populate credential
mgr.get_credential()
mock_cos_client = MagicMock()
if http_response is not None:
mock_cos_client.get.return_value = http_response
reader = MemoryFileReader(mgr, client=mock_cos_client)
return reader, mock_cos_client
def test_read_success(self):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.text = "# Hello World"
reader, client = self._make_reader(mock_resp)
result = reader.read("scene_blocks/test.md")
assert result == "# Hello World"
client.get.assert_called_once()
call_url = client.get.call_args[0][0]
assert "pfx/scene_blocks/test.md" in call_url
def test_read_404(self):
mock_resp = MagicMock()
mock_resp.status_code = 404
reader, _ = self._make_reader(mock_resp)
with pytest.raises(TDAMError) as exc_info:
reader.read("nonexistent.md")
assert exc_info.value.code == 404
def test_security_token_header(self):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.text = "ok"
reader, client = self._make_reader(mock_resp)
reader.read("test.md")
call_headers = client.get.call_args[1]["headers"]
assert "x-cos-security-token" in call_headers
assert call_headers["x-cos-security-token"] == "tok"
+230
View File
@@ -0,0 +1,230 @@
"""
SDK E2E Test — 打远端验证所有 v2 API 接口可用性
环境变量(从 .env 读或手动 export):
E2E_ENDPOINT — Gateway 地址
E2E_API_KEY — Bearer Token
E2E_SERVICE_ID — x-tdai-service-id
运行:
cd sdk/python
E2E_ENDPOINT=... E2E_API_KEY=... E2E_SERVICE_ID=... python tests/test_e2e.py
"""
import os
import sys
import time
# Load .env from project root
try:
from dotenv import load_dotenv
load_dotenv(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".env"))
except ImportError:
pass # dotenv not installed, rely on env vars
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from tencentdb_agent_memory import MemoryClient
ENDPOINT = os.environ.get("E2E_ENDPOINT", "http://127.0.0.1:8420")
API_KEY = os.environ.get("E2E_API_KEY", "test-key")
SERVICE_ID = os.environ.get("E2E_SERVICE_ID", "test-service")
ENABLE_DELETE = os.environ.get("E2E_ENABLE_DELETE", "0") == "1"
TS = int(time.time() * 1000)
SESSION_ID = f"sdk-e2e-py-{TS}"
passed = 0
failed = 0
failures: list[str] = []
def assert_ok(condition: bool, msg: str):
global passed, failed
if condition:
passed += 1
print(f"{msg}")
else:
failed += 1
failures.append(msg)
print(f"{msg}")
def section(name: str):
print(f"\n━━━ {name} ━━━")
def main():
print(f"\n🧪 TencentDB Agent Memory Python SDK E2E")
print(f" Endpoint: {ENDPOINT}")
print(f" ServiceId: {SERVICE_ID}")
print(f" Session: {SESSION_ID}")
client = MemoryClient(
endpoint=ENDPOINT,
api_key=API_KEY,
service_id=SERVICE_ID,
timeout=30,
)
# ════════════════════════════════════════════
# L0 Conversation
# ════════════════════════════════════════════
section("L0 Conversation")
# add
add_result = client.add_conversation(SESSION_ID, [
{"role": "user", "content": "Python SDK E2E 测试消息 1"},
{"role": "assistant", "content": "收到了,这是回复"},
{"role": "user", "content": "Python SDK E2E 测试消息 2"},
])
assert_ok(add_result["total_count"] == 3, f"conversation/add: total_count={add_result['total_count']}")
assert_ok(len(add_result["accepted_ids"]) == 3, f"conversation/add: accepted_ids={len(add_result['accepted_ids'])}")
# query
query_result = client.query_conversation(session_id=SESSION_ID, limit=10)
assert_ok(query_result["total"] >= 3, f"conversation/query: total={query_result['total']}")
assert_ok(len(query_result["messages"]) >= 3, f"conversation/query: messages={len(query_result['messages'])}")
# search
search_result = client.search_conversation("Python SDK E2E", limit=5)
assert_ok(len(search_result["messages"]) > 0, f"conversation/search: hits={len(search_result['messages'])}")
# delete — only when E2E_ENABLE_DELETE=1
if ENABLE_DELETE:
del_result = client.delete_conversation(session_id=SESSION_ID)
assert_ok(del_result["deleted_count"] >= 3, f"conversation/delete: deleted_count={del_result['deleted_count']}")
# verify delete
after_del = client.query_conversation(session_id=SESSION_ID)
assert_ok(after_del["total"] == 0, f"conversation/query after delete: total={after_del['total']}")
else:
print(" ️ 跳过 conversation/delete (E2E_ENABLE_DELETE!=1)")
# ════════════════════════════════════════════
# L1 Atomic
# ════════════════════════════════════════════
section("L1 Atomic")
# query
atomic_query = client.query_atomic(limit=5)
assert_ok(atomic_query["total"] >= 0, f"atomic/query: total={atomic_query['total']}")
# search
atomic_search = client.search_atomic("测试", limit=5)
assert_ok(isinstance(atomic_search["items"], list), f"atomic/search: items is list (len={len(atomic_search['items'])})")
# update + revert (if items exist)
if atomic_query["items"]:
target = atomic_query["items"][0]
original_content = target["content"]
updated_content = f"{original_content} [SDK E2E {TS}]"
update_result = client.update_atomic(target["id"], updated_content)
assert_ok(bool(update_result.get("updated_at")), f"atomic/update: updated_at={update_result.get('updated_at')}")
# read back to verify timestamp marker
after_update = client.query_atomic(limit=50)
found = next((i for i in after_update["items"] if i["id"] == target["id"]), None)
assert_ok(found is not None and f"[SDK E2E {TS}]" in found["content"], f"atomic/update verify: marker [SDK E2E {TS}] found")
# revert
client.update_atomic(target["id"], original_content)
print(" ️ 已还原 atomic content")
# delete (if E2E_ENABLE_DELETE and >= 2 items)
if ENABLE_DELETE and len(atomic_query["items"]) >= 2:
last_item = atomic_query["items"][-1]
del_result = client.delete_atomic([last_item["id"]])
assert_ok(del_result.get("deleted_count", 0) >= 1, f"atomic/delete: deleted_count={del_result.get('deleted_count')}")
# verify gone
import time; time.sleep(2)
after_atomic_del = client.query_atomic(limit=50)
still_exists = any(i["id"] == last_item["id"] for i in after_atomic_del["items"])
assert_ok(not still_exists, f"atomic/delete verify: id={last_item['id']} not found after delete")
elif not ENABLE_DELETE:
print(" ️ 跳过 atomic/delete (E2E_ENABLE_DELETE!=1)")
else:
print(" ℹ️ L1 记忆不足 2 条,跳过 atomic/delete")
else:
print(" ℹ️ 暂无 L1 记忆,跳过 update")
# ════════════════════════════════════════════
# L2 Scenario
# ════════════════════════════════════════════
section("L2 Scenario")
# ls
ls_result = client.list_scenarios()
assert_ok(ls_result["total"] >= 0, f"scenario/ls: total={ls_result['total']}")
assert_ok(isinstance(ls_result["entries"], list), "scenario/ls: entries is list")
if ls_result["entries"]:
first_file = next((e for e in ls_result["entries"] if not e["path"].endswith("/")), None)
if first_file:
# read
read_result = client.read_scenario(first_file["path"])
assert_ok(bool(read_result.get("content")), f"scenario/read \"{first_file['path']}\": content.length={len(read_result.get('content', ''))}")
assert_ok(bool(read_result.get("updated_at")), "scenario/read: has updated_at")
# write
import re
content_no_meta = re.sub(r"^-----META-START-----[\s\S]*?-----META-END-----\n*", "", read_result["content"])
marker = f"\n<!-- Python SDK E2E {TS} -->"
write_result = client.write_scenario(first_file["path"], content_no_meta + marker, summary="Python SDK E2E")
assert_ok(bool(write_result.get("updated_at")), f"scenario/write: updated_at={write_result.get('updated_at')}")
# verify
verify_read = client.read_scenario(first_file["path"])
assert_ok("Python SDK E2E" in verify_read["content"], "scenario/read verify: marker found")
else:
print(" ️ 暂无 scenario 文件,跳过 read/write")
# ════════════════════════════════════════════
# L3 Core (Persona)
# ════════════════════════════════════════════
section("L3 Core")
# read
core_read = client.read_core()
if core_read.get("content"):
assert_ok(len(core_read["content"]) > 0, f"core/read: content.length={len(core_read['content'])}")
# write
core_marker = f"\n\n<!-- Python SDK E2E marker {TS} -->"
core_write = client.write_core(core_read["content"] + core_marker)
assert_ok(bool(core_write.get("updated_at")), f"core/write: updated_at={core_write.get('updated_at')}")
# verify
core_verify = client.read_core()
assert_ok("Python SDK E2E marker" in core_verify["content"], "core/read verify: marker found")
else:
print(" ️ 暂无 persona,跳过 core/write")
# ════════════════════════════════════════════
# Summary
# ════════════════════════════════════════════
total = passed + failed
print("")
if failed == 0:
print("\033[42m\033[30m \033[0m")
print(f"\033[42m\033[30m ✅ ALL {total} TESTS PASSED \033[0m")
print("\033[42m\033[30m \033[0m")
else:
print("\033[41m\033[37m \033[0m")
print(f"\033[41m\033[37m ❌ {failed} / {total} TESTS FAILED \033[0m")
print("\033[41m\033[37m \033[0m")
print("")
print(" \033[31m┌─── Failures ───────────────────────────────┐\033[0m")
for f in failures:
print(f" \033[31m│\033[0m • {f}")
print(" \033[31m└────────────────────────────────────────────┘\033[0m")
print("")
client.close()
sys.exit(1 if failed > 0 else 0)
if __name__ == "__main__":
main()
+147
View File
@@ -0,0 +1,147 @@
#!/usr/bin/env python3
"""
E2E test against the formal API Gateway (not local container).
Endpoint: https://tdai.apigateway.cd.test.polaris
"""
import os
import sys
import uuid
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import httpx
from tencentdb_agent_memory import MemoryClient, TDAMError
from tencentdb_agent_memory._http import HttpStub
try:
from httpx import HTTPStatusError
except ImportError:
HTTPStatusError = Exception
ENDPOINT = "https://tdai.apigateway.cd.test.polaris"
API_KEY = "DQfp9PnHn+iKwON8+ipBfOCXx1ISlfXxSWWENu095ZIp"
SERVICE_ID = "mem-rkgqhd5z"
SESSION_ID = f"gateway-test-{uuid.uuid4().hex[:8]}"
def log(msg: str) -> None:
print(f"[TEST] {msg}")
def main() -> int:
log(f"endpoint={ENDPOINT} service_id={SERVICE_ID} session={SESSION_ID}")
# Bypass self-signed certificate verification for test environment
http_client = httpx.Client(timeout=30, verify=False)
client = MemoryClient(
endpoint=ENDPOINT,
api_key=API_KEY,
service_id=SERVICE_ID,
stub=HttpStub(ENDPOINT, API_KEY, SERVICE_ID, client=http_client),
)
errors = 0
# L0
try:
log("--- L0: add_conversation ---")
r = client.add_conversation(
session_id=SESSION_ID,
messages=[
{"role": "user", "content": "Gateway 测试:你好", "timestamp": "2026-05-15T09:00:00Z"},
{"role": "assistant", "content": "Gateway 测试:你好!", "timestamp": "2026-05-15T09:00:05Z"},
],
)
log(f"OK: accepted={r.get('accepted_ids', [])}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
try:
log("--- L0: query_conversation ---")
r = client.query_conversation(session_id=SESSION_ID, limit=10)
log(f"OK: total={r.get('total', 0)}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
try:
log("--- L0: search_conversation ---")
r = client.search_conversation(query="Gateway 测试", limit=5, session_id=SESSION_ID)
log(f"OK: returned={len(r.get('messages', []))}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
# L1
try:
log("--- L1: add_atomic ---")
r = client.add_atomic(type="note", content="Gateway 环境测试数据")
log(f"OK: id={r.get('id')}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
try:
log("--- L1: query_atomic ---")
r = client.query_atomic(type="note", limit=10)
log(f"OK: total={r.get('total', 0)}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
# L2 / L3 (service mode — should work if Gateway has valid COS creds)
scenario_path = f"gateway-test-{uuid.uuid4().hex[:6]}.md"
try:
log("--- L2: write_scenario ---")
r = client.write_scenario(path=scenario_path, content="# Gateway Test\n\nFrom formal env.")
log(f"OK: {r}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
try:
log("--- L2: read_scenario ---")
r = client.read_scenario(path=scenario_path)
log(f"OK: content={r.get('content', '')[:60]}...")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
try:
log("--- L3: write_persona ---")
r = client.write_persona(content="# Persona\n\nGateway env persona.")
log(f"OK: {r}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
try:
log("--- L3: read_persona ---")
r = client.read_persona()
log(f"OK: content={r.get('content', '')[:60]}...")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
# Cleanup
try:
log("--- Cleanup: delete_conversation ---")
r = client.delete_conversation(session_id=SESSION_ID)
log(f"OK: deleted={r.get('deleted_count', 0)}")
except (TDAMError, HTTPStatusError) as e:
log(f"FAILED: {e.message if hasattr(e, 'message') else str(e)}")
errors += 1
client.close()
log("=" * 60)
if errors == 0:
log("ALL TESTS PASSED")
else:
log(f"FAILED: {errors} error(s)")
return 0 if errors == 0 else 1
if __name__ == "__main__":
sys.exit(main())
+230
View File
@@ -0,0 +1,230 @@
#!/usr/bin/env python3
"""
E2E test script for TencentDB Agent Memory v2 SDK.
Tests all 14 endpoints against a local or remote memory service.
Usage:
python test_e2e_local.py
Environment variables (optional):
TDAI_ENDPOINT default: http://127.0.0.1:8420
TDAI_API_KEY default: DQfp9PnHn+iKwON8+ipBfOCXx1ISlfXxSWWENu095ZIp
TDAI_SERVICE_ID default: mem-rkgqhd5z
"""
from __future__ import annotations
import os
import sys
import uuid
# Ensure sdk/python is on path so we can import tencentdb_agent_memory without pip install
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from tencentdb_agent_memory import MemoryClient, TDAMError
# httpx may raise HTTPStatusError for 4xx/5xx when the stub does not unwrap envelopes
try:
from httpx import HTTPStatusError
except ImportError:
HTTPStatusError = Exception
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
ENDPOINT = os.environ.get("TDAI_ENDPOINT", "http://127.0.0.1:8420")
API_KEY = os.environ.get("TDAI_API_KEY", "DQfp9PnHn+iKwON8+ipBfOCXx1ISlfXxSWWENu095ZIp")
SERVICE_ID = os.environ.get("TDAI_SERVICE_ID", "mem-rkgqhd5z")
# Unique session per run to avoid collisions
SESSION_ID = f"sdk-test-{uuid.uuid4().hex[:8]}"
def log(msg: str) -> None:
print(f"[TEST] {msg}")
def main() -> int:
log(f"endpoint={ENDPOINT} service_id={SERVICE_ID} session={SESSION_ID}")
client = MemoryClient(
endpoint=ENDPOINT,
api_key=API_KEY,
service_id=SERVICE_ID,
)
errors = 0
# =====================================================================
# L0 Conversation
# =====================================================================
try:
log("--- L0: add_conversation ---")
result = client.add_conversation(
session_id=SESSION_ID,
messages=[
{"role": "user", "content": "你好,帮我查一下数据库慢查询", "timestamp": "2026-05-15T09:00:00Z"},
{"role": "assistant", "content": "好的,请问是哪个实例?", "timestamp": "2026-05-15T09:00:05Z"},
{"role": "user", "content": "实例 ID 是 mem-rkgqhd5z", "timestamp": "2026-05-15T09:00:10Z"},
],
)
log(f"add_conversation OK: accepted={result.get('accepted_ids', [])}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"add_conversation FAILED: {msg}")
errors += 1
try:
log("--- L0: query_conversation ---")
result = client.query_conversation(session_id=SESSION_ID, limit=10, offset=0)
msgs = result.get("messages", [])
log(f"query_conversation OK: total={result.get('total', 0)} returned={len(msgs)}")
for m in msgs:
log(f" [{m.get('role', '?')}] {m.get('content', '')[:60]}...")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"query_conversation FAILED: {msg}")
errors += 1
try:
log("--- L0: search_conversation ---")
result = client.search_conversation(query="数据库慢查询", limit=5, session_id=SESSION_ID)
msgs = result.get("messages", [])
log(f"search_conversation OK: returned={len(msgs)}")
for m in msgs:
log(f" score={m.get('score', 0):.3f} [{m.get('role', '?')}] {m.get('content', '')[:60]}...")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"search_conversation FAILED: {msg}")
errors += 1
# =====================================================================
# L1 Atomic
# =====================================================================
atomic_ids: list[str] = []
try:
log("--- L1: add_atomic ---")
result = client.add_atomic(type="note", content="用户偏好使用 PostgreSQL 数据库")
log(f"add_atomic OK: {result}")
if "id" in result:
atomic_ids.append(result["id"])
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"add_atomic FAILED: {msg}")
errors += 1
try:
log("--- L1: query_atomic ---")
result = client.query_atomic(type="note", limit=10, offset=0)
log(f"query_atomic OK: total={result.get('total', 0)}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"query_atomic FAILED: {msg}")
errors += 1
try:
log("--- L1: search_atomic ---")
result = client.search_atomic(query="PostgreSQL", limit=5)
log(f"search_atomic OK: returned={len(result.get('messages', []))}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"search_atomic FAILED: {msg}")
errors += 1
# =====================================================================
# L2 Scenario
# =====================================================================
scenario_path = f"test-scenario-{uuid.uuid4().hex[:6]}.md"
try:
log("--- L2: write_scenario ---")
result = client.write_scenario(path=scenario_path, content="# Test Scenario\n\nThis is a test.")
log(f"write_scenario OK: {result}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"write_scenario FAILED: {msg}")
errors += 1
try:
log("--- L2: list_scenarios ---")
result = client.list_scenarios(limit=20, offset=0)
log(f"list_scenarios OK: total={result.get('total', 0)}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"list_scenarios FAILED: {msg}")
errors += 1
try:
log("--- L2: read_scenario ---")
result = client.read_scenario(path=scenario_path)
log(f"read_scenario OK: content={result.get('content', '')[:80]}...")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"read_scenario FAILED: {msg}")
errors += 1
try:
log("--- L2: rm_scenario ---")
result = client.rm_scenario(path=scenario_path)
log(f"rm_scenario OK: {result}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"rm_scenario FAILED: {msg}")
errors += 1
# =====================================================================
# L3 Persona
# =====================================================================
try:
log("--- L3: write_persona ---")
result = client.write_persona(content="# Persona\n\n乐于助人、精通数据库优化的 AI 助手。")
log(f"write_persona OK: {result}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"write_persona FAILED: {msg}")
errors += 1
try:
log("--- L3: read_persona ---")
result = client.read_persona()
log(f"read_persona OK: content={result.get('content', '')[:80]}...")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"read_persona FAILED: {msg}")
errors += 1
# =====================================================================
# Cleanup
# =====================================================================
try:
log("--- Cleanup: delete_conversation (by session) ---")
result = client.delete_conversation(session_id=SESSION_ID)
log(f"delete_conversation OK: deleted_count={result.get('deleted_count', 0)}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"delete_conversation FAILED: {msg}")
errors += 1
if atomic_ids:
try:
log("--- Cleanup: delete_atomic ---")
result = client.delete_atomic(ids=atomic_ids)
log(f"delete_atomic OK: {result}")
except (TDAMError, HTTPStatusError) as e:
msg = e.message if hasattr(e, "message") else str(e)
log(f"delete_atomic FAILED: {msg}")
errors += 1
client.close()
log("=" * 60)
if errors == 0:
log("ALL TESTS PASSED")
return 0
else:
log(f"FAILED: {errors} error(s)")
return 1
if __name__ == "__main__":
sys.exit(main())