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

package/dist/cjs/bundle.cjs

@@ -2178,7 +2178,10 @@ async function* streamChatCompletions({
       },
       body: JSON.stringify({
         ...rest,
-        stream: true
+        stream: true,
+        stream_options: {
+          "include_usage": true
+        }
       }),
       signal: abortSignal
     },
@@ -2970,6 +2973,8 @@ var ChatNodeImpl = class extends NodeImpl {
             }
           }
           const startTime = Date.now();
+          let usagePromptTokens = -1;
+          let usageCompletionTokens = -1;
           const chunks = streamChatCompletions({
             auth: {
               apiKey: context.settings.openAiKey ?? "",
@@ -2986,6 +2991,10 @@ var ChatNodeImpl = class extends NodeImpl {
             if (!chunk.choices) {
               continue;
             }
+            if (chunk.choices.length == 0 && chunk.usage) {
+              usagePromptTokens = chunk.usage.prompt_tokens;
+              usageCompletionTokens = chunk.usage.completion_tokens;
+            }
             for (const { delta, index } of chunk.choices) {
               if (delta.content != null) {
                 responseChoicesParts[index] ??= [];
@@ -3081,11 +3090,22 @@ var ChatNodeImpl = class 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;
@@ -11184,6 +11204,8 @@ var GraphProcessor = class _GraphProcessor {
   #aborted = false;
   #abortSuccessfully = false;
   #abortError = void 0;
+  #totalRequestTokens = 0;
+  #totalResponseTokens = 0;
   #totalCost = 0;
   #ignoreNodes = void 0;
   #nodeAbortControllers = /* @__PURE__ */ new Map();
@@ -11575,6 +11597,18 @@ var GraphProcessor = 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",
@@ -11886,7 +11920,7 @@ var GraphProcessor = class _GraphProcessor {
     }
   }
   async #processSplitRunNode(node, processId) {
-    var _a;
+    var _a, _b, _c;
     const inputValues = this.#getInputValuesForNode(node);
     if (this.#excludedDueToControlFlow(node, inputValues, processId)) {
       return;
@@ -11917,7 +11951,13 @@ var GraphProcessor = class _GraphProcessor {
               processId,
               (node2, partialOutputs, index) => this.#emitter.emit("partialOutput", { node: node2, outputs: partialOutputs, index, processId })
             );
-            if (((_a = output["cost"]) == null ? void 0 : _a.type) === "number") {
+            if (((_a = output["requestTokens"]) == null ? void 0 : _a.type) === "number") {
+              this.#totalRequestTokens += coerceTypeOptional(output["requestTokens"], "number") ?? 0;
+            }
+            if (((_b = output["responseTokens"]) == null ? void 0 : _b.type) === "number") {
+              this.#totalResponseTokens += coerceTypeOptional(output["responseTokens"], "number") ?? 0;
+            }
+            if (((_c = output["cost"]) == null ? void 0 : _c.type) === "number") {
               this.#totalCost += coerceTypeOptional(output["cost"], "number") ?? 0;
             }
             results.push({ type: "output", output });
@@ -11928,7 +11968,7 @@ var GraphProcessor = class _GraphProcessor {
       } else {
         results = await Promise.all(
           (0, import_lodash_es13.range)(0, splittingAmount).map(async (i) => {
-            var _a2;
+            var _a2, _b2, _c2;
             const inputs = fromEntries(
               entries(inputValues).map(([port, value]) => [
                 port,
@@ -11943,7 +11983,13 @@ var GraphProcessor = class _GraphProcessor {
                 processId,
                 (node2, partialOutputs, index) => this.#emitter.emit("partialOutput", { node: node2, outputs: partialOutputs, index, processId })
               );
-              if (((_a2 = output["cost"]) == null ? void 0 : _a2.type) === "number") {
+              if (((_a2 = output["requestTokens"]) == null ? void 0 : _a2.type) === "number") {
+                this.#totalRequestTokens += coerceTypeOptional(output["requestTokens"], "number") ?? 0;
+              }
+              if (((_b2 = output["responseTokens"]) == null ? void 0 : _b2.type) === "number") {
+                this.#totalResponseTokens += coerceTypeOptional(output["responseTokens"], "number") ?? 0;
+              }
+              if (((_c2 = output["cost"]) == null ? void 0 : _c2.type) === "number") {
                 this.#totalCost += coerceTypeOptional(output["cost"], "number") ?? 0;
               }
               return { type: "output", output };
@@ -11969,6 +12015,14 @@ var GraphProcessor = class _GraphProcessor {
       }, {});
       this.#nodeResults.set(node.id, aggregateResults);
       this.#visitedNodes.add(node.id);
+      this.#totalRequestTokens += (0, import_lodash_es13.sum)(results.map((r) => {
+        var _a2;
+        return coerceTypeOptional((_a2 = r.output) == null ? void 0 : _a2["requestTokens"], "number") ?? 0;
+      }));
+      this.#totalResponseTokens += (0, import_lodash_es13.sum)(results.map((r) => {
+        var _a2;
+        return coerceTypeOptional((_a2 = r.output) == null ? void 0 : _a2["responseTokens"], "number") ?? 0;
+      }));
       this.#totalCost += (0, import_lodash_es13.sum)(results.map((r) => {
         var _a2;
         return coerceTypeOptional((_a2 = r.output) == null ? void 0 : _a2["cost"], "number") ?? 0;
@@ -11979,7 +12033,7 @@ var GraphProcessor = class _GraphProcessor {
     }
   }
   async #processNormalNode(node, processId) {
-    var _a;
+    var _a, _b, _c;
     const inputValues = this.#getInputValuesForNode(node);
     if (this.#excludedDueToControlFlow(node, inputValues, processId)) {
       return;
@@ -11995,7 +12049,13 @@ var GraphProcessor = class _GraphProcessor {
       );
       this.#nodeResults.set(node.id, outputValues);
       this.#visitedNodes.add(node.id);
-      if (((_a = outputValues["cost"]) == null ? void 0 : _a.type) === "number") {
+      if (((_a = outputValues["requestTokens"]) == null ? void 0 : _a.type) === "number") {
+        this.#totalRequestTokens += coerceTypeOptional(outputValues["requestTokens"], "number") ?? 0;
+      }
+      if (((_b = outputValues["responseTokens"]) == null ? void 0 : _b.type) === "number") {
+        this.#totalResponseTokens += coerceTypeOptional(outputValues["responseTokens"], "number") ?? 0;
+      }
+      if (((_c = outputValues["cost"]) == null ? void 0 : _c.type) === "number") {
         this.#totalCost += coerceTypeOptional(outputValues["cost"], "number") ?? 0;
       }
       this.#emitter.emit("nodeFinish", { node, outputs: outputValues, processId });