从零手写 JS 向量数据库与 RAG 引擎：语义分块、余弦相似度与混合检索

// 1. 递归语义分块器：优先按 Markdown 二三级标题、双换行段落切分，保留重叠窗口（Overlap）
export function chunkMarkdown(text, { maxChars = 600, overlapChars = 100 } = {}) {
  const separators = ['\n## ', '\n### ', '\n\n', '\n', '。', ' ']

  function splitRecursively(str, sepIdx = 0) {
    if (str.length <= maxChars || sepIdx >= separators.length) {
      return [str.trim()].filter(Boolean)
    }
    const sep = separators[sepIdx]
    const parts = str.split(sep)
    const chunks = []
    let current = ''

    for (let i = 0; i < parts.length; i++) {
      const piece = i === 0 ? parts[i] : sep + parts[i]
      if ((current + piece).length <= maxChars) {
        current += piece
      } else {
        if (current.trim()) chunks.push(...splitRecursively(current, sepIdx + 1))
        // 保留尾部 overlap 上下文，防止跨块句子断裂
        const tail = current.slice(-overlapChars)
        current = tail + piece
      }
    }
    if (current.trim()) chunks.push(...splitRecursively(current, sepIdx + 1))
    return chunks
  }

  return splitRecursively(String(text || ''))
}

// 2. 基于 Float32Array 的余弦相似度计算
export function cosineSimilarity(vecA, vecB) {
  let dot = 0
  let normA = 0
  let normB = 0
  for (let i = 0; i < vecA.length; i++) {
    const a = vecA[i]
    const b = vecB[i]
    dot += a * b
    normA += a * a
    normB += b * b
  }
  const denom = Math.sqrt(normA) * Math.sqrt(normB)
  return denom === 0 ? 0 : dot / denom
}

// 3. 支持向量 + 关键词 RRF 混合重排的内存检索引擎
export class HybridRagStore {
  constructor({ embedder }) {
    this.embedder = embedder // async (text) => number[]
    this.docs = []           // { id, text, metadata, vector: Float32Array, tokens: Set }
  }

  #tokenize(text) {
    return (
      String(text)
        .toLowerCase()
        .match(/[a-z0-9_]+|[\u4e00-\u9fa5]{2,}/g) || []
    )
  }

  async addDocument(id, text, metadata = {}) {
    const rawVec = await this.embedder(text)
    this.docs.push({
      id,
      text,
      metadata,
      vector: new Float32Array(rawVec),
      tokens: this.#tokenize(text)
    })
  }

  async hybridSearch(query, { topK = 5, rrfK = 60 } = {}) {
    const queryVec = new Float32Array(await this.embedder(query))
    const queryTerms = this.#tokenize(query)

    // 第 1 路：向量余弦相似度排序
    const byVector = this.docs
      .map((doc) => ({ id: doc.id, score: cosineSimilarity(queryVec, doc.vector) }))
      .sort((a, b) => b.score - a.score)

    // 第 2 路：关键词命中覆盖度（轻量 BM25 近似）排序
    const byKeyword = this.docs
      .map((doc) => {
        const tokenSet = new Set(doc.tokens)
        const hits = queryTerms.filter((t) => tokenSet.has(t)).length
        return { id: doc.id, score: queryTerms.length ? hits / queryTerms.length : 0 }
      })
      .sort((a, b) => b.score - a.score)

    // 第 3 步：RRF 倒数秩融合（无需对不同量纲的分数做复杂的归一化）
    const rrfScores = new Map()
    byVector.forEach((item, rank) => {
      rrfScores.set(item.id, (rrfScores.get(item.id) || 0) + 1 / (rrfK + rank + 1))
    })
    byKeyword.forEach((item, rank) => {
      rrfScores.set(item.id, (rrfScores.get(item.id) || 0) + 1 / (rrfK + rank + 1))
    })

    return [...this.docs]
      .map((doc) => ({
        id: doc.id,
        text: doc.text,
        metadata: doc.metadata,
        rrfScore: rrfScores.get(doc.id) || 0
      }))
      .sort((a, b) => b.rrfScore - a.rrfScore)
      .slice(0, topK)
  }
}