harness_zbt/backend/src/zbt/domains/knowledge/service.py

"""知识库应用服务。

它把管理端的“上传 → 索引 → 检索测试 → 发布”串成完整流程，并把
Milvus 语义检索、BM25 关键词检索和 BGE 重排封装在一个稳定入口中。
Agent 只能检索已经发布的知识，草稿和仅完成索引的文档不会进入正式回答。
"""

import logging
from collections.abc import Callable
from dataclasses import replace
from datetime import datetime
from typing import Any

from zbt.core.errors import AppError
from zbt.core.identifiers import new_ulid
from zbt.domains.knowledge.chunking import DocumentChunker
from zbt.domains.knowledge.index import KnowledgeIndex
from zbt.domains.knowledge.models import (
    KnowledgeChunk,
    KnowledgeDocument,
    KnowledgeSearchTest,
)
from zbt.domains.knowledge.repository import KnowledgeRepository
from zbt.domains.knowledge.retrieval import (
    KnowledgeReranker,
    rank_bm25,
    reciprocal_rank_fusion,
)

logger = logging.getLogger(__name__)


class KnowledgeService:
    """协调知识文档状态、索引和检索的应用服务。"""

    def __init__(
        self,
        repository: KnowledgeRepository,
        clock: Callable[[], datetime],
        index: KnowledgeIndex | None = None,
        chunker: DocumentChunker | None = None,
        reranker: KnowledgeReranker | None = None,
        candidate_multiplier: int = 4,
    ) -> None:
        self._repository = repository
        self._clock = clock
        self._index = index
        self._chunker = chunker or DocumentChunker()
        self._reranker = reranker
        self._candidate_multiplier = max(candidate_multiplier, 1)

    def create_document(
        self,
        *,
        title: str,
        document_type: str,
        source_name: str,
        product_code: str | None,
        content: str,
        created_by: str,
    ) -> dict[str, Any]:
        """创建知识草稿。

        这里只把原始文档保存到 MySQL；文档尚未切分，也不会被 Agent 检索到。
        """

        now = self._clock()
        document = KnowledgeDocument(
            id=new_ulid(),
            title=title.strip(),
            document_type=document_type,
            source_name=source_name.strip(),
            product_code=product_code.strip() if product_code else None,
            content=content.strip(),
            status="DRAFT",
            version_no=1,
            chunk_count=0,
            created_by=created_by,
            created_at=now,
            updated_at=now,
        )
        self._repository.save_document(document)
        return self._serialize_document(document)

    def list_documents(self) -> dict[str, Any]:
        documents = sorted(
            self._repository.list_documents(),
            key=lambda item: item.created_at,
            reverse=True,
        )
        items = [self._serialize_document(document) for document in documents]
        return {"items": items, "total": len(items)}

    def index_document(self, document_id: str) -> dict[str, Any]:
        """切分文档并替换其在 Milvus 中的全部向量片段。

        使用“整篇替换”可以避免文档重新索引后残留旧版本片段。
        """

        document = self._repository.get_document(document_id)
        if document is None:
            raise AppError("KNOWLEDGE_DOCUMENT_NOT_FOUND", "未找到知识文档", 404)
        if self._index is None:
            raise AppError(
                "KNOWLEDGE_INDEX_UNAVAILABLE",
                "知识索引服务尚未就绪",
                503,
                retryable=True,
            )
        chunks = self._split_document(document)
        self._index.replace_document(document.id, chunks)
        now = self._clock()
        indexed = replace(
            document,
            status="INDEXED",
            chunk_count=len(chunks),
            indexed_at=now,
            published_at=None,
            updated_at=now,
        )
        self._repository.save_document(indexed)
        return self._serialize_document(indexed)

    def publish_document(self, document_id: str) -> dict[str, Any]:
        """发布已经验证过的知识。

        发布前必须存在本次索引之后产生的有效检索测试记录，防止“写入成功但
        实际搜不到”的知识直接被 Agent 使用。
        """

        document = self._repository.get_document(document_id)
        if document is None:
            raise AppError("KNOWLEDGE_DOCUMENT_NOT_FOUND", "未找到知识文档", 404)
        if document.status != "INDEXED":
            raise AppError(
                "KNOWLEDGE_DOCUMENT_NOT_INDEXED",
                "知识文档完成索引和检索测试后才能发布",
                409,
            )
        effective_tests = [
            search_test
            for search_test in self._repository.list_search_tests(document_id)
            if search_test.passed
            and (
                document.indexed_at is None
                or search_test.created_at >= document.indexed_at
            )
        ]
        if not effective_tests:
            raise AppError(
                "KNOWLEDGE_SEARCH_TEST_REQUIRED",
                "知识文档至少完成一次有结果的检索测试后才能发布",
                409,
            )
        now = self._clock()
        published = replace(
            document,
            status="PUBLISHED",
            published_at=now,
            updated_at=now,
        )
        self._repository.save_document(published)
        return self._serialize_document(published)

    def disable_document(self, document_id: str) -> dict[str, Any]:
        document = self._repository.get_document(document_id)
        if document is None:
            raise AppError("KNOWLEDGE_DOCUMENT_NOT_FOUND", "未找到知识文档", 404)
        if document.status == "INACTIVE":
            return self._serialize_document(document)
        if document.status != "PUBLISHED":
            raise AppError(
                "KNOWLEDGE_DOCUMENT_NOT_PUBLISHED",
                "只有已发布知识才能停用",
                409,
            )
        now = self._clock()
        disabled = replace(document, status="INACTIVE", updated_at=now)
        self._repository.save_document(disabled)
        return self._serialize_document(disabled)

    def search(
        self,
        query: str,
        *,
        limit: int = 5,
        product_code: str | None = None,
        document_types: tuple[str, ...] = (),
    ) -> dict[str, Any]:
        """检索所有符合过滤条件的已发布文档。

        这是 Agent 工具实际调用的入口，状态过滤在服务端完成，模型无法绕过。
        """

        if self._index is None:
            raise AppError(
                "KNOWLEDGE_INDEX_UNAVAILABLE",
                "知识索引服务尚未就绪",
                503,
                retryable=True,
            )
        published_documents = [
            document
            for document in self._repository.list_documents()
            if document.status == "PUBLISHED"
            and (product_code is None or document.product_code == product_code)
            and (
                not document_types
                or document.document_type in set(document_types)
            )
        ]
        return self._search_documents(query, published_documents, limit=limit)

    def test_search(
        self,
        document_id: str,
        query: str,
        *,
        limit: int = 5,
        tested_by: str = "SYSTEM",
    ) -> dict[str, Any]:
        """只针对指定文档执行检索，并持久化本次测试结果。

        该记录既供管理端查看，也是 ``publish_document`` 的发布门禁依据。
        """

        document = self._repository.get_document(document_id)
        if document is None:
            raise AppError("KNOWLEDGE_DOCUMENT_NOT_FOUND", "未找到知识文档", 404)
        if document.status not in {"INDEXED", "PUBLISHED"}:
            raise AppError(
                "KNOWLEDGE_DOCUMENT_NOT_INDEXED",
                "知识文档完成索引后才能测试检索",
                409,
            )
        result = self._search_documents(query, [document], limit=limit)
        now = self._clock()
        search_test = KnowledgeSearchTest(
            id=new_ulid(),
            document_id=document_id,
            query=query.strip(),
            hit_count=int(result["total"]),
            top_score=(
                float(result["items"][0]["score"]) if result["items"] else None
            ),
            passed=bool(result["items"]),
            tested_by=tested_by,
            created_at=now,
        )
        self._repository.save_search_test(search_test)
        return {
            "document_id": document_id,
            "test_record": self._serialize_search_test(search_test),
            **result,
        }

    def list_search_tests(self, document_id: str) -> dict[str, Any]:
        if self._repository.get_document(document_id) is None:
            raise AppError("KNOWLEDGE_DOCUMENT_NOT_FOUND", "未找到知识文档", 404)
        search_tests = sorted(
            self._repository.list_search_tests(document_id),
            key=lambda item: item.created_at,
            reverse=True,
        )
        items = [self._serialize_search_test(item) for item in search_tests]
        return {"document_id": document_id, "items": items, "total": len(items)}

    def _search_documents(
        self,
        query: str,
        documents: list[KnowledgeDocument],
        *,
        limit: int,
    ) -> dict[str, Any]:
        """执行混合检索主链路：向量召回 → BM25 召回 → RRF → 可选重排。"""

        if self._index is None:
            raise AppError(
                "KNOWLEDGE_INDEX_UNAVAILABLE",
                "知识索引服务尚未就绪",
                503,
                retryable=True,
            )
        document_ids = {document.id for document in documents}
        candidate_limit = limit * self._candidate_multiplier
        # 候选数量先放大，给融合和重排阶段留下足够的选择空间。
        dense_hits = self._index.search(
            query.strip(),
            document_ids=document_ids,
            limit=candidate_limit,
        )
        lexical_hits = rank_bm25(
            query,
            (
                chunk
                for document in documents
                for chunk in self._split_document(document)
            ),
            limit=candidate_limit,
        )
        candidates = reciprocal_rank_fusion(
            [dense_hits, lexical_hits],
            limit=candidate_limit,
        )
        hits = candidates[:limit]
        if self._reranker is not None:
            try:
                hits = self._reranker.rerank(query, candidates, limit=limit)
            except Exception:
                # 语义重排只用于提升排序效果，不应成为知识问答的单点故障。
                # 失败时继续使用 Milvus + BM25 的融合结果，并保留完整异常日志。
                logger.exception(
                    "Knowledge reranker failed; falling back to fused retrieval results"
                )
        items = [
            {
                "chunk_id": hit.chunk_id,
                "document_id": hit.document_id,
                "document_version": hit.document_version,
                "ordinal": hit.ordinal,
                "title": hit.title,
                "content": hit.content,
                "document_type": hit.document_type,
                "source_name": hit.source_name,
                "product_code": hit.product_code,
                "score": hit.score,
            }
            for hit in hits
        ]
        return {"items": items, "total": len(items)}

    def _split_document(self, document: KnowledgeDocument) -> list[KnowledgeChunk]:
        """把文档切成带来源、版本和顺序号的可追溯知识片段。"""

        paragraphs = self._chunker.split(document.content)
        if not paragraphs:
            raise AppError("KNOWLEDGE_DOCUMENT_EMPTY", "知识文档没有可索引内容", 422)
        return [
            KnowledgeChunk(
                id=f"{document.id}:{document.version_no}:{ordinal}",
                document_id=document.id,
                document_version=document.version_no,
                ordinal=ordinal,
                title=document.title,
                content=content,
                document_type=document.document_type,
                source_name=document.source_name,
                product_code=document.product_code,
            )
            for ordinal, content in enumerate(paragraphs, start=1)
        ]

    def _serialize_document(self, document: KnowledgeDocument) -> dict[str, Any]:
        search_tests = sorted(
            self._repository.list_search_tests(document.id),
            key=lambda item: item.created_at,
            reverse=True,
        )
        latest_test = search_tests[0] if search_tests else None
        return {
            "document_id": document.id,
            "title": document.title,
            "document_type": document.document_type,
            "source_name": document.source_name,
            "product_code": document.product_code,
            "status": document.status,
            "version_no": document.version_no,
            "chunk_count": document.chunk_count,
            "created_by": document.created_by,
            "created_at": document.created_at,
            "updated_at": document.updated_at,
            "indexed_at": document.indexed_at,
            "published_at": document.published_at,
            "search_test_count": len(search_tests),
            "last_search_test_at": (
                latest_test.created_at if latest_test is not None else None
            ),
            "last_search_test_passed": (
                latest_test.passed if latest_test is not None else None
            ),
            "excerpt": document.content[:160],
        }

    @staticmethod
    def _serialize_search_test(search_test: KnowledgeSearchTest) -> dict[str, Any]:
        return {
            "test_id": search_test.id,
            "document_id": search_test.document_id,
            "query": search_test.query,
            "hit_count": search_test.hit_count,
            "top_score": search_test.top_score,
            "passed": search_test.passed,
            "tested_by": search_test.tested_by,
            "created_at": search_test.created_at,
        }