RAG 数据管道构建：从原始数据到高质量知识库

原始数据 → 接入 → 清洗 → 标准化 → 分块 → 元数据 → 嵌入 → 向量库 → 检索 → 生成
          ──────────── ETL ────────────  ──── 索引 ────

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Iterator
import hashlib

@dataclass
class Document:
    id: str
    content: str
    source: str
    source_type: str  # pdf, html, markdown, database, api
    metadata: dict
    content_hash: str
    updated_at: str

class DataSource(ABC):
    @abstractmethod
    def fetch(self) -> Iterator[Document]:
        pass

class WebPageSource(DataSource):
    def __init__(self, urls: list, cache_dir: str):
        self.urls = urls
        self.cache = cache_dir

    def fetch(self) -> Iterator[Document]:
        for url in self.urls:
            html = self._fetch_with_cache(url)
            text = self._html_to_markdown(html)
            yield Document(
                id=self._url_hash(url),
                content=text,
                source=url,
                source_type="web",
                metadata={"title": self._extract_title(html)},
                content_hash=hashlib.md5(text.encode()).hexdigest(),
                updated_at=datetime.now().isoformat()
            )

class PDFSource(DataSource):
    def fetch(self) -> Iterator[Document]:
        for pdf_path in self.pdf_files:
            pages = self._extract_pdf_pages(pdf_path)
            for i, page_text in enumerate(pages):
                yield Document(
                    id=f"{pdf_path}_p{i+1}",
                    content=page_text,
                    source=pdf_path,
                    source_type="pdf",
                    metadata={"page": i+1, "total_pages": len(pages)},
                    content_hash=hashlib.md5(page_text.encode()).hexdigest(),
                    updated_at=datetime.now().isoformat()
                )

class DatabaseSource(DataSource):
    def __init__(self, conn_str: str, query: str):
        self.conn_str = conn_str
        self.query = query

    def fetch(self) -> Iterator[Document]:
        import psycopg2
        with psycopg2.connect(self.conn_str) as conn:
            with conn.cursor() as cur:
                cur.execute(self.query)
                for row in cur:
                    yield Document(
                        id=f"db_{row[0]}",
                        content=row[1],
                        source=f"database:{row[0]}",
                        source_type="database",
                        metadata={"table": row[2]} if len(row) > 2 else {},
                        content_hash=hashlib.md5(row[1].encode()).hexdigest(),
                        updated_at=datetime.now().isoformat()
                    )

class DocumentCleaner:
    """文档清洗流水线"""

    def clean(self, doc: Document) -> Document:
        steps = [
            self._remove_html_tags,
            self._normalize_whitespace,
            self._remove_boilerplate,
            self._fix_encoding,
            self._normalize_punctuation,
        ]
        for step in steps:
            doc.content = step(doc.content)
        return doc

    def _remove_html_tags(self, text: str) -> str:
        import re
        # 保留代码块中的内容
        code_blocks = re.findall(r'```.*?```', text, re.DOTALL)
        text = re.sub(r'<[^>]+>', '', text)
        # 恢复代码块
        for i, block in enumerate(code_blocks):
            text = text.replace(f'CODE_BLOCK_{i}', block, 1)
        return text

    def _normalize_whitespace(self, text: str) -> str:
        import re
        # 合并连续空行
        text = re.sub(r'\n{3,}', '\n\n', text)
        # 去除行尾空格
        text = '\n'.join(line.rstrip() for line in text.split('\n'))
        # 合并连续空格
        text = re.sub(r' {2,}', ' ', text)
        return text.strip()

    def _remove_boilerplate(self, text: str) -> str:
        """移除样板文字"""
        boilerplate_patterns = [
            r'Cookie.*?声明',
            r'版权所有.*?保留',
            r'本文档仅供参考',
            r'Last updated:.*?$',
        ]
        for pattern in boilerplate_patterns:
            text = re.sub(pattern, '', text, flags=re.IGNORECASE)
        return text

    def _fix_encoding(self, text: str) -> str:
        # 修复常见编码问题
        replacements = {
            '\ufeff': '',       # BOM
            '\u200b': '',       # 零宽空格
            '\u200c': '',       # 零宽非连接符
            '\u00a0': ' ',      # 不间断空格
            ''': "'", ''': "'",
            '"': '"', '"': '"',
            '—': '-', '–': '-',
        }
        for old, new in replacements.items():
            text = text.replace(old, new)
        return text

from typing import List
import tiktoken

class RecursiveChunker:
    def __init__(
        self,
        chunk_size: int = 512,
        chunk_overlap: int = 50,
        model: str = "gpt-4o"
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.encoder = tiktoken.encoding_for_model(model)
        # 分隔符按优先级排序
        self.separators = [
            "\n## ",    # Markdown 二级标题
            "\n### ",   # Markdown 三级标题
            "\n\n",     # 段落
            "\n",       # 行
            "。",       # 中文句号
            ". ",       # 英文句号
            " ",        # 空格
            "",         # 字符级（最后手段）
        ]

    def chunk(self, text: str) -> List[str]:
        return self._recursive_split(text, self.separators)

    def _recursive_split(self, text: str, separators: list) -> List[str]:
        if self._count_tokens(text) <= self.chunk_size:
            return [text] if text.strip() else []

        for i, sep in enumerate(separators):
            if sep == "":
                return self._force_split(text)

            if sep not in text:
                continue

            parts = text.split(sep)
            chunks = []
            current = ""

            for part in parts:
                candidate = current + sep + part if current else part
                if self._count_tokens(candidate) <= self.chunk_size:
                    current = candidate
                else:
                    if current:
                        chunks.append(current)
                    # 如果单个 part 就超长，用更细的分隔符递归
                    if self._count_tokens(part) > self.chunk_size:
                        sub_chunks = self._recursive_split(part, separators[i+1:])
                        chunks.extend(sub_chunks)
                        current = ""
                    else:
                        current = part

            if current:
                chunks.append(current)

            # 添加重叠
            return self._add_overlap(chunks)

        return self._force_split(text)

    def _add_overlap(self, chunks: List[str]) -> List[str]:
        if len(chunks) <= 1 or self.chunk_overlap <= 0:
            return chunks

        result = [chunks[0]]
        for i in range(1, len(chunks)):
            # 从前一个 chunk 末尾取 overlap 个 token
            prev_tokens = self.encoder.encode(chunks[i-1])
            overlap_text = self.encoder.decode(
                prev_tokens[-self.chunk_overlap:]
            )
            result.append(overlap_text + chunks[i])
        return result

    def _count_tokens(self, text: str) -> int:
        return len(self.encoder.encode(text))

    def _force_split(self, text: str) -> List[str]:
        """按 token 强制切割"""
        tokens = self.encoder.encode(text)
        chunks = []
        for i in range(0, len(tokens), self.chunk_size):
            chunk_tokens = tokens[i:i + self.chunk_size]
            chunks.append(self.encoder.decode(chunk_tokens))
        return chunks

class SemanticChunker:
    """基于句子嵌入相似度的语义分块"""

    def __init__(self, model, threshold: float = 0.5):
        self.model = model  # SentenceTransformer
        self.threshold = threshold

    def chunk(self, text: str) -> List[str]:
        # 1. 按句子分割
        sentences = self._split_sentences(text)
        if len(sentences) <= 1:
            return [text]

        # 2. 计算相邻句子嵌入的余弦相似度
        embeddings = self.model.encode(sentences)
        similarities = [
            cosine_similarity([embeddings[i]], [embeddings[i+1]])[0][0]
            for i in range(len(embeddings) - 1)
        ]

        # 3. 在相似度低谷处分割
        chunks = []
        current = [sentences[0]]
        for i, sim in enumerate(similarities):
            if sim < self.threshold:
                chunks.append(" ".join(current))
                current = [sentences[i + 1]]
            else:
                current.append(sentences[i + 1])

        if current:
            chunks.append(" ".join(current))

        return chunks

class MetadataExtractor:
    """从文档中提取结构化元数据"""

    def extract(self, doc: Document) -> dict:
        return {
            **doc.metadata,
            "language": self._detect_language(doc.content),
            "doc_type": self._classify_doc_type(doc.content),
            "key_phrases": self._extract_key_phrases(doc.content),
            "has_code": "```" in doc.content or "def " in doc.content,
            "has_table": "|" in doc.content and "---" in doc.content,
            "reading_time_min": max(1, len(doc.content) // 500),
            "summary": self._generate_summary(doc.content),
        }

    def _detect_language(self, text: str) -> str:
        # 简单启发式
        chinese_ratio = sum(1 for c in text if '\u4e00' <= c <= '\u9fff') / max(len(text), 1)
        if chinese_ratio > 0.3:
            return "zh"
        return "en"

    def _classify_doc_type(self, text: str) -> str:
        if "```" in text:
            return "technical"
        if any(kw in text for kw in ["第", "章", "节"]):
            return "document"
        if "?" in text and len(text) < 500:
            return "faq"
        return "article"

    def _extract_key_phrases(self, text: str) -> list:
        # 使用 TextRank 或 YAKE 提取关键词
        from keybert import KeyBERT
        kw_model = KeyBERT()
        keywords = kw_model.extract_keywords(text, top_n=5)
        return [kw[0] for kw in keywords]

import asyncio
from typing import List, Tuple

class EmbeddingPipeline:
    def __init__(self, model_name: str = "text-embedding-3-large"):
        self.model_name = model_name
        self.client = AsyncOpenAI()
        self.batch_size = 100  # OpenAI 批量限制

    async def embed_chunks(self, chunks: List[str]) -> List[List[float]]:
        # 分批处理
        all_embeddings = []
        for i in range(0, len(chunks), self.batch_size):
            batch = chunks[i:i + self.batch_size]
            response = await self.client.embeddings.create(
                model=self.model_name,
                input=batch
            )
            all_embeddings.extend([d.embedding for d in response.data])
        return all_embeddings

    async def embed_with_retry(self, text: str, max_retries: int = 3):
        for attempt in range(max_retries):
            try:
                response = await self.client.embeddings.create(
                    model=self.model_name,
                    input=text
                )
                return response.data[0].embedding
            except Exception as e:
                if attempt == max_retries - 1:
                    raise
                await asyncio.sleep(2 ** attempt)

class VectorIndexer:
    def __init__(self, pinecone_client):
        self.pinecone = pinecone_client
        self.index_name = "rag-knowledge"

    async def upsert(self, chunks: List[dict]):
        """
        chunks: [{id, content, embedding, metadata}]
        """
        index = self.pinecone.Index(self.index_name)

        # 批量 upsert
        batch_size = 100
        for i in range(0, len(chunks), batch_size):
            batch = chunks[i:i + batch_size]
            vectors = [
                {
                    "id": c["id"],
                    "values": c["embedding"],
                    "metadata": {
                        "content": c["content"][:1000],  # Pinecone metadata 大小限制
                        **c["metadata"]
                    }
                }
                for c in batch
            ]
            await index.upsert(vectors=vectors)

class IncrementalUpdater:
    """只更新变化的部分，避免全量重建"""

    def __init__(self, vector_indexer, storage):
        self.indexer = vector_indexer
        self.storage = storage  # 记录文档状态的数据库

    async def update(self, new_docs: List[Document]):
        added, updated, deleted = 0, 0, 0

        for doc in new_docs:
            existing = await self.storage.get(doc.id)

            if existing is None:
                # 新文档
                await self._index_doc(doc)
                await self.storage.save(doc)
                added += 1

            elif existing.content_hash != doc.content_hash:
                # 内容变更: 删旧建新
                await self._delete_doc_chunks(doc.id)
                await self._index_doc(doc)
                await self.storage.save(doc)
                updated += 1

            # content_hash 相同则跳过

        # 检查删除: 源中不再存在的文档
        all_ids = {d.id for d in new_docs}
        stored_ids = await self.storage.get_all_ids()
        deleted_ids = stored_ids - all_ids
        for doc_id in deleted_ids:
            await self._delete_doc_chunks(doc_id)
            await self.storage.delete(doc_id)
            deleted += 1

        return {"added": added, "updated": updated, "deleted": deleted}

class PipelineMonitor:
    """数据管道质量监控"""

    async def check_quality(self, docs: List[Document]) -> dict:
        checks = {
            "empty_docs": 0,
            "too_short": 0,      # < 50 字符
            "too_long": 0,       # > 100k 字符
            "encoding_issues": 0,
            "duplicate_content": 0,
            "low_density": 0,    # 信息密度低
        }

        contents = [d.content for d in docs]
        seen_hashes = set()

        for doc in docs:
            if not doc.content.strip():
                checks["empty_docs"] += 1
            elif len(doc.content) < 50:
                checks["too_short"] += 1
            elif len(doc.content) > 100000:
                checks["too_long"] += 1

            if doc.content_hash in seen_hashes:
                checks["duplicate_content"] += 1
            seen_hashes.add(doc.content_hash)

            # 信息密度: 非空非标点字符比例
            alpha_ratio = sum(c.isalnum() for c in doc.content) / max(len(doc.content), 1)
            if alpha_ratio < 0.3:
                checks["low_density"] += 1

        # 检索质量测试
        retrieval_score = await self._test_retrieval(docs)

        return {
            "total_docs": len(docs),
            "quality_issues": checks,
            "retrieval_test_score": retrieval_score,
            "timestamp": datetime.now().isoformat()
        }

    async def _test_retrieval(self, docs) -> float:
        """用标准查询测试检索质量"""
        test_queries = load_test_queries()  # 预定义的测试查询
        correct = 0
        for q in test_queries:
            results = await self.search(q["query"], top_k=5)
            if q["expected_doc_id"] in [r["id"] for r in results]:
                correct += 1
        return correct / len(test_queries)