FastGPT 源码：controller.ts 主要定义

文章目录

- FastGPT 源码：controller.ts 主要定义
- - 1. 主要函数和参数定义
  - 2. 参数初始化
  - 3. 内部工具函数
  - 4. 多查询召回实现
  - 5. 主流程执行

FastGPT 源码：controller.ts 主要定义

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

1. 主要函数和参数定义

type SearchDatasetDataProps = {teamId: string;        // 团队IDmodel: string;         // 向量模型similarity?: number;   // 最小相似度阈值limit: number;         // 最大Token限制datasetIds: string[];  // 要搜索的数据集IDssearchMode?: `${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数量限制,最小50if (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);