mirror of
https://github.com/TencentCloud/TencentDB-Agent-Memory
synced 2026-07-10 20:34:30 +00:00
feat: release v1.0.0-beta.1
This commit is contained in:
@@ -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
|
||||
```
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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."""
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user