graphai 0.0.4 → 0.0.5

package/README.md

@@ -1,6 +1,6 @@
 # GraphAI
 
-GraphAI is a TypeScript library, which makes it easy to create AI applications that need to make asynchronous AI API calls multiple times with some dependencies among them, such as giving the answer from one LLM call to another LLM call as a prompt. 
+GraphAI is an asynchronous data flow execution engine, which makes it easy to create AI applications that need to make asynchronous AI API calls multiple times with some dependencies among them, such as giving the answer from one LLM call to another LLM call as a prompt.
 
 You just need to describe dependencies among those API calls in a single graph definition file (in JSON or YAML), create a GraphAI object with it, and run it.
 
@@ -38,7 +38,6 @@ const nodeExecute = async (context: NodeExecuteContext) => {
   const graph = new GraphAI(graph_data, nodeExecute);
   const results = await graph.run();
   return results["taskC"];
-
 ```
 
 

package/lib/graphai.d.ts

@@ -7,7 +7,7 @@ export declare enum NodeState {
 }
 type ResultData<ResultType = Record<string, any>> = ResultType | undefined;
 type ResultDataDictonary<ResultType = Record<string, any>> = Record<string, ResultData<ResultType>>;
-export type NodeDataParams = Record<string, any>;
+export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType;
 type NodeData = {
     inputs: undefined | Array<string>;
     params: NodeDataParams;
@@ -19,14 +19,23 @@ type GraphData = {
     nodes: Record<string, NodeData>;
     concurrency: number;
 };
-export type NodeExecuteContext<ResultType> = {
+type NodeExecuteContext<ResultType, ParamsType> = {
     nodeId: string;
     retry: number;
-    params: NodeDataParams;
+    params: NodeDataParams<ParamsType>;
     payload: ResultDataDictonary<ResultType>;
 };
-type NodeExecute<ResultType> = (context: NodeExecuteContext<ResultType>) => Promise<ResultData<ResultType>>;
-declare class Node<ResultType = Record<string, any>> {
+export type TransactionLog = {
+    nodeId: string;
+    state: NodeState;
+    startTime: undefined | number;
+    endTime: undefined | number;
+    retryCount: number;
+    error: undefined | Error;
+    result?: ResultData;
+};
+export type NodeExecute<ResultType = Record<string, any>, ParamsType = Record<string, any>> = (context: NodeExecuteContext<ResultType, ParamsType>) => Promise<ResultData<ResultType>>;
+declare class Node {
     nodeId: string;
     params: NodeDataParams;
     inputs: Array<string>;
@@ -34,35 +43,41 @@ declare class Node<ResultType = Record<string, any>> {
     waitlist: Set<string>;
     state: NodeState;
     functionName: string;
-    result: ResultData<ResultType>;
+    result: ResultData;
     retryLimit: number;
     retryCount: number;
     transactionId: undefined | number;
     timeout: number;
+    error: undefined | Error;
     private graph;
-    constructor(nodeId: string, data: NodeData, graph: GraphAI<ResultType>);
+    constructor(nodeId: string, data: NodeData, graph: GraphAI);
     asString(): string;
     private retry;
     removePending(nodeId: string): void;
-    payload(): ResultDataDictonary<ResultType>;
+    payload(): ResultDataDictonary<Record<string, any>>;
     pushQueueIfReady(): void;
     execute(): Promise<void>;
 }
-type GraphNodes<ResultType> = Record<string, Node<ResultType>>;
-type NodeExecuteDictonary<ResultType> = Record<string, NodeExecute<ResultType>>;
-export declare class GraphAI<ResultType = Record<string, any>> {
-    nodes: GraphNodes<ResultType>;
-    callbackDictonary: NodeExecuteDictonary<ResultType>;
+type GraphNodes = Record<string, Node>;
+type NodeExecuteDictonary = Record<string, NodeExecute>;
+export declare class GraphAI {
+    nodes: GraphNodes;
+    callbackDictonary: NodeExecuteDictonary;
     private runningNodes;
     private nodeQueue;
     private onComplete;
     private concurrency;
-    constructor(data: GraphData, callbackDictonary: NodeExecuteDictonary<ResultType> | NodeExecute<ResultType>);
-    getCallback(functionName: string): NodeExecute<ResultType>;
+    private logs;
+    constructor(data: GraphData, callbackDictonary: NodeExecuteDictonary | NodeExecute);
+    getCallback(functionName: string): NodeExecute<Record<string, any>, Record<string, any>>;
     asString(): string;
+    results(): ResultDataDictonary<Record<string, any>>;
+    errors(): Record<string, Error>;
     run(): Promise<unknown>;
     private runNode;
-    pushQueue(node: Node<ResultType>): void;
-    removeRunning(node: Node<ResultType>): void;
+    pushQueue(node: Node): void;
+    removeRunning(node: Node): void;
+    appendLog(log: TransactionLog): void;
+    transactionLogs(): TransactionLog[];
 }
 export {};

package/lib/graphai.js

@@ -27,14 +27,16 @@ class Node {
     asString() {
         return `${this.nodeId}: ${this.state} ${[...this.waitlist]}`;
     }
-    retry(state, result) {
+    retry(state, error) {
         if (this.retryCount < this.retryLimit) {
             this.retryCount++;
             this.execute();
         }
         else {
             this.state = state;
-            this.result = result;
+            this.result = undefined;
+            this.error = error;
+            this.transactionId = 0; // This is necessary for timeout case
             this.graph.removeRunning(this);
         }
     }
@@ -54,14 +56,26 @@ class Node {
         }
     }
     async execute() {
+        const log = {
+            nodeId: this.nodeId,
+            retryCount: this.retryCount,
+            state: NodeState.Executing,
+            startTime: Date.now(),
+            endTime: undefined,
+            error: undefined,
+        };
+        this.graph.appendLog(log);
         this.state = NodeState.Executing;
-        const transactionId = Date.now();
+        const transactionId = log.startTime;
         this.transactionId = transactionId;
         if (this.timeout > 0) {
             setTimeout(() => {
                 if (this.state === NodeState.Executing && this.transactionId === transactionId) {
-                    console.log("*** timeout", this.timeout);
-                    this.retry(NodeState.TimedOut, undefined);
+                    console.log(`-- ${this.nodeId}: timeout ${this.timeout}`);
+                    log.error = Error("Timeout");
+                    log.state = NodeState.TimedOut;
+                    log.endTime = Date.now();
+                    this.retry(NodeState.TimedOut, Error("Timeout"));
                 }
             }, this.timeout);
         }
@@ -74,9 +88,12 @@ class Node {
                 payload: this.payload(),
             });
             if (this.transactionId !== transactionId) {
-                console.log("****** transactionId mismatch (success)");
+                console.log(`-- ${this.nodeId}: transactionId mismatch`);
                 return;
             }
+            log.state = NodeState.Completed;
+            log.endTime = Date.now();
+            log.result = result;
             this.state = NodeState.Completed;
             this.result = result;
             this.waitlist.forEach((nodeId) => {
@@ -85,22 +102,34 @@ class Node {
             });
             this.graph.removeRunning(this);
         }
-        catch (e) {
+        catch (error) {
             if (this.transactionId !== transactionId) {
-                console.log("****** transactionId mismatch (failed)");
+                console.log(`-- ${this.nodeId}: transactionId mismatch(error)`);
                 return;
             }
-            this.retry(NodeState.Failed, undefined);
+            log.state = NodeState.Failed;
+            log.endTime = Date.now();
+            if (error instanceof Error) {
+                log.error = error;
+                this.retry(NodeState.Failed, error);
+            }
+            else {
+                console.error(`-- ${this.nodeId}: Unexpecrted error was caught`);
+                log.error = Error("Unknown");
+                this.retry(NodeState.Failed, Error("Unknown"));
+            }
         }
     }
 }
+const defaultConcurrency = 8;
 class GraphAI {
     constructor(data, callbackDictonary) {
+        this.logs = [];
         this.callbackDictonary = typeof callbackDictonary === "function" ? { default: callbackDictonary } : callbackDictonary;
         if (this.callbackDictonary["default"] === undefined) {
             throw new Error("No default function");
         }
-        this.concurrency = data.concurrency ?? 2;
+        this.concurrency = data.concurrency ?? defaultConcurrency;
         this.runningNodes = new Set();
         this.nodeQueue = [];
         this.onComplete = () => { };
@@ -130,6 +159,24 @@ class GraphAI {
         })
             .join("\n");
     }
+    results() {
+        return Object.keys(this.nodes).reduce((results, nodeId) => {
+            const node = this.nodes[nodeId];
+            if (node.result !== undefined) {
+                results[nodeId] = node.result;
+            }
+            return results;
+        }, {});
+    }
+    errors() {
+        return Object.keys(this.nodes).reduce((errors, nodeId) => {
+            const node = this.nodes[nodeId];
+            if (node.error !== undefined) {
+                errors[nodeId] = node.error;
+            }
+            return errors;
+        }, {});
+    }
     async run() {
         // Nodes without pending data should run immediately.
         Object.keys(this.nodes).forEach((nodeId) => {
@@ -138,11 +185,14 @@ class GraphAI {
         });
         return new Promise((resolve, reject) => {
             this.onComplete = () => {
-                const results = Object.keys(this.nodes).reduce((results, nodeId) => {
-                    results[nodeId] = this.nodes[nodeId].result;
-                    return results;
-                }, {});
-                resolve(results);
+                const errors = this.errors();
+                const nodeIds = Object.keys(errors);
+                if (nodeIds.length > 0) {
+                    reject(errors[nodeIds[0]]);
+                }
+                else {
+                    resolve(this.results());
+                }
             };
         });
     }
@@ -170,5 +220,11 @@ class GraphAI {
             this.onComplete();
         }
     }
+    appendLog(log) {
+        this.logs.push(log);
+    }
+    transactionLogs() {
+        return this.logs;
+    }
 }
 exports.GraphAI = GraphAI;

package/package.json

@@ -1,13 +1,14 @@
 {
   "name": "graphai",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "description": "Asynchronous data flow execution engine to make it simple to build LLM apps.",
   "main": "lib/index.js",
   "scripts": {
     "build": "tsc",
     "eslint": "eslint --fix --ext src/**/*.{ts} tests/**/*.ts",
     "format": "prettier --write '{src,tests}/**/*.ts'",
-    "test": "node --test --require ts-node/register ./tests/test_*.ts"
+    "test": "node --test --require ts-node/register ./tests/test_*.ts",
+    "gpt": "npx ts-node ./samples/sample_gpt.ts"
   },
   "repository": {
     "type": "git",
@@ -27,7 +28,7 @@
     "eslint-plugin-import": "^2.25.2",
     "openai": "^4.12.4",
     "prettier": "^3.0.3",
-    "slashgpt": "^0.0.6",
+    "slashgpt": "^0.0.8",
     "ts-node": "^10.9.1",
     "typescript": "^5.2.2"
   },

package/{tests → samples}/sample_gpt.ts

@@ -1,12 +1,12 @@
 import path from "path";
-import { GraphAI, NodeExecuteContext } from "../src/graphai";
+import { GraphAI, NodeExecute } from "../src/graphai";
 import { ChatSession, ChatConfig } from "slashgpt";
-import { readManifestData } from "./file_utils";
+import { readManifestData } from "../tests/file_utils";
 
 const config = new ChatConfig(path.resolve(__dirname));
 
-const testFunction = async (context: NodeExecuteContext<Record<string, string>>) => {
-  console.log("executing", context.nodeId, context.params, context.payload);
+const testFunction: NodeExecute<Record<string, string>> = async (context) => {
+  console.log("executing", context.nodeId, context.params);
   const session = new ChatSession(config, context.params.manifest ?? {});
   const prompt = Object.keys(context.payload).reduce((prompt, key) => {
     return prompt.replace("${" + key + "}", context.payload[key]!["answer"]);
@@ -31,7 +31,7 @@ const test = async (file: string) => {
 };
 
 const main = async () => {
-  await test("/graphs/sample3.yml");
+  await test("/graphs/slash_gpt.yml");
   console.log("COMPLETE 1");
 };
 main();

package/src/graphai.ts

@@ -10,7 +10,7 @@ export enum NodeState {
 type ResultData<ResultType = Record<string, any>> = ResultType | undefined;
 type ResultDataDictonary<ResultType = Record<string, any>> = Record<string, ResultData<ResultType>>;
 
-export type NodeDataParams = Record<string, any>; // App-specific parameters
+export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType; // App-specific parameters
 
 type NodeData = {
   inputs: undefined | Array<string>;
@@ -25,16 +25,28 @@ type GraphData = {
   concurrency: number;
 };
 
-export type NodeExecuteContext<ResultType> = {
+type NodeExecuteContext<ResultType, ParamsType> = {
   nodeId: string;
   retry: number;
-  params: NodeDataParams;
+  params: NodeDataParams<ParamsType>;
   payload: ResultDataDictonary<ResultType>;
 };
 
-type NodeExecute<ResultType> = (context: NodeExecuteContext<ResultType>) => Promise<ResultData<ResultType>>;
+export type TransactionLog = {
+  nodeId: string;
+  state: NodeState;
+  startTime: undefined | number;
+  endTime: undefined | number;
+  retryCount: number;
+  error: undefined | Error;
+  result?: ResultData;
+};
+
+export type NodeExecute<ResultType = Record<string, any>, ParamsType = Record<string, any>> = (
+  context: NodeExecuteContext<ResultType, ParamsType>,
+) => Promise<ResultData<ResultType>>;
 
-class Node<ResultType = Record<string, any>> {
+class Node {
   public nodeId: string;
   public params: NodeDataParams; // App-specific parameters
   public inputs: Array<string>; // List of nodes this node needs data from.
@@ -42,15 +54,16 @@ class Node<ResultType = Record<string, any>> {
   public waitlist: Set<string>; // List of nodes which need data from this node.
   public state: NodeState;
   public functionName: string;
-  public result: ResultData<ResultType>;
+  public result: ResultData;
   public retryLimit: number;
   public retryCount: number;
   public transactionId: undefined | number; // To reject callbacks from timed-out transactions
   public timeout: number; // msec
+  public error: undefined | Error;
 
-  private graph: GraphAI<ResultType>;
+  private graph: GraphAI;
 
-  constructor(nodeId: string, data: NodeData, graph: GraphAI<ResultType>) {
+  constructor(nodeId: string, data: NodeData, graph: GraphAI) {
     this.nodeId = nodeId;
     this.inputs = data.inputs ?? [];
     this.pendings = new Set(this.inputs);
@@ -70,13 +83,15 @@ class Node<ResultType = Record<string, any>> {
     return `${this.nodeId}: ${this.state} ${[...this.waitlist]}`;
   }
 
-  private retry(state: NodeState, result: ResultData<ResultType>) {
+  private retry(state: NodeState, error: Error) {
     if (this.retryCount < this.retryLimit) {
       this.retryCount++;
       this.execute();
     } else {
       this.state = state;
-      this.result = result;
+      this.result = undefined;
+      this.error = error;
+      this.transactionId = 0; // This is necessary for timeout case
       this.graph.removeRunning(this);
     }
   }
@@ -87,7 +102,7 @@ class Node<ResultType = Record<string, any>> {
   }
 
   public payload() {
-    return this.inputs.reduce((results: ResultDataDictonary<ResultType>, nodeId) => {
+    return this.inputs.reduce((results: ResultDataDictonary, nodeId) => {
       results[nodeId] = this.graph.nodes[nodeId].result;
       return results;
     }, {});
@@ -100,15 +115,27 @@ class Node<ResultType = Record<string, any>> {
   }
 
   public async execute() {
+    const log: TransactionLog = {
+      nodeId: this.nodeId,
+      retryCount: this.retryCount,
+      state: NodeState.Executing,
+      startTime: Date.now(),
+      endTime: undefined,
+      error: undefined,
+    };
+    this.graph.appendLog(log);
     this.state = NodeState.Executing;
-    const transactionId = Date.now();
+    const transactionId = log.startTime;
     this.transactionId = transactionId;
 
     if (this.timeout > 0) {
       setTimeout(() => {
         if (this.state === NodeState.Executing && this.transactionId === transactionId) {
-          console.log("*** timeout", this.timeout);
-          this.retry(NodeState.TimedOut, undefined);
+          console.log(`-- ${this.nodeId}: timeout ${this.timeout}`);
+          log.error = Error("Timeout");
+          log.state = NodeState.TimedOut;
+          log.endTime = Date.now();
+          this.retry(NodeState.TimedOut, Error("Timeout"));
         }
       }, this.timeout);
     }
@@ -122,9 +149,12 @@ class Node<ResultType = Record<string, any>> {
         payload: this.payload(),
       });
       if (this.transactionId !== transactionId) {
-        console.log("****** transactionId mismatch (success)");
+        console.log(`-- ${this.nodeId}: transactionId mismatch`);
         return;
       }
+      log.state = NodeState.Completed;
+      log.endTime = Date.now();
+      log.result = result;
       this.state = NodeState.Completed;
       this.result = result;
       this.waitlist.forEach((nodeId) => {
@@ -132,39 +162,51 @@ class Node<ResultType = Record<string, any>> {
         node.removePending(this.nodeId);
       });
       this.graph.removeRunning(this);
-    } catch (e) {
+    } catch (error) {
       if (this.transactionId !== transactionId) {
-        console.log("****** transactionId mismatch (failed)");
+        console.log(`-- ${this.nodeId}: transactionId mismatch(error)`);
         return;
       }
-      this.retry(NodeState.Failed, undefined);
+      log.state = NodeState.Failed;
+      log.endTime = Date.now();
+      if (error instanceof Error) {
+        log.error = error;
+        this.retry(NodeState.Failed, error);
+      } else {
+        console.error(`-- ${this.nodeId}: Unexpecrted error was caught`);
+        log.error = Error("Unknown");
+        this.retry(NodeState.Failed, Error("Unknown"));
+      }
     }
   }
 }
 
-type GraphNodes<ResultType> = Record<string, Node<ResultType>>;
+type GraphNodes = Record<string, Node>;
 
-type NodeExecuteDictonary<ResultType> = Record<string, NodeExecute<ResultType>>;
+type NodeExecuteDictonary = Record<string, NodeExecute>;
 
-export class GraphAI<ResultType = Record<string, any>> {
-  public nodes: GraphNodes<ResultType>;
-  public callbackDictonary: NodeExecuteDictonary<ResultType>;
+const defaultConcurrency = 8;
+
+export class GraphAI {
+  public nodes: GraphNodes;
+  public callbackDictonary: NodeExecuteDictonary;
   private runningNodes: Set<string>;
-  private nodeQueue: Array<Node<ResultType>>;
+  private nodeQueue: Array<Node>;
   private onComplete: () => void;
   private concurrency: number;
+  private logs: Array<TransactionLog> = [];
 
-  constructor(data: GraphData, callbackDictonary: NodeExecuteDictonary<ResultType> | NodeExecute<ResultType>) {
+  constructor(data: GraphData, callbackDictonary: NodeExecuteDictonary | NodeExecute) {
     this.callbackDictonary = typeof callbackDictonary === "function" ? { default: callbackDictonary } : callbackDictonary;
     if (this.callbackDictonary["default"] === undefined) {
       throw new Error("No default function");
     }
-    this.concurrency = data.concurrency ?? 2;
+    this.concurrency = data.concurrency ?? defaultConcurrency;
     this.runningNodes = new Set<string>();
     this.nodeQueue = [];
     this.onComplete = () => {};
-    this.nodes = Object.keys(data.nodes).reduce((nodes: GraphNodes<ResultType>, nodeId: string) => {
-      nodes[nodeId] = new Node<ResultType>(nodeId, data.nodes[nodeId], this);
+    this.nodes = Object.keys(data.nodes).reduce((nodes: GraphNodes, nodeId: string) => {
+      nodes[nodeId] = new Node(nodeId, data.nodes[nodeId], this);
       return nodes;
     }, {});
 
@@ -193,6 +235,26 @@ export class GraphAI<ResultType = Record<string, any>> {
       .join("\n");
   }
 
+  public results() {
+    return Object.keys(this.nodes).reduce((results: ResultDataDictonary, nodeId) => {
+      const node = this.nodes[nodeId];
+      if (node.result !== undefined) {
+        results[nodeId] = node.result;
+      }
+      return results;
+    }, {});
+  }
+
+  public errors() {
+    return Object.keys(this.nodes).reduce((errors: Record<string, Error>, nodeId) => {
+      const node = this.nodes[nodeId];
+      if (node.error !== undefined) {
+        errors[nodeId] = node.error;
+      }
+      return errors;
+    }, {});
+  }
+
   public async run() {
     // Nodes without pending data should run immediately.
     Object.keys(this.nodes).forEach((nodeId) => {
@@ -202,21 +264,23 @@ export class GraphAI<ResultType = Record<string, any>> {
 
     return new Promise((resolve, reject) => {
       this.onComplete = () => {
-        const results = Object.keys(this.nodes).reduce((results: ResultDataDictonary<ResultType>, nodeId) => {
-          results[nodeId] = this.nodes[nodeId].result;
-          return results;
-        }, {});
-        resolve(results);
+        const errors = this.errors();
+        const nodeIds = Object.keys(errors);
+        if (nodeIds.length > 0) {
+          reject(errors[nodeIds[0]]);
+        } else {
+          resolve(this.results());
+        }
       };
     });
   }
 
-  private runNode(node: Node<ResultType>) {
+  private runNode(node: Node) {
     this.runningNodes.add(node.nodeId);
     node.execute();
   }
 
-  public pushQueue(node: Node<ResultType>) {
+  public pushQueue(node: Node) {
     if (this.runningNodes.size < this.concurrency) {
       this.runNode(node);
     } else {
@@ -224,7 +288,7 @@ export class GraphAI<ResultType = Record<string, any>> {
     }
   }
 
-  public removeRunning(node: Node<ResultType>) {
+  public removeRunning(node: Node) {
     this.runningNodes.delete(node.nodeId);
     if (this.nodeQueue.length > 0) {
       const n = this.nodeQueue.shift();
@@ -236,4 +300,12 @@ export class GraphAI<ResultType = Record<string, any>> {
       this.onComplete();
     }
   }
+
+  public appendLog(log: TransactionLog) {
+    this.logs.push(log);
+  }
+
+  public transactionLogs() {
+    return this.logs;
+  }
 }

package/tests/graphs/test_error.yml

@@ -0,0 +1,22 @@
+nodes:
+  node1:
+    params:
+      delay: 500
+  node2:
+    params:
+      delay: 100
+  node3:
+    params:
+      delay: 500
+      fail: true
+    inputs: [node1, node2]
+  node4:
+    timeout: 200
+    retry: 2
+    params:
+      delay: 300
+    inputs: [node3]
+  node5:
+    params:
+      delay: 100
+    inputs: [node4]

package/tests/graphs/test_multiple_functions_1.yml

@@ -19,3 +19,9 @@ nodes:
       delay: 500
     inputs: [node2, node4]
     functionName: test2
+  node6:
+    params:
+      delay: 100
+      number: 10
+    inputs: [node4]
+    functionName: numberTestFunction

package/tests/graphs/test_timeout.yml

@@ -0,0 +1,22 @@
+nodes:
+  node1:
+    params:
+      delay: 500
+  node2:
+    params:
+      delay: 100
+  node3:
+    params:
+      delay: 500
+      fail: true
+    inputs: [node1, node2]
+    retry: 2
+  node4:
+    timeout: 200
+    params:
+      delay: 300
+    inputs: [node3]
+  node5:
+    params:
+      delay: 100
+    inputs: [node4]

package/tests/test_multiple_functions.ts

@@ -1,12 +1,12 @@
 import path from "path";
-import { GraphAI, NodeExecuteContext } from "../src/graphai";
-import { readManifestData } from "../src/file_utils";
+import { GraphAI, NodeExecute } from "../src/graphai";
+import { readManifestData } from "./file_utils";
 import { sleep } from "./utils";
 
 import test from "node:test";
 import assert from "node:assert";
 
-const testFunction1 = async (context: NodeExecuteContext<Record<string, string>>) => {
+const testFunction1: NodeExecute<Record<string, string>> = async (context) => {
   const { nodeId, retry, params } = context;
   console.log("executing", nodeId, params);
 
@@ -15,7 +15,7 @@ const testFunction1 = async (context: NodeExecuteContext<Record<string, string>>
   return result;
 };
 
-const testFunction2 = async (context: NodeExecuteContext<Record<string, string>>) => {
+const testFunction2: NodeExecute<Record<string, string>> = async (context) => {
   const { nodeId, retry, params } = context;
   console.log("executing", nodeId, params);
 
@@ -24,11 +24,20 @@ const testFunction2 = async (context: NodeExecuteContext<Record<string, string>>
   return result;
 };
 
+const numberTestFunction: NodeExecute<Record<string, number>, Record<"number", number>> = async (context) => {
+  const { nodeId, retry, params } = context;
+  console.log("executing", nodeId, params);
+
+  const result = { [nodeId]: params.number };
+  console.log("completing", nodeId, result);
+  return result;
+};
+
 const runTest = async (file: string) => {
   const file_path = path.resolve(__dirname) + file;
   const graph_data = readManifestData(file_path);
 
-  const graph = new GraphAI(graph_data, { default: testFunction1, test2: testFunction2 });
+  const graph = new GraphAI(graph_data, { default: testFunction1, test2: testFunction2, numberTestFunction });
 
   const results = await graph.run();
   console.log(results);
@@ -43,5 +52,6 @@ test("test sample1", async () => {
     node3: { node3: "output 2" },
     node4: { node4: "output 1" },
     node5: { node5: "output 2" },
+    node6: { node6: 10 },
   });
 });

package/tests/test_sample_flow.ts

@@ -1,19 +1,19 @@
 import path from "path";
-import { GraphAI, NodeExecuteContext } from "../src/graphai";
+import { GraphAI, NodeExecute } from "../src/graphai";
 import { readManifestData } from "./file_utils";
 import { sleep } from "./utils";
 
 import test from "node:test";
 import assert from "node:assert";
 
-const testFunction = async (context: NodeExecuteContext<Record<string, string>>) => {
+const testFunction: NodeExecute<Record<string, string>> = async (context) => {
   const { nodeId, retry, params, payload } = context;
-  console.log("executing", nodeId, params);
+  console.log("executing", nodeId);
   await sleep(params.delay / (retry + 1));
 
   if (params.fail && retry < 2) {
     const result = { [nodeId]: "failed" };
-    console.log("failed", nodeId, result, retry);
+    console.log("failed (intentional)", nodeId, retry);
     throw new Error("Intentional Failure");
   } else {
     const result = Object.keys(payload).reduce(
@@ -23,7 +23,7 @@ const testFunction = async (context: NodeExecuteContext<Record<string, string>>)
       },
       { [nodeId]: "output" },
     );
-    console.log("completing", nodeId, result);
+    console.log("completing", nodeId);
     return result;
   }
 };
@@ -34,13 +34,21 @@ const runTest = async (file: string) => {
 
   const graph = new GraphAI(graph_data, testFunction);
 
-  const results = await graph.run();
-  console.log(results);
-  return results;
+  try {
+    const results = await graph.run();
+    console.log(graph.transactionLogs());
+    return results;
+  } catch (error) {
+    if (error instanceof Error) {
+      console.log("Error:", error.message);
+    }
+    console.log(graph.transactionLogs());
+    return graph.results();
+  }
 };
 
-test("test sample1", async () => {
-  const result = await runTest("/graphs/sample1.yml");
+test("test base", async () => {
+  const result = await runTest("/graphs/test_base.yml");
   assert.deepStrictEqual(result, {
     node1: { node1: "output" },
     node2: { node2: "output" },
@@ -50,8 +58,8 @@ test("test sample1", async () => {
   });
 });
 
-test("test sample1", async () => {
-  const result = await runTest("/graphs/sample2.yml");
+test("test retry", async () => {
+  const result = await runTest("/graphs/test_retry.yml");
   assert.deepStrictEqual(result, {
     node1: { node1: "output" },
     node2: { node2: "output" },
@@ -59,5 +67,21 @@ test("test sample1", async () => {
     node4: { node4: "output", node3: "output", node1: "output", node2: "output" },
     node5: { node5: "output", node4: "output", node3: "output", node1: "output", node2: "output" },
   });
-  console.log("COMPLETE 2");
+});
+
+test("test error", async () => {
+  const result = await runTest("/graphs/test_error.yml");
+  assert.deepStrictEqual(result, {
+    node1: { node1: "output" },
+    node2: { node2: "output" },
+  });
+});
+
+test("test timeout", async () => {
+  const result = await runTest("/graphs/test_timeout.yml");
+  assert.deepStrictEqual(result, {
+    node1: { node1: "output" },
+    node2: { node2: "output" },
+    node3: { node3: "output", node1: "output", node2: "output" },
+  });
 });