ai-sdk-graph 0.2.0 → 0.3.1

package/dist/graph.d.ts

@@ -1,10 +1,8 @@
 import type { GraphSDK } from './types';
-import { createUIMessageStream } from 'ai';
 export declare class SuspenseError extends Error {
     readonly data?: unknown;
     constructor(data?: unknown);
 }
-type Writer = Parameters<Parameters<typeof createUIMessageStream>[0]['execute']>[0]['writer'];
 export declare class Graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'> {
     private readonly nodeRegistry;
     private readonly edgeRegistry;
@@ -17,7 +15,7 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     constructor(options?: GraphSDK.GraphOptions<State, NodeKeys>);
     node<NewKey extends string>(id: NewKey, execute: ({ state, writer, suspense, update }: {
         state: () => Readonly<State>;
-        writer: Writer;
+        writer: GraphSDK.Writer;
         suspense: (data?: unknown) => never;
         update: (update: GraphSDK.StateUpdate<State>) => void;
     }) => Promise<void> | void): Graph<State, NodeKeys | NewKey>;
@@ -33,7 +31,7 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
         direction?: 'TB' | 'LR';
     }): string;
     execute(runId: string, initialState: State | ((state: State | undefined) => State)): ReadableStream<import("ai").InferUIMessageChunk<import("ai").UIMessage<unknown, import("ai").UIDataTypes, import("ai").UITools>>>;
-    executeInternal(runId: string, initialState: State, writer: Writer): Promise<State>;
+    executeInternal(runId: string, initialState: State, writer: GraphSDK.Writer): Promise<State>;
     private registerBuiltInNodes;
     private addEdgeToRegistry;
     private createExecutionContext;
@@ -66,4 +64,3 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     private excludeTerminalNodes;
 }
 export declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(options?: GraphSDK.GraphOptions<State, NodeKeys>): Graph<State, NodeKeys>;
-export {};

package/dist/index.d.ts

@@ -1,3 +1,4 @@
 export * from './graph';
 export * from './storage';
 export * from './types';
+export * from './utils';

package/dist/index.js

@@ -407,8 +407,34 @@ class MermaidGenerator {
     });
   }
 }
+// src/utils.ts
+async function consumeAndMergeStream(stream, writer, options) {
+  let resolver;
+  let rejecter;
+  const messages = new Promise((resolve, reject) => {
+    resolver = resolve;
+    rejecter = reject;
+  });
+  try {
+    const uiMessageStream = stream.toUIMessageStream({
+      ...options,
+      onFinish: ({ messages: messages2 }) => {
+        resolver(messages2);
+      },
+      onError: (error) => {
+        rejecter(error);
+        return error instanceof Error ? error.message : "An error occurred.";
+      }
+    });
+    writer.merge(uiMessageStream);
+  } catch (error) {
+    rejecter(error);
+  }
+  return messages;
+}
 export {
   graph,
+  consumeAndMergeStream,
   SuspenseError,
   RedisStorage,
   InMemoryStorage,

package/dist/types.d.ts

@@ -1,4 +1,4 @@
-import type { UIMessageStreamWriter } from 'ai';
+import type { createUIMessageStream, UIMessageStreamWriter } from 'ai';
 export declare namespace GraphSDK {
     type StateUpdate<State> = Partial<State> | ((state: State) => Partial<State>);
     interface SubgraphOptions<ParentState extends Record<string, unknown>, ChildState extends Record<string, unknown>> {
@@ -49,4 +49,5 @@ export declare namespace GraphSDK {
             writer: UIMessageStreamWriter;
         }) => Promise<void> | void;
     }
+    type Writer = Parameters<Parameters<typeof createUIMessageStream>[0]['execute']>[0]['writer'];
 }

package/dist/utils.d.ts

@@ -0,0 +1,3 @@
+import type { StreamTextResult, ToolSet, UIMessage } from 'ai';
+import type { GraphSDK } from './types';
+export declare function consumeAndMergeStream<Stream extends StreamTextResult<ToolSet, any>>(stream: Stream, writer: GraphSDK.Writer, options?: Omit<Parameters<Stream['toUIMessageStream']>[0], 'onFinish' | 'onError'>): Promise<UIMessage<unknown, import("ai").UIDataTypes, import("ai").UITools>[]>;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-sdk-graph",
-  "version": "0.2.0",
+  "version": "0.3.1",
   "description": "Graph-based workflows for the AI SDK",
   "type": "module",
   "main": "dist/index.js",