SDK Demo

"""Memory Python SDK 快速上手 Demo —— C 端 AI 助手对话

这个脚本演示如何用 ``tencentdb-agent-memory-sdk-python`` 把一段普通用户与
AI 健康助手"小番茄"的日常对话（约 10 轮、覆盖 2~3 个生活话题）快速写入
Memory 实例的 L0 层，等待几分钟后再回查 L0/L1/L2/L3 四层记忆数据。

> 注：「小番茄」是一个虚构的 C 端 AI 健康助手产品，定位类似大众熟悉的
>    AI 减脂教练 / 饮食记录助手，仅作为示例对话场景使用。
"""

from __future__ import annotations

import argparse
import sys
import time
from datetime import datetime, timedelta, timezone
from typing import Any

from tencentdb_agent_memory import MemoryClient, TDAMError


# ============================================================================
#                         👇 用户需要修改的配置区域 👇
# ============================================================================
# 把下面三个字符串替换为你自己实例的真实值即可。
# ============================================================================

# Memory 数据面网关地址（控制台 → 实例详情 → 连接信息），形如 https://memory.<region>.example.com
ENDPOINT = "请替换为你的 Memory 实例 Endpoint"

# 数据面 API Key（控制台 → 实例详情 → 连接信息 中获取）
API_KEY = "请替换为你的 API Key"

# Memory 实例 ID，形如 mem-xxxxxxxx
SERVICE_ID = "请替换为你的实例 ID"

# 默认 session_id：同一段对话归属的会话标识，重复运行会"追加"到同一会话
DEFAULT_SESSION_ID = "demo-pm-quickstart-001"

# ============================================================================
#                         👆 用户需要修改的配置区域 👆
# ============================================================================


# ---------------------------------------------------------------------------- #
#  1. 构造 C 端用户和"小番茄" AI 健康助手的 10 轮对话
#     （覆盖 3 个生活话题：减脂目标 / 饮食运动偏好 / 使用习惯）
# ---------------------------------------------------------------------------- #
def build_dialogue() -> list[dict]:
    """生成 10 轮日常对话，共 20 条消息（user + assistant 各 10 条）。"""
    base_time = datetime(2026, 5, 25, 20, 0, 0, tzinfo=timezone.utc)

    # 每个元素是一轮对话：{"user": ..., "assistant": ...}，一共 10 轮
    rounds: list[dict] = [
        # —— 话题 1：减脂目标与计划（4 轮）——
        {
            "user": "我现在 78 公斤，身高 175，想在 3 个月内减到 70 公斤，帮我看看靠不靠谱？",
            "assistant": "你好，按当前数据 BMI 约 25.5，属于轻度超重；3 个月减 8 公斤、平均每周 0.65 公斤，属于健康的减脂速度，是可以达成的目标。我可以帮你制定饮食和训练计划。",
        },
        {
            "user": "OK，那就定 3 个月减到 70 公斤这个目标。我每周能锻炼 4 次，每次 1 小时左右。",
            "assistant": "已记下你的目标：3 个月内从 78 公斤减到 70 公斤，每周训练 4 次、每次 60 分钟。建议安排成 2 次力量训练 + 2 次有氧训练，我会帮你跟进进度。",
        },
        {
            "user": "我比较喜欢跑步和动感单车，不太喜欢撸铁，力量训练能不能换成自重训练？",
            "assistant": "完全可以。你的训练计划调整为：每周 2 次有氧（跑步 / 动感单车任选）+ 2 次自重训练（深蹲、俯卧撑、平板支撑这一类），不安排器械力量。",
        },
        {
            "user": "对了，每周一早上帮我称一次体重，记录下来做趋势图。",
            "assistant": "好的，每周一早上 08:00 我会推送称重提醒，你拍照或手动输入后我会自动汇总成周度体重曲线，方便你看趋势。",
        },

        # —— 话题 2：饮食和运动偏好（3 轮）——
        {
            "user": "我对牛奶过敏，乳糖不耐受那种，喝了会拉肚子，以后推荐食谱别给我加牛奶和奶酪。",
            "assistant": "已记下饮食禁忌：牛奶、奶酪等乳制品（乳糖不耐受）。后续给你的食谱会用无乳糖牛奶或豆浆替代，奶酪类食物也不再推荐。",
        },
        {
            "user": "另外我不吃香菜，看到就反胃，外卖推荐的时候帮我自动过滤掉含香菜的菜。",
            "assistant": "收到，已加入忌口清单：香菜。今后给你推荐外卖菜品时会自动过滤包含香菜的选项。",
        },
        {
            "user": "我习惯早上 6:30 出门跑步，晚上不想再做运动了，安排训练时尽量都放上午。",
            "assistant": "明白，你的训练时间偏好已记下：所有训练统一安排在上午（默认起跑时间 06:30），晚上不再排任何运动课程。",
        },

        # —— 话题 3：使用偏好（3 轮）——
        {
            "user": "App 里所有提醒推送，都改成晚上 10 点之前发，再晚就别推了，会影响我睡觉。",
            "assistant": "好的，推送时段已限定为每日 22:00 之前，22:00~次日 06:30 之间不再推送任何消息。",
        },
        {
            "user": "重量单位我习惯用公斤，距离用公里，别给我用英制的磅和英里。",
            "assistant": "已设置度量单位偏好：体重 / 重量统一用公斤 (kg)，距离统一用公里 (km)，不再出现 lb / mile 等英制单位。",
        },
        {
            "user": "你跟我说话别太正式，像朋友一样就行，不用每句话都'您好''请问'，听着累。",
            "assistant": "明白，以后我会用更轻松的口吻和你聊，去掉敬语和客套话，像朋友那样直接给你建议。",
        },
    ]

    # 把每轮 {user, assistant} 展开成扁平的消息列表，并加上时间戳
    messages: list[dict] = []
    for idx, turn in enumerate(rounds):
        ts_user = base_time + timedelta(minutes=idx * 2)            # 模拟用户每 2 分钟说一句
        ts_ai = ts_user + timedelta(seconds=20)                     # AI 在 20s 后回复
        messages.append({
            "role": "user",
            "content": turn["user"],
            "timestamp": ts_user.strftime("%Y-%m-%dT%H:%M:%SZ"),
        })
        messages.append({
            "role": "assistant",
            "content": turn["assistant"],
            "timestamp": ts_ai.strftime("%Y-%m-%dT%H:%M:%SZ"),
        })
    return messages


# ---------------------------------------------------------------------------- #
#  2. 一些控制台输出的小工具，让 stdout 看起来更友好
# ---------------------------------------------------------------------------- #
def section(title: str) -> None:
    bar = "─" * 70
    print(f"\\n{bar}\\n  {title}\\n{bar}")


def info(msg: str) -> None:
    print(f"  · {msg}")


def ok(msg: str) -> None:
    print(f"  ✓ {msg}")


def warn(msg: str) -> None:
    print(f"  ! {msg}")


def fail(msg: str) -> None:
    print(f"  ✗ {msg}")


def safe_call(label: str, fn) -> Any:
    """统一封装 SDK 调用：成功返回数据、失败返回 None 并友好提示，不抛异常。"""
    try:
        result = fn()
        ok(f"{label} 调用成功")
        return result
    except TDAMError as exc:
        warn(f"{label} 业务返回 code={exc.code}：{exc.message}"
             + (f"（request_id={exc.request_id}）" if exc.request_id else ""))
        return None
    except Exception as exc:
        fail(f"{label} 异常：{type(exc).__name__}: {exc}")
        return None


# ---------------------------------------------------------------------------- #
#  3. 把 10 轮对话写入 L0
# ---------------------------------------------------------------------------- #
def step_write_l0(client: MemoryClient, session_id: str) -> int:
    section(f"Step 1/3  写入 L0 原始对话（session_id = {session_id}）")

    messages = build_dialogue()
    info(f"准备写入 {len(messages) // 2} 轮 / 共 {len(messages)} 条消息")
    info("话题覆盖：① 减脂目标与训练计划　② 饮食运动偏好　③ 使用习惯偏好")

    data = safe_call(
        "L0 add_conversation",
        lambda: client.add_conversation(session_id, messages),
    )
    if not data:
        fail("写入失败，后续步骤无法进行。请检查 ENDPOINT / API_KEY / SERVICE_ID 配置。")
        sys.exit(1)

    accepted = data.get("accepted_ids") or []
    total = data.get("total_count")
    info(f"服务端受理：total_count = {total}")
    info(f"前 3 个 message_id：{accepted[:3]}")
    return total or len(accepted)


# ---------------------------------------------------------------------------- #
#  4. 等待内核异步沉淀（L0 → L1 → L2/L3）
# ---------------------------------------------------------------------------- #
def step_wait(seconds: int) -> None:
    section(f"Step 2/3  等待内核异步沉淀（{seconds} 秒）")

    if seconds <= 0:
        info("已选择不等待，直接进入查看阶段（注意：可能还看不到 L1/L2/L3）")
        return

    info("Memory 内核会异步把 L0 对话蒸馏成 L1 原子记忆、L2 场景文件、L3 核心人格画像。")
    info("一般 3~5 分钟内可以看到结果，本脚本默认等 5 分钟。")

    end_at = time.time() + seconds
    while True:
        remain = int(end_at - time.time())
        if remain <= 0:
            break
        mm, ss = divmod(remain, 60)
        print(f"    ⏳ 剩余 {mm:02d}:{ss:02d} ...", end="\\r", flush=True)
        time.sleep(min(30, remain))
    print(" " * 40, end="\\r")  # 清掉最后一行倒计时
    ok("等待结束，开始查看四层记忆数据")


# ---------------------------------------------------------------------------- #
#  5. 回查 L0 / L1 / L2 / L3 四层记忆
# ---------------------------------------------------------------------------- #
def step_inspect(client: MemoryClient, session_id: str) -> None:
    section("Step 3/3  查看 L0 / L1 / L2 / L3 四层记忆数据")

    # ---- L0：原始对话回查 ----
    print("\\n[L0] 原始对话")
    l0 = safe_call(
        "L0 query_conversation",
        lambda: client.query_conversation(session_id=session_id, limit=100, offset=0),
    )
    if l0:
        # 注：v2 /conversation/query 的列表字段是 messages（不是 items）
        messages = l0.get("messages") or l0.get("items") or []
        info(f"会话 {session_id} 中累计 {len(messages)} 条消息")
        if messages:
            sample = messages[0]
            preview = (sample.get("content") or "")[:50]
            info(f"首条预览：[{sample.get('role')}] {preview}...")

    # ---- L1：原子记忆（按三种 type 分别拉一次）----
    print("\\n[L1] 原子记忆（atomic notes）")
    l1_total = 0
    type_summary: dict[str, int] = {}
    for atom_type in ("episodic", "persona", "instruction"):
        data = safe_call(
            f"L1 query_atomic type={atom_type}",
            lambda t=atom_type: client.query_atomic(type=t, limit=50, offset=0),
        )
        cnt = len(data.get("items") or []) if data else 0
        type_summary[atom_type] = cnt
        l1_total += cnt

    info(f"L1 抽取统计：episodic={type_summary.get('episodic', 0)} / "
         f"persona={type_summary.get('persona', 0)} / "
         f"instruction={type_summary.get('instruction', 0)}（共 {l1_total} 条）")
    if l1_total == 0:
        warn("L1 暂无数据，可能内核还在沉淀中，可以稍后用 --skip-write 再跑一次只观察")

    # ---- L2：场景文件 ----
    print("\\n[L2] 场景文件（scenario files）")
    l2 = safe_call("L2 list_scenarios (root)", lambda: client.list_scenarios(path_prefix=""))
    if l2:
        files = l2.get("items") or l2.get("files") or []
        info(f"根目录下共 {len(files)} 个场景文件")
        for f in files[:5]:
            name = f.get("path") or f.get("name") or f.get("key") or "<unknown>"
            info(f"  - {name}")

    # ---- L3：核心人格画像 ----
    print("\\n[L3] 核心人格画像（core profile）")
    l3 = safe_call("L3 read_core", lambda: client.read_core())
    if l3:
        content = l3.get("content") or l3.get("text") or ""
        if content:
            preview = content[:200].replace("\\n", " ")
            info(f"L3 已生成，长度 {len(content)} 字，前 200 字预览：")
            print(f"      {preview}{'...' if len(content) > 200 else ''}")
        else:
            warn("L3 read 成功但内容为空，可能仍在沉淀中")
    # 注：L3 在实例还没沉淀时常常是 404/业务错误，safe_call 会友好提示并继续，不会中断。


# ---------------------------------------------------------------------------- #
#  main
# ---------------------------------------------------------------------------- #
def main() -> None:
    parser = argparse.ArgumentParser(
        description="Memory Python SDK 快速上手 Demo（C 端 AI 助手 10 轮对话）"
    )
    parser.add_argument("--wait-seconds", type=int, default=300,
                        help="写入后等待秒数，默认 300（5 分钟）")
    parser.add_argument("--no-wait", action="store_true",
                        help="写完立即查看（冒烟用，不推荐生产观察）")
    parser.add_argument("--skip-write", action="store_true",
                        help="跳过写入步骤，直接观察当前 session 的沉淀结果")
    parser.add_argument("--session-id", default=DEFAULT_SESSION_ID,
                        help=f"自定义 session_id（默认 {DEFAULT_SESSION_ID}）")
    args = parser.parse_args()

    # 配置兜底检查：避免用户忘了改默认占位符
    if "请替换" in ENDPOINT or "请替换" in API_KEY or "请替换" in SERVICE_ID:
        fail("检测到 ENDPOINT / API_KEY / SERVICE_ID 还是占位符，请先在脚本顶部填入真实值再运行。")
        sys.exit(2)

    section("Memory Python SDK 快速上手 Demo")
    info(f"实例 ID  ：{SERVICE_ID}")
    info(f"网关地址 ：{ENDPOINT}")
    info(f"会话 ID  ：{args.session_id}")

    # MemoryClient 支持 with 语法，会自动关闭底层 httpx 连接
    with MemoryClient(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        service_id=SERVICE_ID,
        timeout=30,
        verify=False,   # test 网关用自签证书；正式环境删掉这行或改 True
    ) as client:

        # 1) 写入 L0
        if not args.skip_write:
            step_write_l0(client, args.session_id)
        else:
            section("Step 1/3  跳过写入（--skip-write）")
            info("仅观察当前 session 已有沉淀")

        # 2) 等待内核沉淀
        wait = 0 if args.no_wait else args.wait_seconds
        step_wait(wait)

        # 3) 回查四层
        step_inspect(client, args.session_id)

    section("Demo 执行完毕 ✅")
    print("  下一步建议：")
    print("    1. 打开 Memory 控制台 → 实例详情 → 记忆数据，可以可视化看到 L0/L1/L2/L3 全貌")
    print("    2. 如果 L1/L2/L3 还为空，几分钟后再跑一次：")
    print(f"         python {sys.argv[0]} --skip-write --no-wait")
    print("    3. 若要换一个会话演示，加 --session-id <自定义ID> 即可\\n")


if __name__ == "__main__":
    main()

──────────────────────────────────────────────────────────────────────
  Step 3/3  查看 L0 / L1 / L2 / L3 四层记忆数据
──────────────────────────────────────────────────────────────────────

[L0] 原始对话
  ✓ L0 query_conversation 调用成功
  · 会话 demo-pm-quickstart-001 中累计 20 条消息
  · 首条预览：[user] 我现在 78 公斤，身高 175，想在 3 个月内减到...

[L1] 原子记忆（atomic notes）
  · L1 抽取统计：episodic=6 / persona=4 / instruction=5（共 15 条）

[L2] 场景文件（scenario files）
  · 根目录下共 3 个场景文件

[L3] 核心人格画像（core profile）
  · L3 已生成，长度 412 字，前 200 字预览：
    用户小磊，男性上班族，目标在 3 个月内由 78kg 减至 70kg ...

python quickstart.py --skip-write --no-wait

python quickstart.py --session-id my-test-001