mirror of
https://github.com/TencentCloud/TencentDB-Agent-Memory
synced 2026-07-11 04:44:29 +00:00
chrishuan
631 lines
21 KiB
Python
631 lines
21 KiB
Python
"""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})
|
||
|
||
# -- Offload (Ingest + Compact + Query-MMD) ----------------------------
|
||
|
||
def offload_ingest(
|
||
self,
|
||
session_id: str,
|
||
tool_pairs: List[Dict[str, Any]],
|
||
*,
|
||
prompt: Optional[str] = None,
|
||
recent_messages: Optional[List[Dict[str, Any]]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""``POST /v2/offload/ingest`` — 上报工具调用对，触发 L1 异步处理。
|
||
|
||
可 fire-and-forget 使用（忽略返回值）。
|
||
|
||
Parameters
|
||
----------
|
||
session_id : str
|
||
会话 ID。
|
||
tool_pairs : list[dict]
|
||
工具调用对列表，每个元素包含 ``tool_name``、``tool_call_id``、
|
||
``params``、``result``、``timestamp``，可选 ``duration_ms``。
|
||
prompt : str, optional
|
||
最新 user message，用于 L1.5 任务判断。
|
||
recent_messages : list[dict], optional
|
||
近期历史消息列表（``role`` + ``content``），辅助 L1 提取上下文。
|
||
"""
|
||
return self._stub.post(
|
||
f"{_V2}/offload/ingest",
|
||
_strip_none({
|
||
"session_id": session_id,
|
||
"tool_pairs": tool_pairs,
|
||
"prompt": prompt,
|
||
"recent_messages": recent_messages,
|
||
}),
|
||
)
|
||
|
||
def offload_compact(
|
||
self,
|
||
session_id: str,
|
||
messages: List[Dict[str, Any]],
|
||
ratio: float,
|
||
total_tokens: int,
|
||
*,
|
||
context_window: Optional[int] = None,
|
||
message_tokens: Optional[List[int]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""``POST /v2/offload/compact`` — 对 messages 执行服务端上下文压缩。
|
||
|
||
Parameters
|
||
----------
|
||
session_id : str
|
||
会话 ID。
|
||
messages : list[dict]
|
||
当前完整对话消息列表。
|
||
ratio : float
|
||
当前 token 使用比例（已用 / context_window），触发压缩策略判断。
|
||
total_tokens : int
|
||
当前完整上下文的总 token 数（包含 system prompt、tool schemas 等不在
|
||
messages 中的隐性开销）。服务端用于计算 fixed overhead 和校准 token 估算。
|
||
context_window : int, optional
|
||
模型 context window 大小（token 数）。
|
||
message_tokens : list[int], optional
|
||
每条消息对应的 token 数，提供时可跳过服务端估算，提升性能。
|
||
|
||
Returns
|
||
-------
|
||
dict
|
||
``messages``（压缩后消息列表）+ ``report``（压缩报告）。
|
||
"""
|
||
return self._stub.post(
|
||
f"{_V2}/offload/compact",
|
||
_strip_none({
|
||
"session_id": session_id,
|
||
"messages": messages,
|
||
"ratio": ratio,
|
||
"total_tokens": total_tokens,
|
||
"context_window": context_window,
|
||
"message_tokens": message_tokens,
|
||
}),
|
||
)
|
||
|
||
def offload_query_mmd(
|
||
self,
|
||
session_id: str,
|
||
*,
|
||
limit: Optional[int] = None,
|
||
) -> Dict[str, Any]:
|
||
"""``POST /v2/offload/query-mmd`` — 查询 session 的任务流程图（MMD 文件）。
|
||
|
||
Parameters
|
||
----------
|
||
session_id : str
|
||
会话 ID。
|
||
limit : int, optional
|
||
最多返回几个 MMD 文件。``limit=1`` 时走快速路径只返回当前活跃 MMD。
|
||
|
||
Returns
|
||
-------
|
||
dict
|
||
``mmds``（列表，每项含 ``filename``、``content``、``version``）+
|
||
``current_mmd``（当前活跃 MMD 文件名，无则为 ``None``）。
|
||
"""
|
||
return self._stub.post(
|
||
f"{_V2}/offload/query-mmd",
|
||
_strip_none({
|
||
"session_id": session_id,
|
||
"limit": limit,
|
||
}),
|
||
)
|
||
|
||
# -- 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})
|
||
|
||
# -- Offload (Ingest + Compact + Query-MMD) ----------------------------
|
||
|
||
async def offload_ingest(
|
||
self,
|
||
session_id: str,
|
||
tool_pairs: List[Dict[str, Any]],
|
||
*,
|
||
prompt: Optional[str] = None,
|
||
recent_messages: Optional[List[Dict[str, Any]]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""``POST /v2/offload/ingest``（异步）"""
|
||
return await self._stub.post(
|
||
f"{_V2}/offload/ingest",
|
||
_strip_none({
|
||
"session_id": session_id,
|
||
"tool_pairs": tool_pairs,
|
||
"prompt": prompt,
|
||
"recent_messages": recent_messages,
|
||
}),
|
||
)
|
||
|
||
async def offload_compact(
|
||
self,
|
||
session_id: str,
|
||
messages: List[Dict[str, Any]],
|
||
ratio: float,
|
||
total_tokens: int,
|
||
*,
|
||
context_window: Optional[int] = None,
|
||
message_tokens: Optional[List[int]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""``POST /v2/offload/compact``（异步）"""
|
||
return await self._stub.post(
|
||
f"{_V2}/offload/compact",
|
||
_strip_none({
|
||
"session_id": session_id,
|
||
"messages": messages,
|
||
"ratio": ratio,
|
||
"total_tokens": total_tokens,
|
||
"context_window": context_window,
|
||
"message_tokens": message_tokens,
|
||
}),
|
||
)
|
||
|
||
async def offload_query_mmd(
|
||
self,
|
||
session_id: str,
|
||
*,
|
||
limit: Optional[int] = None,
|
||
) -> Dict[str, Any]:
|
||
"""``POST /v2/offload/query-mmd``（异步）"""
|
||
return await self._stub.post(
|
||
f"{_V2}/offload/query-mmd",
|
||
_strip_none({
|
||
"session_id": session_id,
|
||
"limit": limit,
|
||
}),
|
||
)
|
||
|
||
# -- 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()
|