上下文嵌入

使用上下文感知的嵌入提升检索质量

🌐 查看英文原文 | 源码 Notebook

使用上下文检索增强 RAG

Note: For more background information on Contextual Retrieval, including additional performance evaluations on various datasets, we recommend reading our accompanying blog post.

Retrieval Augmented Generation (RAG) enables Claude to leverage your internal knowledge bases, codebases, or any other corpus of documents when providing a response. Enterprises are increasingly building RAG applications to improve workflows in customer support, Q&A over internal company documents, financial & legal analysis, code generation, and much more.

In a separate guide, we walked through setting up a basic retrieval system, demonstrated how to evaluate its performance, and then outlined a few techniques to improve performance. In this guide, we present a technique for improving retrieval performance: Contextual Embeddings.

In traditional RAG, documents are typically split into smaller chunks for efficient retrieval. While this approach works well for many applications, it can lead to problems when individual chunks lack sufficient context. Contextual Embeddings solve this problem by adding relevant context to each chunk before embedding. This method improves the quality of each embedded chunk, allowing for more accurate retrieval and thus better overall performance. Averaged across all data sources we tested, Contextual Embeddings reduced the top-20-chunk retrieval failure rate by 35%.

The same chunk-specific context can also be used with BM25 search to further improve retrieval performance. We introduce this technique in the “Contextual BM25” section. 在这篇指南中，我们将演示如何使用9个代码库的数据集作为我们的知识库，构建和优化一个上下文检索系统。我们将逐步介绍：

设置基本检索管道以建立性能基准。
上下文嵌入：它是什么，为什么它有效，以及为什么提示缓存使其适用于生产场景。
实现上下文嵌入并展示性能改进。
上下文BM25：通过上下文 BM25混合搜索提高性能。
通过重新排序提高性能，

评估指标和数据集：

我们使用了一个预先分块的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知识库和GCP Vertex AI API，此方法可以在任何平台上使用，只需稍作定制。请考虑联系Anthropic或您的AWS/GCP账户团队以获取有关此方面的指导！为了便于在Bedrock上使用此方法，AWS团队为我们提供了代码，您可以使用这些代码来实现一个Lambda函数，为每个文档添加上下文。如果您部署此Lambda函数，您可以在配置Bedrock知识库时将其选为自定义分块选项。您可以在contextual-rag-lambda-function中找到此代码。主要的Lambda函数代码位于lambda_function.py。

环境配置

在开始本指南之前，请确保您具备以下条件：

技术技能： - 中级Python编程 - 对RAG（检索增强生成）的基本理解 - 对向量数据库和嵌入的熟悉 - 基本的命令行操作能力

系统要求： - Python 3.8+ - 安装并运行Docker（可选，用于BM25搜索） - 4GB+可用RAM - 向量数据库需要约5-10 GB磁盘空间

API访问： - Anthropic API密钥（免费层足够） - Voyage AI API密钥 - Cohere API密钥

时间与成本： - 预计完成时间：30-45分钟 - API费用：运行完整数据集约需$5-10

库

我们需要一些库，包括：

anthropic - 用于与Claude交互
voyageai - 用于生成高质量嵌入
cohere - 用于重新排序
elasticsearch - 用于高效的BM25搜索
pandas, numpy, matplotlib, 和 scikit-learn - 用于数据处理和可视化

环境变量

请确保以下环境变量已设置：

请注意，以上翻译保留了原文的Markdown格式、技术术语、代码引用、变量名、函数名、HTML标签和callout语法，并保持了原文的段落结构和换行。 :::- VOYAGE_API_KEY - ANTHROPIC_API_KEY - COHERE_API_KEY

%%capture
!pip install --upgrade anthropic voyageai cohere elasticsearch pandas numpy
```我们提前定义我们的模型名称，以便在发布新模型时更容易更改模型

```markdown
# 模型名称定义

We define our model names up front to make it easier to change models as new models are released.

## 模型名称定义的重要性

We define our model names up front to make it easier to change models as new models are released.

- **原因一**：方便管理和维护。
- **原因二**：提高开发效率。

## 示例

```python
model_name = "RAG-Model-v1"

请确保模型名称遵循一定的命名规范，以便于识别和查找。

```python
model_name = "RAG-Model-v1"

请确保模型名称遵循一定的命名规范，以便于识别和查找。

```python MODEL_NAME = "claude-haiku-4-5"我们将首先初始化我们将用于生成上下文描述的Anthropic客户端。

# 初始化Anthropic客户端以生成上下文描述
``````python
import os

import anthropic

client = anthropic.Anthropic(
    # This is the default and can be omitted
    api_key=os.getenv("ANTHROPIC_API_KEY"),
)
```## 初始化向量数据库类

我们将创建一个VectorDB类来处理嵌入存储和相似性搜索。这个类在我们的RAG管道中扮演三个关键角色：

1. **嵌入生成**：使用Voyage AI的嵌入模型将文本块转换为向量表示
2. **存储与缓存**：将嵌入保存到磁盘以避免重新计算（这可以节省时间和API成本）
3. **相似性搜索**：使用余弦相似性检索给定查询的最相关块

对于本指南，我们使用一个简单的内存向量数据库，并使用pickle序列化。这使得代码易于理解，且不需要外部依赖。该类在生成嵌入后会自动将其保存到磁盘，因此您只需支付一次嵌入成本。

对于生产使用，请考虑托管向量数据库解决方案。

下面的VectorDB类遵循与生产解决方案相同的接口模式，便于以后替换。关键特性包括批量处理（每次128个块）、使用tqdm进行进度跟踪以及查询缓存以加快评估期间的重复搜索。

```python
class VectorDB:
    # ...

注意：本示例代码仅为示意，实际使用时请根据具体需求进行调整。 :::```python 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"])


Now we can use this class to load our dataset

```python
# 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)

Processing chunks: 100%|██████████| 737/737 [00:00<00:00, 985400.72it/s]
Embedding chunks: 100%|██████████| 737/737 [00:42<00:00, 17.28it/s]

Vector database loaded and saved. Total chunks processed: 737
```## 基础 RAG

要开始，我们将使用一种简化的方法来设置一个基本的 RAG 流程。这在业界有时被称为“朴素 RAG”。一个基本的 RAG 流程包括以下 3 个步骤：

1) 按标题分割文档 - 只包含每个子标题的内容

2) 对每个文档进行嵌入

3) 使用余弦相似度检索文档以回答查询```python
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，以查看有多少黄金块出现在检索结果的前列。这为我们提供了一个基准，可以用来衡量改进。

```markdown
# 建立基线性能

现在，让我们通过评估基本的RAG系统来建立我们的基线性能。我们将测试k=5、10和20，以查看有多少黄金块出现在检索结果的前列。这为我们提供了一个基准，可以用来衡量改进。

``````python
results = evaluate_and_display(
    base_db, "data/evaluation_set.jsonl", k_values=[5, 10, 20], db_name="Baseline RAG"
)

============================================================
Evaluation Results: Contextual Embeddings
============================================================

Evaluating Pass@5...

Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 65.26it/s]


Evaluating Pass@10...

Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.87it/s]


Evaluating Pass@20...

Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.72it/s]


============================================================
Metric          Pass Rate       Score          
------------------------------------------------------------
Pass@5          80.92%          0.8092         
Pass@10         87.15%          0.8715         
Pass@20         90.06%          0.9006         
============================================================
```这些结果展示了我们的基线RAG性能。系统在顶部5个结果中成功检索到正确片段的次数为81%，在顶部10个结果中提升至87%，在顶部20个结果中达到90%。

## 上下文 Embedding

使用基本的RAG，当单独嵌入时，单个片段往往缺乏足够的上下文。上下文Embedding通过使用Claude生成一个简短的描述来解决这个问题，该描述将每个片段“定位”在其源文档中。然后我们将片段及其上下文一起嵌入，创建更丰富的向量表示。

对于我们的代码库数据集中的每个片段，我们将其片段及其完整的源文件传递给Claude。Claude生成一个简明扼要的解释，说明该片段包含的内容以及它在整个文件中的位置。这个上下文在嵌入片段之前被添加到片段前面。

### 成本和延迟考虑

**何时产生成本？**上下文化在摄入时发生一次，而不是在每次查询时。与像HyDE（假设文档嵌入）这样的技术不同，它会为每次搜索增加延迟，上下文嵌入在构建您的向量数据库时是一次性成本。提示缓存使其变得可行。由于我们按顺序处理同一文档的所有片段，我们可以利用提示缓存以实现显著的节省。

1. 第一个片段：我们将完整文档写入缓存（支付一小笔溢价）
2. 后续片段：从缓存中读取文档（这些标记的90%折扣）
3. 缓存持续5分钟，足够处理文档中的所有片段

**成本示例**：对于8k标记文档中的800标记片段，以及100标记的生成上下文，总成本为每百万文档标记1.02美元。当您运行下面的代码时，您将在日志中看到缓存节省。

**注意**：一些嵌入模型有固定的输入标记限制。如果您在使用上下文嵌入时看到性能下降，您的上下文化片段可能被截断——考虑使用具有更大上下文窗口的嵌入模型。

---
让我们通过为单个片段生成上下文来查看上下文嵌入的工作示例。我们将使用Claude来创建定位上下文，您还将看到提示缓存指标的实际应用。

```markdown
# 上下文嵌入工作示例

Let's see an example of how contextual embeddings work by generating context for a single chunk. We'll use Claude to create a situating context, and you'll also see the prompt caching metrics in action.

- **步骤 1**: 生成上下文
  - 使用Claude创建定位上下文

- **步骤 2**: 应用提示缓存指标
  - 观察提示缓存指标的实际应用

::: callout-notice
请注意，以下代码块将展示如何使用Python和JSON进行API调用。
:::

```python
# 示例代码
def get_data():
    # 代码逻辑
    pass

# 变量名
data = get_data()

# 函数名
print(data)

{
  "key": "value"
}

SELECT * FROM table;

``````python DOCUMENT_CONTEXT_PROMPT = ““” {doc_content} ““”

CHUNK_CONTEXT_PROMPT = ““” Here is the chunk we want to situate within the whole document {chunk_content}

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}“)

Situated context: This chunk contains the module documentation and initial struct definition for a differential fuzzing executor. It introduces the `DiffExecutor` struct that wraps two executors (primary and secondary) to run them sequentially with the same input, comparing their behavior for differential testing. The chunk establishes the core data structure and imports needed for the differential fuzzing implementation.

Input tokens: 3412 Output tokens: 76 Cache creation input tokens: 0 Cache read input tokens: 0 ```### 构建上下文向量数据库

现在我们已经看到了如何为单个片段生成上下文描述，让我们将其扩展到处理整个数据集。下面的 ContextualVectorDB 类扩展了我们的基本 VectorDB，在摄取过程中实现自动上下文化。

主要功能：

并行处理：使用 ThreadPoolExecutor 同时上下文化多个片段（可配置线程数）
自动提示缓存：按文档顺序处理片段，以最大化缓存命中
Token 跟踪：监控缓存性能并计算实际成本节省
持久化存储：将嵌入和上下文化元数据保存到磁盘

运行时，请注意 token 使用统计信息——您会看到 70-80% 的输入 token 是从缓存中读取的，这证明了提示缓存带来的显著成本节省。在我们的 737 个片段数据集中，这将原本的 ~$15 摄取作业降低到 ~$3。

::: callout-parallel-processing

并行处理是提高数据处理效率的关键。通过使用 ThreadPoolExecutor，我们可以同时上下文化多个片段，从而显著减少处理时间。

::: callout-caching

自动提示缓存是优化性能的另一个重要方面。通过按文档顺序处理片段，我们可以最大化缓存命中，从而减少对原始数据的读取次数。

::: callout-token-tracking

Token 跟踪功能可以帮助我们监控缓存性能，并计算实际成本节省。这对于优化资源使用和降低总体成本至关重要。

::: callout-persistent-storage

持久化存储确保了嵌入和上下文化元数据的安全保存。这样，即使在系统重启后，我们也能恢复之前的工作状态。```python 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"])


```python
# 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)

Processing 737 chunks with 5 threads

Processing chunks: 100%|██████████| 737/737 [05:32<00:00,  2.22it/s]

Contextual Vector database loaded and saved. Total chunks processed: 737
Total input tokens without caching: 1223730
Total output tokens: 58161
Total input tokens written to cache: 176079
Total input tokens read from cache: 2267069
Total input token savings from prompt caching: 61.83% of all input tokens used were read from cache.
Tokens read from cache come at a 90 percent discount!
```这些数字揭示了上下文嵌入中提示缓存的强大功能：

- 我们处理了9个代码库文件中的**737个块**
- **61.83%的输入token**是从缓存中读取的（90%折扣下的2.27M个token）
- 没有缓存，这将花费**~$9.20**的输入token
- 使用缓存后，实际成本降至**~$2.85**（节省69%）

缓存命中率取决于每个文档包含的块的数量。包含更多块的文件从缓存中受益更多，因为我们只需将整个文档写入缓存一次，然后为该文件中的每个块重复读取。这就是为什么按顺序处理文档（而不是随机打乱块）对于最大化缓存效率至关重要。

现在让我们评估这种上下文化如何提高我们的检索性能，与基线相比。

```python
def process_documents():
    # 代码逻辑
    pass

def evaluate_performance():
    # 代码逻辑
    pass

请注意，上述代码仅为示例，实际实现可能有所不同。 :::python results = evaluate_and_display( contextual_db, "data/evaluation_set.jsonl", k_values=[5, 10, 20], db_name="Contextual Embeddings", )

============================================================
Evaluation Results: Contextual Embeddings
============================================================

Evaluating Pass@5...

Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.58it/s]


Evaluating Pass@10...

Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.37it/s]


Evaluating Pass@20...

Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.14it/s]


============================================================
Metric          Pass Rate       Score          
------------------------------------------------------------
Pass@5          88.12%          0.8812         
Pass@10         92.34%          0.9234         
Pass@20         94.29%          0.9429         
============================================================
```通过在嵌入之前为每个块添加上下文，我们在所有 k 值上减少了 **~30-40%** 的检索失败率。这意味着在您检索到的顶级块中，不相关的结果更少，当您将这些块传递给 Claude 进行最终响应生成时，可以得到更好的答案。

改进在 Pass@5 时最为明显，此时精确度最为重要——这表明上下文化的块不仅被检索得更频繁，而且在相关时排名更高。

## 上下文 BM25：混合搜索

仅使用上下文嵌入就将我们的 Pass@10 从 87% 提高到了 92%。通过结合语义搜索和基于关键词的搜索使用 **上下文 BM25**——一种进一步降低检索失败率的混合方法，我们可以进一步提高性能。

### 为什么使用混合搜索？

语义搜索擅长理解意义和上下文，但可能会错过精确的关键词匹配。BM25（一种概率关键词排名算法）擅长找到特定术语，但缺乏语义理解。通过结合两者，我们可以兼得两者之长：

- **语义搜索**：捕捉概念相似性和释义
- **BM25**：捕捉精确术语、函数名和特定短语
- **互反排名融合**：智能地合并来自两个来源的结果

### 什么是 BM25？

BM25 是一种概率排名函数，它通过考虑文档长度和术语饱和度来改进 TF-IDF。它在生产搜索引擎（包括 Elasticsearch）中得到了广泛的应用，因为它在排名关键词相关性方面的有效性。有关技术细节，请参阅 [这篇博客文章](https://www.elastic.co/blog/practical-bm25-part-2-the-bm25-algorithm-and-its-variables)。

我们不仅搜索原始块内容，还搜索我们之前生成的块和上下文描述。这意味着 BM25 可以匹配原始文本或解释性上下文中的关键词。

### 设置：运行 Elasticsearch

在运行以下代码之前，您需要在本地运行 Elasticsearch。最简单的方法是使用 Docker：

::: callout-notice
注意：以下代码示例假设您已经安装了 Elasticsearch 和 Docker。
:::

```python
# 代码示例

# JSON 示例

-- SQL 示例
``````bash
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函数实现了一个三步过程：

1. 获取候选结果：从语义搜索和BM25中获取前150个结果
2. 评分融合：使用加权互逆排名融合方法结合排名
    - 默认：语义搜索占80%权重，BM25占20%
    - 这些权重是可调的——根据您的用例进行实验以优化
3. 返回top-k：融合后选择得分最高的结果

该加权系统允许您根据数据特征在语义理解和关键词精度之间进行平衡。

::: callout-注意
请注意，权重设置可能需要根据您的具体应用场景进行调整。
:::```python
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",
)

Created index: contextual_bm25_index
======================================================================
Evaluation Results: Contextual BM25 Hybrid Search
======================================================================

Evaluating Pass@5...

Pass@5: 100%|██████████| 248/248 [00:05<00:00, 41.79it/s]

Pass@5: 88.86%
Semantic: 54.6% | BM25: 45.4%

Evaluating Pass@10...

Pass@10: 100%|██████████| 248/248 [00:05<00:00, 42.20it/s]

Pass@10: 92.31%
Semantic: 57.6% | BM25: 42.4%

Evaluating Pass@20...

Pass@20: 100%|██████████| 248/248 [00:05<00:00, 42.15it/s]

Pass@20: 95.23%
Semantic: 60.8% | BM25: 39.2%

======================================================================
Metric       Pass Rate    Score        Semantic     BM25        
----------------------------------------------------------------------
Pass@5            88.86%     0.8886       54.6%       45.4%
Pass@10           92.31%     0.9231       57.6%       42.4%
Pass@20           95.23%     0.9523       60.8%       39.2%
======================================================================

Deleted Elasticsearch index: contextual_bm25_index
```## 重新排序（Reranking）

我们在混合搜索（93.21% Pass@10）方面取得了良好的成果，但还有一项技术可以进一步挖掘额外性能：**重新排序**。

### 什么是重新排序？

重新排序是一种两阶段检索方法：

1. **第一阶段 - 广泛检索**：检索比所需更多的候选内容（例如，检索100个片段）
2. **第二阶段 - 精确选择**：使用专门的重新排序模型对这些候选内容进行评分，并仅选择最相关的top-k个

**为什么这种方法有效？** 初始检索方法（嵌入、BM25）针对数百万份文档的快速检索进行了优化。重新排序模型虽然速度较慢，但更准确——它们可以在较小的候选集上进行更深入的分析。这种速度/准确性的权衡在实践中效果良好。

### 我们的重新排序方法

在这个例子中，我们将使用一个更简单的基于上下文嵌入的重新排序管道（而不是完整的混合搜索）。以下是流程：

1. **过度检索**：获取所需结果10倍的数量（例如，当我们只需要10个片段时，检索100个片段）
2. **使用Cohere重新排序**：使用Cohere的`rerank-english-v3.0`模型对所有候选内容进行评分
3. **选择top-k**：仅返回评分最高的结果

重新排序模型可以访问原始片段内容和我们所生成的上下文描述，这为它提供了丰富的信息，以便进行精确的相关性判断。

### 预期性能

添加重新排序可以带来适度但有意义的效果提升：
- **无重新排序**：92.34% Pass@10（仅使用上下文嵌入）
- **有重新排序**：~95% Pass@10（额外增加2-3%）

这可能看起来很小，但在生产系统中，将失败率从7.66%降低到~5%可以显著改善用户体验。权衡是查询延迟——重新排序根据候选集大小增加约100-200ms。```python
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",
)

============================================================
Evaluation Results: Contextual Embeddings + Reranking
============================================================

Evaluating Pass@5 with reranking...

Pass@5: 100%|██████████| 248/248 [01:40<00:00,  2.47it/s]

Pass@5: 92.15%
Average Score: 0.9215

Evaluating Pass@10 with reranking...

Pass@10: 100%|██████████| 248/248 [02:29<00:00,  1.66it/s]

Pass@10: 95.26%
Average Score: 0.9526

Evaluating Pass@20 with reranking...

Pass@20: 100%|██████████| 248/248 [03:03<00:00,  1.35it/s]

``` Pass@20: 97.45% Average Score: 0.9745

============================================================ Metric Pass Rate Score
———————————————————— Pass@5 92.15% 0.9215
Pass@10 95.26% 0.9526
Pass@20 97.45% 0.9745
============================================================ ```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个百分点精度且能承担额外的基础设施或查询成本时，混合搜索和重排序是可用的。

后续步骤和关键要点

我们展示了如何使用上下文嵌入来提高检索性能，然后通过上下文BM25和重排序提供了额外的改进。
此示例使用了代码库，但这些方法也适用于其他数据类型，例如内部公司知识库、金融和法律内容、教育内容等等。
如果您是AWS用户，您可以从contextual-rag-lambda-function中的Lambda函数开始，如果您是GCP用户，您可以启动自己的Cloud Run实例，并遵循类似的模式！

Tip

中文读者提示：本章节的代码和输出为英文原文。如需理解具体实现细节，可参考上方中文导读和代码注释。如有疑问，欢迎在 GitHub Issues 讨论。

使用上下文检索增强 RAG

评估指标和数据集：

补充说明：

目录

环境配置

库

环境变量

后续步骤和关键要点