ai-sdk-graph 0.1.1 → 0.1.2

package/dist/graph.d.ts

@@ -12,7 +12,11 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     private readonly storage;
     private readonly emitter;
     private readonly stateManager;
-    constructor(storage?: GraphSDK.GraphStorage<State, NodeKeys>);
+    private readonly onFinish;
+    constructor(options?: {
+        storage?: GraphSDK.GraphStorage<State, NodeKeys>;
+        onFinish?: (state: State) => Promise<void> | void;
+    });
     node<NewKey extends string>(id: NewKey, execute: ({ state, writer, suspense, update }: {
         state: () => Readonly<State>;
         writer: Writer;
@@ -65,5 +69,8 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     private deduplicateNodes;
     private excludeTerminalNodes;
 }
-export declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(storage?: GraphSDK.GraphStorage<State, NodeKeys>): Graph<State, NodeKeys>;
+export declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(options?: {
+    storage?: GraphSDK.GraphStorage<State, NodeKeys>;
+    onFinish?: (state: State) => Promise<void> | void;
+}): Graph<State, NodeKeys>;
 export {};

package/dist/index.js

@@ -54,9 +54,11 @@ class Graph {
   storage;
   emitter = new NodeEventEmitter;
   stateManager = new StateManager;
-  constructor(storage = new InMemoryStorage) {
-    this.storage = storage;
+  onFinish;
+  constructor(options = {}) {
+    this.storage = options.storage ?? new InMemoryStorage;
     this.registerBuiltInNodes();
+    this.onFinish = options.onFinish ?? ((state) => {});
   }
   node(id, execute) {
     const node = { id, execute };
@@ -144,10 +146,16 @@ class Graph {
     });
   }
   execute(runId, initialState) {
+    let context;
     return createUIMessageStream({
       execute: async ({ writer }) => {
-        const context = await this.createExecutionContext(runId, initialState, writer);
+        context = await this.createExecutionContext(runId, initialState, writer);
         await this.runExecutionLoop(context);
+      },
+      onFinish: async () => {
+        if (context) {
+          await this.onFinish(context.state);
+        }
       }
     });
   }
@@ -337,8 +345,11 @@ class Graph {
     return nodes.filter((node) => node.id !== BUILT_IN_NODES.END);
   }
 }
-function graph(storage) {
-  return new Graph(storage);
+function graph(options = {}) {
+  return new Graph({
+    storage: options.storage,
+    onFinish: options.onFinish
+  });
 }
 
 class NodeEventEmitter {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-sdk-graph",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "Graph-based workflows for the AI SDK",
   "type": "module",
   "main": "dist/index.js",
@@ -18,9 +18,13 @@
   ],
   "scripts": {
     "test": "bun test",
-    "build": "bun build src/index.ts --outdir dist --target node --packages=external && tsc -p tsconfig.build.json"
+    "build": "bun build src/index.ts --outdir dist --target node --packages=external && tsc -p tsconfig.build.json",
+    "changeset": "changeset",
+    "version": "changeset version",
+    "release": "bun run build && changeset publish"
   },
   "devDependencies": {
+    "@changesets/cli": "^2.29.8",
     "@types/bun": "latest"
   },
   "peerDependencies": {