RAG 进阶

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

REWRITE_PROMPT = """你是一个搜索查询优化专家。请将用户的原始问题改写为更适合文档检索的查询。

改写要求：
1. 保留原始问题的核心意图
2. 展开缩写和口语化表达
3. 添加可能相关的同义词/近义词
4. 如果问题有歧义，选择最可能的解释
5. 输出 1-3 个改写版本，每行一个

原始问题：{query}

改写版本："""

def rewrite_query(query: str, llm_model: str = "gpt-4o-mini") -> list[str]:
    """LLM 查询改写，返回多个改写版本"""
    llm = ChatOpenAI(model=llm_model, temperature=0.3)
    prompt = ChatPromptTemplate.from_template(REWRITE_PROMPT)
    chain = prompt | llm
    result = chain.invoke({"query": query})
    return [q.strip() for q in result.content.strip().split("\n") if q.strip()]


def rewrite_and_retrieve(query: str, retriever, top_k: int = 5) -> list[dict]:
    """改写 + 检索 + 结果去重融合"""
    rewrites = rewrite_query(query)
    all_results = []
    seen_texts = set()

    # 对每个改写版本检索，去重融合
    for rewritten_query in [query] + rewrites:
        results = retriever.retrieve(rewritten_query, top_k=top_k)
        for r in results:
            text_key = r["text"][:100]
            if text_key not in seen_texts:
                seen_texts.add(text_key)
                all_results.append(r)

    return all_results[:top_k * 2]

HYDE_PROMPT = """请根据以下问题，生成一段可能包含该问题答案的文档片段。
注意：生成的是"可能的文档内容"，而非直接回答问题。内容应具体、专业。

问题：{query}

文档片段（200字左右）："""

def hyde_search(
    query: str,
    retriever,
    embed_model,
    llm_model: str = "gpt-4o-mini",
    top_k: int = 5,
    n_hypothetical: int = 3,   # 生成多个假设文档取平均，减少 LLM 随机性影响
) -> list[dict]:
    """HyDE 检索：生成假设文档 → Embed → 检索"""
    import numpy as np

    llm = ChatOpenAI(model=llm_model, temperature=0.7)
    prompt = ChatPromptTemplate.from_template(HYDE_PROMPT)
    chain = prompt | llm

    # 生成多个假设文档
    hypothetical_docs = [chain.invoke({"query": query}).content for _ in range(n_hypothetical)]

    # 编码并取平均（平均向量比单个更稳定）
    hyp_embeddings = embed_model.encode(hypothetical_docs, normalize_embeddings=True)
    avg_embedding = hyp_embeddings.mean(axis=0)
    avg_embedding = avg_embedding / np.linalg.norm(avg_embedding)

    # 用平均 embedding 检索
    return retriever.retrieve_by_embedding(avg_embedding, top_k=top_k)

from langchain.retrievers.multi_query import MultiQueryRetriever

def multi_query_retrieval(
    query: str,
    vectorstore,
    top_k: int = 5,
    llm_model: str = "gpt-4o-mini",
) -> list[dict]:
    """
    多路查询检索（LangChain 原生支持）
    自动生成 3 个变体查询 + 原始查询，共 4 路检索
    """
    llm = ChatOpenAI(model=llm_model, temperature=0.7)

    retriever = MultiQueryRetriever.from_llm(
        retriever=vectorstore.as_retriever(search_kwargs={"k": top_k}),
        llm=llm,
        include_original=True,   # 包含原始 query 的检索结果
    )

    docs = retriever.invoke(query)

    # 去重（按文本内容）
    seen_texts = set()
    unique_docs = []
    for doc in docs:
        if doc.page_content not in seen_texts:
            seen_texts.add(doc.page_content)
            unique_docs.append({"text": doc.page_content, "metadata": doc.metadata})

    return unique_docs

STEP_BACK_PROMPT = """你是一个专家，请将以下具体问题抽象为一个更宏观的背景问题。

具体问题：{question}

更宏观的背景问题（一句话）："""

def step_back_retrieval(
    query: str,
    retriever,
    llm_model: str = "gpt-4o-mini",
    top_k: int = 5,
) -> dict:
    """Step-Back 检索：背景知识 + 具体细节双路检索"""
    llm = ChatOpenAI(model=llm_model, temperature=0)
    prompt = ChatPromptTemplate.from_template(STEP_BACK_PROMPT)
    chain = prompt | llm

    step_back_q = chain.invoke({"question": query}).content.strip()
    print(f"原始查询: {query}")
    print(f"后退查询: {step_back_q}")

    specific_results = retriever.retrieve(query, top_k=top_k)
    background_results = retriever.retrieve(step_back_q, top_k=top_k)

    # 融合去重（具体结果优先）
    seen = set()
    unique_results = []
    for r in specific_results + background_results:
        text_key = r["text"][:100]
        if text_key not in seen:
            seen.add(text_key)
            unique_results.append(r)

    return {
        "step_back_query": step_back_q,
        "results": unique_results[:top_k],
    }

from sentence_transformers import CrossEncoder
import time

class CrossEncoderReranker:
    def __init__(
        self,
        model_name: str = "BAAI/bge-reranker-v2-m3",
        device: str = "cuda",
        max_length: int = 512,  # 超长文档用 1024，但延迟翻倍
    ):
        self.model = CrossEncoder(model_name, max_length=max_length, device=device)

    def rerank(
        self,
        query: str,
        documents: list[str],
        top_k: int = 5,
        batch_size: int = 32,
    ) -> list[dict]:
        """Cross-Encoder 批量重排序"""
        if not documents:
            return []

        pairs = [(query, doc) for doc in documents]

        t0 = time.time()
        scores = self.model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
        print(f"Reranked {len(documents)} docs in {(time.time()-t0)*1000:.1f}ms")

        ranked_results = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)

        return [
            {"text": doc, "rerank_score": float(score)}
            for doc, score in ranked_results[:top_k]
        ]

from qdrant_client.models import Filter, FieldCondition, MatchValue, Range
from datetime import datetime

def search_with_filters(
    client: QdrantClient,
    collection_name: str,
    query_vector: list[float],
    department: str = None,
    doc_type: str = None,
    date_after: datetime = None,
    top_k: int = 10,
) -> list[dict]:
    """带元数据过滤的向量检索"""
    conditions = []

    if department:
        conditions.append(
            FieldCondition(key="metadata.department", match=MatchValue(value=department))
        )
    if doc_type:
        conditions.append(
            FieldCondition(key="metadata.doc_type", match=MatchValue(value=doc_type))
        )
    if date_after:
        conditions.append(
            FieldCondition(
                key="metadata.created_at",
                range=Range(gte=date_after.timestamp())
            )
        )

    search_filter = Filter(must=conditions) if conditions else None

    results = client.search(
        collection_name=collection_name,
        query_vector=query_vector,
        query_filter=search_filter,
        limit=top_k,
        with_payload=True,
        search_params={
            "hnsw_ef": 100,   # 查询时动态设置 ef（可按场景调整）
            "exact": False,   # False=HNSW近似检索（快）；True=暴力搜索（超慢但精确）
        },
    )

    return [{"text": r.payload["text"], "score": r.score} for r in results]


def create_payload_indexes(client: QdrantClient, collection_name: str) -> None:
    """
    为过滤字段创建 Payload 索引（必须！）
    未建索引字段过滤 = 全量扫描 = 性能灾难
    """
    # Keyword 索引（用于精确匹配：department、doc_type、status）
    for field in ["metadata.department", "metadata.doc_type", "metadata.status"]:
        client.create_payload_index(
            collection_name=collection_name,
            field_name=field,
            field_schema="keyword",
        )

    # Float 索引（用于范围查询：时间戳、分数）
    client.create_payload_index(
        collection_name=collection_name,
        field_name="metadata.created_at",
        field_schema="float",
    )

    print("✅ Payload 索引创建完成")

from qdrant_client.models import PointStruct, FilterSelector
import hashlib
import uuid

class VectorStoreManager:
    """生产级向量存储管理器，支持完整增量更新生命周期"""

    def __init__(self, client: QdrantClient, collection_name: str, embed_model):
        self.client = client
        self.collection_name = collection_name
        self.embed_model = embed_model

    def _doc_id_to_uuid(self, doc_id: str) -> str:
        """将业务 doc_id 确定性映射为 UUID（Qdrant 要求）"""
        return str(uuid.UUID(bytes=hashlib.md5(doc_id.encode()).digest()))

    def upsert_documents(self, documents: list[dict], batch_size: int = 100) -> None:
        """
        批量 Upsert 文档（存在则更新，不存在则插入）
        documents: [{"doc_id": str, "text": str, "metadata": dict}]
        """
        for i in range(0, len(documents), batch_size):
            batch = documents[i:i + batch_size]
            texts = [d["text"] for d in batch]

            # 批量 Embedding（比逐条快 10x+）
            embeddings = self.embed_model.encode(texts, batch_size=32, normalize_embeddings=True)

            points = [
                PointStruct(
                    id=self._doc_id_to_uuid(doc["doc_id"]),
                    vector=emb.tolist(),
                    payload={
                        "text": doc["text"],
                        "doc_id": doc["doc_id"],
                        "metadata": doc.get("metadata", {}),
                    }
                )
                for doc, emb in zip(batch, embeddings)
            ]

            # wait=False 异步写入（提高吞吐量）；重要数据用 wait=True
            self.client.upsert(
                collection_name=self.collection_name,
                points=points,
                wait=False,
            )

        print(f"✅ Upserted {len(documents)} documents")

    def delete_documents(self, doc_ids: list[str]) -> None:
        """按 doc_id 批量删除"""
        point_ids = [self._doc_id_to_uuid(doc_id) for doc_id in doc_ids]
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=point_ids,
            wait=True,
        )

    def delete_by_metadata(self, field: str, value: str) -> None:
        """
        按元数据批量删除（最常用：删除某文件的所有 chunks）
        例：delete_by_metadata("source_file", "report_2024.pdf")
        """
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=FilterSelector(
                filter=Filter(
                    must=[FieldCondition(key=f"metadata.{field}", match=MatchValue(value=value))]
                )
            ),
            wait=True,
        )

    def incremental_update(
        self,
        new_documents: list[dict],
        updated_documents: list[dict],
        deleted_doc_ids: list[str],
    ) -> dict:
        """
        增量更新：删除 → 更新 → 新增（顺序执行保证一致性）
        """
        stats = {"deleted": 0, "updated": 0, "inserted": 0}

        if deleted_doc_ids:
            self.delete_documents(deleted_doc_ids)
            stats["deleted"] = len(deleted_doc_ids)

        if updated_documents:
            # 先删旧版，再插新版（避免脏读）
            old_ids = [d["doc_id"] for d in updated_documents]
            self.delete_documents(old_ids)
            self.upsert_documents(updated_documents)
            stats["updated"] = len(updated_documents)

        if new_documents:
            self.upsert_documents(new_documents)
            stats["inserted"] = len(new_documents)

        return stats