graphai 0.0.8 → 0.0.9

package/README.md

@@ -42,11 +42,21 @@ const sampleAgentFunction = async (context: AgentFunctionContext) => {
   return results["taskC"];
 ```
 
+## Background
+
+As Andrew Ng has described in his article, "[The batch: Issue 242](https://www.deeplearning.ai/the-batch/issue-242/)", better results can often be achieved by making multiple calls to a Large Language Model (LLM) and allowing it to incrementally build towards a higher-quality output. Dr. Ng refers to this approach as 'agentic workflows.' 
+
+Building applications that employ these workflows, however, is challenging due to the complexities of managing multiple asynchronous API calls, including error handling, timeouts, retries, and logging. 
+
+GraphAI is designed to simplify this process by decoupling the complexity of multiple asynchronous calls from the application's core logic. It enables developers to model these calls and their dependencies within a single Data Flow Graph, enhancing development and debugging processes. 
+
+Furthermore, GraphAI's robust mechanisms for error handling, retry strategies, timeouts, and logging empower developers to concentrate on refining the application logic.
+
 ## Data Flow Graph
 
 A Data Flow Graph (DFG) is a JavaScript object, which defines the flow of data. It is typically described in YAML file and loaded at runtime.
 
-A DFG consists of a collection of 'nodes', which contains a series of nested keys representing individual nodes in the data flow. Each node is identified by a unique key (e.g., node1, node2) and can contain several predefined keys (params, inputs, outputs, retry, timeout, source, agentId) that dictate the node's behavior and its relationship with other nodes.
+A DFG consists of a collection of 'nodes', which contains a series of nested keys representing individual nodes in the data flow. Each node is identified by a unique key (e.g., node1, node2) and can contain several predefined keys (params, inputs, outputs, retry, timeout, source, agentId, result, fork) that dictate the node's behavior and its relationship with other nodes.
 
 ### DFG Structure
 
@@ -67,10 +77,13 @@ An agent function receives two set of parameters via AgentFunctionContext, agent
 
 - 'inputs': An optional list of node identifiers that the current node depends on. This establishes a flow where the current node can only be executed after the completion of the nodes listed under 'inputs'. If this list is empty and the 'source' property is not set, the associated Agent Function will be immediatley executed. 
 - 'source': An optional flag, which specifies if the node is a 'source node' or not. A souce node is a special node, which receives data from either an external code (via GraphAI's injectResult method) or from a 'dispatcher node'.
-- 'outputs': An optinal property, which specifies the mapping from outputId to nodeId. If this property is set, the node become a special node called "dispatcher". A dispatcher node injects result(s) into specified nodes, enabling the dynamic flow of data.
+- 'result': An optional object, which injects the result into the source node (equivalent to calling the injectResult method).
 - 'retry': An optional number, which specifies the maximum number of retries to be made. If the last attempt fails, that return value will be recorded.
 - 'timeout': An optional number, which specifies the maximum waittime in msec. If the associated agent function does not return the value in time, the "Timeout" error will be recorded and the returned value will be discarded. 
 - 'params': An optional parameters to the associated agent function, which are agent specific.
+- 'agentId': An optional parameter, which specifies the id of the agent, when the graph is associated with multiple agents.
+- 'fork': An optional paramter, which specifies the number of concurrent transactions to be created for the current node.
+- 'outputs': An optinal property, which specifies the mapping from outputId to nodeId. If this property is set, the node become a special node called "dispatcher". A dispatcher node injects result(s) into specified nodes, enabling the dynamic flow of data.
 
 ## GraphAI class
 

package/lib/graphai.d.ts

@@ -12,11 +12,13 @@ type ResultDataDictonary<ResultType = Record<string, any>> = Record<string, Resu
 export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType;
 type NodeData = {
     inputs?: Array<string>;
-    params: NodeDataParams;
+    params?: NodeDataParams;
     retry?: number;
     timeout?: number;
     agentId?: string;
+    fork?: number;
     source?: boolean;
+    result?: ResultData;
     outputs?: Record<string, string>;
 };
 export type GraphData = {
@@ -37,6 +39,7 @@ export type TransactionLog = {
 };
 export type AgentFunctionContext<ParamsType, ResultType, PreviousResultType> = {
     nodeId: string;
+    forkIndex?: number;
     retry: number;
     params: NodeDataParams<ParamsType>;
     inputs: Array<PreviousResultType>;
@@ -51,6 +54,8 @@ declare class Node {
     waitlist: Set<string>;
     state: NodeState;
     agentId?: string;
+    fork?: number;
+    forkIndex?: number;
     result: ResultData;
     retryLimit: number;
     retryCount: number;
@@ -60,7 +65,7 @@ declare class Node {
     source: boolean;
     outputs?: Record<string, string>;
     private graph;
-    constructor(nodeId: string, data: NodeData, graph: GraphAI);
+    constructor(nodeId: string, forkIndex: number | undefined, data: NodeData, graph: GraphAI);
     asString(): string;
     private retry;
     removePending(nodeId: string): void;

package/lib/graphai.js

@@ -12,16 +12,18 @@ var NodeState;
     NodeState["Dispatched"] = "dispatched";
 })(NodeState || (exports.NodeState = NodeState = {}));
 class Node {
-    constructor(nodeId, data, graph) {
+    constructor(nodeId, forkIndex, data, graph) {
         this.waitlist = new Set(); // List of nodes which need data from this node.
         this.state = NodeState.Waiting;
         this.result = undefined;
         this.retryCount = 0;
         this.nodeId = nodeId;
+        this.forkIndex = forkIndex;
         this.inputs = data.inputs ?? [];
         this.pendings = new Set(this.inputs);
-        this.params = data.params;
+        this.params = data.params ?? {};
         this.agentId = data.agentId;
+        this.fork = data.fork;
         this.retryLimit = data.retry ?? 0;
         this.timeout = data.timeout;
         this.source = data.source === true;
@@ -62,9 +64,9 @@ class Node {
                 retryCount: this.retryCount,
                 state: NodeState.Injected,
                 startTime: Date.now(),
+                endTime: Date.now(),
                 result,
             };
-            log.endTime = log.startTime;
             this.graph.appendLog(log);
             this.setResult(result, NodeState.Injected);
         }
@@ -82,21 +84,19 @@ class Node {
         });
     }
     async execute() {
+        const results = this.graph.resultsOf(this.inputs);
+        const transactionId = Date.now();
         const log = {
             nodeId: this.nodeId,
             retryCount: this.retryCount,
             state: NodeState.Executing,
-            startTime: Date.now(),
+            startTime: transactionId,
             agentId: this.agentId,
             params: this.params,
+            inputs: results,
         };
-        const results = this.graph.resultsOf(this.inputs);
-        if (results.length > 0) {
-            log.inputs = results;
-        }
         this.graph.appendLog(log);
         this.state = NodeState.Executing;
-        const transactionId = log.startTime;
         this.transactionId = transactionId;
         if (this.timeout && this.timeout > 0) {
             setTimeout(() => {
@@ -116,6 +116,7 @@ class Node {
                 retry: this.retryCount,
                 params: this.params,
                 inputs: results,
+                forkIndex: this.forkIndex,
             });
             if (this.transactionId !== transactionId) {
                 console.log(`-- ${this.nodeId}: transactionId mismatch`);
@@ -169,22 +170,65 @@ class GraphAI {
         this.onComplete = () => {
             console.error("-- SOMETHING IS WRONG: onComplete is called without run()");
         };
+        const nodeId2forkedNodeIds = {};
+        const forkedNodeId2Index = {};
+        // Create node instances from data.nodes
         this.nodes = Object.keys(data.nodes).reduce((nodes, nodeId) => {
-            nodes[nodeId] = new Node(nodeId, data.nodes[nodeId], this);
+            const fork = data.nodes[nodeId].fork;
+            if (fork) {
+                // For fork, change the nodeId and increase the node
+                nodeId2forkedNodeIds[nodeId] = new Array(fork).fill(undefined).map((_, i) => {
+                    const forkedNodeId = `${nodeId}_${i}`;
+                    nodes[forkedNodeId] = new Node(forkedNodeId, i, data.nodes[nodeId], this);
+                    // Data for pending and waiting
+                    forkedNodeId2Index[forkedNodeId] = i;
+                    return forkedNodeId;
+                });
+            }
+            else {
+                nodes[nodeId] = new Node(nodeId, undefined, data.nodes[nodeId], this);
+            }
             return nodes;
         }, {});
-        // Generate the waitlist for each node
+        // Generate the waitlist for each node, and update the pendings in case of forked node.
         Object.keys(this.nodes).forEach((nodeId) => {
             const node = this.nodes[nodeId];
             node.pendings.forEach((pending) => {
-                const node2 = this.nodes[pending];
-                node2.waitlist.add(nodeId);
+                // If the pending(previous) node is forking
+                if (nodeId2forkedNodeIds[pending]) {
+                    //  update node.pending and pending(previous) node.wailtlist
+                    if (node.fork) {
+                        //  1:1 if current nodes are also forking.
+                        const newPendingId = nodeId2forkedNodeIds[pending][forkedNodeId2Index[nodeId]];
+                        this.nodes[newPendingId].waitlist.add(nodeId); // previousNode
+                        node.pendings.add(newPendingId);
+                    }
+                    else {
+                        //  1:n if current node is not forking.
+                        nodeId2forkedNodeIds[pending].forEach((newPendingId) => {
+                            this.nodes[newPendingId].waitlist.add(nodeId); // previousNode
+                            node.pendings.add(newPendingId);
+                        });
+                    }
+                    node.pendings.delete(pending);
+                }
+                else {
+                    this.nodes[pending].waitlist.add(nodeId); // previousNode
+                }
             });
+            node.inputs = Array.from(node.pendings); // for fork.
+        });
+        // If the result property is specified, inject it.
+        // NOTE: This must be done at the end of this constructor
+        Object.keys(data.nodes).forEach((nodeId) => {
+            const result = data.nodes[nodeId].result;
+            if (result) {
+                this.injectResult(nodeId, result);
+            }
         });
     }
     getCallback(_agentId) {
         const agentId = _agentId ?? "_default";
-        console.log(agentId);
         if (this.callbackDictonary[agentId]) {
             return this.callbackDictonary[agentId];
         }

package/lib/graphai_cli.d.ts

@@ -0,0 +1,2 @@
+#!/usr/bin/env node
+export {};

package/lib/graphai_cli.js

@@ -0,0 +1,36 @@
+#!/usr/bin/env node
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const graphai_1 = require("./graphai");
+const fs_1 = __importDefault(require("fs"));
+const path_1 = __importDefault(require("path"));
+const yaml_1 = __importDefault(require("yaml"));
+const testAgent = async (context) => {
+    return {};
+};
+const main = async () => {
+    const file = process.argv[2];
+    if (file === undefined) {
+        console.log("no file");
+        return;
+    }
+    const file_path = path_1.default.resolve(process.cwd() + "/" + file);
+    if (!fs_1.default.existsSync(file_path)) {
+        console.log("no file");
+        return;
+    }
+    try {
+        const graph_data_file = fs_1.default.readFileSync(file_path, "utf8");
+        const graph_data = yaml_1.default.parse(graph_data_file);
+        const graph = new graphai_1.GraphAI(graph_data, { test: testAgent });
+        const results = await graph.run();
+        console.log(results);
+    }
+    catch (e) {
+        console.log("error", e);
+    }
+};
+main();

package/package.json

@@ -1,14 +1,16 @@
 {
   "name": "graphai",
-  "version": "0.0.8",
+  "version": "0.0.9",
   "description": "Asynchronous data flow execution engine to make it simple to build LLM apps.",
   "main": "lib/index.js",
+  "bin": "lib/graphai_cli.js",
   "scripts": {
     "build": "tsc && tsc-alias",
     "eslint": "eslint --fix --ext src/**/*.{ts} tests/**/*.ts",
     "format": "prettier --write '{src,tests,samples}/**/*.ts' .eslintrc.js",
     "test": "node --test  -r tsconfig-paths/register --require ts-node/register ./tests/**/test_*.ts",
-    "gpt": "npx ts-node  -r tsconfig-paths/register ./samples/sample_gpt.ts"
+    "cli": "npx ts-node  -r tsconfig-paths/register ./src/graphai_cli.ts",
+    "sample": "npx ts-node  -r tsconfig-paths/register"
   },
   "repository": {
     "type": "git",

package/samples/agents/slashgpt_agent.ts

@@ -14,7 +14,8 @@ export const slashGPTFuncitons2TextAgent: AgentFunction<
     const { title, description } = r;
     return ["title:", title, "description:", description].join("\n");
   });
-  return result[params.result_key];
+
+  return result[context.forkIndex ?? 0];
 };
 
 export const slashGPTAgent: AgentFunction<

package/samples/interaction.ts

@@ -0,0 +1,49 @@
+import { GraphAI, GraphData } from "@/graphai";
+import * as readline from "readline";
+import path from "path";
+import * as fs from "fs";
+import { testAgent } from "~/agents/agents";
+
+const getUserInput = async (question: string): Promise<string> => {
+  return new Promise((resolve, reject) => {
+    const rl = readline.createInterface({
+      input: process.stdin,
+      output: process.stdout,
+    });
+
+    rl.question(question, (answer) => {
+      rl.close();
+      resolve(answer);
+    });
+  });
+};
+
+const graph_data: GraphData = {
+  nodes: {
+    node1: {
+      source: true,
+    },
+    node2: {
+      inputs: ["node1"],
+    },
+  },
+};
+
+const runAgent = async (query: string) => {
+  console.log("query=", query);
+  graph_data.nodes.node1.result = { query };
+  const graph = new GraphAI(graph_data, testAgent);
+  // graph.injectResult("node1", { query });
+  const result = await graph.run();
+  const log_path = path.resolve(__dirname) + "/../tests/logs/interaction.log";
+  fs.writeFileSync(log_path, JSON.stringify(graph.transactionLogs(), null, 2));
+  console.log(result);
+};
+
+const main = async () => {
+  const query = await getUserInput("Please enter your question: ");
+  await runAgent(query);
+  console.log("COMPLETE 1");
+};
+
+main();

package/samples/sample_co2.ts

@@ -51,6 +51,7 @@ const graph_data = {
       },
     },
     function2prompt0: {
+      fork: 10,
       params: {
         function_data_key: "methods",
         result_key: 0,
@@ -60,6 +61,7 @@ const graph_data = {
     },
     slashGPTAgent0: {
       agentId: "slashGPTAgent",
+      fork: 10,
       params: {
         debug: true,
         manifest: {
@@ -68,26 +70,6 @@ const graph_data = {
       },
       inputs: ["function2prompt0"],
     },
-    // 1
-    function2prompt1: {
-      params: {
-        debug: true,
-        function_data_key: "methods",
-        result_key: 1,
-      },
-      inputs: ["slashGPTAgent"],
-      agentId: "slashGPTFuncitons2TextAgent",
-    },
-    slashGPTAgent1: {
-      agentId: "slashGPTAgent",
-      params: {
-        debug: true,
-        manifest: {
-          prompt: "ユーザの問い合わせにある文章の専門家です。専門家として、ユーザのアイデアに対して実現可能なシナリオを800文字で書いてください。",
-        },
-      },
-      inputs: ["function2prompt1"],
-    },
   },
 };
 const runAgent = async () => {

package/samples/sample_gpt.ts

@@ -1,15 +1,16 @@
 import path from "path";
+import * as fs from "fs";
 import { GraphAI, AgentFunction } from "@/graphai";
 import { ChatSession, ChatConfig, ManifestData } from "slashgpt";
 import { readGraphaiData } from "~/utils/file_utils";
 
 const config = new ChatConfig(path.resolve(__dirname));
 
-const slashGPTAgent: AgentFunction<{ manifest: ManifestData; prompt: string }, { answer: string }> = async (context) => {
+const slashGPTAgent: AgentFunction<{ manifest: ManifestData; prompt: string }, { content: string }> = async (context) => {
   console.log("executing", context.nodeId, context.params);
   const session = new ChatSession(config, context.params.manifest ?? {});
   const prompt = context.inputs.reduce((prompt, input, index) => {
-    return prompt.replace("${" + index + "}", input["answer"]);
+    return prompt.replace("${" + index + "}", input["content"]);
   }, context.params.prompt);
   session.append_user_question(prompt);
 
@@ -18,20 +19,22 @@ const slashGPTAgent: AgentFunction<{ manifest: ManifestData; prompt: string }, {
   if (message === undefined) {
     throw new Error("No message in the history");
   }
-  const result = { answer: message.content };
-  return result;
+  return message;
 };
 
 const runAgent = async (file: string) => {
   const file_path = path.resolve(__dirname) + file;
   const graph_data = readGraphaiData(file_path);
   const graph = new GraphAI(graph_data, slashGPTAgent);
-  const result = await graph.run();
-  console.log(result);
+  const results = (await graph.run()) as Record<string, any>;
+
+  const log_path = path.resolve(__dirname) + "/../tests/logs/" + path.basename(file_path).replace(/\.yml$/, ".log");
+  console.log(log_path);
+  fs.writeFileSync(log_path, JSON.stringify(graph.transactionLogs(), null, 2));
+  console.log(results["node3"]["content"]);
 };
 
 const main = async () => {
   await runAgent("/graphs/slash_gpt.yml");
-  console.log("COMPLETE 1");
 };
 main();

package/src/graphai.ts

@@ -14,11 +14,13 @@ export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType; // Ag
 
 type NodeData = {
   inputs?: Array<string>;
-  params: NodeDataParams;
+  params?: NodeDataParams;
   retry?: number;
   timeout?: number; // msec
   agentId?: string;
+  fork?: number;
   source?: boolean;
+  result?: ResultData; // preset result for source node.
   outputs?: Record<string, string>; // mapping from routeId to nodeId
 };
 
@@ -42,6 +44,7 @@ export type TransactionLog = {
 
 export type AgentFunctionContext<ParamsType, ResultType, PreviousResultType> = {
   nodeId: string;
+  forkIndex?: number;
   retry: number;
   params: NodeDataParams<ParamsType>;
   inputs: Array<PreviousResultType>;
@@ -61,6 +64,8 @@ class Node {
   public waitlist = new Set<string>(); // List of nodes which need data from this node.
   public state = NodeState.Waiting;
   public agentId?: string;
+  public fork?: number;
+  public forkIndex?: number;
   public result: ResultData = undefined;
   public retryLimit: number;
   public retryCount: number = 0;
@@ -72,12 +77,14 @@ class Node {
 
   private graph: GraphAI;
 
-  constructor(nodeId: string, data: NodeData, graph: GraphAI) {
+  constructor(nodeId: string, forkIndex: number | undefined, data: NodeData, graph: GraphAI) {
     this.nodeId = nodeId;
+    this.forkIndex = forkIndex;
     this.inputs = data.inputs ?? [];
     this.pendings = new Set(this.inputs);
-    this.params = data.params;
+    this.params = data.params ?? {};
     this.agentId = data.agentId;
+    this.fork = data.fork;
     this.retryLimit = data.retry ?? 0;
     this.timeout = data.timeout;
     this.source = data.source === true;
@@ -122,9 +129,9 @@ class Node {
         retryCount: this.retryCount,
         state: NodeState.Injected,
         startTime: Date.now(),
+        endTime: Date.now(),
         result,
       };
-      log.endTime = log.startTime;
       this.graph.appendLog(log);
       this.setResult(result, NodeState.Injected);
     } else {
@@ -143,21 +150,19 @@ class Node {
   }
 
   public async execute() {
+    const results = this.graph.resultsOf(this.inputs);
+    const transactionId = Date.now();
     const log: TransactionLog = {
       nodeId: this.nodeId,
       retryCount: this.retryCount,
       state: NodeState.Executing,
-      startTime: Date.now(),
+      startTime: transactionId,
       agentId: this.agentId,
       params: this.params,
+      inputs: results,
     };
-    const results = this.graph.resultsOf(this.inputs);
-    if (results.length > 0) {
-      log.inputs = results;
-    }
     this.graph.appendLog(log);
     this.state = NodeState.Executing;
-    const transactionId = log.startTime;
     this.transactionId = transactionId;
 
     if (this.timeout && this.timeout > 0) {
@@ -179,6 +184,7 @@ class Node {
         retry: this.retryCount,
         params: this.params,
         inputs: results,
+        forkIndex: this.forkIndex,
       });
       if (this.transactionId !== transactionId) {
         console.log(`-- ${this.nodeId}: transactionId mismatch`);
@@ -241,24 +247,66 @@ export class GraphAI {
     this.onComplete = () => {
       console.error("-- SOMETHING IS WRONG: onComplete is called without run()");
     };
+    const nodeId2forkedNodeIds: Record<string, string[]> = {};
+    const forkedNodeId2Index: Record<string, number> = {};
+
+    // Create node instances from data.nodes
     this.nodes = Object.keys(data.nodes).reduce((nodes: GraphNodes, nodeId: string) => {
-      nodes[nodeId] = new Node(nodeId, data.nodes[nodeId], this);
+      const fork = data.nodes[nodeId].fork;
+      if (fork) {
+        // For fork, change the nodeId and increase the node
+        nodeId2forkedNodeIds[nodeId] = new Array(fork).fill(undefined).map((_, i) => {
+          const forkedNodeId = `${nodeId}_${i}`;
+          nodes[forkedNodeId] = new Node(forkedNodeId, i, data.nodes[nodeId], this);
+          // Data for pending and waiting
+          forkedNodeId2Index[forkedNodeId] = i;
+          return forkedNodeId;
+        });
+      } else {
+        nodes[nodeId] = new Node(nodeId, undefined, data.nodes[nodeId], this);
+      }
       return nodes;
     }, {});
 
-    // Generate the waitlist for each node
+    // Generate the waitlist for each node, and update the pendings in case of forked node.
     Object.keys(this.nodes).forEach((nodeId) => {
       const node = this.nodes[nodeId];
       node.pendings.forEach((pending) => {
-        const node2 = this.nodes[pending];
-        node2.waitlist.add(nodeId);
+        // If the pending(previous) node is forking
+        if (nodeId2forkedNodeIds[pending]) {
+          //  update node.pending and pending(previous) node.wailtlist
+          if (node.fork) {
+            //  1:1 if current nodes are also forking.
+            const newPendingId = nodeId2forkedNodeIds[pending][forkedNodeId2Index[nodeId]];
+            this.nodes[newPendingId].waitlist.add(nodeId); // previousNode
+            node.pendings.add(newPendingId);
+          } else {
+            //  1:n if current node is not forking.
+            nodeId2forkedNodeIds[pending].forEach((newPendingId) => {
+              this.nodes[newPendingId].waitlist.add(nodeId); // previousNode
+              node.pendings.add(newPendingId);
+            });
+          }
+          node.pendings.delete(pending);
+        } else {
+          this.nodes[pending].waitlist.add(nodeId); // previousNode
+        }
       });
+      node.inputs = Array.from(node.pendings); // for fork.
+    });
+
+    // If the result property is specified, inject it.
+    // NOTE: This must be done at the end of this constructor
+    Object.keys(data.nodes).forEach((nodeId) => {
+      const result = data.nodes[nodeId].result;
+      if (result) {
+        this.injectResult(nodeId, result);
+      }
     });
   }
 
   public getCallback(_agentId?: string) {
     const agentId = _agentId ?? "_default";
-    console.log(agentId);
     if (this.callbackDictonary[agentId]) {
       return this.callbackDictonary[agentId];
     }

package/src/graphai_cli.ts

@@ -0,0 +1,35 @@
+#!/usr/bin/env node
+
+import { GraphAI, AgentFunction } from "./graphai";
+import fs from "fs";
+import path from "path";
+import YAML from "yaml";
+
+const testAgent: AgentFunction<{ delay: number; fail: boolean }> = async (context) => {
+  return {};
+};
+
+const main = async () => {
+  const file = process.argv[2];
+  if (file === undefined) {
+    console.log("no file");
+    return;
+  }
+  const file_path = path.resolve(process.cwd() + "/" + file);
+  if (!fs.existsSync(file_path)) {
+    console.log("no file");
+    return;
+  }
+  try {
+    const graph_data_file = fs.readFileSync(file_path, "utf8");
+    const graph_data = YAML.parse(graph_data_file);
+
+    const graph = new GraphAI(graph_data, { test: testAgent });
+    const results = await graph.run();
+    console.log(results);
+  } catch (e) {
+    console.log("error", e);
+  }
+};
+
+main();

package/tests/graphai/test_fork.ts

@@ -0,0 +1,106 @@
+import { GraphAI, AgentFunction } from "@/graphai";
+import { testAgent } from "~/agents/agents";
+import { graphDataTestRunner } from "~/utils/runner";
+
+import test from "node:test";
+import assert from "node:assert";
+
+const testAgent1: AgentFunction = async (context) => {
+  const { nodeId, retry, params, inputs } = context;
+  console.log("executing", nodeId, params, inputs);
+
+  const result = {
+    [nodeId]: [nodeId, inputs.map((a) => Object.values(a).flat())]
+      .flat()
+      .filter((a) => !!a)
+      .join(":"),
+  };
+  console.log("completing", nodeId, result);
+  return result;
+};
+
+test("test base", async () => {
+  const forkGraph = {
+    nodes: {
+      node1: {
+        params: {},
+      },
+      node2: {
+        params: {},
+        fork: 10,
+        inputs: ["node1"],
+      },
+      node3: {
+        params: {},
+        // fork: 10,
+        inputs: ["node2"],
+      },
+    },
+  };
+
+  const result = await graphDataTestRunner("fork.yml", forkGraph, testAgent1);
+  // console.log(result);
+  assert.deepStrictEqual(result, {
+    node1: { node1: "node1" },
+    node2_0: { node2_0: "node2_0:node1" },
+    node2_1: { node2_1: "node2_1:node1" },
+    node2_2: { node2_2: "node2_2:node1" },
+    node2_3: { node2_3: "node2_3:node1" },
+    node2_4: { node2_4: "node2_4:node1" },
+    node2_5: { node2_5: "node2_5:node1" },
+    node2_6: { node2_6: "node2_6:node1" },
+    node2_7: { node2_7: "node2_7:node1" },
+    node2_8: { node2_8: "node2_8:node1" },
+    node2_9: { node2_9: "node2_9:node1" },
+    node3: {
+      node3:
+        "node3:node2_0:node1:node2_1:node1:node2_2:node1:node2_3:node1:node2_4:node1:node2_5:node1:node2_6:node1:node2_7:node1:node2_8:node1:node2_9:node1",
+    },
+  });
+});
+
+test("test base", async () => {
+  const forkGraph = {
+    nodes: {
+      node1: {
+        params: {},
+      },
+      node2: {
+        params: {},
+        fork: 10,
+        inputs: ["node1"],
+      },
+      node3: {
+        params: {},
+        fork: 10,
+        inputs: ["node2"],
+      },
+    },
+  };
+
+  const result = await graphDataTestRunner("fork.yml", forkGraph, testAgent1);
+  // console.log(result);
+  assert.deepStrictEqual(result, {
+    node1: { node1: "node1" },
+    node2_0: { node2_0: "node2_0:node1" },
+    node2_1: { node2_1: "node2_1:node1" },
+    node2_2: { node2_2: "node2_2:node1" },
+    node2_3: { node2_3: "node2_3:node1" },
+    node2_4: { node2_4: "node2_4:node1" },
+    node2_5: { node2_5: "node2_5:node1" },
+    node2_6: { node2_6: "node2_6:node1" },
+    node2_7: { node2_7: "node2_7:node1" },
+    node2_8: { node2_8: "node2_8:node1" },
+    node2_9: { node2_9: "node2_9:node1" },
+    node3_0: { node3_0: "node3_0:node2_0:node1" },
+    node3_1: { node3_1: "node3_1:node2_1:node1" },
+    node3_2: { node3_2: "node3_2:node2_2:node1" },
+    node3_3: { node3_3: "node3_3:node2_3:node1" },
+    node3_4: { node3_4: "node3_4:node2_4:node1" },
+    node3_5: { node3_5: "node3_5:node2_5:node1" },
+    node3_6: { node3_6: "node3_6:node2_6:node1" },
+    node3_7: { node3_7: "node3_7:node2_7:node1" },
+    node3_8: { node3_8: "node3_8:node2_8:node1" },
+    node3_9: { node3_9: "node3_9:node2_9:node1" },
+  });
+});

package/tests/graphai/test_sample_flow.ts

@@ -60,13 +60,12 @@ test("test source", async () => {
 test("test source2", async () => {
   const result = await fileTestRunner("/graphs/test_source2.yml", testAgent, (graph: GraphAI) => {
     graph.injectResult("node1", { node1: "injected" });
-    graph.injectResult("node2", { node2: "injected" });
   });
   assert.deepStrictEqual(result, {
     node1: { node1: "injected" },
-    node2: { node2: "injected" },
-    node3: { node3: "output", node1: "injected", node2: "injected" },
-    node4: { node4: "output", node3: "output", node1: "injected", node2: "injected" },
-    node5: { node5: "output", node4: "output", node3: "output", node1: "injected", node2: "injected" },
+    node2: { node2: "preset" },
+    node3: { node3: "output", node1: "injected", node2: "preset" },
+    node4: { node4: "output", node3: "output", node1: "injected", node2: "preset" },
+    node5: { node5: "output", node4: "output", node3: "output", node1: "injected", node2: "preset" },
   });
 });

package/tests/graphs/test_cli.yaml

@@ -0,0 +1,4 @@
+nodes:
+  node1:
+    agentId: test
+    

package/tests/graphs/test_source2.yml

@@ -3,6 +3,8 @@ nodes:
     source: true
   node2:
     source: true
+    result:
+      node2: preset
   node3:
     params:
       delay: 500