vent_ai_platform/tests/test_knowledge_hub.py

from __future__ import annotations

from typing import Any

from fastapi.testclient import TestClient

from app.main import app
from app.services.knowledge_hub_service import (
    KnowledgeHubClient,
    _extract_section_blocks,
    _extract_section_ranges,
    reset_knowledge_hub_state,
)

NOTES = [
    {
        "note_id": "raw/papers/2025-safety-regulations",
        "file_path": "raw/papers/2025-safety-regulations.md",
        "title": "煤矿安全规程",
        "aliases": ["2025版安全规程"],
        "tags": ["regulation/national", "safety/gas"],
        "created": "2026-01-01",
        "updated": "2026-09-01",
        "chunk_count": 12,
        "confidence": "high",
        "links": [],
    },
    {
        "note_id": "concepts/局部通风与掘进通风",
        "file_path": "concepts/局部通风与掘进通风.md",
        "title": "局部通风与掘进通风",
        "aliases": ["局部通风机"],
        "tags": ["ventilation/local-fan", "ventilation/airflow"],
        "created": "2026-02-01",
        "updated": "2026-08-01",
        "chunk_count": 4,
        "confidence": "high",
        "links": [],
    },
    {
        "note_id": "queries/瓦斯超限处置分析",
        "file_path": "queries/瓦斯超限处置分析.md",
        "title": "瓦斯超限处置分析",
        "aliases": [],
        "tags": ["workflow/query", "safety/gas"],
        "created": "2026-03-01",
        "updated": "2026-07-01",
        "chunk_count": 3,
        "confidence": "medium",
        "links": [],
    },
    {
        "note_id": "raw/cases/瓦斯超限事故案例",
        "file_path": "raw/cases/瓦斯超限事故案例.md",
        "title": "瓦斯超限事故案例",
        "aliases": [],
        "tags": ["source/case", "safety/gas"],
        "created": "2026-04-01",
        "updated": "2026-06-01",
        "chunk_count": 2,
        "confidence": "high",
        "links": [],
    },
]


async def fake_request(
    self: KnowledgeHubClient,
    method: str,
    path: str,
    *,
    params: dict[str, Any] | None = None,
    json: dict[str, Any] | None = None,
) -> Any:
    del self, method, json
    if path == "/api/admin/status":
        return {
            "status": "ok",
            "mode": "retrieval-only",
            "stats": {
                "total_notes": len(NOTES),
                "total_chunks": 21,
                "last_indexed": "2026-10-13T10:00:00+08:00",
            },
            "embedding_model": "disabled",
        }
    if path == "/api/tags":
        return {"tags": [{"tag": "safety/gas", "count": 3}]}
    if path == "/api/notes":
        return {
            "total": len(NOTES),
            "page": 1,
            "page_size": 200,
            "notes": NOTES,
        }
    if path == "/api/search":
        query = str((params or {}).get("q") or "")
        if "瓦斯" not in query:
            return {"query": query, "result": [], "results": []}
        return {
            "query": query,
            "mode": "keyword",
            "total": 1,
            "results": [
                {
                    "chunk_id": "raw/cases/瓦斯超限事故案例:0",
                    "note_id": "raw/cases/瓦斯超限事故案例",
                    "title": "瓦斯超限事故案例",
                    "file_path": "raw/cases/瓦斯超限事故案例.md",
                    "heading": "处置要求",
                    "content": "瓦斯超限后应当立即停止作业、切断电源、撤出人员，并检查通风系统。",
                    "score": 0.82,
                }
            ],
        }
    if path == "/api/note_full":
        return {
            "note_id": "raw/papers/2025-safety-regulations",
            "full_content": "# 煤矿安全规程\n瓦斯超限后应当停止作业。",
            "chunks": [
                {
                    "chunk_id": "raw/papers/2025-safety-regulations:0",
                    "heading": "瓦斯管理",
                    "content": "瓦斯超限后应当停止作业。",
                }
            ],
        }
    if path.startswith("/api/notes/"):
        return {
            **NOTES[0],
            "backlinks": [
                {
                    "note_id": "concepts/局部通风与掘进通风",
                    "title": "局部通风与掘进通风",
                }
            ],
        }
    raise AssertionError(f"unexpected request: {path}")


def test_knowledge_hub_page_endpoints(monkeypatch) -> None:
    monkeypatch.setattr(KnowledgeHubClient, "_request", fake_request)
    reset_knowledge_hub_state()
    try:
        with TestClient(app) as client:
            overview = client.get("/api/v1/knowledge/overview")
            assert overview.status_code == 200
            overview_data = overview.json()
            assert overview_data["total_notes"] == 4
            assert overview_data["total_chunks"] == 21
            assert overview_data["service"]["semantic_enabled"] is False
            counts = {item["key"]: item["count"] for item in overview_data["category_counts"]}
            assert counts == {
                "standards": 1,
                "cases": 1,
                "expert": 1,
                "management": 1,
            }

            catalog = client.get(
                "/api/v1/knowledge/catalog",
                params={"category": "standards", "page": 1, "page_size": 20},
            )
            assert catalog.status_code == 200
            assert catalog.json()["total"] == 1
            assert catalog.json()["items"][0]["title"] == "煤矿安全规程"

            answer = client.post(
                "/api/v1/knowledge/ask",
                json={"question": "瓦斯超限如何处置？", "top_k": 3},
            )
            assert answer.status_code == 200
            answer_data = answer.json()
            assert answer_data["answer_status"] == "extractive"
            assert answer_data["search_mode"] == "keyword"
            assert answer_data["sources"][0]["title"] == "瓦斯超限事故案例"
            assert "停止作业" in answer_data["answer"]

            async def fake_generated_answer(question, candidates):
                del question
                assert candidates
                return (
                    "结论\n瓦斯超限后应立即停止作业。[1]\n\n操作建议\n1. 切断电源并撤出人员。",
                    ["切断电源并撤出人员。"],
                    "pytest-rag-model",
                    128,
                )

            monkeypatch.setattr(
                "app.services.knowledge_hub_service._generated_answer",
                fake_generated_answer,
            )
            rag_answer = client.post(
                "/api/v1/knowledge/rag",
                json={"question": "瓦斯超限如何处置？", "top_k": 3},
            )
            assert rag_answer.status_code == 200
            rag_data = rag_answer.json()
            assert rag_data["answer_status"] == "generated"
            assert rag_data["answer_model"] == "pytest-rag-model"
            assert rag_data["generation_enabled"] is True
            assert rag_data["context_chars"] == 128
            assert rag_data["sources"][0]["citation_index"] == 1

            async def fake_stream_deltas(prompt, *, system_prompt):
                del prompt, system_prompt
                yield "thinking", "先确认瓦斯超限，再给出处置步骤。"
                for part in ["瓦斯超限后应立即停止作业。[1]", "\n切断电源并撤出人员。"]:
                    yield "answer", part

            monkeypatch.setattr(
                "app.services.knowledge_hub_service._stream_model_deltas",
                fake_stream_deltas,
            )
            streamed = client.post(
                "/api/v1/knowledge/rag/stream",
                json={"question": "瓦斯超限如何处置？", "top_k": 3},
            )
            assert streamed.status_code == 200
            assert "event: meta" in streamed.text
            assert "event: thinking" in streamed.text
            assert "event: delta" in streamed.text
            assert "event: done" in streamed.text
            assert "停止作业" in streamed.text

            streamed_get = client.get(
                "/api/v1/knowledge/rag/stream",
                params={"question": "瓦斯超限如何处置？", "top_k": 3},
            )
            assert streamed_get.status_code == 200
            assert "event: meta" in streamed_get.text
            assert "event: thinking" in streamed_get.text
            assert "event: delta" in streamed_get.text
            assert "event: done" in streamed_get.text

            with client.websocket_connect("/api/v1/knowledge/rag/ws") as websocket:
                websocket.send_json({"question": "瓦斯超限如何处置？", "top_k": 3})
                websocket_events = []
                while True:
                    message = websocket.receive_json()
                    websocket_events.append(message)
                    if message["event"] == "done":
                        break
            assert [item["event"] for item in websocket_events][:3] == [
                "meta",
                "thinking",
                "delta",
            ]
            assert websocket_events[-1]["data"]["answer_status"] == "generated"

            chat_answer = client.post(
                "/api/v1/knowledge/rag",
                json={"question": "你好"},
            )
            assert chat_answer.status_code == 200
            chat_data = chat_answer.json()
            assert chat_data["answer_status"] == "conversational"
            assert chat_data["search_mode"] == "conversation"
            assert chat_data["sources"] == []
            assert "通防知识助手" in chat_data["answer"]

            graph = client.get("/api/v1/knowledge/graph")
            assert graph.status_code == 200
            assert len(graph.json()["nodes"]) == 5
            assert graph.json()["nodes"][0]["node_type"] == "risk"

            detail = client.get("/api/v1/knowledge/catalog/raw/papers/2025-safety-regulations")
            assert detail.status_code == 200
            assert detail.json()["title"] == "煤矿安全规程"
            assert "瓦斯超限" in detail.json()["full_content"]
            assert detail.json()["chunks"][0]["heading"] == "瓦斯管理"
    finally:
        reset_knowledge_hub_state()


def test_section_range_expansion() -> None:
    assert _extract_section_ranges("应按 5.2.2—5.2.6 分别计算") == [
        ("5.2.2", "5.2.6")
    ]
    assert _extract_section_ranges("应按 5.2.2 一 5.2.6 分别计算") == [
        ("5.2.2", "5.2.6")
    ]
    content = "\n".join(
        [
            "5.2.1 采区需风量",
            "5.2.2 采煤工作面需风量",
            "5.2.2.1 按瓦斯涌出量计算",
            "5.2.6 备用工作面需风量",
            "5.3 矿井总需风量",
        ]
    )
    block = _extract_section_blocks(content, "5.2.2", "5.2.6")
    assert "5.2.2 采煤工作面需风量" in block
    assert "5.2.2.1 按瓦斯涌出量计算" in block
    assert "5.2.6 备用工作面需风量" in block
    assert "5.3 矿井总需风量" not in block