deerflow-code/offline-backend-20260512/backend/tests/test_llmwiki_wiki_retrieval.py

from __future__ import annotations

from types import SimpleNamespace

import pytest
from openpyxl import load_workbook

from app.gateway.routers.llmwiki import _build_wiki_excel
from deerflow.integrations.weknora.local_index.retrieval import WikiRetrievalService, is_fully_vectorized


class _Store:
    def __init__(self, states):
        self.states = states

    async def list_index_status(self, mapping_ids=None):
        wanted = set(mapping_ids or [])
        return [state for state in self.states if not wanted or state["knowledge_base_mapping_id"] in wanted]


class _WikiClient:
    def __init__(self):
        self.list_calls = []
        self.raw_search_called = False

    async def list_wiki_pages(self, knowledge_base_id, *, page=1, page_size=50, query=""):
        self.list_calls.append((knowledge_base_id, query))
        return {"pages": [{"slug": "guides/install", "title": "安装"}], "total": 1, "total_pages": 1}

    async def get_wiki_page(self, knowledge_base_id, slug):
        return {
            "id": f"{knowledge_base_id}:{slug}",
            "slug": slug,
            "title": "安装",
            "summary": "Wiki 摘要",
            "content": "# 安装\n这是处理后的 Wiki 正文。",
            "status": "published",
            "category_path": ["指南"],
        }

    async def search(self, query, knowledge_base_ids):
        self.raw_search_called = True
        raise AssertionError("raw knowledge search must never be called")


class _VectorSearch:
    def __init__(self):
        self.embedding = SimpleNamespace(fingerprint="fp-current")
        self.calls = []

    async def search(self, query, mappings, **kwargs):
        self.calls.append([mapping["id"] for mapping in mappings])
        mapping = mappings[0]
        return {
            "results": [
                {
                    "kind": "wiki_page",
                    "id": "local-page",
                    "wiki_page_id": "local-page",
                    "wiki_slug": "local/page",
                    "knowledge_base_id": mapping["id"],
                    "title": "本地页面",
                    "content": "本地 Wiki",
                    "matched_sections": [],
                    "score": 0.95,
                }
            ],
            "warnings": [],
        }


def _ready_state(mapping_id: str):
    return {
        "knowledge_base_mapping_id": mapping_id,
        "state": "idle",
        "last_success_at": "2026-08-24T00:00:00Z",
        "embedding_fingerprint": "fp-current",
        "local_page_count": 1,
        "ready_page_count": 1,
        "failed_page_count": 0,
        "vector_count": 3,
    }


def test_fully_vectorized_requires_completed_full_consistent_scan() -> None:
    state = _ready_state("kb")
    assert is_fully_vectorized(state, "fp-current")
    assert not is_fully_vectorized({**state, "state": "partial"}, "fp-current")
    assert not is_fully_vectorized({**state, "ready_page_count": 0}, "fp-current")
    assert not is_fully_vectorized(state, "different-fingerprint")


@pytest.mark.asyncio
async def test_unvectorized_library_uses_only_wiki_page_search_api() -> None:
    client = _WikiClient()
    service = WikiRetrievalService(_Store([]), vector_search=_VectorSearch(), client=client)
    mapping = {"id": "local-kb", "weknora_id": "remote-kb", "name": "测试库"}

    result = await service.search("安装", [mapping])

    assert result["mode"] == "wiki_api"
    assert result["fallback_knowledge_base_ids"] == ["local-kb"]
    assert result["results"][0]["kind"] == "wiki_page"
    assert result["results"][0]["content"] == "# 安装\n这是处理后的 Wiki 正文。"
    assert "chunk_id" not in result["results"][0]
    assert client.list_calls == [("remote-kb", "安装")]
    assert not client.raw_search_called


@pytest.mark.asyncio
async def test_mixed_search_vectors_complete_library_and_falls_back_incomplete_library() -> None:
    client = _WikiClient()
    vector = _VectorSearch()
    service = WikiRetrievalService(_Store([_ready_state("ready")]), vector_search=vector, client=client)
    mappings = [
        {"id": "ready", "weknora_id": "remote-ready", "name": "已完成"},
        {"id": "pending", "weknora_id": "remote-pending", "name": "未完成"},
    ]

    result = await service.search("安装", mappings, top_k_pages=8)

    assert result["mode"] == "mixed"
    assert vector.calls == [["ready"]]
    assert client.list_calls == [("remote-pending", "安装")]
    assert {item["knowledge_base_id"] for item in result["results"]} == {"ready", "pending"}


def test_excel_export_contains_current_directory_and_every_parent_directory() -> None:
    long_content = "正文" * 20_000
    pages = [
        {"id": "1", "slug": "root", "title": "根专题", "content": "根", "status": "published"},
        {"id": "2", "slug": "root/child", "title": "子专题", "parent_slug": "root", "content": "子", "status": "published"},
        {
            "id": "3",
            "slug": "root/child/article",
            "title": "文章",
            "parent_slug": "root/child",
            "content": long_content,
            "status": "published",
        },
    ]

    workbook = load_workbook(_build_wiki_excel("知识库", pages), read_only=True)
    sheet = workbook["Wiki 全量导出"]
    rows = list(sheet.iter_rows(values_only=True))
    headings = list(rows[0])
    article = dict(zip(headings, rows[3], strict=True))

    assert article["目录完整路径"] == "根专题 / 子专题"
    assert article["所在目录"] == "子专题"
    assert article["父级目录完整路径"] == "根专题"
    assert article["第1级目录"] == "根专题"
    assert article["第2级目录"] == "子专题"
    assert article["正文"] + article["正文（续1）"] == long_content