import numpy as np
import matplotlib.pyplot as plt

def simulate_cosine_distribution(dim=1536, n_pairs=100000):
    """模拟高维空间中随机向量对的 cosine 相似度分布"""
    # 生成随机向量对
    u = np.random.randn(n_pairs, dim)
    v = np.random.randn(n_pairs, dim)
    # 归一化
    u_norm = u / np.linalg.norm(u, axis=1, keepdims=True)
    v_norm = v / np.linalg.norm(v, axis=1, keepdims=True)
    # cosine 相似度
    cos_sim = np.sum(u_norm * v_norm, axis=1)
    return cos_sim

# 不同维度下的 cosine 相似度标准差
for d in [128, 256, 512, 768, 1024, 1536, 3072]:
    sims = simulate_cosine_distribution(dim=d, n_pairs=50000)
    print(f"dim={d:4d}: mean={np.mean(sims):.6f}, std={np.std(sims):.6f}, "
          f"99%区间=[{np.percentile(sims, 0.5):.4f}, {np.percentile(sims, 99.5):.4f}]")

dim= 128: mean=0.000132, std=0.088384, 99%区间=[-0.2239, 0.2248]
dim= 256: mean=0.000025, std=0.062417, 99%区间=[-0.1598, 0.1595]
dim= 512: mean=0.000532, std=0.044161, 99%区间=[-0.1138, 0.1136]
dim= 768: mean=0.000059, std=0.036108, 99%区间=[-0.0938, 0.0929]
dim=1024: mean=0.000003, std=0.031198, 99%区间=[-0.0805, 0.0805]
dim=1536: mean=0.000078, std=0.025511, 99%区间=[-0.0651, 0.0656]
dim=3072: mean=0.000032, std=0.018038, 99%区间=[-0.0463, 0.0464]

# Embedding 微调的最小闭环：对比学习 + Hard Negative Mining
from sentence_transformers import SentenceTransformer, losses, InputExample
from torch.utils.data import DataLoader
import torch

class HardNegativeMiningDataset:
    """
    生产级 Hard Negative Mining：从检索结果中自动挖掘"被模型误判为高相似度但实际不相关"的负样本
    """
    def __init__(self, model: SentenceTransformer, corpus: list[str],
                 queries: list[str], relevant_docs: dict[str, list[int]],
                 top_k_hard_neg: int = 5):
        self.model = model
        self.corpus = corpus
        self.queries = queries
        self.relevant_docs = relevant_docs  # {query_idx: [relevant_doc_indices]}
        self.top_k_hard_neg = top_k_hard_neg

    def mine_hard_negatives(self) -> list[InputExample]:
        """基于当前模型的检索结果挖掘 Hard Negatives"""
        corpus_emb = self.model.encode(self.corpus, show_progress_bar=True,
                                        convert_to_tensor=True)
        query_emb = self.model.encode(self.queries, show_progress_bar=True,
                                       convert_to_tensor=True)

        examples = []
        for q_idx, q_emb in enumerate(query_emb):
            # 检索 Top-50
            cos_scores = torch.nn.functional.cosine_similarity(
                q_emb.unsqueeze(0), corpus_emb
            )
            _, top_indices = torch.topk(cos_scores, k=50)

            # 正例：标注的相关文档
            pos_docs = set(self.relevant_docs.get(q_idx, []))
            # Hard Negative：Top-50 中非相关的文档（被模型误判为高相似度）
            hard_neg = [i for i in top_indices.tolist()
                        if i not in pos_docs][:self.top_k_hard_neg]

            for pos_idx in pos_docs:
                examples.append(InputExample(
                    texts=[self.queries[q_idx], self.corpus[pos_idx],
                           self.corpus[hard_neg[0]] if hard_neg else
                           self.corpus[np.random.randint(0, len(self.corpus))]],
                    label=1.0  # Triplet loss 不需要 label
                ))
        return examples

# 训练配置
def finetune_embedding(base_model_name: str, train_examples: list,
                       output_path: str, epochs: int = 3):
    model = SentenceTransformer(base_model_name)
    train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32)
    # MultipleNegativesRankingLoss：InfoNCE 的变体，利用 batch 内其他样本作为负例
    train_loss = losses.MultipleNegativesRankingLoss(model)
    model.fit(
        train_objectives=[(train_dataloader, train_loss)],
        epochs=epochs,
        warmup_steps=int(len(train_dataloader) * 0.1),
        output_path=output_path,
        show_progress_bar=True
    )
    return model

检索的候选集大小 = (n_probe / n_list) × N

当 n_list=4096, n_probe=64, N=10M:
  候选集 = 64/4096 × 10M ≈ 156K  → 搜索空间压缩 98.4%

原始存储：N × d × 4 bytes = 10M × 1024 × 4 = 40 GB
PQ 存储：  N × m × 1 byte  = 10M × 128  × 1 = 1.28 GB  ← 压缩比 31:1
其中 m = d / 子空间维度，通常每个子空间 8 维，所以 m = 1024 / 8 = 128 个子空间

import faiss
import numpy as np
import time

def benchmark_ivf_pq_params(dim=1024, n_total=5_000_000, n_query=1000,
                             k=10, ground_truth_indices=None):
    """
    网格搜索 IVF-PQ 最优参数组合，输出 Recall@10 与 QPS 的 Pareto 前沿
    """
    # 生成模拟数据（建议替换为真实向量）
    np.random.seed(42)
    data = np.random.randn(n_total, dim).astype(np.float32)
    queries = np.random.randn(n_query, dim).astype(np.float32)

    # 如果无 ground truth，用暴力检索生成
    if ground_truth_indices is None:
        print("Computing ground truth via brute-force...")
        index_flat = faiss.IndexFlatIP(dim)
        index_flat.add(data)
        _, ground_truth_indices = index_flat.search(queries, k)

    results = []
    for nlist in [512, 1024, 2048, 4096, 8192]:
        for m in [64, 128, 256]:  # 子空间数：总维度/子空间维度
            nbits = 8  # 每个子空间用 8-bit 量化（256 个聚类中心）

            # 构建 IVF-PQ 索引（L2 距离；若需要 cosine，请先对 data 做 L2 归一化）
            quantizer = faiss.IndexFlatL2(dim)
            index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, nbits)
            index.train(data[:500000])  # 用 10% 数据训练
            index.add(data)
            index.nprobe = max(8, nlist // 64)  # probe 数设为合理的默认值

            # 测试 QPS
            start = time.perf_counter()
            _, indices = index.search(queries, k)
            elapsed = time.perf_counter() - start
            qps = n_query / elapsed

            # 计算 Recall
            recall_at_k = np.mean([
                len(set(result) & set(gt)) / k
                for result, gt in zip(indices, ground_truth_indices)
            ])

            # 计算内存占用
            code_size = m * nbits / 8  # bytes per vector
            memory_gb = n_total * code_size / 1e9

            results.append({
                'nlist': nlist, 'm': m, 'nprobe': index.nprobe,
                'recall@10': round(recall_at_k, 4),
                'qps': round(qps, 0), 'latency_ms': round(1000/qps, 1),
                'memory_gb': round(memory_gb, 2),
                'compression_ratio': round(4*dim / code_size, 1)
            })

    # 按 Recall 降序排列，找 Pareto 最优
    results.sort(key=lambda x: x['recall@10'], reverse=True)
    return results

# 运行基准（数据量大时需要 GPU 加速，这里给出可运行的演示版）
benchmark_results = benchmark_ivf_pq_params()
for r in benchmark_results[:10]:
    print(f"nlist={r['nlist']:>5}, m={r['m']:>3}, "
          f"Recall@10={r['recall@10']:.3f}, "
          f"QPS={r['qps']:.0f}, latency={r['latency_ms']:.1f}ms, "
          f"内存={r['memory_gb']:.1f}GB, 压缩比={r['compression_ratio']:.0f}x")

import hnswlib
import numpy as np
import time

def hnsw_benchmark(dim=1024, n_total=1_000_000, k=10):
    """HNSW 索引的性能基准"""
    data = np.random.randn(n_total, dim).astype(np.float32)
    queries = np.random.randn(100, dim).astype(np.float32)

    for M in [8, 16, 32, 64]:
        # 初始化索引
        index = hnswlib.Index(space='cosine', dim=dim)
        index.init_index(
            max_elements=n_total,
            ef_construction=M * 25,  # ef_construction = M × 25 是经验公式
            M=M
        )

        # 构建索引
        start = time.perf_counter()
        index.add_items(data)
        build_time = time.perf_counter() - start

        # 检索基准
        for ef in [16, 32, 64, 128, 256]:
            index.set_ef(ef)
            start = time.perf_counter()
            labels, distances = index.knn_query(queries, k=k)
            latency = (time.perf_counter() - start) / 100 * 1000  # ms per query

            print(f"M={M:>2}, ef={ef:>3}: 构建={build_time:.0f}s, "
                  f"延迟={latency:.2f}ms")

hnsw_benchmark()

内存充裕（向量存储 < 10% 可用内存）？
├── 是 → 需要最高 Recall（>0.97）？
│   ├── 是 → HNSW (M=64, ef=256)
│   └── 否 → HNSW (M=16, ef=64)
└── 否 → 数据规模（N）？
    ├── N < 100 万 → 可以接受 IVF_FLAT（无 PQ 量化，精度无损）
    ├── 100 万 < N < 1 亿 → IVF-PQ (nlist=4√N, m=128/d_subspace=8)
    └── N > 1 亿 → 考虑 DiskANN 或 ScaNN（存储与计算分离）

chunk_size=512, chunk_overlap=50, separators=["\n\n", "\n", "。", ".", " ", ""]

import numpy as np
from sentence_transformers import SentenceTransformer

class SemanticChunker:
    """
    基于 Embedding 相似度梯度的自适应语义分块器

    原理：遍历文档的每个句子，计算当前句子与前一句的 Embedding 余弦相似度。
    当相似度低于阈值时，在当前位置切分。这避免了在语义连贯的段落中间硬切。
    """
    def __init__(self, model: SentenceTransformer,
                 similarity_threshold: float = 0.65,
                 min_chunk_chars: int = 200,
                 max_chunk_chars: int = 1500):
        self.model = model
        self.threshold = similarity_threshold
        self.min_chars = min_chunk_chars
        self.max_chars = max_chunk_chars

    def _split_sentences(self, text: str) -> list[str]:
        """按句号、问号、感叹号、换行拆分句子，保留分隔符"""
        import re
        sentences = re.split(r'(?<=[。！？\n])\s*', text)
        return [s.strip() for s in sentences if s.strip()]

    def chunk(self, text: str) -> list[dict]:
        sentences = self._split_sentences(text)
        if len(sentences) <= 1:
            return [{"text": text, "start": 0, "end": len(text)}]

        # 批量编码所有句子
        embeddings = self.model.encode(sentences, show_progress_bar=False)
        embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

        # 计算相邻句子的余弦相似度
        similarities = []
        for i in range(1, len(embeddings)):
            sim = np.dot(embeddings[i-1], embeddings[i])
            similarities.append(float(sim))

        # 在局部相似度极小值处切分
        chunks = []
        current_chunk = [sentences[0]]
        current_start = 0
        sentence_starts = self._get_sentence_starts(text, sentences)

        for i in range(1, len(sentences)):
            current_len = sum(len(s) for s in current_chunk)
            sim = similarities[i-1]

            # 切分条件：相似度低于阈值 AND（超过最小长度 OR 后面还有低相似度）
            should_split = (
                sim < self.threshold and
                (current_len >= self.min_chars or
                 (i+1 < len(similarities) and similarities[i] < self.threshold))
            )

            if should_split or current_len + len(sentences[i]) > self.max_chars:
                chunk_text = ''.join(current_chunk)
                chunks.append({
                    "text": chunk_text,
                    "start": current_start,
                    "end": current_start + len(chunk_text)
                })
                current_chunk = [sentences[i]]
                current_start = sentence_starts[i]
            else:
                current_chunk.append(sentences[i])

        # 最后一个 chunk
        if current_chunk:
            chunk_text = ''.join(current_chunk)
            chunks.append({
                "text": chunk_text,
                "start": current_start,
                "end": current_start + len(chunk_text)
            })

        return chunks

    def _get_sentence_starts(self, text: str, sentences: list[str]) -> list[int]:
        """获取每个句子在原文中的起始位置"""
        starts = []
        pos = 0
        for s in sentences:
            idx = text.find(s, pos)
            starts.append(idx if idx >= 0 else pos)
            pos = starts[-1] + len(s)
        return starts


# 对比实验：固定分块 vs 语义分块
def chunking_comparison():
    sample_doc = """
    向量数据库是专门用于存储和检索高维向量的数据库系统。它通过近似最近邻搜索算法，
    在海量向量中快速找到与查询向量最相似的 Top-K 个结果。常见的 ANN 算法包括 HNSW、IVF-PQ 等。

    在生产环境中，向量数据库的选型需要综合考虑数据规模、查询延迟、内存开销和运维成本。
    对于百万级以下的数据量，Qdrant 是一个性价比极高的选择。对于亿级以上的数据量，
    Milvus 的分布式架构提供了更好的水平扩展能力。

    Embedding 模型将文本映射到高维向量空间，语义相近的文本在向量空间中距离较近。
    目前中文场景表现最好的开源模型是 BGE-M3 和 Stella-base-zh-v3。
    但模型选型必须基于目标领域的实际评测数据，而非通用榜单排名。
    """

    # 固定分块
    fixed_chunks = [
        sample_doc[i:i+200]
        for i in range(0, len(sample_doc), 200)
    ]
    print("=== 固定分块 (200 chars) ===")
    for i, chunk in enumerate(fixed_chunks):
        print(f"Chunk {i}: '{chunk[:80]}...'")

    # 语义分块
    model = SentenceTransformer('BAAI/bge-small-zh-v1.5')
    chunker = SemanticChunker(model, similarity_threshold=0.65)
    sem_chunks = chunker.chunk(sample_doc)
    print("\n=== 语义分块 ===")
    for i, chunk in enumerate(sem_chunks):
        print(f"Chunk {i} (start={chunk['start']}, end={chunk['end']}): "
              f"'{chunk['text'][:80]}...'")

chunking_comparison()

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class MultiGranularityChunk:
    chunk_id: str
    text: str
    embedding: Optional[np.ndarray] = None
    parent_id: Optional[str] = None
    granularity: str = "fine"  # "fine" | "coarse" | "summary"
    metadata: dict = field(default_factory=dict)

class MultiGranularityIndexer:
    """
    多粒度索引构建器

    流程：
    1. 文档 → 细粒度 chunk（精细检索单元）
    2. 细粒度 chunk → 聚合为粗粒度 parent chunk（上下文窗口）
    3. 粗粒度 parent chunk → 生成摘要（全局概览）

    检索时：
    - 先用 query 在细粒度索引中检索 Top-K
    - 对每个命中的细粒度 chunk，展开其 parent chunk 作为上下文
    - 可选：如果 query 与 summary 的相似度很高，直接使用 summary 作为补充上下文
    """
    def __init__(self, embed_model: SentenceTransformer,
                 fine_chunk_size: int = 256,
                 coarse_chunk_multiplier: int = 4):
        self.embed_model = embed_model
        self.fine_size = fine_chunk_size
        self.coarse_size = fine_chunk_size * coarse_chunk_multiplier

    def index_document(self, doc_id: str, text: str) -> dict:
        fine_chunks = self._create_fine_chunks(doc_id, text)
        coarse_chunks = self._aggregate_coarse_chunks(fine_chunks)
        summaries = self._generate_summaries(coarse_chunks)

        # 编码所有粒度的 chunk
        all_texts = (
            [c.text for c in fine_chunks] +
            [c.text for c in coarse_chunks] +
            [c.text for c in summaries]
        )
        embeddings = self.embed_model.encode(all_texts)

        # 分配 embedding
        offset = 0
        for chunk in fine_chunks:
            chunk.embedding = embeddings[offset]
            offset += 1
        for chunk in coarse_chunks:
            chunk.embedding = embeddings[offset]
            offset += 1
        for chunk in summaries:
            chunk.embedding = embeddings[offset]
            offset += 1

        return {
            "doc_id": doc_id,
            "fine_chunks": fine_chunks,
            "coarse_chunks": coarse_chunks,
            "summaries": summaries
        }

    def _create_fine_chunks(self, doc_id: str, text: str):
        """创建细粒度 chunk（128-256 tokens）"""
        # 简化实现；生产环境使用前面的 SemanticChunker
        chunks = []
        for i, start in enumerate(range(0, len(text), self.fine_size)):
            chunks.append(MultiGranularityChunk(
                chunk_id=f"{doc_id}_fine_{i}",
                text=text[start:start+self.fine_size],
                granularity="fine",
                parent_id=f"{doc_id}_coarse_{i // 4}"
            ))
        return chunks

    def _aggregate_coarse_chunks(self, fine_chunks: list):
        """将连续的细粒度 chunk 聚合为粗粒度 parent chunk"""
        coarse = {}
        for fc in fine_chunks:
            pid = fc.parent_id
            if pid not in coarse:
                coarse[pid] = MultiGranularityChunk(
                    chunk_id=pid, text="", granularity="coarse"
                )
            coarse[pid].text += fc.text + " "
        return list(coarse.values())

    def _generate_summaries(self, coarse_chunks: list):
        """为每个粗粒度 chunk 生成摘要（简化版用 LLM，这里用占位）"""
        summaries = []
        for cc in coarse_chunks:
            # 生产环境：调用 LLM 生成摘要
            summary_text = f"[摘要] {cc.text[:200]}..."
            summaries.append(MultiGranularityChunk(
                chunk_id=f"{cc.chunk_id}_summary",
                text=summary_text,
                granularity="summary"
            ))
        return summaries

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

class HybridRetriever:
    """
    混合检索器：Dense (Embedding) + Sparse (BM25) → RRF 融合

    为什么不直接用分数加权？
    Dense 的 cosine 相似度 (0.3-0.95) 和 BM25 分数 (0-200+) 的量纲完全不同，
    直接线性加权需要对每个排序器的分数做归一化（min-max / Z-score），
    但归一化参数依赖于查询分布，难以稳定。RRF 只依赖排名，天然量纲无关。
    """
    def __init__(self, dense_model: SentenceTransformer, corpus: list[str]):
        self.dense_model = dense_model
        self.corpus = corpus

        # 构建 BM25 索引
        tokenized_corpus = [doc.split() for doc in corpus]
        self.bm25 = BM25Okapi(tokenized_corpus)

        # 预编码 Dense 向量
        self.dense_embeddings = dense_model.encode(
            corpus, show_progress_bar=True, normalize_embeddings=True
        )

    def search(self, query: str, top_k: int = 10, rrf_k: int = 60) -> list[dict]:
        # Dense 检索
        q_emb = self.dense_model.encode(
            [query], normalize_embeddings=True
        )[0]
        dense_scores = np.dot(self.dense_embeddings, q_emb)
        dense_ranking = np.argsort(-dense_scores)[:top_k * 3]  # 取 3× 候选

        # Sparse 检索
        tokenized_query = query.split()
        bm25_scores = np.array(self.bm25.get_scores(tokenized_query))
        sparse_ranking = np.argsort(-bm25_scores)[:top_k * 3]

        # RRF 融合
        rrf_scores = {}
        for rank, idx in enumerate(dense_ranking):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1)
        for rank, idx in enumerate(sparse_ranking):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1)

        # 按 RRF 分数排序
        sorted_items = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
        results = []
        for idx, score in sorted_items[:top_k]:
            results.append({
                "doc_id": idx,
                "text": self.corpus[idx],
                "rrf_score": score,
                "dense_score": float(dense_scores[idx]),
                "bm25_score": float(bm25_scores[idx])
            })
        return results

def adaptive_rrf(query: str, dense_ranking, sparse_ranking,
                 top_k: int = 10, rrf_k: int = 60):
    """
    自适应 RRF：根据查询特征动态调整融合权重

    判定逻辑：
    - 查询包含数字/编号/专有名词 → BM25 权重 0.7
    - 查询为自然语言长句 → Dense 权重 0.7
    - 中性查询 → 等权重
    """
    import re
    keyword_indicators = bool(re.search(r'[0-9]{2,}|第[一二三四五六七八九十]|[A-Z]{2,}', query))
    long_query = len(query) > 50

    if keyword_indicators and not long_query:
        w_dense, w_sparse = 0.3, 0.7
    elif long_query and not keyword_indicators:
        w_dense, w_sparse = 0.7, 0.3
    else:
        w_dense, w_sparse = 0.5, 0.5

    rrf_scores = {}
    for rank, idx in enumerate(dense_ranking):
        rrf_scores[idx] = rrf_scores.get(idx, 0) + w_dense / (rrf_k + rank + 1)
    for rank, idx in enumerate(sparse_ranking):
        rrf_scores[idx] = rrf_scores.get(idx, 0) + w_sparse / (rrf_k + rank + 1)

    sorted_items = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
    return [idx for idx, _ in sorted_items[:top_k]]

from FlagEmbedding import FlagReranker
import numpy as np
import time

class RerankerPipeline:
    """
    两阶段检索+重排序 Pipeline

    阶段 1：Bi-Encoder 粗排 → Top-N 候选（N=50-100）
    阶段 2：Cross-Encoder 精排 → Top-K 最终结果（K=5-10）

    延迟预算分配：
    - 粗排：<10ms（向量库索引检索）
    - 精排：N × 单对延迟（例如 20 × 8ms = 160ms）
    - 总计：<200ms（满足实时交互体验）
    """
    def __init__(self, reranker_model: str = "BAAI/bge-reranker-v2-m3",
                 use_fp16: bool = True, batch_size: int = 16):
        self.reranker = FlagReranker(
            reranker_model,
            use_fp16=use_fp16
        )
        self.batch_size = batch_size

    def rerank(self, query: str, candidates: list[dict],
               top_k: int = 10) -> list[dict]:
        """
        对候选文档重排序

        Args:
            query: 用户查询
            candidates: 候选文档列表，每项含 'text' 字段
            top_k: 返回的最终文档数

        Returns:
            按相关性分数降序排列的文档列表
        """
        if len(candidates) <= top_k:
            # 候选集太小，不需要重排序
            return candidates

        # 构建 query-document pairs
        pairs = [[query, doc["text"]] for doc in candidates]

        # 批量推理（批量处理比逐对处理快 3-5 倍）
        all_scores = []
        for i in range(0, len(pairs), self.batch_size):
            batch = pairs[i:i + self.batch_size]
            scores = self.reranker.compute_score(batch, normalize=True)
            if isinstance(scores, float):
                scores = [scores]
            all_scores.extend(scores)

        # 按分数排序并附加 rerank_score
        for doc, score in zip(candidates, all_scores):
            doc["rerank_score"] = float(score) if score else 0.0

        candidates.sort(key=lambda x: x.get("rerank_score", 0), reverse=True)
        return candidates[:top_k]


# 完整的两阶段检索管道
def full_rag_retrieval_pipeline(
    query: str,
    hybrid_retriever: HybridRetriever,
    reranker: RerankerPipeline,
    top_k_candidates: int = 50,
    top_k_final: int = 5
) -> list[dict]:
    """
    RAG 完整检索管道：混合检索 → 重排序 → Top-K 结果

    延迟分解（典型值）：
    1. 混合检索（粗排）：8ms
    2. Cross-Encoder 重排序（50 candidates）：约 40-60ms（jina-reranker-v2, batch_size=16，含数据搬运开销）
    3. 总计：~50-70ms
    """
    # 阶段 1：混合检索粗排
    candidates = hybrid_retriever.search(query, top_k=top_k_candidates)
    print(f"[Pipeline] 混合检索返回 {len(candidates)} 候选文档")

    # 阶段 2：Cross-Encoder 精排
    start = time.perf_counter()
    results = reranker.rerank(query, candidates, top_k=top_k_final)
    elapsed = time.perf_counter() - start
    print(f"[Pipeline] 重排序完成: {len(candidates)} → {len(results)}, "
          f"耗时 {elapsed*1000:.0f}ms")

    return results

def asymmetric_truncation(query: str, document: str, max_tokens: int = 7168) -> str:
    """
    非对称截断：在文档中找到与 query 最相关的片段

    对于长文档，Cross-Encoder 的注意力矩阵是 O((L_q + L_d)²)，
    截断文档长度对延迟的影响是平方级的。
    """
    doc_words = document.split()
    if len(doc_words) <= max_tokens:
        return document

    # 滑窗计算每个窗口与 query 的重合度
    query_words = set(query.lower().split())
    best_window, best_score = "", 0
    window_size = max_tokens

    for i in range(0, len(doc_words) - window_size, window_size // 2):
        window = doc_words[i:i + window_size]
        window_set = set(w.lower() for w in window)
        overlap = len(query_words & window_set)

        # 加分项：窗口包含标题/关键短语
        if any(kw in " ".join(window).lower() for kw in ['定义', '说明', '总结', '结论']):
            overlap += 3

        if overlap > best_score:
            best_score = overlap
            best_window = " ".join(window)

    return best_window if best_window else " ".join(doc_words[:max_tokens])

用户 Query
    │
    ├── 1. Query 改写（可选）
    │   ├── 意图分类（事实查询 / 解释性查询 / 对比查询）
    │   ├── 查询扩展（HyDE：生成假设答案再检索）
    │   └── 多语言翻译（中英文检索统一）
    │
    ├── 2. 多路检索（并发）
    │   ├── Dense Retrieval (bge-m3, HNSW 索引)
    │   ├── Sparse Retrieval (BM25, Elasticsearch)
    │   └── 结构化过滤（元数据/标签/时间范围）
    │
    ├── 3. 融合（自适应 RRF）
    │   └── Top-100 候选集
    │
    ├── 4. 粗排重排序（jina-reranker-v2）
    │   └── Top-30 候选集
    │
    ├── 5. 精排重排序（bge-reranker-v2-minicpm）
    │   └── Top-5 最终文档
    │
    ├── 6. 上下文组装
    │   ├── 多粒度展开（fine chunk → parent chunk）
    │   ├── 去重与合并（MinHash LSH 去重）
    │   └── Token 预算控制（prompt 总长度限制）
    │
    └── 7. LLM 生成
        ├── System Prompt + RAG Context + User Query
        └── 引用溯源（返回引用文档 ID 和原文位置）

from scipy.spatial.distance import jensenshannon
from scipy.stats import wasserstein_distance
import numpy as np
from collections import deque

class RAGMonitor:
    """
    RAG 生产环境核心监控模块

    三个关键指标：
    1. 检索精度代理指标（无需人工标注）
    2. Embedding 分布漂移检测
    3. 反馈闭环数据收集
    """
    def __init__(self, window_size: int = 1000):
        self.retrieval_window = deque(maxlen=window_size)
        self.feedback_pool = deque(maxlen=10000)

    def log_retrieval(self, query_embedding: np.ndarray,
                      top_k_embeddings: np.ndarray):
        """
        记录检索的 Top-K 余弦相似度分布

        监测异常：如果 Top-1 相似度突然整体下降（>2σ），
        可能是索引损坏、Embedding 模型退化或数据漂移。
        """
        cos_sims = np.dot(top_k_embeddings, query_embedding)
        self.retrieval_window.append({
            "top1_sim": float(cos_sims[0]),
            "top5_avg_sim": float(np.mean(cos_sims[:5])),
            "top10_avg_sim": float(np.mean(cos_sims))
        })

    def check_degradation(self) -> dict:
        """退化检测：当前滑动窗口 vs 历史基线"""
        if len(self.retrieval_window) < 100:
            return {"status": "insufficient_data"}

        recent = [r["top1_sim"] for r in list(self.retrieval_window)[-100:]]
        baseline = [r["top1_sim"] for r in list(self.retrieval_window)[:500]]

        mean_recent = np.mean(recent)
        mean_baseline = np.mean(baseline)
        std_baseline = np.std(baseline)

        z_score = (mean_recent - mean_baseline) / (std_baseline + 1e-8)
        degradation_detected = z_score < -2.0

        return {
            "status": "degraded" if degradation_detected else "healthy",
            "z_score": round(z_score, 3),
            "mean_recent_top1": round(mean_recent, 4),
            "mean_baseline_top1": round(mean_baseline, 4),
            "drift_ratio": round(mean_recent / mean_baseline, 4)
        }

    def check_embedding_drift(self, current_embeddings: np.ndarray,
                              baseline_embeddings: np.ndarray,
                              n_bins: int = 50) -> float:
        """检测 Embedding 分布漂移（按维度聚合的 JS 散度）"""
        drift_scores = []
        for dim in range(min(current_embeddings.shape[1], baseline_embeddings.shape[1])):
            hist_curr, _ = np.histogram(current_embeddings[:, dim], bins=n_bins, density=True)
            hist_base, _ = np.histogram(baseline_embeddings[:, dim], bins=n_bins, density=True)
            js_div = jensenshannon(hist_curr + 0.001, hist_base + 0.001)
            drift_scores.append(js_div)
        return float(np.mean(drift_scores))

    def collect_feedback(self, query: str, top_docs: list[str],
                          llm_response: str, user_rating: int,
                          cited_doc_indices: list[int]):
        """
        收集用户反馈用于持续优化

        user_rating: 1-5，<3 表示不满意
        cited_doc_indices: LLM 回答中引用的文档索引
        """
        self.feedback_pool.append({
            "query": query,
            "top_docs": top_docs,
            "response": llm_response,
            "rating": user_rating,
            "cited_docs": cited_doc_indices
        })

        # 提取 Hard Negatives：用户点踩时，被检索到但排在正确文档之前的文档
        if user_rating < 3 and cited_doc_indices:
            correct_idx = cited_doc_indices[0]
            hard_negatives = []
            for i in range(correct_idx):
                if i not in cited_doc_indices:
                    hard_negatives.append({
                        "query": query,
                        "positive_doc": top_docs[correct_idx],
                        "negative_doc": top_docs[i]
                    })
            # 这些 hard_negatives 可直接用于 Reranker/Embedding 微调
            return hard_negatives
        return []