agentframe/tests/test_core.py

"""AgentFrame 核心测试 (离线, 无需 API key)"""
import os
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from agentframe.config import AgentFrameConfig
from agentframe.core.engine import ContextEngine
from agentframe.embed.provider import HashEmbedding


def make_engine() -> ContextEngine:
    cfg = AgentFrameConfig.from_env()
    cfg.llm.provider = "mock"   # 离线
    return ContextEngine(cfg)


def test_ingest_and_retrieve():
    eng = make_engine()
    eng.ingest("KV 缓存压缩测试数据", ["kv"])
    eng.ingest("注意力分数不等于任务重要性", ["agent"])
    eng.ingest("Minecraft 1.20.4 逆向完成", ["mc"])
    assert len(eng.agent.chunk_meta) == 3
    r = eng.ask("KV 压缩多少倍？", chat=False)
    assert len(r.retrieved) >= 1
    print("✅ test_ingest_and_retrieve")


def test_similar_text_retrieval():
    eng = make_engine()
    eng.ingest("吸收式 MLA 缓存 576 维潜在向量", ["method"])
    eng.ingest("HiLS 分层软max 端到端块选择", ["method"])
    eng.ingest("今天天气很好适合出去玩", ["life"])
    # 查询与第一条相关
    r = eng.ask("MLA 潜在向量维度是多少？", chat=False)
    tops = [cid for cid, _ in r.retrieved]
    assert 0 in tops, f"期望命中 chunk_0, 实际 {tops}"
    print(f"✅ test_similar_text_retrieval (top: {tops[:3]})")


def test_forget_curve():
    eng = make_engine()
    eng.ingest("A", ["x"])
    eng.ingest("B", ["x"])
    eng.ingest("C", ["x"])
    # 时间推进 (默认半衰期 100, 500 轮后 decay=2^-5=0.031)
    for _ in range(500):
        eng.now += 1
    # 阈值设 0.1: strength = 0.5*0.031 + 0 ≈ 0.016 < 0.1 → 应遗忘
    victims = eng.forget(0.1)
    assert len(victims) >= 2, f"长时间不访问应遗忘, 实际 {len(victims)}"
    print(f"✅ test_forget_curve (遗忘 {len(victims)}/3)")


def test_save_load():
    eng = make_engine()
    eng.ingest("持久化测试内容", ["test"])
    with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
        path = f.name
    eng.save(path)
    eng2 = make_engine()
    ok = eng2.load(path)
    assert ok
    assert len(eng2.agent.chunk_meta) == 1
    os.unlink(path)
    print("✅ test_save_load")


def test_hash_embedding_deterministic():
    emb = HashEmbedding(dim=576)
    v1 = emb.embed("同一句话")
    v2 = emb.embed("同一句话")
    v3 = emb.embed("完全不同的话")
    assert (v1 == v2).all()
    sim_same = float(v1 @ v2)
    sim_diff = float(v1 @ v3)
    assert sim_same > sim_diff, f"{sim_same} vs {sim_diff}"
    print(f"✅ test_hash_embedding (同句相似度 {sim_same:.3f} > 异句 {sim_diff:.3f})")


def test_tool_exec():
    eng = make_engine()
    out = eng._exec_tool("print(6*7)")
    assert "42" in out
    bad = eng._exec_tool("print(undefined_var)")
    assert "Traceback" in bad or "Error" in bad
    print("✅ test_tool_exec")


def test_lfru_hysteresis():
    """LFRU 滞回驱逐: 历史热块获得保护, 信用随时间衰减 (colibrì #441/#497)"""
    import numpy as np
    from agentframe.core.quad import KVPager, CompressedKV

    def mk(cid, max_heat, acc):
        return CompressedKV(
            chunk_id=cid, latent=np.zeros(576, dtype=np.float32),
            quant_bits=4, heat=max_heat, size_bytes=1000,
            importance=0.5, access_count=acc, max_heat=max_heat,
            last_access=0)

    # 冷却 200 轮 (2 个半衰期): 历史热块应被保护
    p = KVPager(vram_limit_mb=0.006, ram_limit_mb=64)
    for kv in [mk(0, 0.95, 0), mk(1, 0.4, 0), mk(2, 0.5, 0),
               mk(3, 0.5, 0), mk(4, 0.5, 0)]:
        p.vram[kv.chunk_id] = kv
        p.vram_used += kv.size_bytes
    victim = p.evict(200.0, "vram")
    assert victim == 1, f"历史热块0应被保护, 实际驱逐 {victim}"

    # 完全冷透 (2000 轮): 信用衰减殆尽, 照常驱逐
    p2 = KVPager(vram_limit_mb=0.006, ram_limit_mb=64)
    for kv in [mk(0, 0.95, 0), mk(1, 0.4, 0), mk(2, 0.5, 0),
               mk(3, 0.5, 0), mk(4, 0.5, 0)]:
        p2.vram[kv.chunk_id] = kv
        p2.vram_used += kv.size_bytes
    assert p2.evict(2000.0, "vram") is not None
    print("✅ test_lfru_hysteresis (滞回保护 + 信用衰减)")


def test_tool_safety():
    """工具执行安全: 危险命令/超长代码应被拦截"""
    eng = make_engine()
    out = eng._exec_tool("import os; os.system('rm -rf /tmp/x')")
    assert "安全拦截" in out, out
    out2 = eng._exec_tool("import subprocess; subprocess.run(['shutdown'])")
    assert "安全拦截" in out2, out2
    print("✅ test_tool_safety (危险命令拦截)")


def test_api_auth():
    """API 认证: token 模式 + 回环模式"""
    import os
    from agentframe.api.server import AgentFrameAPI
    from agentframe.config import AgentFrameConfig
    os.environ["AGENTFRAME_STATE_DIR"] = "/tmp/af_test_auth"
    cfg = AgentFrameConfig.from_env()
    cfg.api.token = "test-token"
    api = AgentFrameAPI(cfg)
    client = api.app.test_client()
    assert client.get("/v1/sessions").status_code == 401
    ok = client.get("/v1/sessions",
                    headers={"Authorization": "Bearer test-token"})
    assert ok.status_code == 200
    print("✅ test_api_auth (Bearer 认证)")


def test_couple_prefetch():
    """跨轮共现预取: 学习共现 → 预测 → disk→RAM 提升 (colibrì couple 移植)"""
    import numpy as np
    from agentframe.core.couple import CouplePrefetcher
    from agentframe.core.quad import KVPager, CompressedKV

    def mk(cid):
        return CompressedKV(chunk_id=cid, latent=np.zeros(576, dtype=np.float32),
                            quant_bits=4, heat=0.1, size_bytes=1000,
                            importance=0.5)

    # 1. 共现学习: 轮1检索{0,1}, 轮2检索{1,2}, 轮3检索{2,3}
    cp = CouplePrefetcher(top_k=8)
    cp.record([0, 1])
    cp.record([1, 2])
    cp.record([2, 3])
    # 轮2后: (0,2)共现1, (1,2)共现1; 轮3后: (1,2)再+1, (1,3)共现1, (2,3)共现1
    assert cp.cooccur[0][2] == 1
    assert cp.cooccur[1][2] == 2  # 轮2(1在prev,2在cur) + 轮3(1在prev,2在cur)
    assert cp.cooccur[2][3] == 1
    assert cp.cooccur[1][3] == 1
    # 2. 预测: 当前检索{1} → 预测 2 (共现1)
    pred = cp.predict([1])
    assert 2 in pred, f"应预测到块2, 实际 {pred}"
    assert 1 not in pred, "预测不应包含当前集合中的块"
    print(f"✅ test_couple_prefetch (学习+预测: {pred[:4]})")

    # 3. 预取落地: 块2在 disk, prefetch 后应到 RAM (不进 VRAM)
    p = KVPager(vram_limit_mb=64, ram_limit_mb=0.006)  # 6KB RAM
    for cid in (0, 1, 3):
        p.disk[cid] = mk(cid)
    p.disk[2] = mk(2)  # 目标块在 disk
    moved = p.prefetch([2], now=1.0)
    assert 2 in moved
    assert 2 in p.ram and 2 not in p.vram and 2 not in p.disk
    print("✅ test_couple_prefetch (预取 disk→RAM, 不占 VRAM)")

    # 4. RAM 满时腾位: 塞满 RAM 后预取新块
    p2 = KVPager(vram_limit_mb=64, ram_limit_mb=0.005)  # 5KB RAM (最多5块)
    for cid in range(5):
        p2.ram[cid] = mk(cid)
        p2.ram_used += 1000
    p2.disk[99] = mk(99)
    moved2 = p2.prefetch([99], now=1.0)
    assert 99 in p2.ram
    assert len(p2.ram) == 5  # 腾位后仍不超限
    print("✅ test_couple_prefetch (RAM 满时 LFRU 腾位)")


def test_topk_protection():
    """Top-K 保护 (colibrì 接入): 路由命中块标记高精度, 检索不翻转"""
    import numpy as np
    from agentframe.core.quad import (AbsorbedMLA, ReversibleQuantizer,
                                      LandmarkRouter)

    # 1. 保护标记: 命中后 size 变 1152B (16bit), 未命中保持 352B (4bit)
    m = AbsorbedMLA(n_layers=27, quant_bits=4, n_ch=16)
    rng = np.random.default_rng(42)
    latents = np.tanh(rng.normal(0, 1, (20, 576)).astype(np.float32))
    for L in latents:
        m.encode(L)
    assert all(not kv.protected for kv in m.chunks.values())
    m.protect_topk(0)
    m.protect_topk(1)
    assert m.chunks[0].protected and m.chunks[0].size_bytes == 1152
    assert m.chunks[2].protected is False and m.chunks[2].size_bytes == 352
    print("✅ test_topk_protection (16bit 保护标记 + 大小区分)")

    # 2. 检索一致性: 原始 latent 构建摘要 vs INT4 解包摘要
    router = LandmarkRouter(top_k=8, seed=42)
    summaries_orig = {b: router.build_summary(latents[b].reshape(1, -1))
                      for b in range(20)}
    flip_orig = 0
    for q in range(10):
        qv = latents[rng.integers(0, 20)]
        s1 = router.route(qv, qv, summaries_orig)
        flip_orig += 0  # 无损路径不翻转
    # 3. 解包路径翻转率 (对照): 应显著高于无损路径
    flip_deq = 0
    summaries_deq = {}
    for b in range(20):
        q4, sc, _ = ReversibleQuantizer.quantize_int4(latents[b], n_ch=16)
        deq = ReversibleQuantizer.dequant_int4(q4, sc, 576)
        k, bias = router.build_summary(deq.reshape(1, -1))
        summaries_deq[b] = (k, bias)
    for q in range(10):
        qv = latents[rng.integers(0, 20)]
        s_orig = router.route(qv, qv, summaries_orig)
        s_deq = router.route(qv, qv, summaries_deq)
        if s_orig.chunk_ids != s_deq.chunk_ids:
            flip_deq += 1
    print(f"  ✅ 无损路径 0 翻转 | INT4 解包路径 {flip_deq}/10 翻转")
    assert flip_deq >= 1, "对照: INT4 解包路径应存在翻转 (证明保护必要性)"
    print("✅ test_topk_protection (无损检索 0 翻转, 对照解包路径有翻转)")


def test_int4_packing():
    """真 INT4 打包 (colibrì quant.h 移植): 往返精度 + 真实压缩比"""
    import numpy as np
    from agentframe.core.quad import ReversibleQuantizer, AbsorbedMLA

    rng = np.random.default_rng(42)
    latent = np.tanh(rng.normal(0, 1, 576).astype(np.float32))
    q4, scales, size = ReversibleQuantizer.quantize_int4(latent, n_ch=16)
    deq = ReversibleQuantizer.dequant_int4(q4, scales, 576)
    # 往返精度: 余弦相似度 > 0.99
    sim = float(deq @ latent) / (np.linalg.norm(deq) * np.linalg.norm(latent))
    assert sim > 0.99, f"INT4 往返相似度 {sim:.4f}"
    # 真实压缩: 352B vs 原始 2304B
    assert size == 288 + 64, f"打包大小 {size} != 352"
    # 每 token 压缩比对齐 L40S 实测 28.4x
    m = AbsorbedMLA(n_layers=27, quant_bits=4, n_ch=16)
    ratio = 270 * 1024 / m.bytes_per_token()
    assert ratio > 25, f"压缩比 {ratio:.1f}x 不足"
    print(f"✅ test_int4_packing (cos={sim:.4f}, {ratio:.1f}x)")


def test_incremental_persist():
    """增量持久化 (colibrì kv_persist 移植): append + crash-safe 恢复"""
    import tempfile, os
    from agentframe.memory.incremental import IncrementalKVStore

    path = tempfile.mktemp(suffix=".kv")
    store = IncrementalKVStore(path)
    # 追加 3 条
    import numpy as np
    n1 = store.append(0, np.zeros(288, dtype=np.uint8), np.ones(16, dtype=np.float32),
                      np.zeros(576, dtype=np.float32), 4, 352, {"text": "A"})
    n2 = store.append(1, np.ones(288, dtype=np.uint8), np.ones(16, dtype=np.float32) * 2,
                      np.ones(576, dtype=np.float32), 4, 352, {"text": "B"})
    n3 = store.append(2, np.zeros(288, dtype=np.uint8), np.ones(16, dtype=np.float32),
                      np.zeros(576, dtype=np.float32), 4, 352, {"text": "C"})
    assert n3 == 3
    # 模拟崩溃: 写入半行垃圾
    with open(path, "a") as f:
        f.write('{"magic": "AFKV1", "chunk_id": 99, "q4": "zz')
    recs = store.load()
    assert len(recs) == 3, f"坏行应被跳过, 实际 {len(recs)}"
    assert recs[1]["meta"]["text"] == "B"
    print("✅ test_incremental_persist (append + crash-safe 坏行跳过)")
    os.unlink(path)


def test_prefix_reuse():
    """查询前缀复用 (colibrì kv_prefix 移植)"""
    eng = make_engine()
    eng.ingest("KV 压缩 29 倍", ["kv"])
    eng.ingest("LFRU 滞回驱逐", ["method"])
    eng.ingest("Couple 预取", ["method"])
    # 第一轮: 正常检索
    r1 = eng.ask("KV 压缩是多少倍", chat=False)
    first = r1.retrieved
    # 第二轮: 相同前缀 → 复用
    r2 = eng.ask("KV 压缩是多少倍呢", chat=False)
    assert eng._prefix_hits == 1, "应触发前缀复用"
    assert [c for c, _ in r2.retrieved] == [c for c, _ in first]
    # 第三轮: 完全不同 → 不复用
    r3 = eng.ask("今天天气怎么样", chat=False)
    assert eng._prefix_hits == 1, "不同 query 不应复用"
    print("✅ test_prefix_reuse (前缀命中复用, 不同 query 不复用)")


def test_incremental_engine():
    """引擎级增量恢复: ingest → enable → load_incremental 重建"""
    import tempfile, os
    path = tempfile.mktemp(suffix=".kv")
    eng = make_engine()
    eng.enable_incremental(path)
    eng.ingest("增量持久化内容 A", ["a"])
    eng.ingest("增量持久化内容 B", ["b"])
    # 新引擎从日志恢复
    eng2 = make_engine()
    n = eng2.load_incremental(path)
    assert n == 2, f"应恢复 2 块, 实际 {n}"
    r = eng2.ask("增量持久化", chat=False)
    assert len(r.retrieved) >= 1
    os.unlink(path)
    print("✅ test_incremental_engine (引擎级增量恢复 + 检索)")


if __name__ == "__main__":
    test_ingest_and_retrieve()
    test_similar_text_retrieval()
    test_forget_curve()
    test_save_load()
    test_hash_embedding_deterministic()
    test_tool_exec()
    test_lfru_hysteresis()
    test_tool_safety()
    test_api_auth()
    test_couple_prefetch()
    test_int4_packing()
    test_topk_protection()
    test_incremental_persist()
    test_prefix_reuse()
    test_incremental_engine()
    print("\n🎉 全部核心测试通过!")