rag/knowledge/base.py

"""
知识库管理器 - 基础模块

包含：
- 配置常量
- 数据类定义
- 辅助函数
- BM25Index 类
"""

import os
import json
import pickle
import logging
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from rank_bm25 import BM25Okapi
import jieba

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# ==================== 配置常量 ====================

# 向量存储基础路径（位于 knowledge/vector_store/）
VECTOR_STORE_BASE_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "vector_store"
)

# 向量库基础路径（ChromaDB 数据存储）
CHROMA_DB_BASE_PATH = os.path.join(VECTOR_STORE_BASE_PATH, "chroma")

# BM25 索引基础路径
BM25_INDEX_BASE_PATH = os.path.join(VECTOR_STORE_BASE_PATH, "bm25")

# 向量库元数据文件
KB_METADATA_FILE = "kb_metadata.json"

# 预定义的公开知识库名称
PUBLIC_KB_NAME = "public_kb"

# 默认部门列表
DEFAULT_DEPARTMENTS = ["finance", "hr", "tech", "operation", "marketing"]

# 部门名称映射
DEPARTMENT_NAME_MAP = {
    "财务部": "finance", "财务": "finance",
    "人事部": "hr", "人事": "hr", "人力资源部": "hr", "人力资源": "hr",
    "技术部": "tech", "技术": "tech", "研发部": "tech", "研发": "tech",
    "运营部": "operation", "运营": "operation",
    "市场部": "marketing", "市场": "marketing",
    "法务部": "legal", "法务": "legal",
    "行政部": "admin", "行政": "admin",
    "finance": "finance", "hr": "hr", "tech": "tech",
    "operation": "operation", "marketing": "marketing", "legal": "legal", "admin": "admin",
}


# ==================== 数据结构 ====================

@dataclass
class CollectionInfo:
    """向量库信息"""
    name: str
    display_name: str
    document_count: int = 0
    created_at: str = ""
    department: str = ""
    description: str = ""


@dataclass
class SearchResult:
    """检索结果"""
    ids: List[str]
    documents: List[str]
    metadatas: List[dict]
    distances: List[float]
    collection_name: str = ""


# ==================== 辅助函数 ====================

def _get_doc_type(filename: str) -> str:
    """
    根据文件扩展名判断文档类型

    Args:
        filename: 文件名，包含扩展名

    Returns:
        文档类型字符串: 'pdf' | 'word' | 'excel' | 'ppt' | 'other'

    Example:
        >>> _get_doc_type("report.pdf")
        'pdf'
        >>> _get_doc_type("data.xlsx")
        'excel'
    """
    ext = Path(filename).suffix.lower()
    type_map = {
        '.pdf': 'pdf',
        '.docx': 'word', '.doc': 'word',
        '.xlsx': 'excel', '.xls': 'excel',
        '.pptx': 'ppt', '.ppt': 'ppt',
    }
    return type_map.get(ext, 'other')


def _extract_figure_number(caption: str, section: str = '') -> str:
    """
    从 caption 或 section 中提取图号（增强版）

    支持格式：
    - 图2.4, 图2-4
    - Fig.2.4, Fig 2.4, Figure 2.4
    - （图2）

    Args:
        caption: 图片标题/说明文字
        section: 章节信息（可选）

    Returns:
        图号字符串，如 "2.4"；未找到返回空字符串

    Example:
        >>> _extract_figure_number("图2.4 系统架构图")
        '2.4'
        >>> _extract_figure_number("参见Figure 3.1所示")
        '3.1'
    """
    import re
    text = f"{caption} {section}"
    patterns = [
        r'图\s*(\d+[\.\-]\d+)',
        r'Fig\.?\s*(\d+[\.\-]\d+)',
        r'Figure\s*(\d+[\.\-]\d+)',
        r'[（(]\s*图\s*(\d+)\s*[)）]',
    ]
    for pattern in patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            return match.group(1).replace('-', '.')
    return ""


def normalize_department_name(department: str) -> str:
    """
    将部门名称标准化为英文标识

    支持中文部门名（如"财务部"）和英文标识（如"finance"），
    返回符合 ChromaDB 命名规范的英文标识。

    Args:
        department: 原始部门名称（中文或英文）

    Returns:
        标准化的英文标识；无法识别时返回空字符串

    Example:
        >>> normalize_department_name("财务部")
        'finance'
        >>> normalize_department_name("tech")
        'tech'
    """
    if not department:
        return ""
    if department in DEPARTMENT_NAME_MAP:
        return DEPARTMENT_NAME_MAP[department]
    if department.replace("_", "").replace("-", "").isalnum() and department.isascii():
        return department.lower()
    logger.warning(f"无法识别的部门名称: {department}")
    return ""


def _extract_section(section_path: str, max_levels: int = 3) -> str:
    """
    动态截断章节路径，保留核心+末尾层级

    用于在向量检索时提供简洁的章节上下文，
    避免过长的章节路径影响语义匹配。

    Args:
        section_path: 完整章节路径，用 '>' 分隔
        max_levels: 最多保留的层级数（默认3级）

    Returns:
        截断后的章节路径

    Example:
        >>> _extract_section("第一章 > 1.1 概述 > 1.1.1 背景 > 1.1.1.1 详细说明")
        '1.1 概述 > 1.1.1 背景 > 1.1.1.1 详细说明'
    """
    parts = [p.strip() for p in section_path.split('>') if p.strip()]
    if len(parts) <= max_levels:
        return ' > '.join(parts)
    return ' > '.join(parts[-max_levels:])


def _build_semantic_content_for_text(chunk, page_info: dict, doc_type: str) -> str:
    """
    构建语义增强内容（文本类型）

    将原始文本切片转换为适合向量检索的语义增强格式，
    包含标题、章节上下文和正文内容。

    PDF 和 Word 差异化处理：
    - PDF: 有 text_level、bbox，标题识别准确
    - Word: 无 text_level、bbox，依赖启发式识别

    Args:
        chunk: 文档切片对象，含 title、content、text_level 属性
        page_info: 页面信息字典，含 section_path、section 等
        doc_type: 文档类型 ('pdf' | 'word' | 'excel' | 'ppt')

    Returns:
        语义增强后的内容字符串，格式为：
        标题
        主题：章节路径
        正文内容
    """
    parts = []
    title = getattr(chunk, 'title', '') or ''
    if isinstance(title, list):
        title = ' '.join(str(t) for t in title if t) or ''
    if not isinstance(title, str):
        title = ''
    text_level = getattr(chunk, 'text_level', 0)

    if title and title.strip() and text_level > 0:
        parts.append(title.strip())

    section = page_info.get('section_path', '') or page_info.get('section', '')
    if isinstance(section, list):
        section = ' > '.join(str(s) for s in section if s) or ''
    if not isinstance(section, str):
        section = ''
    if section and section.strip():
        section = _extract_section(section, max_levels=3)
        parts.append(f"主题：{section}")

    content = chunk.content if hasattr(chunk, 'content') else page_info.get('text', '')
    if isinstance(content, list):
        content = '\n'.join(str(item) for item in content)
    parts.append(content)
    return "\n".join(parts)


def _build_semantic_content_for_table(table_md: str, page_info: dict, chunk, doc_type: str) -> str:
    """
    构建语义增强内容（表格类型）

    将 Markdown 表格转换为包含语义摘要的增强格式，
    提取表头、字段、行数和示例数据，提升向量检索命中率。

    Args:
        table_md: 表格的 Markdown 内容
        page_info: 页面信息字典，含 section_path 等
        chunk: 文档切片对象，含 title 属性
        doc_type: 文档类型 ('pdf' | 'word' | 'excel' | 'ppt')

    Returns:
        语义增强后的内容字符串，包含：
        主题：章节路径
        表格：标题
        字段：表头列表
        描述：行数和字段概要
        示例：首行数据示例
    """
    parts = []
    section = page_info.get('section_path', '') or page_info.get('section', '')
    if isinstance(section, list):
        section = ' > '.join(str(s) for s in section if s) or ''
    if not isinstance(section, str):
        section = ''
    if section and section.strip():
        section = _extract_section(section, max_levels=3)
        parts.append(f"主题：{section}")

    caption = getattr(chunk, 'title', '') or ''
    if isinstance(caption, list):
        caption = ' '.join(str(c) for c in caption if c) or ''
    if caption and caption.strip() and caption != "表格":
        parts.append(f"表格：{caption.strip()}")

    # 检测是否是 HTML 格式，如果是则转换为 Markdown
    if '<table' in table_md.lower():
        try:
            from parsers.mineru_parser import html_table_to_markdown
            table_md = html_table_to_markdown(table_md)
        except Exception as e:
            logger.warning(f"HTML 表格转换失败: {e}")

    lines = table_md.split('\n')
    headers = []
    for line in lines:
        if line.startswith('|') and '---' not in line:
            headers = [h.strip() for h in line.split('|') if h.strip()]
            if headers and len(headers) > 1:
                parts.append(f"字段：{', '.join(headers)}")
            break

    row_count = len([l for l in lines if l.startswith('|') and '---' not in l])
    if headers:
        parts.append(f"描述：该表包含{row_count}行数据，记录各{', '.join(headers[:3])}信息")
    else:
        parts.append(f"描述：该表包含{row_count}行数据")

    for line in lines:
        if line.startswith('|') and '---' not in line and headers:
            cells = [c.strip() for c in line.split('|') if c.strip()]
            if cells and cells != headers:
                example_parts = []
                for i, h in enumerate(headers[:2]):
                    if i < len(cells):
                        example_parts.append(f"{h}={cells[i]}")
                if example_parts:
                    parts.append(f"示例：{', '.join(example_parts)}")
                break

    # 添加完整表格内容（Markdown 格式）
    if lines and any(l.startswith('|') for l in lines):
        parts.append("")  # 空行分隔
        parts.append("表格内容：")
        parts.extend(lines)

    return "\n".join(parts)


# ==================== BM25 索引管理 ====================

class BM25Index:
    """
    BM25 关键词检索索引

    基于 rank_bm25.BM25Okapi 实现，使用 jieba 进行中文分词。
    支持文档的添加、检索、持久化和加载。

    Attributes:
        bm25: BM25Okapi 索引实例
        ids: 文档 ID 列表
        documents: 文档内容列表
        metadatas: 文档元数据列表

    Example:
        >>> index = BM25Index()
        >>> index.add_documents(
        ...     ids=["doc1", "doc2"],
        ...     documents=["财务报销流程", "请假审批制度"],
        ...     metadatas=[{"source": "a.pdf"}, {"source": "b.pdf"}]
        ... )
        >>> ids, docs, metas, scores = index.search("报销", top_k=5)
    """

    def __init__(self) -> None:
        """初始化空的 BM25 索引"""
        self.bm25: Optional[BM25Okapi] = None
        self.ids: List[str] = []
        self.documents: List[str] = []
        self.metadatas: List[dict] = []

    def tokenize(self, text: str) -> List[str]:
        """
        使用 jieba 对文本进行分词

        Args:
            text: 待分词的文本

        Returns:
            分词结果列表
        """
        return list(jieba.cut(text))

    def add_documents(self, ids: List[str], documents: List[str], metadatas: List[dict]) -> None:
        """
        添加文档到索引（追加模式，自动去重）

        如果 ID 已存在则更新对应文档，否则追加新文档。
        添加后自动重建 BM25 索引。

        Args:
            ids: 文档 ID 列表
            documents: 文档内容列表
            metadatas: 文档元数据列表
        """
        # 建立已有 ID -> 索引位置 的映射，用于去重
        existing_map = {doc_id: idx for idx, doc_id in enumerate(self.ids)}

        for i, doc_id in enumerate(ids):
            if doc_id in existing_map:
                # 更新已有文档
                pos = existing_map[doc_id]
                self.documents[pos] = documents[i]
                self.metadatas[pos] = metadatas[i]
            else:
                # 追加新文档
                existing_map[doc_id] = len(self.ids)
                self.ids.append(doc_id)
                self.documents.append(documents[i])
                self.metadatas.append(metadatas[i])

        if self.documents:
            tokenized = [self.tokenize(doc) for doc in self.documents]
            self.bm25 = BM25Okapi(tokenized)

    def search(self, query: str, top_k: int = 10) -> Tuple[List[str], List[str], List[dict], List[float]]:
        """
        检索与查询最相关的文档

        Args:
            query: 查询文本
            top_k: 返回的最大文档数

        Returns:
            元组 (ids, documents, metadatas, scores)：
            - ids: 文档 ID 列表
            - documents: 文档内容列表
            - metadatas: 文档元数据列表
            - scores: BM25 分数列表
        """
        if not self.bm25 or not self.documents:
            return [], [], [], []
        tokenized_query = self.tokenize(query)
        scores = self.bm25.get_scores(tokenized_query)
        top_indices = np.argsort(scores)[::-1][:top_k]
        return (
            [self.ids[i] for i in top_indices],
            [self.documents[i] for i in top_indices],
            [self.metadatas[i] for i in top_indices],
            [float(scores[i]) for i in top_indices]
        )

    def save(self, filepath: str) -> None:
        """
        持久化索引到文件

        Args:
            filepath: 保存路径（.pkl 文件）
        """
        data = {'ids': self.ids, 'documents': self.documents, 'metadatas': self.metadatas}
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'wb') as f:
            pickle.dump(data, f)

    def load(self, filepath: str) -> bool:
        """
        从文件加载索引

        Args:
            filepath: 索引文件路径（.pkl 文件）

        Returns:
            加载成功返回 True，失败返回 False
        """
        if not os.path.exists(filepath):
            return False
        try:
            with open(filepath, 'rb') as f:
                data = pickle.load(f)
            self.ids = data.get('ids', [])
            self.documents = data.get('documents', [])
            self.metadatas = data.get('metadatas', [])
            if self.documents:
                tokenized = [self.tokenize(doc) for doc in self.documents]
                self.bm25 = BM25Okapi(tokenized)
            return True
        except Exception as e:
            logger.error(f"加载 BM25 索引失败: {e}")
            return False

    def clear(self) -> None:
        """清空索引数据"""
        self.bm25 = None
        self.ids = []
        self.documents = []
        self.metadatas = []