File size: 2,394 Bytes
aec3094 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import type {
NodeParameterValueType,
INodeParameterResourceLocator,
IRunData,
INodeExecutionData,
} from 'n8n-workflow';
import { NodeConnectionTypes } from 'n8n-workflow';
type TokenUsageValues = {
completionTokens: number;
promptTokens: number;
totalTokens: number;
};
type TokenUsageInfo = Record<`${string}__${number}` | 'total', TokenUsageValues>;
function isRlcValue(value: NodeParameterValueType): value is INodeParameterResourceLocator {
return Boolean(
typeof value === 'object' && value && 'value' in value && '__rl' in value && value.__rl,
);
}
export function checkNodeParameterNotEmpty(value: NodeParameterValueType) {
if (value === undefined || value === null || value === '') {
return false;
}
if (isRlcValue(value)) {
return checkNodeParameterNotEmpty(value.value);
}
return true;
}
export function extractTokenUsage(executionRunData: IRunData) {
const result: TokenUsageInfo = {
total: {
completionTokens: 0,
promptTokens: 0,
totalTokens: 0,
},
};
const extractFromNode = (nodeName: string, nodeData: INodeExecutionData, index: number) => {
function isValidTokenInfo(data: unknown): data is TokenUsageValues {
return (
typeof data === 'object' &&
data !== null &&
'completionTokens' in data &&
'promptTokens' in data &&
'totalTokens' in data &&
typeof data.completionTokens === 'number' &&
typeof data.promptTokens === 'number' &&
typeof data.totalTokens === 'number'
);
}
const tokenInfo = nodeData.json?.tokenUsage ?? nodeData.json?.tokenUsageEstimate;
if (tokenInfo && isValidTokenInfo(tokenInfo)) {
result[`${nodeName}__${index}`] = {
completionTokens: tokenInfo.completionTokens,
promptTokens: tokenInfo.promptTokens,
totalTokens: tokenInfo.totalTokens,
};
result.total.completionTokens += tokenInfo.completionTokens;
result.total.promptTokens += tokenInfo.promptTokens;
result.total.totalTokens += tokenInfo.totalTokens;
}
};
for (const [nodeName, nodeData] of Object.entries(executionRunData)) {
if (nodeData[0]?.data?.[NodeConnectionTypes.AiLanguageModel]) {
for (const [index, node] of nodeData.entries()) {
const modelNodeExecutionData = node.data?.[NodeConnectionTypes.AiLanguageModel]?.[0]?.[0];
if (modelNodeExecutionData) {
extractFromNode(nodeName, modelNodeExecutionData, index);
}
}
}
}
return result;
}
|