当前位置：首页 > article >正文

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

article 2025/3/6 19:48:51

文章目录

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

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

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

datasetSearchResultConcat

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

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数量。

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;
};