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

package/dist/cjs/bundle.cjs

@@ -21882,12 +21882,16 @@ var dotProductSimilarity = (a, b) => {
 };
 
 // src/api/streaming.ts
+function nodeMatches(spec, event) {
+  var _a;
+  return (spec == null ? void 0 : spec.includes(event.node.id)) || (spec == null ? void 0 : spec.includes((_a = event.node.data) == null ? void 0 : _a.id)) || (spec == null ? void 0 : spec.includes(event.node.title));
+}
 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))) {
+      if (spec.partialOutputs === true || nodeMatches(spec.partialOutputs, event)) {
         const currentOutput = coerceType(event.outputs["response"], "string");
         const delta = currentOutput.slice(previousIndexes.get(event.node.id) ?? 0);
         yield {
@@ -21900,9 +21904,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") {
@@ -21913,7 +21934,7 @@ async function* getProcessorEvents(processor, spec) {
         };
       }
     } else if (event.type === "nodeStart") {
-      if (spec.nodeStart === true || ((_c = spec.nodeStart) == null ? void 0 : _c.includes(event.node.id)) || ((_d = spec.nodeStart) == null ? void 0 : _d.includes(event.node.title))) {
+      if (spec.nodeStart === true || nodeMatches(spec.nodeStart, event)) {
         yield {
           type: "nodeStart",
           inputs: event.inputs,
@@ -21922,7 +21943,13 @@ async function* getProcessorEvents(processor, spec) {
         };
       }
     } else if (event.type === "nodeFinish") {
-      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))) {
+      if (spec.exposeUsage) {
+        const usage = event.outputs["usage"];
+        if (usage !== void 0) {
+          usages.push(usage);
+        }
+      }
+      if (spec.nodeFinish === true || nodeMatches(spec.nodeFinish, event)) {
         yield {
           type: "nodeFinish",
           outputs: event.outputs,
@@ -21955,19 +21982,53 @@ 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) {
+      var _a;
       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 === "nodeFinish") {
+            const value = (_a = event.outputs["valueOutput"]) == null ? void 0 : _a.value;
+            controller.enqueue(`data: ${JSON.stringify(value)}
+
+`);
+          } 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();
           }
         }