mirror of
https://github.com/TencentCloud/TencentDB-Agent-Memory
synced 2026-07-11 12:54:29 +00:00
feat: release v1.0.0-beta.1
This commit is contained in:
chrishuan
@@ -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."""
|
||||
Reference in New Issue
Block a user