ai-sdk-graph 0.4.0 → 0.5.0

package/dist/compiled-graph.d.ts

@@ -0,0 +1,50 @@
+import type { GraphSDK } from './types';
+import type { Graph } from './graph';
+export declare class SuspenseError extends Error {
+    readonly data?: unknown;
+    constructor(data?: unknown);
+}
+export declare class CompiledGraph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'> {
+    private readonly nodeRegistry;
+    private readonly edgeRegistry;
+    private readonly subgraphRegistry;
+    private readonly storage;
+    private readonly emitter;
+    private readonly stateManager;
+    private readonly onFinish;
+    private readonly onStart;
+    constructor(nodeRegistry: ReadonlyMap<NodeKeys, GraphSDK.Node<State, NodeKeys>>, edgeRegistry: ReadonlyMap<NodeKeys, GraphSDK.Edge<State, NodeKeys>[]>, subgraphRegistry: ReadonlyMap<NodeKeys, {
+        subgraph: Graph<any, any>;
+        options: GraphSDK.SubgraphOptions<State, any>;
+    }>, options?: GraphSDK.CompileOptions<State, NodeKeys>);
+    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: GraphSDK.Writer): Promise<State>;
+    private createExecutionContext;
+    private runExecutionLoop;
+    private runExecutionLoopInternal;
+    private executeWithStrategy;
+    private resumeWithStrategy;
+    private executeNodesWithStrategy;
+    private handleBatchResultWithStrategy;
+    private hasSuspendedNodes;
+    private hasNodesToExecute;
+    private restoreCheckpoint;
+    private isValidCheckpoint;
+    private hasNodeIds;
+    private hasAtLeastOneNode;
+    private restoreFromCheckpoint;
+    private createFreshExecution;
+    private persistCheckpoint;
+    private executeBatch;
+    private executeSingleNode;
+    private createNodeExecutionParams;
+    private executeSubgraphNode;
+    private generateSubgraphRunId;
+    private createSuspenseFunction;
+    private resolveNodeIds;
+    private computeNextNodes;
+    private findSuccessors;
+    private resolveEdgeTarget;
+    private deduplicateNodes;
+    private excludeTerminalNodes;
+}

package/dist/graph.d.ts

@@ -1,18 +1,11 @@
 import type { GraphSDK } from './types';
-export declare class SuspenseError extends Error {
-    readonly data?: unknown;
-    constructor(data?: unknown);
-}
+import { CompiledGraph } from './compiled-graph';
+export { SuspenseError } from './compiled-graph';
 export declare class Graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'> {
     private readonly nodeRegistry;
     private readonly edgeRegistry;
     private readonly subgraphRegistry;
-    private readonly storage;
-    private readonly emitter;
-    private readonly stateManager;
-    private readonly onFinish;
-    private readonly onStart;
-    constructor(options?: GraphSDK.GraphOptions<State, NodeKeys>);
+    constructor();
     node<NewKey extends string>(id: NewKey, execute: ({ state, writer, suspense, update }: {
         state: () => Readonly<State>;
         writer: GraphSDK.Writer;
@@ -27,40 +20,11 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
         subgraph: Graph<any, any>;
         options: GraphSDK.SubgraphOptions<State, any>;
     }>;
+    compile(options?: GraphSDK.CompileOptions<State, NodeKeys>): CompiledGraph<State, NodeKeys>;
     toMermaid(options?: {
         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: GraphSDK.Writer): Promise<State>;
     private registerBuiltInNodes;
     private addEdgeToRegistry;
-    private createExecutionContext;
-    private runExecutionLoop;
-    private runExecutionLoopInternal;
-    private executeWithStrategy;
-    private resumeWithStrategy;
-    private executeNodesWithStrategy;
-    private handleBatchResultWithStrategy;
-    private hasSuspendedNodes;
-    private hasNodesToExecute;
-    private restoreCheckpoint;
-    private isValidCheckpoint;
-    private hasNodeIds;
-    private hasAtLeastOneNode;
-    private restoreFromCheckpoint;
-    private createFreshExecution;
-    private persistCheckpoint;
-    private executeBatch;
-    private executeSingleNode;
-    private createNodeExecutionParams;
-    private executeSubgraphNode;
-    private generateSubgraphRunId;
-    private createSuspenseFunction;
-    private resolveNodeIds;
-    private computeNextNodes;
-    private findSuccessors;
-    private resolveEdgeTarget;
-    private deduplicateNodes;
-    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 declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(): Graph<State, NodeKeys>;

package/dist/index.d.ts

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

package/dist/index.js

@@ -31,7 +31,7 @@ class RedisStorage {
   }
 }
 
-// src/graph.ts
+// src/compiled-graph.ts
 import { createUIMessageStream } from "ai";
 var BUILT_IN_NODES = {
   START: "START",
@@ -47,53 +47,23 @@ class SuspenseError extends Error {
   }
 }
 
-class Graph {
-  nodeRegistry = new Map;
-  edgeRegistry = new Map;
-  subgraphRegistry = new Map;
+class CompiledGraph {
+  nodeRegistry;
+  edgeRegistry;
+  subgraphRegistry;
   storage;
   emitter = new NodeEventEmitter;
   stateManager = new StateManager;
   onFinish;
   onStart;
-  constructor(options = {}) {
+  constructor(nodeRegistry, edgeRegistry, subgraphRegistry, options = {}) {
+    this.nodeRegistry = nodeRegistry;
+    this.edgeRegistry = edgeRegistry;
+    this.subgraphRegistry = subgraphRegistry;
     this.storage = options.storage ?? new InMemoryStorage;
-    this.registerBuiltInNodes();
     this.onFinish = options.onFinish ?? (() => {});
     this.onStart = options.onStart ?? (() => {});
   }
-  node(id, execute) {
-    const node = { id, execute };
-    this.nodeRegistry.set(node.id, node);
-    return this;
-  }
-  edge(from, to) {
-    const edge = { from, to };
-    this.addEdgeToRegistry(edge);
-    return this;
-  }
-  graph(id, subgraph, options) {
-    this.subgraphRegistry.set(id, { subgraph, options });
-    const node = {
-      id,
-      execute: async () => {}
-    };
-    this.nodeRegistry.set(node.id, node);
-    return this;
-  }
-  get nodes() {
-    return this.nodeRegistry;
-  }
-  get edges() {
-    return this.edgeRegistry;
-  }
-  get subgraphs() {
-    return this.subgraphRegistry;
-  }
-  toMermaid(options) {
-    const generator = new MermaidGenerator(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry);
-    return generator.generate(options);
-  }
   execute(runId, initialState) {
     let context;
     return createUIMessageStream({
@@ -118,15 +88,6 @@ class Graph {
     await this.runExecutionLoopInternal(context);
     return context.state;
   }
-  registerBuiltInNodes() {
-    this.node(BUILT_IN_NODES.START, () => {});
-    this.node(BUILT_IN_NODES.END, () => {});
-  }
-  addEdgeToRegistry(edge) {
-    const existingEdges = this.edgeRegistry.get(edge.from) ?? [];
-    existingEdges.push(edge);
-    this.edgeRegistry.set(edge.from, existingEdges);
-  }
   async createExecutionContext(runId, initialState, writer) {
     const { context, firstTime } = await this.restoreCheckpoint(runId, initialState, writer);
     return { context: { ...context, runId, writer }, firstTime };
@@ -255,7 +216,8 @@ class Graph {
     this.emitter.emitStart(context.writer, node.id);
     const subgraphRunId = this.generateSubgraphRunId(context.runId, node.id);
     try {
-      const childFinalState = await subgraph.executeInternal(subgraphRunId, options.input(context.state), context.writer);
+      const childRunner = new CompiledGraph(subgraph.nodes, subgraph.edges, subgraph.subgraphs, { storage: this.storage });
+      const childFinalState = await childRunner.executeInternal(subgraphRunId, options.input(context.state), context.writer);
       const parentUpdate = options.output(childFinalState, context.state);
       context.state = this.stateManager.apply(context.state, parentUpdate, context.writer);
       this.emitter.emitEnd(context.writer, node.id);
@@ -299,9 +261,6 @@ class Graph {
     return nodes.filter((node) => node.id !== BUILT_IN_NODES.END);
   }
 }
-function graph(options = {}) {
-  return new Graph(options);
-}
 
 class NodeEventEmitter {
   emitStart(writer, nodeId) {
@@ -339,6 +298,63 @@ class StateManager {
   }
 }
 
+// src/graph.ts
+class Graph {
+  nodeRegistry = new Map;
+  edgeRegistry = new Map;
+  subgraphRegistry = new Map;
+  constructor() {
+    this.registerBuiltInNodes();
+  }
+  node(id, execute) {
+    const node = { id, execute };
+    this.nodeRegistry.set(node.id, node);
+    return this;
+  }
+  edge(from, to) {
+    const edge = { from, to };
+    this.addEdgeToRegistry(edge);
+    return this;
+  }
+  graph(id, subgraph, options) {
+    this.subgraphRegistry.set(id, { subgraph, options });
+    const node = {
+      id,
+      execute: async () => {}
+    };
+    this.nodeRegistry.set(node.id, node);
+    return this;
+  }
+  get nodes() {
+    return this.nodeRegistry;
+  }
+  get edges() {
+    return this.edgeRegistry;
+  }
+  get subgraphs() {
+    return this.subgraphRegistry;
+  }
+  compile(options = {}) {
+    return new CompiledGraph(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry, options);
+  }
+  toMermaid(options) {
+    const generator = new MermaidGenerator(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry);
+    return generator.generate(options);
+  }
+  registerBuiltInNodes() {
+    this.node("START", () => {});
+    this.node("END", () => {});
+  }
+  addEdgeToRegistry(edge) {
+    const existingEdges = this.edgeRegistry.get(edge.from) ?? [];
+    existingEdges.push(edge);
+    this.edgeRegistry.set(edge.from, existingEdges);
+  }
+}
+function graph() {
+  return new Graph;
+}
+
 class MermaidGenerator {
   nodeRegistry;
   edgeRegistry;
@@ -456,5 +472,6 @@ export {
   RedisStorage,
   InMemoryStorage,
   Graph,
-  GRAPH_DATA_PART_TYPES
+  GRAPH_DATA_PART_TYPES,
+  CompiledGraph
 };

package/dist/types.d.ts

@@ -49,5 +49,15 @@ export declare namespace GraphSDK {
             writer: UIMessageStreamWriter;
         }) => Promise<void> | void;
     }
+    interface CompileOptions<State extends Record<string, unknown>, NodeKeys extends string> {
+        storage?: GraphStorage<State, NodeKeys>;
+        onFinish?: (args: {
+            state: State;
+        }) => Promise<void> | void;
+        onStart?: (args: {
+            state: State;
+            writer: UIMessageStreamWriter;
+        }) => Promise<void> | void;
+    }
     type Writer = Parameters<Parameters<typeof createUIMessageStream>[0]['execute']>[0]['writer'];
 }

package/package.json

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