@alpic80/rivet-core 1.19.1-aidon.1 → 1.19.1-aidon.2

package/dist/esm/model/GraphProcessor.js

@@ -64,6 +64,8 @@ export class GraphProcessor {
     #aborted = false;
     #abortSuccessfully = false;
     #abortError = undefined;
+    #totalRequestTokens = 0;
+    #totalResponseTokens = 0;
     #totalCost = 0;
     #ignoreNodes = undefined;
     #nodeAbortControllers = new Map();
@@ -480,6 +482,18 @@ export class GraphProcessor {
                 }
                 throw error;
             }
+            if (this.#graphOutputs['requestTokens'] == null) {
+                this.#graphOutputs['requestTokens'] = {
+                    type: 'number',
+                    value: this.#totalRequestTokens,
+                };
+            }
+            if (this.#graphOutputs['responseTokens'] == null) {
+                this.#graphOutputs['responseTokens'] = {
+                    type: 'number',
+                    value: this.#totalResponseTokens,
+                };
+            }
             if (this.#graphOutputs['cost'] == null) {
                 this.#graphOutputs['cost'] = {
                     type: 'number',
@@ -811,6 +825,12 @@ export class GraphProcessor {
                     ]));
                     try {
                         const output = await this.#processNodeWithInputData(node, inputs, i, processId, (node, partialOutputs, index) => this.#emitter.emit('partialOutput', { node, outputs: partialOutputs, index, processId }));
+                        if (output['requestTokens']?.type === 'number') {
+                            this.#totalRequestTokens += coerceTypeOptional(output['requestTokens'], 'number') ?? 0;
+                        }
+                        if (output['responseTokens']?.type === 'number') {
+                            this.#totalResponseTokens += coerceTypeOptional(output['responseTokens'], 'number') ?? 0;
+                        }
                         if (output['cost']?.type === 'number') {
                             this.#totalCost += coerceTypeOptional(output['cost'], 'number') ?? 0;
                         }
@@ -829,6 +849,12 @@ export class GraphProcessor {
                     ]));
                     try {
                         const output = await this.#processNodeWithInputData(node, inputs, i, processId, (node, partialOutputs, index) => this.#emitter.emit('partialOutput', { node, outputs: partialOutputs, index, processId }));
+                        if (output['requestTokens']?.type === 'number') {
+                            this.#totalRequestTokens += coerceTypeOptional(output['requestTokens'], 'number') ?? 0;
+                        }
+                        if (output['responseTokens']?.type === 'number') {
+                            this.#totalResponseTokens += coerceTypeOptional(output['responseTokens'], 'number') ?? 0;
+                        }
                         if (output['cost']?.type === 'number') {
                             this.#totalCost += coerceTypeOptional(output['cost'], 'number') ?? 0;
                         }
@@ -858,6 +884,8 @@ export class GraphProcessor {
             }, {});
             this.#nodeResults.set(node.id, aggregateResults);
             this.#visitedNodes.add(node.id);
+            this.#totalRequestTokens += sum(results.map((r) => coerceTypeOptional(r.output?.['requestTokens'], 'number') ?? 0));
+            this.#totalResponseTokens += sum(results.map((r) => coerceTypeOptional(r.output?.['responseTokens'], 'number') ?? 0));
             this.#totalCost += sum(results.map((r) => coerceTypeOptional(r.output?.['cost'], 'number') ?? 0));
             this.#emitter.emit('nodeFinish', { node, outputs: aggregateResults, processId });
         }
@@ -875,6 +903,12 @@ export class GraphProcessor {
             const outputValues = await this.#processNodeWithInputData(node, inputValues, 0, processId, (node, partialOutputs, index) => this.#emitter.emit('partialOutput', { node, outputs: partialOutputs, index, processId }));
             this.#nodeResults.set(node.id, outputValues);
             this.#visitedNodes.add(node.id);
+            if (outputValues['requestTokens']?.type === 'number') {
+                this.#totalRequestTokens += coerceTypeOptional(outputValues['requestTokens'], 'number') ?? 0;
+            }
+            if (outputValues['responseTokens']?.type === 'number') {
+                this.#totalResponseTokens += coerceTypeOptional(outputValues['responseTokens'], 'number') ?? 0;
+            }
             if (outputValues['cost']?.type === 'number') {
                 this.#totalCost += coerceTypeOptional(outputValues['cost'], 'number') ?? 0;
             }

package/dist/esm/model/nodes/ChatNode.js

@@ -728,6 +728,8 @@ export class ChatNodeImpl extends NodeImpl {
                     }
                 }
                 const startTime = Date.now();
+                let usagePromptTokens = -1;
+                let usageCompletionTokens = -1;
                 const chunks = streamChatCompletions({
                     auth: {
                         apiKey: context.settings.openAiKey ?? '',
@@ -746,6 +748,10 @@ export class ChatNodeImpl extends NodeImpl {
                         // Could be error for some reason 🤷‍♂️ but ignoring has worked for me so far.
                         continue;
                     }
+                    if (chunk.choices.length == 0 && chunk.usage) { //capture the usage info
+                        usagePromptTokens = chunk.usage.prompt_tokens;
+                        usageCompletionTokens = chunk.usage.completion_tokens;
+                    }
                     for (const { delta, index } of chunk.choices) {
                         if (delta.content != null) {
                             responseChoicesParts[index] ??= [];
@@ -847,11 +853,22 @@ export class ChatNodeImpl extends NodeImpl {
                     throw new Error('No response from OpenAI');
                 }
                 output['in-messages'] = { type: 'chat-message[]', value: messages };
-                output['requestTokens'] = { type: 'number', value: tokenCount * (numberOfChoices ?? 1) };
+                let finalTokenCount = tokenCount * (numberOfChoices ?? 1);
                 let responseTokenCount = 0;
                 for (const choiceParts of responseChoicesParts) {
-                    responseTokenCount += await context.tokenizer.getTokenCountForString(choiceParts.join(), tokenizerInfo);
+                    responseTokenCount += await context.tokenizer.getTokenCountForString(choiceParts.join(''), tokenizerInfo);
+                }
+                if (usagePromptTokens != -1 && usageCompletionTokens != -1) {
+                    if (finalTokenCount != usagePromptTokens) {
+                        console.log(`calculated token count:${finalTokenCount}, usage:${usagePromptTokens}`);
+                        finalTokenCount = usagePromptTokens;
+                    }
+                    if (responseTokenCount != usageCompletionTokens) {
+                        console.log(`calculated response token count:${responseTokenCount}, usage:${usageCompletionTokens}`);
+                        responseTokenCount = usageCompletionTokens;
+                    }
                 }
+                output['requestTokens'] = { type: 'number', value: finalTokenCount };
                 output['responseTokens'] = { type: 'number', value: responseTokenCount };
                 const promptCostPerThousand = model in openaiModels ? openaiModels[model].cost.prompt : 0;
                 const completionCostPerThousand = model in openaiModels ? openaiModels[model].cost.completion : 0;

package/dist/esm/utils/openai.js

@@ -169,6 +169,7 @@ export class OpenAIError extends Error {
         this.name = 'OpenAIError';
     }
 }
+;
 export async function* streamChatCompletions({ endpoint, auth, signal, headers, timeout, ...rest }) {
     const abortSignal = signal ?? new AbortController().signal;
     const response = await fetchEventSource(endpoint, {
@@ -182,6 +183,9 @@ export async function* streamChatCompletions({ endpoint, auth, signal, headers,
         body: JSON.stringify({
             ...rest,
             stream: true,
+            stream_options: {
+                "include_usage": true
+            },
         }),
         signal: abortSignal,
     }, timeout ?? DEFAULT_CHAT_NODE_TIMEOUT);

package/dist/types/utils/openai.d.ts

@@ -311,6 +311,7 @@ export type ChatCompletionChunk = {
     created: number;
     model: string;
     choices?: ChatCompletionChunkChoice[];
+    usage?: CompletionUsage;
 };
 export type GptFunctionCall = {
     name: string;
@@ -320,6 +321,11 @@ export type GptFunctionCallDelta = {
     name?: string;
     arguments?: string;
 };
+export interface CompletionUsage {
+    completion_tokens: number;
+    prompt_tokens: number;
+    total_tokens: number;
+}
 export type ChatCompletionChunkChoice = {
     index: number;
     message_index: number;

package/package.json

@@ -2,7 +2,7 @@
   "name": "@alpic80/rivet-core",
   "license": "MIT",
   "repository": "https://github.com/castortech/rivet",
-  "version": "1.19.1-aidon.1",
+  "version": "1.19.1-aidon.2",
   "packageManager": "yarn@3.5.0",
   "main": "dist/cjs/bundle.cjs",
   "module": "dist/esm/index.js",