@alpic80/rivet-core 1.24.0-aidon.3 → 1.24.0-aidon.4

package/dist/cjs/bundle.cjs

@@ -21885,6 +21885,7 @@ var dotProductSimilarity = (a, b) => {
 async function* getProcessorEvents(processor, spec) {
   var _a, _b, _c, _d, _e, _f;
   const previousIndexes = /* @__PURE__ */ new Map();
+  const usages = [];
   for await (const event of processor.events()) {
     if (event.type === "partialOutput") {
       if (spec.partialOutputs === true || ((_a = spec.partialOutputs) == null ? void 0 : _a.includes(event.node.id)) || ((_b = spec.partialOutputs) == null ? void 0 : _b.includes(event.node.title))) {
@@ -21900,9 +21901,26 @@ async function* getProcessorEvents(processor, spec) {
       }
     } else if (event.type === "done") {
       if (spec.done) {
+        const results = event.results;
+        if (!spec.exposeCost) {
+          delete results.cost;
+        }
+        if (!spec.exposeUsage) {
+          delete results.requestTokens;
+          delete results.responseTokens;
+        } else if (usages.length) {
+          const usageOutput = {
+            type: "any[]",
+            value: usages
+          };
+          results["usages"] = usageOutput;
+        }
+        if (spec.removeFinalOutput) {
+          delete results.output;
+        }
         yield {
           type: "done",
-          graphOutput: event.results
+          graphOutput: results
         };
       }
     } else if (event.type === "error") {
@@ -21922,6 +21940,12 @@ async function* getProcessorEvents(processor, spec) {
         };
       }
     } else if (event.type === "nodeFinish") {
+      if (spec.exposeUsage) {
+        const usage = event.outputs["usage"];
+        if (usage !== void 0) {
+          usages.push(usage);
+        }
+      }
       if (spec.nodeFinish === true || ((_e = spec.nodeFinish) == null ? void 0 : _e.includes(event.node.id)) || ((_f = spec.nodeFinish) == null ? void 0 : _f.includes(event.node.title))) {
         yield {
           type: "nodeFinish",
@@ -21955,19 +21979,47 @@ data: ${JSON.stringify(data)}
     }
   });
 }
-function getSingleNodeStream(processor, nodeIdOrTitle) {
+function getSingleNodeStream(processor, arg) {
+  let spec;
+  if (typeof arg === "string") {
+    const nodeIdOrTitle = arg;
+    spec = {
+      partialOutputs: [nodeIdOrTitle],
+      nodeFinish: [nodeIdOrTitle]
+    };
+  } else {
+    spec = arg;
+  }
   return new ReadableStream({
     async start(controller) {
       try {
-        for await (const event of getProcessorEvents(processor, {
-          partialOutputs: [nodeIdOrTitle],
-          nodeFinish: [nodeIdOrTitle]
-        })) {
-          if (event.type === "partialOutput" && (event.nodeId === nodeIdOrTitle || event.nodeTitle === nodeIdOrTitle)) {
+        for await (const event of getProcessorEvents(processor, spec)) {
+          if (event.type === "partialOutput") {
             controller.enqueue(`data: ${JSON.stringify(event.delta)}
 
 `);
-          } else if (event.type === "nodeFinish" && (event.nodeId === nodeIdOrTitle || event.nodeTitle === nodeIdOrTitle)) {
+          } else if (event.type === "error") {
+            controller.enqueue(`error: ${JSON.stringify(event.error)}
+
+`);
+          } else if (event.type === "done") {
+            if (spec.done) {
+              const results = event.graphOutput;
+              if (!spec.exposeCost) {
+                delete results.cost;
+              }
+              if (!spec.exposeUsage) {
+                delete results.requestTokens;
+                delete results.responseTokens;
+                delete results.usages;
+              }
+              if (spec.removeFinalOutput) {
+                delete results.output;
+              }
+              controller.enqueue(`graphOutput: ${JSON.stringify(results)}
+
+`);
+            }
             controller.close();
           }
         }