ChatGPT 检索插件,支持通过自然语言查询对个人或组织文档进行语义搜索和检索
基于 ChatGPT Retrieval Plugin 实现三步检索优化,将在现有架构上添加:relevance reranking、冗余移除、task-aware filters。
该插件实现了检索增强生成(RAG)系统,使大型语言模型能够访问和推理外部文档。其工作原理是:使用 OpenAI 的嵌入模型将文档转换为向量嵌入,将这些嵌入存储在向量数据库中,然后根据与用户查询的语义相似性检索相关文档片段。
步骤 1:扩展 Metadata 支持任务感知过滤
首先,在 models/models.py 中扩展 DocumentMetadata,添加更丰富的过滤字段:
classDocumentMetadata(BaseModel):
source: Optional[Source] = None
source_id: Optional[str] = None
url: Optional[str] = None
created_at: Optional[str] = None
author: Optional[str] = None
# 新增字段
updated_at: Optional[str] = None# 最后更新时间,关键!
region: Optional[str] = None# 区域:EU, US, CN 等
product_version: Optional[str] = None# 产品版本
document_type: Optional[str] = None# 文档类型:policy, faq, spec 等
department: Optional[str] = None# 部门
同样更新 DocumentMetadataFilter:
classDocumentMetadataFilter(BaseModel):
document_id: Optional[str] = None
source: Optional[Source] = None
source_id: Optional[str] = None
author: Optional[str] = None
start_date: Optional[str] = None
end_date: Optional[str] = None
# 新增过滤字段
updated_after: Optional[str] = None# 只检索指定日期后更新的文档
region: Optional[str] = None
product_version: Optional[str] = None
document_type: Optional[str] = None
department: Optional[str] = None
步骤 2:创建 Reranking 服务
创建新文件 services/reranker.py:
from typing importList
from sentence_transformers import CrossEncoder
import numpy as np
from models.models import DocumentChunkWithScore
classReranker:
def__init__(self, model_name: str = "BAAI/bge-reranker-v2-m3"):
"""
初始化 Cross-Encoder 用于相关性重排序
比纯向量相似度慢,但准确率高得多
"""
self.model = CrossEncoder(model_name)
asyncdefrerank(
self,
query: str,
documents: List[DocumentChunkWithScore],
top_k: int = 5
) -> List[DocumentChunkWithScore]:
"""
使用 cross-encoder 重新排序文档
将 query 与每个文档实际配对计算相关性
"""
iflen(documents) <= top_k:
return documents
# 准备 (query, document) 对
pairs = [(query, doc.text) for doc in documents]
# Cross-encoder 预测相关性分数
scores = self.model.predict(pairs)
# 根据分数重新排序
indexed_scores = list(enumerate(scores))
indexed_scores.sort(key=lambda x: x[1], reverse=True)
# 返回前 top_k 个
reranked_docs = []
for idx, score in indexed_scores[:top_k]:
doc = documents[idx]
# 用 reranker 的分数替换向量相似度分数
doc.score = float(score)
reranked_docs.append(doc)
return reranked_docs
步骤 3:创建去重服务
创建新文件 services/deduplicator.py:
from typing importList
from sentence_transformers import SentenceTransformer
import numpy as np
from models.models import DocumentChunkWithScore
classDeduplicator:
def__init__(self, model_name: str = "all-MiniLM-L6-v2", threshold: float = 0.9):
"""
初始化去重器
threshold: cosine similarity 阈值,超过则视为重复
"""
self.model = SentenceTransformer(model_name)
self.threshold = threshold
asyncdefdeduplicate(self, documents: List[DocumentChunkWithScore]) -> List[DocumentChunkWithScore]:
"""
移除冗余文档
使用 embeddings 计算相似度,超过阈值则删除
"""
iflen(documents) <= 1:
return documents
# 提取文本并编码
texts = [doc.text for doc in documents]
embeddings = self.model.encode(texts)
# 归一化
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
# 计算相似度矩阵
similarity_matrix = np.dot(embeddings, embeddings.T)
# 去重:如果与已有文档相似度过高,则跳过
unique_docs = []
unique_indices = set()
for i, doc inenumerate(documents):
if i in unique_indices:
continue
# 保留这个文档
unique_docs.append(doc)
# 标记与它相似的所有其他文档
similar_indices = np.where(similarity_matrix[i] > self.threshold)[0]
for idx in similar_indices:
if idx != i: # 不标记自己
unique_indices.add(idx)
return unique_docs
步骤 4:增强 DataStore 查询方法
在 datastore/datastore.py 中,修改 query 方法以集成优化层:
from services.reranker import Reranker
from services.deduplicator import Deduplicator
classDataStore(ABC):
def__init__(self):
# 初始化优化组件
self.reranker = Reranker()
self.deduplicator = Deduplicator()
asyncdefquery(
self,
queries: List[Query],
enable_rerank: bool = True,
enable_dedup: bool = True,
rerank_top_k: int = 5
) -> List[QueryResult]:
"""
增强的查询方法,支持三步优化
"""
# 第一步:初始向量检索(返回 top 50)
initial_top_k = 50
queries_with_top_k = [
QueryWithEmbedding(
query=q.query,
filter=q.filter,
top_k=initial_top_k,
embedding=awaitself._get_embedding(q.query) # 你需要实现这个
)
for q in queries
]
# 获取初始结果
query_results = awaitself._query(queries_with_top_k)
# 对每个查询结果应用优化
optimized_results = []
for result in query_results:
documents = result.results
# 第二步:相关性重排序
if enable_rerank andlen(documents) > rerank_top_k:
documents = awaitself.reranker.rerank(
query=result.query,
documents=documents,
top_k=rerank_top_k
)
# 第三步:冗余移除
if enable_dedup:
documents = awaitself.deduplicator.deduplicate(documents)
optimized_results.append(
QueryResult(query=result.query, results=documents)
)
return optimized_results
asyncdef_get_embedding(self, text: str) -> List[float]:
"""获取文本的 embedding"""
from services.openai import get_embeddings
embeddings = get_embeddings([text])
return embeddings[0]
步骤 5:实际使用示例
from models.models import Query, DocumentMetadataFilter
# 用户提问:"Summarize the latest refund policy changes for EU customers."
query = Query(
query="Summarize the latest refund policy changes for EU customers.",
filter=DocumentMetadataFilter(
region="EU", # 只检索 EU 文档
document_type="policy", # 只检索政策文档
updated_after="2025-01-01"# 只检索 2025 年后更新的文档
)
)
# 执行优化查询
results = await datastore.query(
queries=[query],
enable_rerank=True,
enable_dedup=True,
rerank_top_k=5
)
# 结果:只有 3-5 个高度相关、最新、不冗余的 chunks
for result in results:
print(f"Query: {result.query}")
print(f"Found {len(result.results)} optimized chunks:")
for doc in result.results:
print(f" - Score: {doc.score:.3f}")
print(f" Region: {doc.metadata.region}")
print(f" Updated: {doc.metadata.updated_at}")
print(f" Text: {doc.text[:100]}...")
步骤 6:Dockerfile 依赖更新
在 Dockerfile 中添加必要的依赖:
# 添加 reranking 和去重的依赖
RUN pip install sentence-transformers torch
关键收益总结
-
1. Token 节省 20-40%:从 50 个 chunks 减少到 3-5 个 -
2. 准确率提升 15-30%:通过 reranking 和过滤 -
3. 可调试性:清楚知道模型看了哪些 context -
4. 避免幻觉:移除相互冲突的旧文档
夜雨聆风
