利用上下文检索增强RAG：提升检索性能的实用指南
本文介绍如何通过上下文嵌入、上下文BM25和重排序技术显著提升RAG系统的检索性能。传统RAG中，文档分块后孤立嵌入常导致上下文缺失，影响检索准确性。上下文嵌入利用Claude为每个块生成文档级上下文描述，再嵌入，平均降低35%的检索失败率。结合BM25混合搜索和重排序，Pass@10可从87%提升至95%以上。文章提供完整代码实现、环境配置、性能评估及生产优化建议，适合有编程基础的开发者学习AI工程实践。
技术背景
RAG（检索增强生成）是当前企业构建AI应用的核心技术，通过结合内部知识库与LLM生成能力，实现客户支持、文档问答、代码生成等场景。然而，传统RAG在文档分块后，每个块独立嵌入，缺乏全局上下文，导致检索结果不准确，尤其当块内容依赖文档其他部分时。例如，代码库中一个函数块可能依赖其他文件或上下文，孤立嵌入难以捕捉其语义。上下文检索技术通过为每个块生成上下文描述，解决此问题，提升检索质量。适用场景包括代码库、法律文档、金融报告等长文档，主要限制是增加摄取成本和延迟，但通过提示缓存可缓解。
核心原理
上下文嵌入的核心原理是：在嵌入前，利用LLM（如Claude）为每个块生成一段简短的上下文描述，描述该块在源文档中的位置和内容概要，然后将该描述与原始块拼接，再进行嵌入。这样，嵌入向量不仅包含块本身的信息，还包含文档级上下文，使得语义更丰富，检索更准确。上下文BM25则是在BM25关键词搜索中，同时索引原始块和上下文描述，使关键词匹配能命中上下文中的术语。重排序采用两阶段检索：先广泛检索候选（如top100），再用专门的重排序模型（如Cohere rerank）对候选进行精细评分，选出最相关的top-k。数据流为：文档分块→生成上下文→嵌入/索引→查询时检索（语义+BM25）→融合→重排序→返回结果。
架构分析
系统架构分为数据摄取和查询两个阶段。摄取阶段：文档分块后，通过Claude API生成上下文描述（利用提示缓存降低成本），然后使用Voyage AI嵌入模型将上下文+块内容转为向量，存入向量数据库（如内存VectorDB或生产级向量库），同时将原始块和上下文描述索引到Elasticsearch（用于BM25）。查询阶段：用户查询先通过向量数据库进行语义检索（余弦相似度），同时通过Elasticsearch进行BM25检索，然后使用加权倒数排名融合（RRF）合并结果，可选地使用Cohere重排序模型对融合结果进行重排，最终返回top-k块。组件关系：VectorDB类负责嵌入存储和搜索，ContextualVectorDB扩展了上下文生成，ElasticsearchBM25负责关键词索引，evaluate函数评估性能。设计决策：使用提示缓存降低摄取成本，使用批量嵌入提高效率，使用线程池并行生成上下文，使用加权融合平衡语义和关键词。
完整实战教程
注意：有关上下文检索的更多背景信息，包括在不同数据集上的额外性能评估，我们建议阅读随附的博客文章。
检索增强生成（RAG）使Claude能够在提供响应时利用您的内部知识库、代码库或任何其他文档语料库。企业越来越多地构建RAG应用，以改善客户支持、内部公司文档问答、金融和法律分析、代码生成等场景的工作流程。
在另一份指南中，我们介绍了如何设置基本的检索系统，演示了如何评估其性能，并概述了几种改进性能的技术。在本指南中，我们介绍一种改进检索性能的技术：上下文嵌入（Contextual Embeddings）。
在传统的RAG中，文档通常被分割成较小的块以便高效检索。虽然这种方法在许多应用中效果良好，但当单个块缺乏足够上下文时可能会导致问题。上下文嵌入通过在嵌入之前为每个块添加上下文来解决这个问题。这种方法提高了每个嵌入块的质量，从而实现更准确的检索，进而提升整体性能。在我们测试的所有数据源中，上下文嵌入将前20块检索失败率平均降低了35%。
相同的块特定上下文也可以与BM25搜索结合使用，以进一步提高检索性能。我们在“上下文BM25”部分介绍了这一技术。
在本指南中，我们将演示如何使用9个代码库的数据集作为知识库来构建和优化上下文检索系统。我们将逐步介绍：
1) 设置基本的检索流程以建立性能基线。
2) 上下文嵌入：它是什么，为什么有效，以及提示缓存如何使其适用于生产环境。
3) 实现上下文嵌入并展示性能改进。
4) 上下文BM25：使用*上下文*BM25混合搜索提高性能。
5) 通过重排序提高性能。
评估指标与数据集：
我们使用一个预先分块的9个代码库数据集——所有代码库都根据基本的字符分割机制进行了分块。我们的评估数据集包含248个查询——每个查询都包含一个“黄金块”。我们将使用一个名为Pass@k的指标来评估性能。Pass@k检查每个查询检索到的前k个文档中是否存在“黄金文档”。在这种情况下，上下文嵌入帮助我们将Pass@10性能从约87%提高到约95%。
您可以在data/codebase_chunks.json中找到代码文件及其块，在data/evaluation_set.jsonl中找到评估数据集。
附加说明：
提示缓存有助于在使用这种检索方法时管理成本。此功能目前可在Anthropic的第一方API上使用，并即将在AWS Bedrock和GCP Vertex的第三方合作伙伴环境中推出。我们知道许多客户在构建RAG解决方案时使用AWS Knowledge Bases和GCP Vertex AI API，此方法可通过一些定制在任一平台上使用。考虑联系Anthropic或您的AWS/GCP客户团队以获取指导！
为了更容易在Bedrock上使用此方法，AWS团队为我们提供了代码，您可以使用这些代码实现一个Lambda函数，为每个文档添加上下文。如果您部署此Lambda函数，您可以在配置Bedrock Knowledge Base时将其选择为自定义分块选项。您可以在contextual-rag-lambda-function中找到此代码。主要的lambda函数代码位于lambda_function.py中。
目录
1) 设置
2) 基本RAG
3) 上下文嵌入
4) 上下文BM25
5) 重排序
环境设置
在开始本指南之前，请确保你具备：
技术技能：
- 中级 Python 编程能力
- 对 RAG（检索增强生成）的基本理解
- 熟悉向量数据库和嵌入
- 基本的命令行操作能力
系统要求：
- Python 3.8 或更高版本
- 已安装并运行 Docker（可选，用于 BM25 搜索）
- 4GB 以上可用内存
- 约 5-10 GB 磁盘空间用于向量数据库
API 访问权限：
时间与成本：
- 预计完成时间：30-45 分钟
- API 成本：运行完整数据集约需 5-10 美元
所需库
我们需要几个库，包括：
1) anthropic - 用于与 Claude 交互
2) voyageai - 用于生成高质量嵌入
3) cohere - 用于重排序
4) elasticsearch - 用于高性能 BM25 搜索
5) pandas、numpy、matplotlib 和 scikit-learn - 用于数据处理和可视化
环境变量
请确保设置以下环境变量：
- VOYAGE_API_KEY
- ANTHROPIC_API_KEY
- COHERE_API_KEY
%%capture
!pip install --upgrade anthropic voyageai cohere elasticsearch pandas numpy
我们预先定义模型名称，以便在新模型发布时更容易更换模型。
MODEL_NAME = "claude-haiku-4-5"
我们首先初始化 Anthropic 客户端，用于生成上下文描述。
import os
import anthropic
client = anthropic.Anthropic(
# This is the default and can be omitted
api_key=os.getenv("ANTHROPIC_API_KEY"),
)
初始化向量数据库类
我们将创建一个 VectorDB 类来处理嵌入存储和相似性搜索。这个类在我们的 RAG 流程中承担三个关键功能：
- 嵌入生成：使用 Voyage AI 的嵌入模型将文本块转换为向量表示
- 存储与缓存：将嵌入保存到磁盘，避免重复计算（节省时间和 API 成本）
- 相似性搜索：使用余弦相似度为给定查询检索最相关的块
在本指南中，我们使用一个简单的内存向量数据库，并通过 pickle 序列化。这使得代码易于理解，且无需外部依赖。该类在生成嵌入后自动将其保存到磁盘，因此您只需支付一次嵌入成本。
对于生产环境，请考虑使用托管的向量数据库解决方案。
下面的 VectorDB 类遵循与生产解决方案相同的接口模式，便于以后替换。关键特性包括批量处理（每次 128 个块）、使用 tqdm 进行进度跟踪，以及查询缓存以加速评估期间的重复搜索。
import json
import pickle
from typing import Any
import numpy as np
import voyageai
from tqdm import tqdm
class VectorDB:
def __init__(self, name: str, api_key=None):
if api_key is None:
api_key = os.getenv("VOYAGE_API_KEY")
self.client = voyageai.Client(api_key=api_key)
self.name = name
self.embeddings = []
self.metadata = []
self.query_cache = {}
self.db_path = f"./data/{name}/vector_db.pkl"
def load_data(self, dataset: list[dict[str, Any]]):
if self.embeddings and self.metadata:
print("Vector database is already loaded. Skipping data loading.")
return
if os.path.exists(self.db_path):
print("Loading vector database from disk.")
self.load_db()
return
texts_to_embed = []
metadata = []
total_chunks = sum(len(doc["chunks"]) for doc in dataset)
with tqdm(total=total_chunks, desc="Processing chunks") as pbar:
for doc in dataset:
for chunk in doc["chunks"]:
texts_to_embed.append(chunk["content"])
metadata.append(
{
"doc_id": doc["doc_id"],
"original_uuid": doc["original_uuid"],
"chunk_id": chunk["chunk_id"],
"original_index": chunk["original_index"],
"content": chunk["content"],
}
)
pbar.update(1)
self._embed_and_store(texts_to_embed, metadata)
self.save_db()
print(f"Vector database loaded and saved. Total chunks processed: {len(texts_to_embed)}")
def _embed_and_store(self, texts: list[str], data: list[dict[str, Any]]):
batch_size = 128
with tqdm(total=len(texts), desc="Embedding chunks") as pbar:
result = []
for i in range(0, len(texts), batch_size):
batch = texts[i : i + batch_size]
batch_result = self.client.embed(batch, model="voyage-2").embeddings
result.extend(batch_result)
pbar.update(len(batch))
self.embeddings = result
self.metadata = data
def search(self, query: str, k: int = 20) -> list[dict[str, Any]]:
if query in self.query_cache:
query_embedding = self.query_cache[query]
else:
query_embedding = self.client.embed([query], model="voyage-2").embeddings[0]
self.query_cache[query] = query_embedding
if not self.embeddings:
raise ValueError("No data loaded in the vector database.")
similarities = np.dot(self.embeddings, query_embedding)
top_indices = np.argsort(similarities)[::-1][:k]
top_results = []
for idx in top_indices:
result = {
"metadata": self.metadata[idx],
"similarity": float(similarities[idx]),
}
top_results.append(result)
return top_results
def save_db(self):
data = {
"embeddings": self.embeddings,
"metadata": self.metadata,
"query_cache": json.dumps(self.query_cache),
}
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
with open(self.db_path, "wb") as file:
pickle.dump(data, file)
def load_db(self):
if not os.path.exists(self.db_path):
raise ValueError(
"Vector database file not found. Use load_data to create a new database."
)
with open(self.db_path, "rb") as file:
data = pickle.load(file)
self.embeddings = data["embeddings"]
self.metadata = data["metadata"]
self.query_cache = json.loads(data["query_cache"])
现在我们可以使用这个类来加载我们的数据集。
# Load your transformed dataset
with open("data/codebase_chunks.json") as f:
transformed_dataset = json.load(f)
# Initialize the VectorDB
base_db = VectorDB("base_db")
# Load and process the data
base_db.load_data(transformed_dataset)
基础 RAG
为了开始，我们将使用一种非常基础的方法来搭建一个简单的 RAG 流程。这通常被业内称为“朴素 RAG”。一个基础的 RAG 流程包括以下三个步骤：
1) 按标题对文档进行分块——仅包含每个子标题下的内容
2) 对每个文档进行嵌入
3) 使用余弦相似度检索文档，以回答查询
import json
from collections.abc import Callable
from typing import Any
from tqdm import tqdm
def load_jsonl(file_path: str) -> list[dict[str, Any]]:
"""Load JSONL file and return a list of dictionaries."""
with open(file_path) as file:
return [json.loads(line) for line in file]
def evaluate_retrieval(
queries: list[dict[str, Any]], retrieval_function: Callable, db, k: int = 20
) -> dict[str, float]:
total_score = 0
total_queries = len(queries)
for query_item in tqdm(queries, desc="Evaluating retrieval"):
query = query_item["query"]
golden_chunk_uuids = query_item["golden_chunk_uuids"]
# Find all golden chunk contents
golden_contents = []
for doc_uuid, chunk_index in golden_chunk_uuids:
golden_doc = next(
(doc for doc in query_item["golden_documents"] if doc["uuid"] == doc_uuid), None
)
if not golden_doc:
print(f"Warning: Golden document not found for UUID {doc_uuid}")
continue
golden_chunk = next(
(chunk for chunk in golden_doc["chunks"] if chunk["index"] == chunk_index), None
)
if not golden_chunk:
print(
f"Warning: Golden chunk not found for index {chunk_index} in document {doc_uuid}"
)
continue
golden_contents.append(golden_chunk["content"].strip())
if not golden_contents:
print(f"Warning: No golden contents found for query: {query}")
continue
retrieved_docs = retrieval_function(query, db, k=k)
# Count how many golden chunks are in the top k retrieved documents
chunks_found = 0
for golden_content in golden_contents:
for doc in retrieved_docs[:k]:
retrieved_content = (
doc["metadata"]
.get("original_content", doc["metadata"].get("content", ""))
.strip()
)
if retrieved_content == golden_content:
chunks_found += 1
break
query_score = chunks_found / len(golden_contents)
total_score += query_score
average_score = total_score / total_queries
pass_at_n = average_score * 100
return {"pass_at_n": pass_at_n, "average_score": average_score, "total_queries": total_queries}
def retrieve_base(query: str, db, k: int = 20) -> list[dict[str, Any]]:
"""
Retrieve relevant documents using either VectorDB or ContextualVectorDB.
:param query: The query string
:param db: The VectorDB or ContextualVectorDB instance
:param k: Number of top results to retrieve
:return: List of retrieved documents
"""
return db.search(query, k=k)
def evaluate_db(db, original_jsonl_path: str, k):
# Load the original JSONL data for queries and ground truth
original_data = load_jsonl(original_jsonl_path)
# Evaluate retrieval
results = evaluate_retrieval(original_data, retrieve_base, db, k)
return results
def evaluate_and_display(db, jsonl_path: str, k_values: list[int] = None, db_name: str = ""):
"""
Evaluate retrieval performance across multiple k values and display formatted results.
Args:
db: Vector database instance (VectorDB or ContextualVectorDB)
jsonl_path: Path to evaluation dataset
k_values: List of k values to evaluate (default: [5, 10, 20])
db_name: Optional name for the database being evaluated
Returns:
Dict mapping k values to their results
"""
if k_values is None:
k_values = [5, 10, 20]
results = {}
print(f"{'=' * 60}")
if db_name:
print(f"Evaluation Results: {db_name}")
else:
print("Evaluation Results")
print(f"{'=' * 60}\n")
for k in k_values:
print(f"Evaluating Pass@{k}...")
results[k] = evaluate_db(db, jsonl_path, k)
print() # Add spacing between evaluations
# Print summary table
print(f"{'=' * 60}")
print(f"{'Metric':<15} {'Pass Rate':<15} {'Score':<15}")
print(f"{'-' * 60}")
for k in k_values:
pass_rate = f"{results[k]['pass_at_n']:.2f}%"
score = f"{results[k]['average_score']:.4f}"
print(f"{'Pass@' + str(k):<15} {pass_rate:<15} {score:<15}")
print(f"{'=' * 60}\n")
return results
现在，让我们通过评估基础RAG系统来建立基线性能。我们将测试k=5、10和20，看看有多少黄金块出现在检索结果的前列。这为我们提供了一个衡量改进的基准。
results = evaluate_and_display(
base_db, "data/evaluation_set.jsonl", k_values=[5, 10, 20], db_name="Baseline RAG"
)
这些结果展示了我们基线RAG的性能。系统在排名前5的结果中成功检索到正确块的概率为81%，在前10名中提升至87%，在前20名中达到90%。
上下文嵌入
在基础 RAG 中，单独的块在孤立嵌入时往往缺乏足够的上下文。上下文嵌入通过使用 Claude 生成简短描述来解决这个问题，该描述将每个块“定位”在其源文档中。然后我们将块与这个上下文一起嵌入，创建更丰富的向量表示。
对于代码库数据集中的每个块，我们将块及其完整源文件传递给 Claude。Claude 生成一个简洁的解释，说明块包含什么以及它在整个文件中的位置。这个上下文在嵌入之前被前置到块中。
成本和延迟考虑
这个成本何时发生？ 上下文化在摄取时发生一次，而不是在每次查询时。与 HyDE（假设文档嵌入）等为每次搜索增加延迟的技术不同，上下文嵌入在构建向量数据库时是一次性成本。提示缓存使这变得实用。由于我们按顺序处理同一文档的所有块，我们可以利用提示缓存来节省大量成本。
- 第一个块：我们将完整文档写入缓存（支付少量额外费用）
- 后续块：从缓存读取文档（这些令牌享受 90% 折扣）
- 缓存持续 5 分钟，足以处理文档中的所有块
成本示例：对于 8k 令牌文档中的 800 令牌块，生成 100 令牌上下文，每百万文档令牌的总成本为 1.02 美元。当你运行下面的代码时，你会在日志中看到缓存节省。
注意： 一些嵌入模型有固定的输入令牌限制。如果你使用上下文嵌入时性能变差，你的上下文块可能被截断——考虑使用具有更大上下文窗口的嵌入模型。
让我们通过为单个块生成上下文来了解上下文嵌入的工作原理。我们将使用Claude创建一个定位上下文，同时你也会看到提示缓存的指标在实际中如何运作。
DOCUMENT_CONTEXT_PROMPT = """
<document>
{doc_content}
</document>
"""
CHUNK_CONTEXT_PROMPT = """
Here is the chunk we want to situate within the whole document
<chunk>
{chunk_content}
</chunk>
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.
Answer only with the succinct context and nothing else.
"""
def situate_context(doc: str, chunk: str) -> str:
response = client.messages.create(
model=MODEL_NAME,
max_tokens=1024,
temperature=0.0,
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": DOCUMENT_CONTEXT_PROMPT.format(doc_content=doc),
"cache_control": {
"type": "ephemeral"
}, # we will make use of prompt caching for the full documents
},
{
"type": "text",
"text": CHUNK_CONTEXT_PROMPT.format(chunk_content=chunk),
},
],
}
],
)
return response
jsonl_data = load_jsonl("data/evaluation_set.jsonl")
# Example usage
doc_content = jsonl_data[0]["golden_documents"][0]["content"]
chunk_content = jsonl_data[0]["golden_chunks"][0]["content"]
response = situate_context(doc_content, chunk_content)
print(f"Situated context: {response.content[0].text}")
print("-" * 10)
# Print cache performance metrics
print(f"Input tokens: {response.usage.input_tokens}")
print(f"Output tokens: {response.usage.output_tokens}")
print(f"Cache creation input tokens: {response.usage.cache_creation_input_tokens}")
print(f"Cache read input tokens: {response.usage.cache_read_input_tokens}")
构建上下文向量数据库
现在我们已经了解了如何为单个块生成上下文描述，接下来将其扩展到处理整个数据集。下面的 ContextualVectorDB 类在摄取过程中自动添加上下文，扩展了我们的基础 VectorDB。
关键特性：
- 并行处理：使用 ThreadPoolExecutor 同时为多个块添加上下文（可配置线程数）
- 自动提示缓存：按文档逐个处理块，以最大化缓存命中率
- 令牌跟踪：监控缓存性能并计算实际成本节省
- 持久化存储：将嵌入和上下文元数据保存到磁盘
运行此代码时，请注意令牌使用统计信息——您会看到 70-80% 的输入令牌是从缓存中读取的，这证明了提示缓存带来的巨大成本节省。在我们的 737 块数据集中，这将原本约 15 美元的摄取成本降低到约 3 美元。
import json
import os
import pickle
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any
import anthropic
import numpy as np
import voyageai
from tqdm import tqdm
class ContextualVectorDB:
def __init__(self, name: str, voyage_api_key=None, anthropic_api_key=None):
if voyage_api_key is None:
voyage_api_key = os.getenv("VOYAGE_API_KEY")
if anthropic_api_key is None:
anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
self.voyage_client = voyageai.Client(api_key=voyage_api_key)
self.anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)
self.name = name
self.embeddings = []
self.metadata = []
self.query_cache = {}
self.db_path = f"./data/{name}/contextual_vector_db.pkl"
self.token_counts = {"input": 0, "output": 0, "cache_read": 0, "cache_creation": 0}
self.token_lock = threading.Lock()
def situate_context(self, doc: str, chunk: str) -> tuple[str, Any]:
DOCUMENT_CONTEXT_PROMPT = """
<document>
{doc_content}
</document>
"""
CHUNK_CONTEXT_PROMPT = """
Here is the chunk we want to situate within the whole document
<chunk>
{chunk_content}
</chunk>
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.
Answer only with the succinct context and nothing else.
"""
response = self.anthropic_client.messages.create(
model=MODEL_NAME,
max_tokens=1000,
temperature=0.0,
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": DOCUMENT_CONTEXT_PROMPT.format(doc_content=doc),
"cache_control": {
"type": "ephemeral"
}, # we will make use of prompt caching for the full documents
},
{
"type": "text",
"text": CHUNK_CONTEXT_PROMPT.format(chunk_content=chunk),
},
],
},
],
extra_headers={"anthropic-beta": "prompt-caching-2024-07-31"},
)
return response.content[0].text, response.usage
def load_data(self, dataset: list[dict[str, Any]], parallel_threads: int = 1):
if self.embeddings and self.metadata:
print("Vector database is already loaded. Skipping data loading.")
return
if os.path.exists(self.db_path):
print("Loading vector database from disk.")
self.load_db()
return
texts_to_embed = []
metadata = []
total_chunks = sum(len(doc["chunks"]) for doc in dataset)
def process_chunk(doc, chunk):
# for each chunk, produce the context
contextualized_text, usage = self.situate_context(doc["content"], chunk["content"])
with self.token_lock:
self.token_counts["input"] += usage.input_tokens
self.token_counts["output"] += usage.output_tokens
self.token_counts["cache_read"] += usage.cache_read_input_tokens
self.token_counts["cache_creation"] += usage.cache_creation_input_tokens
return {
# append the context to the original text chunk
"text_to_embed": f"{contextualized_text}\n\n{chunk['content']}",
"metadata": {
"doc_id": doc["doc_id"],
"original_uuid": doc["original_uuid"],
"chunk_id": chunk["chunk_id"],
"original_index": chunk["original_index"],
"original_content": chunk["content"],
"contextualized_content": contextualized_text,
},
}
print(f"Processing {total_chunks} chunks with {parallel_threads} threads")
with ThreadPoolExecutor(max_workers=parallel_threads) as executor:
futures = []
for doc in dataset:
for chunk in doc["chunks"]:
futures.append(executor.submit(process_chunk, doc, chunk))
for future in tqdm(as_completed(futures), total=total_chunks, desc="Processing chunks"):
result = future.result()
texts_to_embed.append(result["text_to_embed"])
metadata.append(result["metadata"])
self._embed_and_store(texts_to_embed, metadata)
self.save_db()
# logging token usage
print(
f"Contextual Vector database loaded and saved. Total chunks processed: {len(texts_to_embed)}"
)
print(f"Total input tokens without caching: {self.token_counts['input']}")
print(f"Total output tokens: {self.token_counts['output']}")
print(f"Total input tokens written to cache: {self.token_counts['cache_creation']}")
print(f"Total input tokens read from cache: {self.token_counts['cache_read']}")
total_tokens = (
self.token_counts["input"]
+ self.token_counts["cache_read"]
+ self.token_counts["cache_creation"]
)
savings_percentage = (
(self.token_counts["cache_read"] / total_tokens) * 100 if total_tokens > 0 else 0
)
print(
f"Total input token savings from prompt caching: {savings_percentage:.2f}% of all input tokens used were read from cache."
)
print("Tokens read from cache come at a 90 percent discount!")
# we use voyage AI here for embeddings. Read more here: https://docs.voyageai.com/docs/embeddings
def _embed_and_store(self, texts: list[str], data: list[dict[str, Any]]):
batch_size = 128
result = [
self.voyage_client.embed(texts[i : i + batch_size], model="voyage-2").embeddings
for i in range(0, len(texts), batch_size)
]
self.embeddings = [embedding for batch in result for embedding in batch]
self.metadata = data
def search(self, query: str, k: int = 20) -> list[dict[str, Any]]:
if query in self.query_cache:
query_embedding = self.query_cache[query]
else:
query_embedding = self.voyage_client.embed([query], model="voyage-2").embeddings[0]
self.query_cache[query] = query_embedding
if not self.embeddings:
raise ValueError("No data loaded in the vector database.")
similarities = np.dot(self.embeddings, query_embedding)
top_indices = np.argsort(similarities)[::-1][:k]
top_results = []
for idx in top_indices:
result = {
"metadata": self.metadata[idx],
"similarity": float(similarities[idx]),
}
top_results.append(result)
return top_results
def save_db(self):
data = {
"embeddings": self.embeddings,
"metadata": self.metadata,
"query_cache": json.dumps(self.query_cache),
}
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
with open(self.db_path, "wb") as file:
pickle.dump(data, file)
def load_db(self):
if not os.path.exists(self.db_path):
raise ValueError(
"Vector database file not found. Use load_data to create a new database."
)
with open(self.db_path, "rb") as file:
data = pickle.load(file)
self.embeddings = data["embeddings"]
self.metadata = data["metadata"]
self.query_cache = json.loads(data["query_cache"])
# Load the transformed dataset
with open("data/codebase_chunks.json") as f:
transformed_dataset = json.load(f)
# Initialize the ContextualVectorDB
contextual_db = ContextualVectorDB("my_contextual_db")
# Load and process the data
# note: consider increasing the number of parallel threads to run this faster, or reducing the number of parallel threads if concerned about hitting your API rate limit
contextual_db.load_data(transformed_dataset, parallel_threads=5)
这些数字揭示了提示缓存对上下文嵌入的强大作用：
- 我们处理了 737 个块，分布在 9 个代码库文件中
- 61.83% 的输入令牌 从缓存中读取（227 万令牌享受 90% 折扣）
- 如果没有缓存，输入令牌成本约为 9.20 美元
- 使用缓存后，实际成本降至 约 2.85 美元（节省 69%）
缓存命中率取决于每个文档包含的块数。包含更多块的文件从缓存中获益更多，因为我们将整个文档写入缓存一次，然后针对该文件中的每个块重复读取。这就是为什么按顺序处理文档（而不是随机打乱块）对于最大化缓存效率至关重要。
现在，让我们评估这种上下文化相比基线在多大程度上提升了我们的检索性能。
results = evaluate_and_display(
contextual_db,
"data/evaluation_set.jsonl",
k_values=[5, 10, 20],
db_name="Contextual Embeddings",
)
通过在嵌入前为每个分块添加上下文，我们在所有 k 值下将检索失败率降低了 约 30-40%。这意味着在检索到的顶部块中不相关结果更少，当您将这些块传递给 Claude 生成最终回答时，答案质量会更高。
这种改进在 Pass@5 上最为显著，此时精确度最为关键——这表明带有上下文的块不仅被更频繁地检索到，而且在相关时排名也更高。
上下文BM25：混合搜索
仅使用上下文嵌入就将我们的Pass@10从87%提升到了92%。通过将语义搜索与基于关键词的搜索相结合，使用上下文BM25——一种混合方法，可以进一步降低检索失败率，从而将性能提升到更高水平。
为什么使用混合搜索？
语义搜索擅长理解含义和上下文，但可能会错过精确的关键词匹配。BM25（一种概率关键词排序算法）擅长查找特定术语，但缺乏语义理解。通过结合两者，我们可以兼得两者的优点：
- 语义搜索：捕捉概念相似性和释义
- BM25：捕捉精确术语、函数名和特定短语
- 倒数排名融合：智能地合并来自两个来源的结果
什么是BM25？
BM25是一种概率排序函数，通过考虑文档长度和词频饱和度来改进TF-IDF。它因其在关键词相关性排序方面的有效性而被广泛应用于生产搜索引擎（包括Elasticsearch）。有关技术细节，请参阅这篇博客文章。
我们不仅搜索原始块内容，还搜索我们之前生成的块*和*上下文描述。这意味着BM25可以匹配原始文本或解释性上下文中的关键词。
设置：运行Elasticsearch
在运行以下代码之前，您需要在本地运行Elasticsearch。最简单的方法是使用Docker：
docker run -d --name elasticsearch -p 9200:9200 -p 9300:9300 \
-e "discovery.type=single-node" \
-e "xpack.security.enabled=false" \
elasticsearch:9.2.0
故障排查：
- 验证它是否在运行：docker ps | grep elasticsearch
- 如果端口 9200 被占用：docker stop elasticsearch && docker rm elasticsearch
- 如果出现问题，检查日志：docker logs elasticsearch
混合搜索的工作原理
下面的 retrieve_advanced 函数实现了一个三步过程：
- 检索候选：从语义搜索和 BM25 中分别获取前 150 个结果。
- 分数融合：使用加权倒数排名融合来合并排名。
- 默认：语义搜索权重为 80%，BM25 为 20%。
- 这些权重是可调的——请根据您的用例进行实验以优化。
- 返回 top-k：融合后选择得分最高的结果。
加权系统允许您根据数据特征在语义理解和关键词精确度之间进行平衡。
import json
import os
from typing import Any
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
from tqdm import tqdm
class ElasticsearchBM25:
def __init__(self, index_name: str = "contextual_bm25_index"):
self.es_client = Elasticsearch("http://localhost:9200")
self.index_name = index_name
self.create_index()
def create_index(self):
index_settings = {
"settings": {
"analysis": {"analyzer": {"default": {"type": "english"}}},
"similarity": {"default": {"type": "BM25"}},
"index.queries.cache.enabled": False,
},
"mappings": {
"properties": {
"content": {"type": "text", "analyzer": "english"},
"contextualized_content": {"type": "text", "analyzer": "english"},
"doc_id": {"type": "keyword", "index": False},
"chunk_id": {"type": "keyword", "index": False},
"original_index": {"type": "integer", "index": False},
}
},
}
# Change this line - remove 'body=' parameter
if not self.es_client.indices.exists(index=self.index_name):
self.es_client.indices.create(
index=self.index_name,
settings=index_settings["settings"],
mappings=index_settings["mappings"],
)
print(f"Created index: {self.index_name}")
def index_documents(self, documents: list[dict[str, Any]]):
actions = [
{
"_index": self.index_name,
"_source": {
"content": doc["original_content"],
"contextualized_content": doc["contextualized_content"],
"doc_id": doc["doc_id"],
"chunk_id": doc["chunk_id"],
"original_index": doc["original_index"],
},
}
for doc in documents
]
success, _ = bulk(self.es_client, actions)
self.es_client.indices.refresh(index=self.index_name)
return success
def search(self, query: str, k: int = 20) -> list[dict[str, Any]]:
self.es_client.indices.refresh(index=self.index_name)
# Change this - remove 'body=' and pass query directly
response = self.es_client.search(
index=self.index_name,
query={
"multi_match": {
"query": query,
"fields": ["content", "contextualized_content"],
}
},
size=k,
)
return [
{
"doc_id": hit["_source"]["doc_id"],
"original_index": hit["_source"]["original_index"],
"content": hit["_source"]["content"],
"contextualized_content": hit["_source"]["contextualized_content"],
"score": hit["_score"],
}
for hit in response["hits"]["hits"]
]
def create_elasticsearch_bm25_index(db: ContextualVectorDB):
es_bm25 = ElasticsearchBM25()
es_bm25.index_documents(db.metadata)
return es_bm25
def retrieve_advanced(
query: str,
db: ContextualVectorDB,
es_bm25: ElasticsearchBM25,
k: int,
semantic_weight: float = 0.8,
bm25_weight: float = 0.2,
):
num_chunks_to_recall = 150
# Semantic search
semantic_results = db.search(query, k=num_chunks_to_recall)
ranked_chunk_ids = [
(result["metadata"]["doc_id"], result["metadata"]["original_index"])
for result in semantic_results
]
# BM25 search using Elasticsearch
bm25_results = es_bm25.search(query, k=num_chunks_to_recall)
ranked_bm25_chunk_ids = [
(result["doc_id"], result["original_index"]) for result in bm25_results
]
# Combine results
chunk_ids = list(set(ranked_chunk_ids + ranked_bm25_chunk_ids))
chunk_id_to_score = {}
# Initial scoring with weights
for chunk_id in chunk_ids:
score = 0
if chunk_id in ranked_chunk_ids:
index = ranked_chunk_ids.index(chunk_id)
score += semantic_weight * (1 / (index + 1)) # Weighted 1/n scoring for semantic
if chunk_id in ranked_bm25_chunk_ids:
index = ranked_bm25_chunk_ids.index(chunk_id)
score += bm25_weight * (1 / (index + 1)) # Weighted 1/n scoring for BM25
chunk_id_to_score[chunk_id] = score
# Sort chunk IDs by their scores in descending order
sorted_chunk_ids = sorted(
chunk_id_to_score.keys(), key=lambda x: (chunk_id_to_score[x], x[0], x[1]), reverse=True
)
# Assign new scores based on the sorted order
for index, chunk_id in enumerate(sorted_chunk_ids):
chunk_id_to_score[chunk_id] = 1 / (index + 1)
# Prepare the final results
final_results = []
semantic_count = 0
bm25_count = 0
for chunk_id in sorted_chunk_ids[:k]:
chunk_metadata = next(
chunk
for chunk in db.metadata
if chunk["doc_id"] == chunk_id[0] and chunk["original_index"] == chunk_id[1]
)
is_from_semantic = chunk_id in ranked_chunk_ids
is_from_bm25 = chunk_id in ranked_bm25_chunk_ids
final_results.append(
{
"chunk": chunk_metadata,
"score": chunk_id_to_score[chunk_id],
"from_semantic": is_from_semantic,
"from_bm25": is_from_bm25,
}
)
if is_from_semantic and not is_from_bm25:
semantic_count += 1
elif is_from_bm25 and not is_from_semantic:
bm25_count += 1
else: # it's in both
semantic_count += 0.5
bm25_count += 0.5
return final_results, semantic_count, bm25_count
def evaluate_db_advanced(
db: ContextualVectorDB,
original_jsonl_path: str,
k_values: list[int] = None,
db_name: str = "Hybrid Search",
):
"""
Evaluate hybrid search (semantic + BM25) at multiple k values with formatted results.
Args:
db: ContextualVectorDB instance
original_jsonl_path: Path to evaluation dataset
k_values: List of k values to evaluate (default: [5, 10, 20])
db_name: Name for the evaluation display
Returns:
Dict mapping k values to their results and source breakdowns
"""
if k_values is None:
k_values = [5, 10, 20]
original_data = load_jsonl(original_jsonl_path)
es_bm25 = create_elasticsearch_bm25_index(db)
results = {}
print(f"{'=' * 70}")
print(f"Evaluation Results: {db_name}")
print(f"{'=' * 70}\n")
try:
# Warm-up queries
warm_up_queries = original_data[:10]
for query_item in warm_up_queries:
_ = retrieve_advanced(query_item["query"], db, es_bm25, k_values[0])
for k in k_values:
print(f"Evaluating Pass@{k}...")
total_score = 0
total_semantic_count = 0
total_bm25_count = 0
total_results = 0
for query_item in tqdm(original_data, desc=f"Pass@{k}"):
query = query_item["query"]
golden_chunk_uuids = query_item["golden_chunk_uuids"]
golden_contents = []
for doc_uuid, chunk_index in golden_chunk_uuids:
golden_doc = next(
(doc for doc in query_item["golden_documents"] if doc["uuid"] == doc_uuid),
None,
)
if golden_doc:
golden_chunk = next(
(
chunk
for chunk in golden_doc["chunks"]
if chunk["index"] == chunk_index
),
None,
)
if golden_chunk:
golden_contents.append(golden_chunk["content"].strip())
if not golden_contents:
continue
retrieved_docs, semantic_count, bm25_count = retrieve_advanced(
query, db, es_bm25, k
)
chunks_found = 0
for golden_content in golden_contents:
for doc in retrieved_docs[:k]:
retrieved_content = doc["chunk"]["original_content"].strip()
if retrieved_content == golden_content:
chunks_found += 1
break
query_score = chunks_found / len(golden_contents)
total_score += query_score
total_semantic_count += semantic_count
total_bm25_count += bm25_count
total_results += len(retrieved_docs)
total_queries = len(original_data)
average_score = total_score / total_queries
pass_at_n = average_score * 100
semantic_percentage = (
(total_semantic_count / total_results) * 100 if total_results > 0 else 0
)
bm25_percentage = (total_bm25_count / total_results) * 100 if total_results > 0 else 0
results[k] = {
"pass_at_n": pass_at_n,
"average_score": average_score,
"total_queries": total_queries,
"semantic_percentage": semantic_percentage,
"bm25_percentage": bm25_percentage,
}
print(f"Pass@{k}: {pass_at_n:.2f}%")
print(f"Semantic: {semantic_percentage:.1f}% | BM25: {bm25_percentage:.1f}%\n")
# Print summary table
print(f"{'=' * 70}")
print(f"{'Metric':<12} {'Pass Rate':<12} {'Score':<12} {'Semantic':<12} {'BM25':<12}")
print(f"{'-' * 70}")
for k in k_values:
r = results[k]
print(
f"{'Pass@' + str(k):<12} {r['pass_at_n']:>10.2f}% {r['average_score']:>10.4f} "
f"{r['semantic_percentage']:>10.1f}% {r['bm25_percentage']:>10.1f}%"
)
print(f"{'=' * 70}\n")
return results
finally:
# Delete the Elasticsearch index
if es_bm25.es_client.indices.exists(index=es_bm25.index_name):
es_bm25.es_client.indices.delete(index=es_bm25.index_name)
print(f"Deleted Elasticsearch index: {es_bm25.index_name}")
results = evaluate_db_advanced(
contextual_db,
"data/evaluation_set.jsonl",
k_values=[5, 10, 20],
db_name="Contextual BM25 Hybrid Search",
)
重排序
通过混合搜索，我们已经取得了不错的结果（Pass@10 为 93.21%），但还有一种技术可以进一步提升性能：重排序。
什么是重排序？
重排序是一种两阶段的检索方法：
- 第一阶段 - 广泛检索：通过检索比实际需要更多的候选结果来扩大范围（例如，检索 100 个文本块）
- 第二阶段 - 精确选择：使用专门的重排序模型对这些候选结果进行评分，并只选择最相关的 top-k 个结果
为什么这种方法有效？ 初始检索方法（嵌入、BM25）针对数百万文档的检索速度进行了优化。重排序模型速度较慢但更准确——它们可以对较小的候选集进行更深入的分析。这形成了一种速度与准确性的权衡，在实践中效果很好。
我们的重排序方法
在本例中，我们将使用一个更简单的重排序流程，该流程仅基于上下文嵌入（而非完整的混合搜索）。流程如下：
- 过度检索：获取比所需多 10 倍的结果（例如，当我们需要 10 个结果时，检索 100 个文本块）
- 使用 Cohere 重排序：使用 Cohere 的
rerank-english-v3.0模型对所有候选结果进行评分 - 选择 top-k：仅返回得分最高的结果
重排序模型可以访问原始文本块内容和生成的上下文描述，这为其提供了丰富的信息来进行精确的相关性判断。
预期性能
添加重排序可以带来适度但显著的改进：
- 无重排序：Pass@10 为 92.34%（仅上下文嵌入）
- 有重排序：Pass@10 约为 95%（额外提升 2-3%）
这可能看起来提升不大，但在生产系统中，将失败率从 7.66% 降低到约 5% 可以显著改善用户体验。代价是查询延迟——根据候选集的大小，重排序每次查询会增加约 100-200 毫秒的延迟。
import json
from collections.abc import Callable
from typing import Any
import cohere
from tqdm import tqdm
def evaluate_db_rerank(
db, original_jsonl_path: str, k_values: list[int] = None, db_name: str = "Reranking"
):
"""
Evaluate reranking performance at multiple k values with formatted results.
Args:
db: ContextualVectorDB instance
original_jsonl_path: Path to evaluation dataset
k_values: List of k values to evaluate (default: [5, 10, 20])
db_name: Name for the evaluation display
Returns:
Dict mapping k values to their results
"""
if k_values is None:
k_values = [5, 10, 20]
original_data = load_jsonl(original_jsonl_path)
co = cohere.Client(os.getenv("COHERE_API_KEY"))
results = {}
print(f"{'=' * 60}")
print(f"Evaluation Results: {db_name}")
print(f"{'=' * 60}\n")
for k in k_values:
print(f"Evaluating Pass@{k} with reranking...")
total_score = 0
total_queries = len(original_data)
for query_item in tqdm(original_data, desc=f"Pass@{k}"):
query = query_item["query"]
golden_chunk_uuids = query_item["golden_chunk_uuids"]
# Find golden contents
golden_contents = []
for doc_uuid, chunk_index in golden_chunk_uuids:
golden_doc = next(
(doc for doc in query_item["golden_documents"] if doc["uuid"] == doc_uuid), None
)
if golden_doc:
golden_chunk = next(
(chunk for chunk in golden_doc["chunks"] if chunk["index"] == chunk_index),
None,
)
if golden_chunk:
golden_contents.append(golden_chunk["content"].strip())
if not golden_contents:
continue
# Retrieve and rerank
semantic_results = db.search(query, k=k * 10)
# Prepare documents for reranking
documents = [
f"{res['metadata']['original_content']}\n\nContext: {res['metadata']['contextualized_content']}"
for res in semantic_results
]
# Rerank
rerank_response = co.rerank(
model="rerank-english-v3.0", query=query, documents=documents, top_n=k
)
time.sleep(0.1) # Rate limiting
# Get final results
retrieved_docs = []
for r in rerank_response.results:
original_result = semantic_results[r.index]
retrieved_docs.append(
{"chunk": original_result["metadata"], "score": r.relevance_score}
)
# Check if golden chunks are in results
chunks_found = 0
for golden_content in golden_contents:
for doc in retrieved_docs[:k]:
retrieved_content = doc["chunk"]["original_content"].strip()
if retrieved_content == golden_content:
chunks_found += 1
break
query_score = chunks_found / len(golden_contents)
total_score += query_score
average_score = total_score / total_queries
pass_at_n = average_score * 100
results[k] = {
"pass_at_n": pass_at_n,
"average_score": average_score,
"total_queries": total_queries,
}
print(f"Pass@{k}: {pass_at_n:.2f}%")
print(f"Average Score: {average_score:.4f}\n")
# Print summary table
print(f"{'=' * 60}")
print(f"{'Metric':<15} {'Pass Rate':<15} {'Score':<15}")
print(f"{'-' * 60}")
for k in k_values:
pass_rate = f"{results[k]['pass_at_n']:.2f}%"
score = f"{results[k]['average_score']:.4f}"
print(f"{'Pass@' + str(k):<15} {pass_rate:<15} {score:<15}")
print(f"{'=' * 60}\n")
return results
results = evaluate_db_rerank(
contextual_db,
"data/evaluation_set.jsonl",
k_values=[5, 10, 20],
db_name="Contextual Embeddings + Reranking",
)
重排序带来了最强的结果，几乎消除了检索失败。让我们看看每种技术是如何在前一种技术的基础上实现这一改进的。
从我们基线 RAG 系统的 87% Pass@10 开始，通过系统性地应用高级检索技术，我们已经攀升到超过 95%。每种方法都针对不同的弱点：上下文嵌入解决了“孤立分块”问题，混合搜索捕获了嵌入可能遗漏的特定关键词查询，而重排序则应用更复杂的相关性评分来优化最终选择。
| 方法 | Pass@5 | Pass@10 | Pass@20 |
|----------|--------|---------|---------|
| 基线 RAG | 80.92% | 87.15% | 90.06% |
| + 上下文嵌入 | 88.12% | 92.34% | 94.29% |
| + 混合搜索 (BM25) | 86.43% | 93.21% | 94.99% |
| + 重排序 | 92.15% | 95.26% | 97.45% |
关键要点：
- 上下文嵌入提供了最大的单一改进（+5-7 个百分点），验证了向分块添加文档级上下文可以显著提高检索质量。仅此技术就能让你达到最优性能的 90%。
- 重排序实现了最高的绝对性能，达到 95.26% 的 Pass@10——这意味着对于 95% 的查询，正确的分块出现在前 10 个结果中。与基线 RAG 相比，这代表了 47% 的检索失败减少（失败率从 12.85% 降至 4.74%）。
- 权衡很重要：每种技术都增加了复杂性和成本：
- 上下文嵌入：一次性摄取成本（对于此数据集，使用提示缓存约 3 美元）
- 混合搜索：需要 Elasticsearch 基础设施和维护
- 重排序：增加 100-200 毫秒的查询延迟和每次查询的 API 成本（每次查询约 0.002 美元）
- 根据你的需求选择方法：
- 高容量、成本敏感：仅使用上下文嵌入（92% Pass@10，无每次查询成本）
- 最大准确性、延迟容忍：完整的重排序管道（95% Pass@10，最佳精度）
- 平衡的生产系统：混合搜索，无需每次查询成本即可获得强大性能（93% Pass@10）
对于大多数生产 RAG 系统，上下文嵌入提供了最佳的性能成本比，仅需一次性摄取成本即可实现 92% 的 Pass@10。当你需要额外的 2-3 个百分点精度并且可以承担额外的基础设施或查询成本时，可以使用混合搜索和重排序。
后续步骤和关键要点
1) 我们演示了如何使用上下文嵌入来提高检索性能，然后通过上下文 BM25 和重排序提供了额外的改进。
2) 此示例使用了代码库，但这些方法也适用于其他数据类型，例如公司内部知识库、金融和法律内容、教育内容等等。
3) 如果你是 AWS 用户，可以从 contextual-rag-lambda-function 中的 Lambda 函数开始；如果你是 GCP 用户，可以启动自己的 Cloud Run 实例并遵循类似的模式！
常见问题
常见错误：1) 未设置环境变量导致API调用失败，需确保VOYAGE_API_KEY、ANTHROPIC_API_KEY、COHERE_API_KEY已设置。2) Elasticsearch未启动或端口冲突，使用docker run命令启动，检查docker ps和日志。3) 嵌入模型输入限制导致上下文截断，若性能下降，考虑使用更大上下文窗口的嵌入模型。4) 并行线程过多导致API限流，可减少parallel_threads。5) 缓存未生效，确保按文档顺序处理块，并设置cache_control。排查方法：检查API响应状态码，查看日志中的token使用统计，验证向量数据库加载是否成功。生产优化：使用托管向量数据库（如Pinecone、Weaviate），使用异步处理，优化批量大小，使用缓存策略，监控延迟和成本。
企业应用扩展
本文方法可扩展至其他数据类型，如企业内部知识库、金融报告、法律文档等。对于AWS用户，可使用提供的Lambda函数集成Bedrock Knowledge Base；GCP用户可部署Cloud Run。未来可探索更先进的上下文生成策略，如动态上下文长度调整、多级上下文（章节级、文档级），以及结合知识图谱增强检索。此外，可尝试不同的融合算法（如基于学习的排序）和重排序模型（如Cross-Encoder）。生产环境中，建议结合用户反馈进行持续优化，并监控检索质量指标。
总结
本文介绍如何通过上下文嵌入、上下文BM25和重排序技术显著提升RAG系统的检索性能。传统RAG中，文档分块后孤立嵌入常导致上下文缺失，影响检索准确性。上下文嵌入利用Claude为每个块生成文档级上下文描述，再嵌入，平均降低35%的检索失败率。结合BM25混合搜索和重排序，Pass@10可从87%提升至95%以上。文章提供完整代码实现、环境配置、性能评估及生产优化建议，适合有编程基础的开发者学习AI工程实践。
- --
内容来源与验证
- 内容等级：PREMIUM_A
- 质量等级：S_CANDIDATE
- 难度：中级
- 验证等级：V1
- 原文作者：Anthropic
- 原文日期：2024-09-13
- 来源许可：MIT
- 来源版本：
06a851f84b4e9919b78b9655845f9900f6d5a9ab - 查看原文
站内还有更多完整实战教程
加入会员，解锁全部教程、完整代码与持续更新内容。
加入会员