mirror of
https://github.com/TencentCloud/TencentDB-Agent-Memory
synced 2026-07-10 20:34:30 +00:00
chrishuan
231 lines
8.5 KiB
Python
231 lines
8.5 KiB
Python
#!/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())
|