FastGPT 源码：controller.ts 主要定义

FastGPT 源码：controller.ts 主要定义

controller.ts 中的核心搜索实现。

1. 主要函数和参数定义

type SearchDatasetDataProps = {
  teamId: string;        // 团队ID
  model: string;         // 向量模型
  similarity?: number;   // 最小相似度阈值
  limit: number;         // 最大Token限制
  datasetIds: string[];  // 要搜索的数据集IDs
  searchMode?: `${DatasetSearchModeEnum}`; // 搜索模式(embedding/全文/混合)
  usingReRank?: boolean; // 是否使用重排序
  reRankQuery: string;   // 重排序查询
  queries: string[];     // 搜索查询列表
};

2. 参数初始化

export async function searchDatasetData(props: SearchDatasetDataProps) {
  let {
    teamId,
    reRankQuery,
    queries,
    model,
    similarity = 0,
    limit: maxTokens,
    searchMode = DatasetSearchModeEnum.embedding,  // 默认使用向量搜索
    usingReRank = false,
    datasetIds = []
  } = props;

  // 初始化搜索模式
  searchMode = DatasetSearchModeMap[searchMode] ? searchMode : DatasetSearchModeEnum.embedding;
  // 检查是否可以使用重排序
  usingReRank = usingReRank && global.reRankModels.length > 0;

  // token数量限制,最小50
  if (maxTokens < 50) {
    maxTokens = 1500;
  }
}

3. 内部工具函数

// 计算不同搜索模式下的召回限制
const countRecallLimit = () => {
  if (searchMode === DatasetSearchModeEnum.embedding) {
    return { embeddingLimit: 150, fullTextLimit: 0 };  // 纯向量搜索
  }
  if (searchMode === DatasetSearchModeEnum.fullTextRecall) {
    return { embeddingLimit: 0, fullTextLimit: 150 };  // 纯全文搜索
  }
  return { embeddingLimit: 100, fullTextLimit: 80 };   // 混合搜索
};

// 向量检索实现
const embeddingRecall = async ({ query, limit }) => {
  // 1. 获取查询文本的向量
  const { vectors, tokens } = await getVectorsByText({...});
  
  // 2. 从向量库检索
  const { results } = await recallFromVectorStore({...});
  
  // 3. 获取完整的QA数据
  const dataList = await MongoDatasetData.find({...});
  
  // 4. 格式化结果
  return { embeddingRecallResults: formatResult, tokens };
};

// 全文检索实现
const fullTextRecall = async ({ query, limit }) => {
  // 使用MongoDB的全文索引搜索
  let searchResults = await Promise.all(
    datasetIds.map(id => MongoDatasetData.find({
      $text: { $search: jiebaSplit({ text: query }) }
    }))
  );
  
  return { fullTextRecallResults, tokenLen: 0 };
};

// 重排序实现
const reRankSearchResult = async ({ data, query }) => {
  // 调用重排序模型
  const results = await reRankRecall({
    query,
    documents: data.map(item => ({
      id: item.id,
      text: `${item.q}\n${item.a}`
    }))
  });
  
  // 合并重排序分数
  return mergeResult;
};

4. 多查询召回实现

const multiQueryRecall = async ({ embeddingLimit, fullTextLimit }) => {
  // 并行执行每个query的向量和全文检索
  await Promise.all(
    queries.map(async (query) => {
      const [{ tokens, embeddingRecallResults }, { fullTextRecallResults }] = 
        await Promise.all([
          embeddingRecall({ query, limit: embeddingLimit }),
          fullTextRecall({ query, limit: fullTextLimit })
        ]);
      // 收集结果...
    })
  );

  // 使用RRF合并多个查询的结果
  const rrfEmbRecall = datasetSearchResultConcat(
    embeddingRecallResList.map(list => ({ k: 60, list }))
  );
  const rrfFTRecall = datasetSearchResultConcat(
    fullTextRecallResList.map(list => ({ k: 60, list }))
  );
};

5. 主流程执行

/* main step */
// 1. 获取召回限制
const { embeddingLimit, fullTextLimit } = countRecallLimit();

// 2. 执行多查询召回
const { embeddingRecallResults, fullTextRecallResults } = 
  await multiQueryRecall({...});

// 3. 执行重排序
const reRankResults = await (async () => {
  if (!usingReRank) return [];
  // 合并向量和全文结果
  const concatRecallResults = embeddingRecallResults.concat(
    fullTextRecallResults.filter(item => !set.has(item.id))
  );
  // 去重并重排序
  return reRankSearchResult({...});
})();

// 4. 最终RRF合并
const rrfConcatResults = datasetSearchResultConcat([
  { k: 60, list: embeddingRecallResults },
  { k: 60, list: fullTextRecallResults },
  { k: 58, list: reRankResults }
]);

// 5. 结果过滤
// - 去重
// - 相似度过滤
// - Token限制过滤
const filterSameDataResults = rrfConcatResults.filter(...);
const scoreFilter = filterByScore(filterSameDataResults);
const finalResults = filterResultsByMaxTokens(scoreFilter, maxTokens);

实现完整覆盖如下流程:

1. 多种搜索模式(向量/全文/混合)
2. 多查询并行检索
3. 查询结果RRF合并
4. 重排序优化
5. 结果过滤和限制

整个过程保证了搜索的全面性(多种召回方式)和准确性(重排序和过滤)。