aiquery2/api/app/utils/fastgptResponseData.js

'use strict';

const toTokenNumber = (value) => Math.max(0, Number.parseInt(value, 10) || 0);
const toText = (value, max = 128) => String(value ?? '').trim().slice(0, max);

const unwrapResponseData = (value = {}) => {
   const body = value?._body && typeof value._body === 'object' ? value._body : value;
   if (body && typeof body === 'object' && Object.prototype.hasOwnProperty.call(body, 'responseData')) {
      return body.responseData;
   }
   return body;
};

const collectResponseDataNodes = (value) => {
   const root = unwrapResponseData(value);
   if (!root || typeof root !== 'object') return [];

   const nodes = [];
   const seen = new WeakSet();
   const visit = (item, path = 'root', depth = 0) => {
      if (!item || typeof item !== 'object' || depth > 12 || seen.has(item)) return;
      seen.add(item);

      const inputTokens = toTokenNumber(item.inputTokens ?? item.input_tokens ?? item.toolCallInputTokens);
      const outputTokens = toTokenNumber(item.outputTokens ?? item.output_tokens ?? item.toolCallOutputTokens);
      const totalTokens = toTokenNumber(item.totalTokens ?? item.total_tokens) || inputTokens + outputTokens;
      const model = toText(item.model || item.modelName || item.aiModel);
      const embeddingTokens = toTokenNumber(item.embeddingTokens ?? item.embedding_tokens);
      const embeddingModel = toText(item.embeddingModel);
      const quotes = Array.isArray(item.quoteList) ? item.quoteList : (Array.isArray(item.cites) ? item.cites : []);
      const quote = quotes[0] || {};
      const knowledgeBaseId = toText(item.datasetId || item.knowledgeBaseId || item.collectionId || quote.datasetId);
      const knowledgeBaseName = toText(item.datasetName || item.knowledgeBaseName || quote.datasetName || quote.collection?.name);
      const moduleType = toText(item.moduleType).toLowerCase();
      const isDatasetSearchNode = moduleType === 'datasetsearchnode' || moduleType === 'datasetsearch';
      const ragQueryCount = isDatasetSearchNode ? 1 : 0;
      const ragHitCount = toTokenNumber(item.hitCount ?? item.matchCount) || quotes.length;
      const nodeId = toText(item.id || item.nodeId || path);

      if (totalTokens || model) {
         nodes.push({
            kind: 'model', nodeId, path, model, inputTokens, outputTokens, totalTokens,
            knowledgeBaseId, knowledgeBaseName, ragQueryCount, ragHitCount,
         });
      } else if (ragQueryCount) {
         nodes.push({ kind: 'rag', nodeId, path, knowledgeBaseId, knowledgeBaseName, ragQueryCount, ragHitCount });
      }
      if (embeddingTokens) {
         nodes.push({
            kind: 'embedding', nodeId: `${nodeId}:embedding`, path, model: embeddingModel,
            inputTokens: embeddingTokens, outputTokens: 0, totalTokens: embeddingTokens,
         });
      }

      Object.entries(item).forEach(([key, child]) => {
         if (child && typeof child === 'object') visit(child, `${path}.${key}`, depth + 1);
      });
   };
   visit(root);
   return nodes;
};

const summarizeResponseDataUsage = (value) => collectResponseDataNodes(value).reduce((usage, node) => {
   if (node.kind === 'embedding') usage.embeddingTokens += node.totalTokens;
   if (node.kind === 'model') {
      usage.inputTokens += node.inputTokens;
      usage.outputTokens += node.outputTokens;
   }
   return usage;
}, { inputTokens: 0, outputTokens: 0, embeddingTokens: 0, totalTokens: 0 });

const getResponseDataUsage = (value) => {
   const usage = summarizeResponseDataUsage(value);
   return { ...usage, totalTokens: usage.inputTokens + usage.outputTokens + usage.embeddingTokens };
};

module.exports = {
   collectResponseDataNodes,
   getResponseDataUsage,
   unwrapResponseData,
};