RAG 推理性能优化：Embedding 缓存 + 向量检索加速的工程方案

RAG 推理性能优化Embedding 缓存 向量检索加速的工程方案一、知识库才 5000 条怎么检索就花了 1.8 秒产品团队抱怨 Agent 响应太慢全链路追踪发现用户问题 Embedding 化耗时 400ms向量检索耗时 1100ms上下文拼接 200ms模型推理反而只要 2 秒。检索速度竟然快追上推理速度了——这显然是工程问题不是模型问题。对于中小规模的知识库1 万-10 万条向量检索本身不应该成为瓶颈。大部分延迟来自两个地方Embedding API 的网络调用和向量数据库的查询优化不足。这两者都有成熟的优化手段。二、RAG 推理性能优化的四层架构RAG 的性能优化可以按收益从高到低分为第一层Embedding 缓存。相同或相似文本的 Embedding 只计算一次。对于高频查询Embedding 缓存可以消除 90% 的 API 调用。第二层索引优化。向量索引类型FLAT、IVF、HNSW的选择直接决定了检索延迟。中大规模用 HNSW小规模用 FLAT 都无所谓。第三层查询优化。预过滤 向量检索的组合、分片并行检索、结果数量限制。第四层硬件优化。GPU 加速 Embedding 计算、SSD 存储加速、网络就近部署。flowchart LR Q[用户查询] -- C1{Embedding 缓存} C1 --|命中| V[跳过计算] C1 --|未命中| E[调用 Embedding API] E -- V V -- C2{检索缓存} C2 --|命中| D[复用检索结果] C2 --|未命中| S{向量检索} S --|预过滤| F[元数据过滤: 日期/标签] F -- I[向量索引] I -- I1[HNSW 图搜索] I -- I2[IVF 聚类剪枝] I1 -- R[Top-K 结果] I2 -- R R -- D D -- LLM[模型推理] LLM -- RES[最终回答]三、Go 实现缓存 批量 预过滤的检索优化package ragoptimizer import ( context crypto/md5 encoding/hex fmt sync time ) // 数据结构 // EmbeddingCache Embedding 向量缓存 type EmbeddingCache struct { cache map[string][]float32 // 文本哈希 → 向量 mu sync.RWMutex hits int64 misses int64 } func NewEmbeddingCache() *EmbeddingCache { return EmbeddingCache{ cache: make(map[string][]float32), } } func (ec *EmbeddingCache) Get(text string) ([]float32, bool) { ec.mu.RLock() defer ec.mu.RUnlock() key : hashText(text) vec, ok : ec.cache[key] if ok { ec.hits } return vec, ok } func (ec *EmbeddingCache) Set(text string, vec []float32) { ec.mu.Lock() defer ec.mu.Unlock() key : hashText(text) if _, exists : ec.cache[key]; !exists { ec.misses } ec.cache[key] vec } // RetrievalCache 检索结果缓存 type RetrievalCache struct { cache map[string]*CachedResult mu sync.RWMutex } type CachedResult struct { Documents []Document ExpiresAt time.Time } func (rc *RetrievalCache) IsExpired(r *CachedResult) bool { return time.Now().After(r.ExpiresAt) } // Embedding 服务客户端带重试 批量 超时 // EmbeddingClient Embedding 服务客户端 type EmbeddingClient struct { endpoint string httpClient *httpClientWithTimeout batchSize int // 批量处理大小 mu sync.Mutex } // GetEmbedding 获取单条文本的向量走缓存优先 func (ec *EmbeddingClient) GetEmbedding( ctx context.Context, text string, cache *EmbeddingCache, ) ([]float32, error) { // 第一步检查缓存 if cached, ok : cache.Get(text); ok { return cached, nil } // 第二步调用 Embedding API vec, err : ec.callEmbeddingAPI(ctx, []string{text}) if err ! nil { return nil, fmt.Errorf(embedding api 调用失败: %w, err) } // 第三步写入缓存 cache.Set(text, vec[0]) return vec[0], nil } // BatchGetEmbeddings 批量获取多条文本的向量 func (ec *EmbeddingClient) BatchGetEmbeddings( ctx context.Context, texts []string, cache *EmbeddingCache, ) ([][]float32, error) { if len(texts) 0 { return nil, nil } results : make([][]float32, len(texts)) var uncachedTexts []string var uncachedIndices []int // 分离已缓存和未缓存的文本 for i, text : range texts { if vec, ok : cache.Get(text); ok { results[i] vec } else { uncachedTexts append(uncachedTexts, text) uncachedIndices append(uncachedIndices, i) } } if len(uncachedTexts) 0 { return results, nil // 全部命中缓存 } // 分批调用 API控制每批大小防止超时 for start : 0; start len(uncachedTexts); start ec.batchSize { end : start ec.batchSize if end len(uncachedTexts) { end len(uncachedTexts) } batch : uncachedTexts[start:end] vecs, err : ec.callEmbeddingAPI(ctx, batch) if err ! nil { return nil, fmt.Errorf(批量 embedding 失败 (batch %d-%d): %w, start, end, err) } for j, vec : range vecs { idx : uncachedIndices[startj] results[idx] vec cache.Set(uncachedTexts[startj], vec) } } return results, nil } // callEmbeddingAPI 实际调用 Embedding API func (ec *EmbeddingClient) callEmbeddingAPI(ctx context.Context, texts []string) ([][]float32, error) { // 重试逻辑 const maxRetries 2 var lastErr error for attempt : 0; attempt maxRetries; attempt { if attempt 0 { time.Sleep(time.Duration(attempt*200) * time.Millisecond) } resp, err : ec.httpClient.Post(ctx, ec.endpoint, map[string]interface{}{ input: texts, model: text-embedding-ada-002, }) if err ! nil { lastErr err continue } // 解析响应提取向量具体实现略 _ resp return make([][]float32, len(texts)), nil } return nil, fmt.Errorf(embedding api 重试耗尽: %w, lastErr) } // 向量检索优化器 // SearchOptions 检索选项 type SearchOptions struct { TopK int // 返回结果数 MinScore float32 // 最低相似度 PreFilters map[string]string // 预过滤条件标签、日期等 UseApproximate bool // 是否使用近似检索HNSW/IVF } // OptimizedRetriever 优化后的检索器 type OptimizedRetriever struct { embedClient *EmbeddingClient embedCache *EmbeddingCache resultCache *RetrievalCache vectorStore VectorStore } func NewOptimizedRetriever( embedClient *EmbeddingClient, vectorStore VectorStore, ) *OptimizedRetriever { return OptimizedRetriever{ embedClient: embedClient, embedCache: NewEmbeddingCache(), resultCache: RetrievalCache{ cache: make(map[string]*CachedResult), }, vectorStore: vectorStore, } } // Search 一次完整的 RAG 检索流程带所有优化 func (r *OptimizedRetriever) Search( ctx context.Context, query string, opts SearchOptions, ) ([]Document, error) { startTime : time.Now() // 优化 1使用默认值避免零值导致异常 if opts.TopK 0 { opts.TopK 5 } // 优化 2检查检索缓存 cacheKey : r.makeCacheKey(query, opts) r.resultCache.mu.RLock() if cached, ok : r.resultCache.cache[cacheKey]; ok !r.resultCache.IsExpired(cached) { r.resultCache.mu.RUnlock() return cached.Documents, nil } r.resultCache.mu.RUnlock() // 优化 3获取查询向量从缓存或 API queryVec, err : r.embedClient.GetEmbedding(ctx, query, r.embedCache) if err ! nil { return nil, fmt.Errorf(获取查询向量失败: %w, err) } // 优化 4预过滤 向量检索 var docs []Document if len(opts.PreFilters) 0 { // 先按元数据过滤再做向量检索缩小检索范围 docs, err r.vectorStore.SearchWithFilter( ctx, queryVec, opts.TopK, opts.PreFilters, opts.UseApproximate, ) } else { // 无预过滤直接向量检索 docs, err r.vectorStore.Search( ctx, queryVec, opts.TopK, opts.UseApproximate, ) } if err ! nil { return nil, fmt.Errorf(向量检索失败: %w, err) } // 优化 5过滤低相似度结果 if opts.MinScore 0 { filtered : make([]Document, 0, len(docs)) for _, d : range docs { if d.Score opts.MinScore { filtered append(filtered, d) } } docs filtered } // 优化 6存储检索结果缓存 r.resultCache.mu.Lock() r.resultCache.cache[cacheKey] CachedResult{ Documents: docs, ExpiresAt: time.Now().Add(5 * time.Minute), } r.resultCache.mu.Unlock() elapsed : time.Since(startTime) fmt.Printf([RAG] 检索完成: %d 结果, 耗时 %v\n, len(docs), elapsed) return docs, nil } func (r *OptimizedRetriever) makeCacheKey(query string, opts SearchOptions) string { data : fmt.Sprintf(%s|%d|%.2f, query, opts.TopK, opts.MinScore) h : md5.Sum([]byte(data)) return hex.EncodeToString(h[:]) } // 接口定义 type Document struct { ID string Content string Score float32 } type VectorStore interface { Search(ctx context.Context, vec []float32, topK int, approximate bool) ([]Document, error) SearchWithFilter(ctx context.Context, vec []float32, topK int, filters map[string]string, approximate bool) ([]Document, error) } type httpClientWithTimeout struct{} func (h *httpClientWithTimeout) Post(ctx context.Context, url string, body interface{}) (interface{}, error) { return nil, nil } func hashText(text string) string { h : md5.Sum([]byte(text)) return hex.EncodeToString(h[:]) } // 缓存预热 // WarmUpCache 缓存预热预加载热门查询的向量 func (r *OptimizedRetriever) WarmUpCache(ctx context.Context, hotQueries []string) error { fmt.Printf([RAG] 开始预热缓存: %d 条查询\n, len(hotQueries)) _, err : r.embedClient.BatchGetEmbeddings(ctx, hotQueries, r.embedCache) if err ! nil { return fmt.Errorf(缓存预热失败: %w, err) } fmt.Printf([RAG] 缓存预热完成: hits待统计\n) return nil }四、性能优化的边界与陷阱不要过早优化。1000 条文档的知识库用 FLAT 索引就够了上 HNSW 只会增加内存开销。先用最简单的方案跑起来pprof 告诉你哪里慢了再优化哪里。缓存失效时机。知识库更新后检索缓存需要失效。最简单做法是按知识库版本号做缓存命名空间更新后新查询用新命名空间旧缓存自然过期。批量处理的复杂度。批量 Embedding 能节省网络开销但也增加了单次请求失败的爆炸半径。一批 100 条文本API 调用失败意味着 100 条全需要重试。batch size 建议从 20 开始调。预过滤的顺序很重要。先做元数据过滤再做向量检索缩小范围大多数情况下比先检索再过滤要快。但如果过滤条件是全文包含某个关键词那可能先检索再过滤更快。五、总结RAG 检索性能优化的核心策略缓存优先Embedding 缓存 检索结果缓存 批量调用减少网络往返 预过滤缩小检索范围 近似索引加速检索。实施顺序先上缓存效果立竿见影再调批量大小和索引类型投入产出比高。监控几个核心指标缓存命中率、检索 P95 延迟、Embedding API 调用次数——这三个指标直接决定了优化的方向。