lambdagent/tests/test_sample_kb.py

"""
tests/test_sample_kb.py — 示例知识库种子 + 平台 KB 工具桥接（KB 调用链核查修复）。

覆盖：
1. ensure_sample_kb: 建库/登记/索引、幂等、无租户时延迟。
2. _make_platform_kb_tools: KBSearch 查到挂接库内容；KBList 列出库；
   无挂接时返回 {}（不注入，保留 lambdagent 内置行为）。
3. compiler 工具合并: subAgents 的 call_* 与宿主注入的自定义工具共存
   （原来注入任何 tools 都会让 subAgents 整段跳过）。
"""
from __future__ import annotations

import json
import os

import pytest

os.environ.setdefault("AGENTPAAS_DATABASE_URL", "sqlite:///:memory:")
os.environ.setdefault("AGENTPAAS_TESTING", "1")


@pytest.fixture()
def fresh_db(monkeypatch):
    """独立内存 DB + 一个活跃租户。"""
    from agentpaas.db.models import Database, gen_id, now_utc
    import agentpaas.db.session as _session_mod

    prev = _session_mod._db
    _session_mod._db = Database("sqlite:///:memory:")
    db = _session_mod._db
    tid = gen_id("tn_")
    db.execute(
        "INSERT INTO tenants (id, name, plan, status, created_at) "
        "VALUES (?, 'test', 'free', 'active', ?)", (tid, now_utc()))
    db.commit()
    yield db, tid
    _session_mod._db = prev


def test_sample_kb_created_and_indexed(fresh_db, tmp_path):
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME

    db, tid = fresh_db
    ensure_sample_kb(str(tmp_path))

    kb = db.fetchone("SELECT * FROM knowledge_bases WHERE name = ?", (SAMPLE_KB_NAME,))
    assert kb, "示例知识库应被创建"
    assert kb["tenant_id"] == tid

    # 文件已拷贝并登记
    files = db.fetchall("SELECT * FROM kb_files WHERE kb_id = ?", (kb["id"],))
    assert len(files) >= 4, files
    names = {f["file_name"] for f in files}
    assert "知识库使用指南.md" in names

    # pageindex 已生成且可检索
    idx = os.path.join(kb["root_dir"], "rag_page_index.json")
    assert os.path.isfile(idx)
    entries = json.loads(open(idx, encoding="utf-8").read())
    assert len(entries) > 10  # 按 ## 小节切分

    from agentpaas.api.v1.knowledge import _page_index_search
    hits = _page_index_search(kb["root_dir"], "顺序表 插入 删除", top_k=3)
    assert hits and "顺序" in (hits[0].get("text", "") + hits[0].get("summary", ""))


def test_sample_kb_idempotent(fresh_db, tmp_path):
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME

    db, _ = fresh_db
    ensure_sample_kb(str(tmp_path))
    ensure_sample_kb(str(tmp_path))  # 再跑一次
    rows = db.fetchall("SELECT id FROM knowledge_bases WHERE name = ?", (SAMPLE_KB_NAME,))
    assert len(rows) == 1, "幂等：不应重复创建"


def test_sample_kb_defers_without_tenant(tmp_path, monkeypatch):
    from agentpaas.db.models import Database
    import agentpaas.db.session as _session_mod
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME

    prev = _session_mod._db
    _session_mod._db = Database("sqlite:///:memory:")  # 无租户
    try:
        ensure_sample_kb(str(tmp_path))
        kb = _session_mod._db.fetchone(
            "SELECT id FROM knowledge_bases WHERE name = ?", (SAMPLE_KB_NAME,))
        assert kb is None, "无租户时应延迟到下次启动"
    finally:
        _session_mod._db = prev


# ── 平台 KB 工具桥接 ────────────────────────────────────────────────────────

def test_platform_kb_tools_search_and_list(fresh_db, tmp_path):
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME
    from agentpaas.api.v1.agents import _make_platform_kb_tools

    db, tid = fresh_db
    ensure_sample_kb(str(tmp_path))
    kb = db.fetchone("SELECT * FROM knowledge_bases WHERE name = ?", (SAMPLE_KB_NAME,))

    tools = _make_platform_kb_tools(db, [kb["id"]], tid, search_mode="pageindex")
    assert set(tools) == {"KBSearch", "KBList"}

    # KBSearch: JSON 入参与裸字符串都支持
    out = tools["KBSearch"]('{"query": "循环队列 队满"}')
    assert "知识库检索结果" in out and "队" in out
    out2 = tools["KBSearch"]("栈 后进先出")
    assert "知识库检索结果" in out2

    # 查不到 → NO_MATCH 提示而非空串
    miss = tools["KBSearch"]("量子引力 弦论")
    assert "NO_MATCH" in miss or "知识库检索结果" in miss

    # KBList 列出挂接库与文件数
    lst = tools["KBList"]("")
    assert SAMPLE_KB_NAME in lst and "文件" in lst


def test_platform_kb_tools_empty_when_no_kb(fresh_db):
    from agentpaas.api.v1.agents import _make_platform_kb_tools
    db, tid = fresh_db
    assert _make_platform_kb_tools(db, [], tid) == {}


def test_platform_kb_tools_tenant_scoped(fresh_db, tmp_path):
    """别的租户挂了这个 kb_id 也查不到（audit #30 语义在工具层同样成立）。"""
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME
    from agentpaas.api.v1.agents import _make_platform_kb_tools

    db, tid = fresh_db
    ensure_sample_kb(str(tmp_path))
    kb = db.fetchone("SELECT * FROM knowledge_bases WHERE name = ?", (SAMPLE_KB_NAME,))

    tools = _make_platform_kb_tools(db, [kb["id"]], "tn_other_tenant", "pageindex")
    out = tools["KBSearch"]("顺序表")
    assert "NO_MATCH" in out  # 跨租户被静默跳过 → 查无结果
    assert SAMPLE_KB_NAME not in tools["KBList"]("")


# ── compiler 工具合并（subAgents 与注入工具共存）────────────────────────────

def test_compiler_merges_injected_tools_with_subagents():
    from lambdagent.fromconfig.compiler import build_agent

    cfg = {
        "name": "orch", "type": "react",
        "model": {"name": "ollama/qwen2.5:7b"},
        "systemPrompt": "test",
        "react": {"maxSteps": 3},
        "subAgents": {
            "helper": {"inline": {"type": "simple",
                                  "model": {"name": "ollama/qwen2.5:7b"},
                                  "systemPrompt": "sub"},
                       "tool": "call_helper"},
        },
        "mcp": {"localTools": ["KBSearch", "terminate"]},
    }
    sentinel = {"value": None}

    def fake_kb_search(x):
        sentinel["value"] = x
        return "[平台KB] hit"

    term = build_agent(cfg, {"tools": {"KBSearch": fake_kb_search}})
    assert term is not None
    # 注入的 KBSearch 生效（而非 lambdagent 内置实现）
    from lambdagent.fromconfig.compiler import _compile_tools  # noqa: F401
    # 通过重新编译工具表来断言合并语义
    tools = {}
    from lambdagent.fromconfig import compiler as _c
    merged = {** _c._compile_sub_agents(cfg, cfg["subAgents"], {}),
              **{"KBSearch": fake_kb_search}}
    assert "call_helper" in merged and merged["KBSearch"] is fake_kb_search


def test_page_index_search_chinese_query(fresh_db, tmp_path):
    """中文检索回归：原 \\W+ 分词把整句中文当一个 token → 全 0 分乱序。
    现在 CJK bigram 切分，长中文问句应命中正确小节。"""
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME
    from agentpaas.api.v1.knowledge import _page_index_search

    db, _ = fresh_db
    ensure_sample_kb(str(tmp_path))
    kb = db.fetchone("SELECT * FROM knowledge_bases WHERE name = ?", (SAMPLE_KB_NAME,))

    hits = _page_index_search(kb["root_dir"], "循环队列怎么判断队满？讲义里怎么说的", 3)
    assert hits and hits[0]["score"] > 0, "中文长句必须有非零得分"
    assert "栈与队列" in hits[0]["source"], f"应命中第3章, got {hits[0]['source']}"

    # 英文 query 不受影响
    hits_en = _page_index_search(kb["root_dir"], "ADT stack push pop", 3)
    assert hits_en and hits_en[0]["score"] > 0


def test_resolve_tenant_prefers_busiest(fresh_db, tmp_path):
    """多租户老库：示例库应归属活跃智能体最多的租户（浏览器会话所在地），
    而非最老租户或 config.json 租户。"""
    from agentpaas.db.models import gen_id, now_utc
    from agentpaas.engine.sample_kb import ensure_sample_kb, SAMPLE_KB_NAME

    db, tid_old = fresh_db
    # 第二个租户，拥有 2 个活跃智能体（模拟日常使用租户）
    tid_busy = gen_id("tn_")
    now = now_utc()
    db.execute("INSERT INTO tenants (id, name, plan, status, created_at) "
               "VALUES (?, 'busy', 'free', 'active', ?)", (tid_busy, now))
    for i in range(2):
        db.execute(
            "INSERT INTO agents (id, tenant_id, name, current_version, status, "
            "created_at, updated_at) VALUES (?, ?, ?, 1, 'active', ?, ?)",
            (gen_id("ag_"), tid_busy, f"a{i}", now, now))
    db.commit()

    ensure_sample_kb(str(tmp_path))
    kb = db.fetchone("SELECT tenant_id FROM knowledge_bases WHERE name = ?",
                     (SAMPLE_KB_NAME,))
    assert kb["tenant_id"] == tid_busy, "应归属智能体最多的租户"