14.12-混合检索

class BM25 {
  private k1 = 1.5
  private b = 0.75
  private avgDocLength = 0
  private docCount = 0
  private docFreqs: Map<string, number> = new Map()  // 每个词出现在多少文档
  private docLengths: number[] = []
  private docs: Array<Map<string, number>> = []  // 每个文档的词频
 
  fit(documents: string[]) {
    this.docCount = documents.length
    this.docLengths = []
    this.docFreqs.clear()
    this.docs = []
 
    let totalLength = 0
 
    for (const doc of documents) {
      const words = this.tokenize(doc)
      const tf = new Map<string, number>()
 
      for (const word of words) {
        tf.set(word, (tf.get(word) ?? 0) + 1)
        this.docFreqs.set(word, (this.docFreqs.get(word) ?? 0) + 1)
      }
 
      this.docs.push(tf)
      this.docLengths.push(words.length)
      totalLength += words.length
    }
 
    this.avgDocLength = totalLength / this.docCount
  }
 
  search(query: string, topK: number): Array<{ index: number; score: number }> {
    const queryWords = this.tokenize(query)
    const scores: Array<{ index: number; score: number }> = []
 
    for (let i = 0; i < this.docCount; i++) {
      let score = 0
 
      for (const word of queryWords) {
        const tf = this.docs[i].get(word) ?? 0
        const df = this.docFreqs.get(word) ?? 0
        const idf = Math.log((this.docCount - df + 0.5) / (df + 0.5) + 1)
 
        const numerator = tf * (this.k1 + 1)
        const denominator = tf + this.k1 * (1 - this.b + this.b * this.docLengths[i] / this.avgDocLength)
 
        score += idf * numerator / denominator
      }
 
      if (score > 0) {
        scores.push({ index: i, score })
      }
    }
 
    scores.sort((a, b) => b.score - a.score)
    return scores.slice(0, topK)
  }
 
  private tokenize(text: string): string[] {
    // 简单分词——实际项目用 jieba 或 ICU
    return text.toLowerCase().split(/[\s,.;:!?]+/).filter(Boolean)
  }
}

// SPLADE 模型会做词扩展——即使文档里没有「返款」这个词，
// 如果语义上相关，SPLADE 也会在「返款」对应的维度上给出非零权重
const sparseVector = await spladeModel.encode("退款政策说明")
// { indices: [142, 892, 2051, ...], values: [0.82, 0.45, 0.67, ...] }

function rrfFusion(
  resultsList: Array<Array<{ id: string; score: number }>>,
  k = 60
): Array<{ id: string; score: number }> {
  const scores = new Map<string, number>()
 
  for (const results of resultsList) {
    for (let rank = 0; rank < results.length; rank++) {
      const { id } = results[rank]
      const current = scores.get(id) ?? 0
      scores.set(id, current + 1 / (k + rank + 1))  // rank 从 1 开始
    }
  }
 
  return Array.from(scores.entries())
    .map(([id, score]) => ({ id, score }))
    .sort((a, b) => b.score - a.score)
}

function weightedFusion(
  vectorResults: Array<{ id: string; score: number }>,
  keywordResults: Array<{ id: string; score: number }>,
  vectorWeight = 0.7,
  keywordWeight = 0.3
): Array<{ id: string; score: number }> {
  // 归一化分数到 [0, 1]
  const normalizeVector = normalizeScores(vectorResults)
  const normalizeKeyword = normalizeScores(keywordResults)
 
  const scores = new Map<string, number>()
 
  for (const { id, score } of normalizeVector) {
    scores.set(id, (scores.get(id) ?? 0) + vectorWeight * score)
  }
  for (const { id, score } of normalizeKeyword) {
    scores.set(id, (scores.get(id) ?? 0) + keywordWeight * score)
  }
 
  return Array.from(scores.entries())
    .map(([id, score]) => ({ id, score }))
    .sort((a, b) => b.score - a.score)
}
 
function normalizeScores(results: Array<{ id: string; score: number }>): Array<{ id: string; score: number }> {
  if (results.length === 0) return []
 
  const max = Math.max(...results.map((r) => r.score))
  const min = Math.min(...results.map((r) => r.score))
 
  if (max === min) {
    return results.map((r) => ({ ...r, score: 1 }))
  }
 
  return results.map((r) => ({ ...r, score: (r.score - min) / (max - min) }))
}

// 写入时同时存稠密和稀疏向量
await client.upsert('documents', {
  points: [
    {
      id: 'doc-1-chunk-0',
      vector: {
        dense: denseEmbedding,           // 768 维稠密向量
        'text-sparse': sparseVector,     // 稀疏向量（BM25 或 SPLADE）
      },
      payload: { text: '...' },
    },
  ],
})
 
// 混合查询
const results = await client.query('documents', {
  prefetch: [
    {
      query: denseQueryVector,
      using: 'dense',
      limit: 20,
    },
    {
      query: sparseQueryVector,
      using: 'text-sparse',
      limit: 20,
    },
  ],
  query: { fusion: 'rrf' },  // RRF 融合
  limit: 5,
})

-- 混合查询
SELECT
  id,
  content,
  1 - (embedding <=> $1) AS vector_score,
  ts_rank(content_tsv, plainto_tsquery('simple', $2)) AS keyword_score,
  -- 加权融合
  0.7 * (1 - (embedding <=> $1))
  + 0.3 * ts_rank(content_tsv, plainto_tsquery('simple', $2))
  AS combined_score
FROM document_chunks
WHERE content_tsv @@ plainto_tsquery('simple', $2)
   OR 1 - (embedding <=> $1) >= 0.8
ORDER BY combined_score DESC
LIMIT 5

async function hybridSearch(
  queryText: string,
  queryVector: number[],
  options: SearchOptions
): Promise<SearchResult[]> {
  // 并行执行两路检索
  const [vectorResults, keywordResults] = await Promise.all([
    vectorDB.search(queryVector, { topK: 20, ...options }),
    bm25Search(queryText, { topK: 20, ...options }),
  ])
 
  // RRF 融合
  return rrfFusion([vectorResults, keywordResults]).slice(0, options.topK ?? 5)
}