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

package/dist/esm/api/streaming.js

@@ -3,6 +3,7 @@ import { coerceType } from '../utils/coerceType.js';
 /** A simplified way to listen and stream processor events, including filtering. */
 export async function* getProcessorEvents(processor, spec) {
     const previousIndexes = new Map();
+    const usages = [];
     for await (const event of processor.events()) {
         if (event.type === 'partialOutput') {
             if (spec.partialOutputs === true ||
@@ -21,9 +22,27 @@ export 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,
                 };
             }
         }
@@ -48,6 +67,12 @@ export async function* getProcessorEvents(processor, spec) {
             }
         }
         else if (event.type === 'nodeFinish') {
+            if (spec.exposeUsage) {
+                const usage = event.outputs['usage'];
+                if (usage !== undefined) {
+                    usages.push(usage);
+                }
+            }
             if (spec.nodeFinish === true ||
                 spec.nodeFinish?.includes(event.node.id) ||
                 spec.nodeFinish?.includes(event.node.title)) {
@@ -90,19 +115,44 @@ spec) {
         },
     });
 }
-export function getSingleNodeStream(processor, nodeIdOrTitle) {
+export 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') { //nodeIdOrTitle filter managed by spec
                         controller.enqueue(`data: ${JSON.stringify(event.delta)}\n\n`);
                     }
-                    else if (event.type === 'nodeFinish' &&
-                        (event.nodeId === nodeIdOrTitle || event.nodeTitle === nodeIdOrTitle)) {
+                    else if (event.type === 'error') {
+                        controller.enqueue(`error: ${JSON.stringify(event.error)}\n\n`);
+                    }
+                    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)}\n\n`);
+                        }
                         controller.close();
                     }
                 }

package/dist/types/api/streaming.d.ts

@@ -6,6 +6,12 @@ export type RivetEventStreamFilterSpec = {
     done?: boolean;
     /** If the graph errors, send an error event? */
     error?: boolean;
+    /** Expose the cost of the graph run in the response */
+    exposeCost?: boolean;
+    /** Expose the token usage of the graph run in the response */
+    exposeUsage?: boolean;
+    /** Whether to remove the final output when done is true */
+    removeFinalOutput?: boolean;
     /** Stream node start events for the specified node IDs or node titles. */
     nodeStart?: string[] | true;
     /** Stream node finish events for the specified nodeIDs or node titles. */
@@ -54,3 +60,4 @@ export declare function getProcessorSSEStream(processor: GraphProcessor,
 /** The spec for what you're streaming to the client */
 spec: RivetEventStreamFilterSpec): ReadableStream<Uint8Array<ArrayBufferLike>>;
 export declare function getSingleNodeStream(processor: GraphProcessor, nodeIdOrTitle: string): ReadableStream<string>;
+export declare function getSingleNodeStream(processor: GraphProcessor, spec: RivetEventStreamFilterSpec): ReadableStream<string>;

package/package.json

@@ -2,7 +2,7 @@
   "name": "@alpic80/rivet-core",
   "license": "MIT",
   "repository": "https://github.com/castortech/rivet",
-  "version": "1.24.0-aidon.3",
+  "version": "1.24.0-aidon.4",
   "packageManager": "yarn@3.5.0",
   "main": "dist/cjs/bundle.cjs",
   "module": "dist/esm/index.js",
@@ -48,6 +48,7 @@
     "@google-cloud/vertexai": "^0.1.3",
     "@google/generative-ai": "^0.24.0",
     "@huggingface/inference": "^2.6.4",
+    "@ironclad/rivet-core": "npm:@alpic80/rivet-core@1.24.0-aidon.4",
     "assemblyai": "^4.6.0",
     "autoevals": "^0.0.26",
     "cron-parser": "^4.9.0",