FastGPT 源码：utils.ts 中的 RRF 实现

文章目录

- [FastGPT 源码：utils.ts 中的 RRF 实现](#FastGPT 源码：utils.ts 中的 RRF 实现)
- - datasetSearchResultConcat
  - filterSearchResultsByMaxChars
  - 关键点说明
  - - [1. RRF算法实现:](#1. RRF算法实现:)
    - [2. 分数合并策略](#2. 分数合并策略)
    - [3. Token限制实现](#3. Token限制实现)

FastGPT 源码：utils.ts 中的 RRF 实现

utils.ts 中的核心函数实现。

datasetSearchResultConcat

这是一个实现 RRF(Reciprocal Rank Fusion) 算法的函数,用于合并多个搜索结果。

typescript 复制代码

export const datasetSearchResultConcat = (
  // 入参是一个数组,每个元素包含 k值 和 搜索结果列表
  arr: { k: number; list: SearchDataResponseItemType[] }[]
): SearchDataResponseItemType[] => {
  // 过滤掉空结果
  arr = arr.filter((item) => item.list.length > 0);

  // 处理边界情况
  if (arr.length === 0) return [];
  if (arr.length === 1) return arr[0].list;

  // 用Map存储合并结果,key是文档id,value是文档信息和RRF分数
  const map = new Map<string, SearchDataResponseItemType & { rrfScore: number }>();

  // RRF算法实现
  arr.forEach((item) => {
    const k = item.k;  // RRF算法的k参数

    item.list.forEach((data, index) => {
      const rank = index + 1;  // 当前item在结果中的排名
      const score = 1 / (k + rank);  // RRF分数计算公式

      const record = map.get(data.id);
      if (record) {  // 如果该文档已存在
        // 合并分数逻辑:同类型分数取最大值
        const concatScore = [...record.score];
        for (const dataItem of data.score) {
          const sameScore = concatScore.find((item) => item.type === dataItem.type);
          if (sameScore) {
            sameScore.value = Math.max(sameScore.value, dataItem.value);
          } else {
            concatScore.push(dataItem);
          }
        }

        // 更新文档信息和累加RRF分数
        map.set(data.id, {
          ...record,
          score: concatScore,
          rrfScore: record.rrfScore + score  // RRF分数累加
        });
      } else {  // 新文档
        map.set(data.id, {
          ...data,
          rrfScore: score
        });
      }
    });
  });

  // 根据RRF分数排序
  const mapArray = Array.from(map.values());
  const results = mapArray.sort((a, b) => b.rrfScore - a.rrfScore);

  // 格式化最终结果
  return results.map((item, index) => {
    // 将RRF分数添加到score数组中
    const rrfScore = item.score.find((item) => item.type === SearchScoreTypeEnum.rrf);
    if (rrfScore) {
      rrfScore.value = item.rrfScore;
      rrfScore.index = index;
    } else {
      item.score.push({
        type: SearchScoreTypeEnum.rrf,
        value: item.rrfScore,
        index
      });
    }

    // 删除临时的rrfScore字段
    delete item.rrfScore;
    return item;
  });
};

filterSearchResultsByMaxChars

这个函数用于限制搜索结果的总token数量。

typescript 复制代码

export const filterSearchResultsByMaxChars = (
  list: SearchDataResponseItemType[],  // 搜索结果列表
  maxTokens: number                    // 最大token限制
) => {
  const results: SearchDataResponseItemType[] = [];
  let totalTokens = 0;

  // 遍历结果列表
  for (let i = 0; i < list.length; i++) {
    const item = list[i];
    // 计算当前项的token数(问题+答案)
    totalTokens += countPromptTokens(item.q + item.a);
    
    // 如果超过最大限制500,停止添加
    if (totalTokens > maxTokens + 500) {
      break;
    }
    results.push(item);
    // 如果超过最大限制,停止添加
    if (totalTokens > maxTokens) {
      break;
    }
  }

  // 如果没有结果,至少返回第一条
  return results.length === 0 ? list.slice(0, 1) : results;
};

关键点说明

1. RRF算法实现:

使用公式 1/(k + rank) 计算每个结果的分数
k 值通常为 60,用于平滑排名差异
同一文档的RRF分数会累加

2. 分数合并策略

同类型的分数(如 embedding 分数)取最大值
RRF 分数采用累加方式
最终结果包含所有类型的分数

3. Token限制实现

累计计算 QA 对的 token 数
允许超过限制 500 个 token
确保至少返回一条结果