在 PostgreSQL 中利用上下文检索优化 RAG
2025-05-30 · 发布于
工程实践 · 2141 字 · 5 分钟 · 阅读原文
Keming
在 PostgreSQL 中结合向量搜索与关键词 BM25 评分实现上下文检索（Contextual Retrieval），全面提升大语言模型的生成质量。
在标准的检索增强生成（RAG）方案中，一种常见做法是将长篇文档切分为更小、更易于管理的文本块（chunk）。这样做主要是为了结合句子嵌入（sentence embedding）与关键词 BM25 评分来提升检索效率，使系统能够快速定位并提取出相关的信息片段。虽然这种方法在许多场景下行之有效，但它也带来了一个显著的挑战：单个文本块很可能会丢失其原本所处的上下文语境。
不妨以历史文献集合为例。如果用户询问某个特定历史事件，而检索出的文本块写着：“条约正式签署，敌对状态宣告结束。”
脱离了前后的文本，这段分块就缺失了至关重要的背景信息：签署的是哪一份条约 、缔约双方是谁 ，以及这一重大事件发生在何时何地 。如果仅依赖这样一个脱离上下文的片段，提供给大语言模型（LLM）的信息将是残缺甚至可能产生误导的，从而妨碍其针对该历史事件生成全面、准确的回答。
这个例子清楚地说明：虽然分块（chunking）有助于提升检索效率，但它可能会无意间剥离赋予信息片段具体含义和相关性的必要上下文。这种上下文缺失会严重影响 RAG 系统中 LLM 生成回答的质量与准确性。
为了解决这一问题，Anthropic 探索并提出了一种名为上下文检索（Contextual Retrieval）
简而言之，可以将它视为利用整篇文档的上下文对分块进行增强（chunk augmentation）。
本文将带你一步步借助上下文信息来提升检索准确率。我们使用了与 Anthropic 教程中相同的提示词（prompt） 、数据集和评估指标
我们主要使用 Gemini 模型来生成分块上下文和文本嵌入，其他模型也同样适用。嵌入质量对向量检索指标影响很大，建议参考 MTEB 排行榜 挑选最契合你业务场景的模型。在关键词分词方面，我们使用默认的 BERT 分词器，你也可以尝试其他分词方案 。
重排序（rerank）并非本文的重点；如果你想深入了解重排序方法，欢迎阅读我们关于 Cross-Encoder 重排序 以及词元级延迟交互重排序 的专题文章。顺便提一句，由于我们已经支持多向量 MaxSim 索引 ，词元级延迟交互重排序现在可以更加高效地执行。
准备环境
我们在 VectorChord-suite
docker run --rm -d --name vdb -e POSTGRES_PASSWORD = postgres -p 5432:5432 ghcr.io/tensorchord/vchord-suite:pg17-20250414 我们将使用 vechord
pip install 'vechord[gemini]' 让我们先从基础的 RAG 流程开始。
from vechord.spec import (
ForeignKey ,
Keyword ,
PrimaryKeyAutoIncrease ,
Table ,
UniqueIndex ,
Vector ,
)
DenseVector = Vector [ 768 ]
class Document ( Table , kw_only = True ):
uid : Optional [ PrimaryKeyAutoIncrease ] = None
uuid : Annotated [ str , UniqueIndex ()]
content : str
class Chunk ( Table , kw_only = True ):
uid : Optional [ PrimaryKeyAutoIncrease ] = None
doc_uuid : Annotated [ str , ForeignKey [ Document . uuid ]]
index : int
content : str
vector : DenseVector
keyword : Keyword
class Query ( Table , kw_only = True ):
uid : Optional [ PrimaryKeyAutoIncrease ] = None
content : str
answer : str
doc_uuids : list [ str ]
chunk_index : list [ int ]
vector : DenseVector 上述定义已经自动包含了向量嵌入索引和关键词索引。
加载数据集的代码如下：
from vechord.embedding import GeminiDenseEmbedding
from vechord.registry import VechordRegistry
emb = GeminiDenseEmbedding ()
vr = VechordRegistry ( "anthropic" , "postgresql://postgres:[email protected] :5432/" )
vr . register ([ Document , Chunk , Query ])
def load_data ( filepath : str ):
with open ( filepath , "r" , encoding = "utf-8" ) as f :
docs = json . load ( f )
for doc in docs :
vr . insert (
Document (
uuid = doc [ "original_uuid" ],
content = doc [ "content" ],
)
)
for chunk in doc [ "chunks" ]:
vr . insert (
Chunk (
doc_uuid = doc [ "original_uuid" ],
index = chunk [ "original_index" ],
content = chunk [ "content" ],
vector = emb . vectorize_chunk ( chunk [ "content" ]),
keyword = Keyword ( chunk [ "content" ]),
)
)
def load_query ( filepath : str ):
queries = []
with open ( filepath , "r" , encoding = "utf-8" ) as f :
for line in f :
query = json . loads ( line )
queries . append (
Query (
content = query [ "query" ],
answer = query [ "answer" ],
doc_uuids = [ x [ 0 ] for x in query [ "golden_chunk_uuids" ]],
chunk_index = [ x [ 1 ] for x in query [ "golden_chunk_uuids" ]],
vector = emb . vectorize_query ( query [ "query" ]),
)
)
vr . copy_bulk ( queries )
评估
现在我们已经具备了基础 RAG 流程所需的全部组件。接下来定义评估指标 Pass@k，它表示检索出的 Top-K 文本块中包含了多少比例的标准答案（ground truth）文本块。
def evaluate ( topk = 5 , search_func = vector_search ):
print ( f "TopK= { topk } , search by: { search_func . __name__ } " )
queries : list [ Query ] = vr . select_by ( Query . partial_init ())
total_score = 0
start = perf_counter ()
for query in queries :
chunks : list [ Chunk ] = search_func ( query , topk )
count = 0
for doc_uuid , chunk_index in zip (
query . doc_uuids , query . chunk_index , strict = True
):
for chunk in chunks :
if chunk . doc_uuid == doc_uuid and chunk . index == chunk_index :
count += 1
break
score = count / len ( query . doc_uuids )
total_score += score
print (
f "Pass@ { topk } : { total_score / len ( queries ) : .4f } , total queries: { len ( queries ) } , QPS: { len ( queries ) / ( perf_counter () - start ) : .3f } "
) 我们可以尝试不同的检索策略，例如向量搜索、关键词搜索、结合融合（fusion）或重排序（rerank）的混合搜索。
这些检索策略可定义如下：
from vechord.rerank import CohereReranker , ReciprocalRankFusion
def vector_search ( query : Query , topk : int ) -> list [ Chunk ]:
return vr . search_by_vector ( Chunk , query . vector , topk = topk )
def keyword_search ( query : Query , topk : int ) -> list [ Chunk ]:
return vr . search_by_keyword ( Chunk , query . content , topk = topk )
def hybrid_search_fuse ( query : Query , topk : int ) -> list [ Chunk ]:
rrf = ReciprocalRankFusion ()
return rrf . fuse ([ vector_search ( query , topk ), keyword_search ( query , topk )])[: topk ]
def hybrid_search_rerank ( query : Query , topk : int , boost = 3 ) -> list [ Chunk ]:
ranker = CohereReranker ()
vecs = vector_search ( query , topk * boost )
keys = keyword_search ( query , topk * boost )
chunks = list ({ chunk . uid : chunk for chunk in vecs + keys } . values ())
indices = ranker . rerank ( query . content , [ chunk . content for chunk in chunks ])
return [ chunks [ i ] for i in indices [: topk ]] 上下文检索（Contextual Retrieval）
LLM 使用如下提示词生成上下文信息：
prompt = (
"<document> \n {whole_document} \n </document>"
"Here is the chunk we want to situate within the whole document \n "
"<chunk> \n {chunk} \n </chunk> \n "
"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."
) 上面的提示词是专为通用检索任务设计的。针对你的具体业务场景，定制更专门的提示词往往能取得更好的效果。
这一功能已经内置在 vechord 库中。
GeminiAugmenter 同样支持类似 Claude 的提示词缓存（prompt caching）功能，但它要求缓存内容至少达到 32768 个 token，而本实验所用的数据集未达到该长度阈值。
from vechord.augment import GeminiAugmenter
class ContextualChunk ( Table , kw_only = True ):
uid : Optional [ PrimaryKeyAutoIncrease ] = None
doc_uuid : Annotated [ str , ForeignKey [ Document . uuid ]]
index : int
content : str
context : str
vector : DenseVector
keyword : Keyword
def load_contextual_chunks ( filepath : str ):
augmenter = GeminiAugmenter ()
with open ( filepath , "r" , encoding = "utf-8" ) as f :
docs = json . load ( f )
for doc in docs :
augmenter . reset ( doc [ "content" ])
chunks = doc [ "chunks" ]
augments = augmenter . augment_context ([ chunk [ "content" ] for chunk in chunks ])
if len ( augments ) != len ( chunks ):
print ( f "augments length not match for uuid: { doc [ 'original_uuid' ] } , { len ( augments ) } != { len ( chunks ) } " )
for chunk , context in zip ( chunks , augments , strict = False ):
contextual_content = f " { chunk [ 'content' ] } \n\n { context } "
vr . insert (
ContextualChunk (
doc_uuid = doc [ "original_uuid" ],
index = chunk [ "original_index" ],
content = chunk [ "content" ],
context = context ,
vector = emb . vectorize_chunk ( contextual_content ),
keyword = Keyword ( contextual_content ),
)
) 随后，我们可以将检索策略扩展为如下形式：
def vector_contextual_search ( query : Query , topk : int ) -> list [ ContextualChunk ]:
return vr . search_by_vector ( ContextualChunk , query . vector , topk = topk )
def keyword_contextual_search ( query : Query , topk : int ) -> list [ ContextualChunk ]:
return vr . search_by_keyword ( ContextualChunk , query . content , topk = topk )
def hybrid_contextual_search_fuse ( query : Query , topk : int ) -> list [ ContextualChunk ]:
rrf = ReciprocalRankFusion ()
return rrf . fuse (
[ vector_contextual_search ( query , topk ), keyword_contextual_search ( query , topk )]
)[: topk ]
def hybrid_contextual_search_rerank (
query : Query , topk : int , boost = 3
) -> list [ ContextualChunk ]:
ranker = CohereReranker ()
vecs = vector_contextual_search ( query , topk * boost )
keys = keyword_contextual_search ( query , topk * boost )
chunks = list ({ chunk . uid : chunk for chunk in vecs + keys } . values ())
indices = ranker . rerank (
query . content , [ f " { chunk . content } \n { chunk . context } " for chunk in chunks ]
)
return [ chunks [ i ] for i in indices [: topk ]]
基准测试结果
所有代码均可在我们的 vechord
