"""
=============================================================================
第 11 讲 · 01：MemGPT —— 把 LLM 当作操作系统，在"内存"和"磁盘"之间搬数据
=============================================================================

对应讲义：lectures/lecture11_memory.md 第 2 节（Packer et al. 2023）

-----------------------------------------------------------------------------
【设定：一段很长的多次会话】
-----------------------------------------------------------------------------
用户和助手聊了很多次（session）。每次会话里，用户顺口提到几个关于自己的事实（"我的猫叫豆豆""我在杭州工作"），
其余是闲聊。有些事实后来会**更新**（"我搬到上海了"）—— 正确答案是最新的那个。
最后问几个"深层记忆"问题（Deep Memory Retrieval）：关于很久以前提到过的事实。

上下文窗口只能放 WINDOW 条消息。三种做法：
    截断（FIFO）   ：只保留最近的 WINDOW 条消息。事实不在窗口里，模型只能猜（必错）。
    递归摘要       ：窗口满了，就把最旧的一半消息压缩成摘要，和旧的摘要合并。每压缩一次，
                     摘要里的每个事实有 SUMMARY_KEEP 的概率被保留下来（摘要是有损的，而且损失会累积）；
                     更新的事实也可能保留了旧的值。
    MemGPT        ：分层的记忆 ——
                     · 主上下文 = 系统指令 + 核心记忆（working context，最多 CORE_SLOTS 个事实）+ 最近的消息（FIFO）；
                     · 归档存储（archival storage）：容量不限的外部数据库。
                     模型通过**函数调用**管理记忆：看到重要的事实，写进核心记忆或归档存储；被问到时，先查核心记忆，
                     没有就调用 archival_search 去查。窗口用到 70% 时，系统发出**内存压力警告**，
                     提醒模型把重要的信息保存下来（警告后写入的概率更高）；用到 100% 时，系统把最旧的一半消息
                     移出窗口（flush），窗口回到 50%，然后又慢慢变满 —— 所以只有 70%~100% 这一段时间里有警告。
                     （真实的 MemGPT 还会把移出的消息放进 recall storage，可以用 conversation_search 检索；
                     这里省略了它，只看"模型主动保存"的作用。把它加回来是讲义的练习。）
【重要】写入、检索、摘要保留的概率都是**设定**的（见下面的常数），演示的是机制，不是真实系统的数值。

运行时间：不到 1 秒（只用标准库）。
"""

import random

WINDOW = 40                 # 上下文窗口能放的消息数
MSGS_PER_SESSION = 12       # 每次会话的消息数（其中 2 条是事实）
FACTS_PER_SESSION = 2
SUMMARY_KEEP = 0.85         # 每次压缩，摘要里的每个事实被保留下来的概率
CORE_SLOTS = 8              # 核心记忆最多放几个事实
WRITE_P, WRITE_P_WARNED = 0.6, 0.95   # 模型主动把一个事实写进归档存储的概率（平时 / 收到内存压力警告后）
SEARCH_P = 0.9              # 归档检索找到正确记录的概率
LATEST_P = 0.9              # 检索到同一个属性的多个版本时，模型挑对"最新那个"的概率
UPDATE_P = 0.25             # 一个新提到的事实是"更新已有的属性"的概率

ATTRIBUTES = [f"属性{i:02d}" for i in range(30)]    # 例如"猫的名字""工作的城市"……


def make_conversation(rng, n_sessions):
    """
    生成一段多次会话的对话：消息列表，每条是 ("chat", None) 或 ("fact", (属性, 值, 时间))。

    【返回】(messages, truth)：truth 是每个属性**最新**的值（上帝视角，只用来评估）
    """
    messages, truth, t = [], {}, 0
    for _ in range(n_sessions):
        fact_pos = set(rng.sample(range(MSGS_PER_SESSION), FACTS_PER_SESSION))
        for j in range(MSGS_PER_SESSION):
            if j in fact_pos:
                if truth and rng.random() < UPDATE_P:
                    attr = rng.choice(sorted(truth))                # 更新一个已有的属性
                else:
                    unused = [a for a in ATTRIBUTES if a not in truth]
                    attr = rng.choice(unused) if unused else rng.choice(ATTRIBUTES)
                value = rng.randint(100, 999)
                truth[attr] = value
                messages.append(("fact", (attr, value, t)))
            else:
                messages.append(("chat", None))
            t += 1
    return messages, truth


def questions(rng, truth, n=6):
    """从提到过的属性里挑 n 个来提问（不够就全部）。"""
    attrs = sorted(truth)
    return rng.sample(attrs, min(n, len(attrs)))


def answer_truncation(messages, attr):
    """截断：只看最近 WINDOW 条消息里，这个属性最新的值；不在窗口里就答不出来（返回 None）。"""
    for kind, payload in reversed(messages[-WINDOW:]):
        if kind == "fact" and payload[0] == attr:
            return payload[1]
    return None


def answer_summary(messages, attr, rng):
    """
    递归摘要：模拟一路压缩下来的摘要 + 当前窗口。

    【思路】每当窗口满了，把最旧的一半消息"压缩"进摘要：摘要里的事实（包括刚压进去的）各自以 SUMMARY_KEEP 的概率
        存活。更新时，新值进入摘要会覆盖旧值；但如果新值在压缩时丢了，摘要里可能还留着旧值（答错）。
    """
    summary, window = {}, []
    for msg in messages:
        window.append(msg)
        if len(window) > WINDOW:
            old, window = window[:WINDOW // 2], window[WINDOW // 2:]
            for kind, payload in old:
                if kind == "fact":
                    summary[payload[0]] = payload[1]               # 新的值进入摘要（覆盖旧值）
            summary = {a: v for a, v in summary.items() if rng.random() < SUMMARY_KEEP}
    for kind, payload in reversed(window):
        if kind == "fact" and payload[0] == attr:
            return payload[1]
    return summary.get(attr)


def answer_memgpt(messages, attr, rng, pressure_warning=True):
    """
    MemGPT：核心记忆 + 归档存储 + FIFO 窗口。

    【思路】按时间顺序处理消息：
        · fill 记录窗口里现在有几条消息：每来一条加 1；超过 WINDOW 时 flush，回到 WINDOW // 2。
        · 每个事实：模型以一定概率把它写进归档存储（WRITE_P；如果此时窗口已经用到 70% 以上、并且有内存压力警告，
          概率提高到 WRITE_P_WARNED）；同时写进核心记忆（核心记忆满了就挤掉最久没更新的那个）。
        · 回答时：先看窗口里的 fill 条消息；再看核心记忆；再调用 archival_search：以 SEARCH_P 的概率找到这个属性的记录，
          如果有多个版本，以 LATEST_P 的概率挑对最新的那个。
    【边界】fill 只在 WINDOW // 2 到 WINDOW 之间来回变化（对话开头除外）。和"截断"不同，MemGPT 的窗口
        刚 flush 完时只有一半是满的 —— 它用这一半的空间换来了给核心记忆和函数调用留出的余地。
    """
    core, archive, fill = {}, [], 0
    for kind, payload in messages:
        fill += 1
        if fill > WINDOW:
            fill = WINDOW // 2                                     # flush：最旧的一半消息被移出窗口
        if kind != "fact":
            continue
        a, v, t = payload
        p_write = WRITE_P_WARNED if (pressure_warning and fill >= 0.7 * WINDOW) else WRITE_P
        if rng.random() < p_write:
            archive.append((a, v, t))
        core[a] = (v, t)
        if len(core) > CORE_SLOTS:
            oldest = min(core, key=lambda k: core[k][1])
            del core[oldest]
    for kind, payload in reversed(messages[len(messages) - fill:]):   # 窗口里最近的 fill 条消息
        if kind == "fact" and payload[0] == attr:
            return payload[1]
    if attr in core:
        return core[attr][0]
    if rng.random() >= SEARCH_P:
        return None
    records = sorted((r for r in archive if r[0] == attr), key=lambda r: r[2])
    if not records:
        return None
    if len(records) > 1 and rng.random() >= LATEST_P:
        return records[rng.randrange(len(records) - 1)][1]         # 挑成了一个旧版本
    return records[-1][1]


def main():
    lengths = (2, 5, 10, 20, 40)
    methods = (("截断（只保留最近 40 条）", lambda m, a, r: answer_truncation(m, a)),
               ("递归摘要", answer_summary),
               ("MemGPT（无内存压力警告）", lambda m, a, r: answer_memgpt(m, a, r, False)),
               ("MemGPT", answer_memgpt))
    print("=" * 78)
    print(f"深层记忆问题的准确率 vs 对话长度（每次会话 {MSGS_PER_SESSION} 条消息，窗口 {WINDOW} 条；每个长度 100 段对话）")
    print("=" * 78)
    print(f"{'':28}" + "".join(f"{n} 次会话".rjust(10) for n in lengths))
    for name, fn in methods:
        row = []
        for n in lengths:
            rng = random.Random(240 + n)
            hits = []
            for _ in range(100):
                msgs, truth = make_conversation(rng, n)
                for attr in questions(rng, truth):
                    hits.append(fn(msgs, attr, rng) == truth[attr])
            row.append(sum(hits) / len(hits))
        width = sum(2 if ord(ch) > 127 else 1 for ch in name)
        print(name + " " * max(1, 28 - width) + "".join(f"{x:10.3f}" for x in row))

    print("\n观察：")
    print("  · 截断：对话一长，早期的事实就被挤出窗口，准确率随对话长度断崖式下降。")
    print("  · 递归摘要：能撑得更久，但每压缩一次就丢一点，损失会累积 —— 对话越长，最早的事实被压缩的次数越多，")
    print("    越可能已经丢了；更新过的事实还可能留着旧值。")
    print("  · MemGPT：下降得慢得多。从 10 次到 40 次会话，它只从 0.843 降到 0.757，递归摘要从 0.843 降到 0.497。")
    print("    剩下的下降来自三处：核心记忆只放得下 8 个事实；不是每个事实都被存进了归档；属性被更新得越多，")
    print("    检索时挑到旧版本的机会越多。上限由\"存了没有\"和\"能不能检索到、挑对最新版本\"决定，而不是由窗口大小决定。")
    print("  · 内存压力警告的作用很大：没有它，模型只有 60% 的概率主动保存事实，MemGPT 在 10 次会话时")
    print("    甚至不如递归摘要（0.758 vs 0.843）；有了它（只在窗口用到 70%~100% 时出现），10 次会话以上的长度都高 0.08~0.16。")
    print("    \"会管理记忆\"不只是有存储，还要知道**什么时候**该存。")
    print("    MemGPT 论文在 Deep Memory Retrieval 上的结果：GPT-4 直接用对话历史 32.1%，加上 MemGPT 92.5%。\n")


if __name__ == "__main__":
    main()