AI RAG实战指南：从基础检索到自适应反思

def split_documents(docs, chunk_size=512, overlap=50):
    chunks = []
    for doc in docs:
        for i in range(0, len(doc), chunk_size - overlap):
            chunks.append(doc[i:i + chunk_size])
    return chunks

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", "。", "！", "？", ";", "，", " ", ""]
)

import time

def batch_embeddings(texts, batch_size=10):
    """批量生成向量，处理速率限制"""
    embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        try:
            response = client.embeddings.create(
                input=batch,
                model="text-embedding-3-small"
            )
            embeddings.extend([np.array(d.embedding) for d in response.data])
        except Exception as e:
            print(f"Batch {i} failed: {e}")
            time.sleep(5)
        time.sleep(0.1)  # 避免触发速率限制
    return embeddings

import faiss
import numpy as np

# 假设有 N 个 1024 维向量
embeddings = np.array([doc.embedding for doc in documents]).astype('float32')

# IndexFlatL2：精确搜索，简单但慢
index_flat = faiss.IndexFlatL2(1024)
index_flat.add(embeddings)

# IndexIVFFlat：聚类后搜索，更快
n_clusters = 1000  # 一般是文档数的 sqrt
index_ivf = faiss.IndexIVFFlat(faiss.IndexFlatL2(1024), 1024, n_clusters)
index_ivf.train(embeddings)
index_ivf.add(embeddings)
index_ivf.nprobe = 20  # 查询时探查的簇数

def rebuild_index():
    # 备份
    backup_dir = f"./chroma_db_backup_{datetime.now().strftime('%Y%m%d')}"
    shutil.copytree("./chroma_db", backup_dir)

    # 获取所有数据并重建
    collection = client.get_collection("knowledge_base")
    all_data = collection.get(include=['documents', 'metadatas', 'embeddings'])

    client.delete_collection("knowledge_base")
    collection = client.create_collection("knowledge_base")
    collection.add(
        documents=all_data['documents'],
        metadatas=all_data['metadatas'],
        embeddings=all_data['embeddings'],
        ids=[f"doc_{i}" for i in range(len(all_data['documents']))]
    )

from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_community.document_compressors.cross_encoder import CrossEncoderReranker

def create_hybrid_retriever(vectorstore, documents, top_k=10):
    # 向量检索
    vector_retriever = vectorstore.as_retriever(search_kwargs={"k": top_k * 2})

    # BM25 关键词检索
    bm25_retriever = BM25Retriever.from_documents(documents)
    bm25_retriever.k = top_k * 2

    # 组合两种检索（权重可调）
    ensemble_retriever = EnsembleRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        weights=[0.7, 0.3]  # 语义权重高一些
    )

    # 交叉编码器重排序
    cross_encoder = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
    compressor = CrossEncoderReranker(compressor=cross_encoder, top_n=top_k)

    return ContextualCompressionRetriever(
        base_compressor=compressor,
        base_retriever=ensemble_retriever
    )

prompt = """你是一个专业的知识库助手。请基于以下上下文回答问题。

上下文：
{context}

问题：{question}

请用中文回答，并在回答中用【来源：文档名】的形式标注信息来源。
如果上下文中没有相关信息，请直接说"我没有找到相关信息"，不要编造答案。

答案："""

from langchain.chains import RetrievalQAWithSourcesChain

chain = RetrievalQAWithSourcesChain.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt_template}
)

result = chain.invoke({"question": "公司的休假政策是什么？"})
print(result["answer"])
print("来源：", result["source_documents"])

def route_query(question: str) -> QueryType:
    """根据问题类型决定路由"""
    prompt = f"""
    分析以下问题的类型：
    问题：{question}

    类型说明：
    - factual: 事实性查询，需要精确信息
    - analytical: 分析性查询，需要综合多个信息
    - creative: 创造性查询，需要发散思维

    返回类型和置信度(0-1)。
    """
    result = llm.predict(prompt, response_model=QueryType)
    return result

class BaseRetriever(ABC):
    @abstractmethod
    def retrieve(self, query: str, top_k: int = 5) -> list[dict]:
        pass

class HybridRetriever(BaseRetriever):
    def __init__(self, retrievers: list, weights: list):
        self.retrievers = retrievers
        self.weights = weights

    def retrieve(self, query: str, top_k: int = 5) -> list[dict]:
        all_results = []
        for retriever, weight in zip(self.retrievers, self.weights):
            results = retriever.retrieve(query, top_k=top_k)
            for result in results:
                result['score'] *= weight
            all_results.extend(results)
        return self.merge_results(all_results)[:top_k * 2]

class Reranker:
    def __init__(self, model_name="BAAI/bge-reranker-base"):
        from sentence_transformers import CrossEncoder
        self.model = CrossEncoder(model_name)

    def rerank(self, query, documents, top_k=5):
        pairs = [(query, doc['content']) for doc in documents]
        scores = self.model.predict(pairs)

        for doc, score in zip(documents, scores):
            doc['rerank_score'] = score

        documents.sort(key=lambda x: x['rerank_score'], reverse=True)
        return documents[:top_k]

class ContextBuilder:
    def __init__(self, max_tokens=3000):
        self.max_tokens = max_tokens

    def build(self, query, documents):
        # 按相关性排序
        documents = sorted(documents, key=lambda x: x.get('rerank_score', 0), reverse=True)

        # 前 3 个文档完整保留
        context_parts = []
        current_tokens = 0

        for doc in documents[:3]:
            content = f"## 来源：{doc.get('source', '未知')}\n{doc['content']}\n"
            doc_tokens = len(content) // 4
            if current_tokens + doc_tokens > self.max_tokens * 0.8:
                break
            context_parts.append(content)
            current_tokens += doc_tokens

        # 剩余文档只保留摘要
        for doc in documents[3:]:
            if current_tokens >= self.max_tokens:
                break
            summary = self._summarize(doc['content'])
            context_parts.append(f"## 来源：{doc.get('source', '未知')}\n{summary}\n")
            current_tokens += len(summary) // 4

        return "\n".join(context_parts)

def extract_entities_relations(chunk_text):
    prompt = f"""
从以下文本中提取实体和关系，按照给定的类型约束输出 JSON。

文本：{chunk_text}

实体类型：服务、组件、配置、错误、API、产品、版本
关系类型：调用、依赖、导致、修复、包含、兼容

输出格式：
{{
  "entities": [{{"text": "...", "type": "..."}}],
  "relations": [{{"source": "...", "target": "...", "relation": "..."}}]
}}
"""
    response = llm.generate(prompt, temperature=0.1)  # 低温度保证稳定
    return parse_json(response)

import networkx as nx

def retrieve_context(G, entities, max_hops=2):
    """从问题实体出发，做多跳检索"""
    context_nodes = set(entities)
    for entity in entities:
        if G.has_node(entity):
            for node, distance in nx.single_source_shortest_path_length(
                G, entity, cutoff=max_hops
            ).items():
                context_nodes.add(node)

    subgraph = G.subgraph(context_nodes).copy()
    return subgraph

class ChainHopRAG:
    def __init__(self, vectorstore, max_hops=3):
        self.vectorstore = vectorstore
        self.max_hops = max_hops
        self.llm = ChatOpenAI(temperature=0)

    def run(self, question):
        collected_info = []
        current_queries = [question]

        for hop in range(self.max_hops):
            # 检索当前查询
            for query in current_queries:
                docs = self.vectorstore.similarity_search(query, k=3)
                collected_info.extend(docs)

            # 判断是否继续
            if hop < self.max_hops - 1:
                response = self.llm.invoke(
                    self.query_gen_prompt.format(
                        question=question,
                        collected_info="\n".join([d.page_content for d in collected_info[-6:]])
                    )
                )

                if "不需要进一步检索" in response.content:
                    break

                current_queries = self._extract_queries(response.content)
                if not current_queries:
                    break

        # 生成最终答案
        return self.llm.invoke(
            self.answer_prompt.format(
                question=question,
                collected_info="\n".join([d.page_content for d in collected_info])
            )
        ).content

def should_use_multi_hop(question):
    """判断问题是否适合多跳"""
    keywords = ["对比", "差异", "为什么", "关系", "影响"]
    return any(keyword in question for keyword in keywords)

def hierarchical_search(query, vector_db, top_k=5):
    query_embedding = embedding_model.encode(query)

    # 1. 判断查询意图
    intent = classify_query_intent(query)

    # 2. 根据意图选择搜索层级
    if intent == "overview":
        search_levels = [2]  # 章节级
    elif intent == "detail":
        search_levels = [4]  # 段落级
    else:
        search_levels = [2, 3, 4]  # 多层级

    # 3. 在目标层级搜索
    results = vector_db.search(
        query_embedding,
        top_k=top_k * 2,
        filter={'level': {'$in': search_levels}}
    )

    # 4. 去重和扩展（必要时补充父节点）
    results = deduplicate_and_expand(results, vector_db)

    return results[:top_k]

def should_retrieve_with_rules(question, llm):
    """带规则的检索判断"""
    # 关键词触发规则
    force_retrieve_keywords = ['最新', '今年', '2024', '2025', '具体数据', '详细']
    no_retrieve_keywords = ['编程', '代码', '算法', '解释', '是什么']

    if any(kw in question for kw in force_retrieve_keywords):
        return True

    if any(kw in question for kw in no_retrieve_keywords):
        return False

    # 其他情况由 LLM 判断
    prompt = f"问题: {question}\n判断是否需要检索外部信息，只回答 yes 或 no。"
    response = llm.generate(prompt)
    return "yes" in response.lower()

def adaptive_retrieval(question, vector_db, llm, max_iterations=3):
    all_docs = []
    search_terms = [question]

    for iteration in range(max_iterations):
        docs = vector_db.search(search_terms[-1], top_k=5)
        all_docs.extend(docs)

        # 评估检索结果
        evaluation = evaluate_retrieval(question, docs, llm)

        # 质量足够，停止
        if evaluation["relevance"] >= 7 and evaluation["completeness"] >= 7:
            break

        # 相关性低，调整搜索词
        if evaluation["relevance"] < 5:
            search_terms.append(generate_alternative_query(question, llm))
        # 信息不足，扩大范围
        elif evaluation["completeness"] < 7:
            search_terms.append(expand_query(question, llm))

    return remove_duplicates(all_docs)[:10]

import hashlib
from functools import lru_cache

# 答案缓存
def query_with_cache(question):
    cache_key = hashlib.md5(question.encode()).hexdigest()

    if cache_key in cache:
        print("命中缓存")
        return cache[cache_key]

    result = chain.invoke({"question": question})
    cache[cache_key] = result
    return result

def route_rag_strategy(question):
    """根据问题类型选择 RAG 策略"""
    # 简单事实查询：单体 RAG
    if is_simple_factual(question):
        return "monolithic"

    # 需要跨文档推理：多跳 RAG
    if any(kw in question for kw in ["对比", "差异", "为什么", "关系"]):
        return "multi_hop"

    # 需要理解实体关系：图谱 RAG
    if any(kw in question for kw in ["调用链", "依赖", "导致", "影响"]):
        return "graph"

    # 默认：混合检索 RAG
    return "hybrid"