lambdagent/tests/test_assistant.py

"""
tests/test_assistant.py — 工作台 AI 助手入口（方案 A+B 后端）。

覆盖：
- GET /status/recent-runs: 跨智能体最近运行、运行中排最前、租户隔离。
- POST /assistant/route: 唯一智能体直返、分类命中、拿不准回退、
  分类器异常降级、choice 越界保护（注入 fake 分类器，零 LLM 调用）。
"""
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 api_client(tmp_path, monkeypatch):
    import secrets
    from fastapi.testclient import TestClient
    from agentpaas.api.app import app
    from agentpaas.api.middleware.auth import hash_key
    from agentpaas.config import settings
    from agentpaas.db.models import Database, gen_id, now_utc
    import agentpaas.db.session as _session_mod

    monkeypatch.setattr(settings, "workspace_base", str(tmp_path / "Workspace"))
    prev = _session_mod._db
    _session_mod._db = Database("sqlite:///:memory:")
    db = _session_mod._db
    tid, uid = gen_id("tn_"), gen_id("usr_")
    raw_key = f"ap_{secrets.token_hex(16)}"
    now = now_utc()
    db.execute("INSERT INTO tenants (id, name, plan, status, created_at) "
               "VALUES (?, 'test', 'free', 'active', ?)", (tid, now))
    db.execute("INSERT INTO users (id, tenant_id, email, role, created_at) "
               "VALUES (?, ?, '', 'admin', ?)", (uid, tid, now))
    db.execute(
        "INSERT INTO api_keys (id, tenant_id, user_id, key_hash, key_prefix, name, "
        "scopes, rate_limit, status, created_at) "
        "VALUES (?, ?, ?, ?, ?, 'test', ?, 600, 'active', ?)",
        (gen_id("key_"), tid, uid, hash_key(raw_key), raw_key[:8],
         json.dumps(["agents:*"]), now))
    db.commit()
    with TestClient(app) as client:
        yield client, raw_key, tid, db
    _session_mod._db = prev


def _auth(k):
    return {"Authorization": f"Bearer {k}"}


def _mk_agent(client, key, name, desc=""):
    cfg = {"name": name, "type": "simple",
           "model": {"name": "ollama/qwen2.5:7b"}, "systemPrompt": "t"}
    r = client.post("/api/v1/agents", headers=_auth(key),
                    json={"name": name, "description": desc, "config": cfg})
    assert r.status_code in (200, 201), r.text[:200]
    return r.json()["agent_id"]


# ── recent-runs ──────────────────────────────────────────────────────────────

def test_recent_runs_running_first_and_truncated(api_client):
    from agentpaas.db.models import gen_id, now_utc
    client, key, tid, db = api_client
    aid = _mk_agent(client, key, "甲")
    now = now_utc()
    for i, status in enumerate(["completed", "running", "failed"]):
        db.execute(
            "INSERT INTO runs (id, agent_id, agent_version, tenant_id, input, "
            "status, created_at) VALUES (?, ?, 1, ?, ?, ?, ?)",
            (gen_id("run_"), aid, tid, f"输入{i}" + "长" * 200, status, now))
    db.commit()

    r = client.get("/api/v1/status/recent-runs?limit=10", headers=_auth(key))
    assert r.status_code == 200
    items = r.json()["items"]
    assert len(items) == 3
    assert items[0]["status"] == "running"          # 运行中排最前
    assert items[0]["agent_name"] == "甲"
    assert len(items[0]["input_preview"]) <= 120    # 截断


def test_recent_runs_tenant_scoped(api_client):
    from agentpaas.db.models import gen_id, now_utc
    client, key, tid, db = api_client
    # 别的租户的 run 不可见
    other_t = gen_id("tn_")
    db.execute("INSERT INTO tenants (id, name, plan, status, created_at) "
               "VALUES (?, 'o', 'free', 'active', ?)", (other_t, now_utc()))
    other_a = gen_id("ag_")
    db.execute("INSERT INTO agents (id, tenant_id, name, current_version, status, "
               "created_at, updated_at) VALUES (?, ?, 'x', 1, 'active', ?, ?)",
               (other_a, other_t, now_utc(), now_utc()))
    db.execute("INSERT INTO runs (id, agent_id, agent_version, tenant_id, input, "
               "status, created_at) VALUES (?, ?, 1, ?, 'i', 'completed', ?)",
               (gen_id("run_"), other_a, other_t, now_utc()))
    db.commit()
    r = client.get("/api/v1/status/recent-runs", headers=_auth(key))
    assert r.json()["items"] == []


# ── assistant/route ──────────────────────────────────────────────────────────

def test_route_single_agent_no_llm(api_client):
    client, key, *_ = api_client
    aid = _mk_agent(client, key, "唯一助手")
    r = client.post("/api/v1/assistant/route", headers=_auth(key),
                    json={"input": "随便什么任务"})
    assert r.status_code == 200
    d = r.json()
    assert d["matched"] is True and d["agent_id"] == aid


def test_route_classifier_picks(api_client, monkeypatch):
    from agentpaas.api.v1 import assistant as mod
    client, key, *_ = api_client
    _mk_agent(client, key, "审稿助手", "论文评审")
    aid_exam = _mk_agent(client, key, "试卷助手", "出试卷与题库")

    monkeypatch.setattr(mod, "_pick_classifier_model",
                        lambda: {"provider": "fake", "name": "f"})
    seen = {}

    def fake_classify(user_input, candidates, model):
        seen["input"] = user_input
        seen["names"] = [c["name"] for c in candidates]
        # 选「试卷助手」
        idx = next(i for i, c in enumerate(candidates, 1) if c["name"] == "试卷助手")
        return idx, "出题类请求", {"input_tokens": 50, "output_tokens": 10}

    monkeypatch.setattr(mod, "_classify_fn", fake_classify)
    r = client.post("/api/v1/assistant/route", headers=_auth(key),
                    json={"input": "给数据结构出一份期末卷"})
    d = r.json()
    assert d["matched"] is True and d["agent_id"] == aid_exam
    assert d["reason"] == "出题类请求"
    assert "试卷助手" in seen["names"] and "期末卷" in seen["input"]


def test_route_unsure_and_out_of_range(api_client, monkeypatch):
    from agentpaas.api.v1 import assistant as mod
    client, key, *_ = api_client
    _mk_agent(client, key, "甲")
    _mk_agent(client, key, "乙")
    monkeypatch.setattr(mod, "_pick_classifier_model",
                        lambda: {"provider": "fake", "name": "f"})
    # choice=0 拿不准
    monkeypatch.setattr(mod, "_classify_fn", lambda *a: (0, "都不像", None))
    assert client.post("/api/v1/assistant/route", headers=_auth(key),
                       json={"input": "x"}).json()["matched"] is False
    # choice 越界
    monkeypatch.setattr(mod, "_classify_fn", lambda *a: (99, "", None))
    assert client.post("/api/v1/assistant/route", headers=_auth(key),
                       json={"input": "x"}).json()["matched"] is False
    # 分类器抛异常 → 降级 matched=false 而非 500
    def boom(*a):
        raise RuntimeError("provider down")
    monkeypatch.setattr(mod, "_classify_fn", boom)
    r = client.post("/api/v1/assistant/route", headers=_auth(key), json={"input": "x"})
    assert r.status_code == 200 and r.json()["matched"] is False


def test_route_no_classifier_model(api_client, monkeypatch):
    from agentpaas.api.v1 import assistant as mod
    client, key, *_ = api_client
    _mk_agent(client, key, "甲")
    _mk_agent(client, key, "乙")
    monkeypatch.setattr(mod, "_pick_classifier_model", lambda: None)
    r = client.post("/api/v1/assistant/route", headers=_auth(key), json={"input": "x"})
    assert r.json()["matched"] is False