graphai 0.0.7 → 0.0.8

package/README.md

@@ -4,7 +4,7 @@
 
 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 data flow graph (typically in YAML), create a GraphAI object with it, and run it.
+You just need to describe dependencies among those API calls in a single data flow graph, create a GraphAI object with that graph, and run it.
 
 Here is an example:
 
@@ -27,7 +27,7 @@ const sampleAgentFunction = async (context: AgentFunctionContext) => {
   const { 
     nodeId, // taskA, taskB or taskC 
     params, // agent-specific parameters specified in the graph definition file
-    payload // for taskC, { taskA: resultA, taskB: resultB }
+    inputs // inputs from previous nodes
   } = context;
   // Agent-specific code (such as calling OpenAI's chat.completions API)
   ...
@@ -46,7 +46,7 @@ const sampleAgentFunction = async (context: AgentFunctionContext) => {
 
 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, retry, timeout, source, dispatch, 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) that dictate the node's behavior and its relationship with other nodes.
 
 ### DFG Structure
 
@@ -61,19 +61,16 @@ An agent is an abstract object, which takes some inputs and generates an output
 
 An agent function is a TypeScript function, which implements an agent. A DFG is associated one or more agent functions. If the DFG is associated with multiple agent functions, each node needs to be associated only one of them (either explicitly with 'agentId' or implicitly to the default Agent function).
 
-An agent function receives two set of parameters via AgentFunctionContext, agent specific parameters specified in the DFG and input data came from other nodes (payload).
-
-A regular agent function returns the data (type: ```Record<string, any>```), but a dispatcher agent function returns the date with outputID(s) (type: ```Record<string: Record<string, any>>```). 
+An agent function receives two set of parameters via AgentFunctionContext, agent specific parameters specified in the DFG and input data came from other nodes.
 
 ## Node Structure
 
 - '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'.
-- 'dispatch': An optinal property (type: ```Record<string, string>```) which indicates that the current node is a 'dispatcher node'. A dispatcher node is associated is a dispaher-style agent function, which has multiple possible 'outputs'.
+- '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.
 - '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.
-- 'payloadMapping': An optional property (type: ```Record<string, string>```), which maps input nodeIds to agent specific inputIDs.
 
 ## GraphAI class
 

package/lib/graphai.d.ts

@@ -13,12 +13,11 @@ export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType;
 type NodeData = {
     inputs?: Array<string>;
     params: NodeDataParams;
-    payloadMapping?: Record<string, string>;
     retry?: number;
     timeout?: number;
     agentId?: string;
     source?: boolean;
-    dispatch?: Record<string, string>;
+    outputs?: Record<string, string>;
 };
 export type GraphData = {
     nodes: Record<string, NodeData>;
@@ -32,7 +31,7 @@ export type TransactionLog = {
     retryCount: number;
     agentId?: string;
     params?: NodeDataParams;
-    payload?: ResultDataDictonary<ResultData>;
+    inputs?: Array<ResultData>;
     errorMessage?: string;
     result?: ResultData;
 };
@@ -40,7 +39,7 @@ export type AgentFunctionContext<ParamsType, ResultType, PreviousResultType> = {
     nodeId: string;
     retry: number;
     params: NodeDataParams<ParamsType>;
-    payload: ResultDataDictonary<PreviousResultType>;
+    inputs: Array<PreviousResultType>;
 };
 export type AgentFunction<ParamsType = Record<string, any>, ResultType = Record<string, any>, PreviousResultType = Record<string, any>> = (context: AgentFunctionContext<ParamsType, ResultType, PreviousResultType>) => Promise<ResultData<ResultType>>;
 export type AgentFunctionDictonary = Record<string, AgentFunction<any, any, any>>;
@@ -48,11 +47,10 @@ declare class Node {
     nodeId: string;
     params: NodeDataParams;
     inputs: Array<string>;
-    payloadMapping: Record<string, string>;
     pendings: Set<string>;
     waitlist: Set<string>;
     state: NodeState;
-    agentId: string;
+    agentId?: string;
     result: ResultData;
     retryLimit: number;
     retryCount: number;
@@ -60,13 +58,12 @@ declare class Node {
     timeout?: number;
     error?: Error;
     source: boolean;
-    dispatch?: Record<string, string>;
+    outputs?: Record<string, string>;
     private graph;
     constructor(nodeId: string, data: NodeData, graph: GraphAI);
     asString(): string;
     private retry;
     removePending(nodeId: string): void;
-    payload(): ResultDataDictonary<Record<string, any>>;
     pushQueueIfReady(): void;
     injectResult(result: ResultData): void;
     private setResult;
@@ -83,7 +80,7 @@ export declare class GraphAI {
     private concurrency;
     private logs;
     constructor(data: GraphData, callbackDictonary: AgentFunctionDictonary | AgentFunction<any, any, any>);
-    getCallback(agentId: string): AgentFunction<any, any, any>;
+    getCallback(_agentId?: string): AgentFunction<any, any, any>;
     asString(): string;
     results(): ResultDataDictonary<Record<string, any>>;
     errors(): Record<string, Error>;
@@ -94,5 +91,6 @@ export declare class GraphAI {
     appendLog(log: TransactionLog): void;
     transactionLogs(): TransactionLog[];
     injectResult(nodeId: string, result: ResultData): void;
+    resultsOf(nodeIds: Array<string>): ResultData<Record<string, any>>[];
 }
 export {};

package/lib/graphai.js

@@ -19,14 +19,13 @@ class Node {
         this.retryCount = 0;
         this.nodeId = nodeId;
         this.inputs = data.inputs ?? [];
-        this.payloadMapping = data.payloadMapping ?? {};
         this.pendings = new Set(this.inputs);
         this.params = data.params;
-        this.agentId = data.agentId ?? "default";
+        this.agentId = data.agentId;
         this.retryLimit = data.retry ?? 0;
         this.timeout = data.timeout;
         this.source = data.source === true;
-        this.dispatch = data.dispatch;
+        this.outputs = data.outputs;
         this.graph = graph;
     }
     asString() {
@@ -51,17 +50,6 @@ class Node {
             this.pushQueueIfReady();
         }
     }
-    payload() {
-        return this.inputs.reduce((results, nodeId) => {
-            if (this.payloadMapping && this.payloadMapping[nodeId]) {
-                results[this.payloadMapping[nodeId]] = this.graph.nodes[nodeId].result;
-            }
-            else {
-                results[nodeId] = this.graph.nodes[nodeId].result;
-            }
-            return results;
-        }, {});
-    }
     pushQueueIfReady() {
         if (this.pendings.size === 0 && !this.source) {
             this.graph.pushQueue(this);
@@ -74,6 +62,7 @@ class Node {
                 retryCount: this.retryCount,
                 state: NodeState.Injected,
                 startTime: Date.now(),
+                result,
             };
             log.endTime = log.startTime;
             this.graph.appendLog(log);
@@ -93,7 +82,6 @@ class Node {
         });
     }
     async execute() {
-        const payload = this.payload();
         const log = {
             nodeId: this.nodeId,
             retryCount: this.retryCount,
@@ -101,8 +89,11 @@ class Node {
             startTime: Date.now(),
             agentId: this.agentId,
             params: this.params,
-            payload,
         };
+        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;
@@ -124,7 +115,7 @@ class Node {
                 nodeId: this.nodeId,
                 retry: this.retryCount,
                 params: this.params,
-                payload,
+                inputs: results,
             });
             if (this.transactionId !== transactionId) {
                 console.log(`-- ${this.nodeId}: transactionId mismatch`);
@@ -132,10 +123,10 @@ class Node {
             }
             log.endTime = Date.now();
             log.result = result;
-            const dispatch = this.dispatch;
-            if (dispatch !== undefined) {
+            const outputs = this.outputs;
+            if (outputs !== undefined) {
                 Object.keys(result).forEach((outputId) => {
-                    const nodeId = dispatch[outputId];
+                    const nodeId = outputs[outputId];
                     this.graph.injectResult(nodeId, result[outputId]);
                 });
                 log.state = NodeState.Dispatched;
@@ -173,10 +164,7 @@ class GraphAI {
         this.runningNodes = new Set();
         this.nodeQueue = [];
         this.logs = [];
-        this.callbackDictonary = typeof callbackDictonary === "function" ? { default: callbackDictonary } : callbackDictonary;
-        if (this.callbackDictonary["default"] === undefined) {
-            throw new Error("No default function");
-        }
+        this.callbackDictonary = typeof callbackDictonary === "function" ? { _default: callbackDictonary } : callbackDictonary;
         this.concurrency = data.concurrency ?? defaultConcurrency;
         this.onComplete = () => {
             console.error("-- SOMETHING IS WRONG: onComplete is called without run()");
@@ -194,11 +182,13 @@ class GraphAI {
             });
         });
     }
-    getCallback(agentId) {
-        if (agentId && this.callbackDictonary[agentId]) {
+    getCallback(_agentId) {
+        const agentId = _agentId ?? "_default";
+        console.log(agentId);
+        if (this.callbackDictonary[agentId]) {
             return this.callbackDictonary[agentId];
         }
-        return this.callbackDictonary["default"];
+        throw new Error("No agent: " + agentId);
     }
     asString() {
         return Object.keys(this.nodes)
@@ -288,5 +278,10 @@ class GraphAI {
             console.error("-- Invalid nodeId", nodeId);
         }
     }
+    resultsOf(nodeIds) {
+        return nodeIds.map((nodeId) => {
+            return this.nodes[nodeId].result;
+        });
+    }
 }
 exports.GraphAI = GraphAI;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "graphai",
-  "version": "0.0.7",
+  "version": "0.0.8",
   "description": "Asynchronous data flow execution engine to make it simple to build LLM apps.",
   "main": "lib/index.js",
   "scripts": {

package/samples/agents/arxiv_agent.ts

@@ -35,12 +35,11 @@ export const arxivAgent: AgentFunction<{ keywords: string[]; limit: number }, ar
 };
 
 export const arxiv2TextAgent: AgentFunction<{}, string, string[]> = async (context) => {
-  const result = (context?.payload?.inputData || [])
+  const result = (context.inputs[0] || [])
     .map((r: any) => {
       const { id, title, summary } = r;
       return ["id:", id, "title:", title, "summary:", summary].join("\n");
     })
     .join("\n\n\n");
-
   return result;
 };

package/samples/agents/parroting_agent.ts

@@ -0,0 +1,5 @@
+import { AgentFunction } from "@/graphai";
+
+export const parrotingAgent: AgentFunction = async (context) => {
+  return context.params;
+};

package/samples/agents/slashgpt_agent.ts

@@ -4,18 +4,46 @@ import { ChatSession, ChatConfig, ManifestData } from "slashgpt";
 
 const config = new ChatConfig(path.resolve(__dirname));
 
-export const slashGPTAgent: AgentFunction<{ manifest: ManifestData; prompt: string }, { answer: string }> = async (context) => {
-  console.log("executing", context.nodeId, context);
-  const session = new ChatSession(config, context.params?.manifest ?? {});
+export const slashGPTFuncitons2TextAgent: AgentFunction<
+  { function_data_key: string; result_key: number },
+  string,
+  { function_data: { [key: string]: string[] } }
+> = async (context) => {
+  const { params } = context;
+  const result = (context?.inputs[0].function_data[params.function_data_key] || []).map((r: any) => {
+    const { title, description } = r;
+    return ["title:", title, "description:", description].join("\n");
+  });
+  return result[params.result_key];
+};
+
+export const slashGPTAgent: AgentFunction<
+  { manifest: ManifestData; prompt: string; function_result?: boolean; debug?: boolean },
+  { answer: string },
+  string
+> = async (context) => {
+  const { params } = context;
+  if (params.debug) {
+    console.log("executing", context.nodeId, context);
+  }
+  const session = new ChatSession(config, params?.manifest ?? {});
 
-  const prompt = [context.params?.prompt, context.payload.inputData].join("\n\n");
+  const prompt = params?.prompt ?? [context.inputs].filter((a) => a !== undefined).join("\n\n");
+  // console.log(prompt);
   session.append_user_question(prompt);
 
   await session.call_loop(() => {});
-  const message = session.history.last_message();
+
+  // console.log(session.history)
+  const message = (() => {
+    if (params.function_result) {
+      return session.history.messages().find((m) => m.role === "function_result");
+    }
+    return session.history.last_message();
+  })();
   if (message === undefined) {
     throw new Error("No message in the history");
   }
-  const result = { answer: message.content };
+  const result = { answer: message.content, function_data: message.function_data };
   return result;
 };

package/samples/graphs/arxiv.yml

@@ -5,17 +5,13 @@ nodes:
         - llm
         - gpt
       limit: 10
-    functionName: arxivAgent
+    agentId: arxivAgent
   arxiv2TextAgent:
     inputs: [searchArxiv]
-    functionName: arxiv2TextAgent
-    payloadMapping:
-      searchArxiv: inputData
+    agentId: arxiv2TextAgent
   slashGPTAgent:
     inputs: [arxiv2TextAgent]
-    payloadMapping:
-      arxiv2TextAgent: inputData
-    functionName: slashGPTAgent
+    agentId: slashGPTAgent
     params:
       prompt: |
         与えられたそれぞれの論文の要点をまとめ、以下の項目で日本語で出力せよ。それぞれの項目は最大でも180文字以内に要約せよ。

package/samples/graphs/slash_gpt.yml

@@ -5,8 +5,15 @@ nodes:
   node2:
     inputs: [node1]
     params:
-      prompt: Please evaluate following business ideas. ${node1}
+      prompt: |
+        Please evaluate following business ideas.
+        ${0}
   node3:
     inputs: [node1, node2]
     params:
-      prompt: Please pick the winner of this business idea contest. ${node1} ${node2}
+      prompt:  |
+        Please pick the winner of this business idea contest.
+        [ideas]
+        ${0}
+        [evalutations]
+        ${1}

package/samples/sample_co2.ts

@@ -0,0 +1,103 @@
+import { GraphAI, AgentFunction } from "@/graphai";
+import { readGraphaiData } from "~/utils/file_utils";
+
+import { slashGPTAgent, slashGPTFuncitons2TextAgent } from "./agents/slashgpt_agent";
+
+const graph_data = {
+  nodes: {
+    slashGPTAgent: {
+      agentId: "slashGPTAgent",
+      params: {
+        function_result: true,
+        prompt: "世界で協力してco2を減らす方法を教えて下さい",
+        manifest: {
+          prompt: "あなたは世界経済の専門家です。ユーザの問い合わせについて考え、10この結果をfunctionsの結果に返してください。",
+          skip_function_result: true,
+          actions: {
+            your_ideas: {
+              type: "message_template",
+              message: "dummy",
+            },
+          },
+          functions: [
+            {
+              name: "your_ideas",
+              description: "Your answer to a user's inquiry",
+              parameters: {
+                type: "object",
+                properties: {
+                  methods: {
+                    type: "array",
+                    items: {
+                      type: "object",
+                      properties: {
+                        title: {
+                          type: "string",
+                          description: "title",
+                        },
+                        description: {
+                          type: "string",
+                          description: "description",
+                        },
+                      },
+                      required: ["title", "description"],
+                    },
+                  },
+                },
+              },
+            },
+          ],
+        },
+      },
+    },
+    function2prompt0: {
+      params: {
+        function_data_key: "methods",
+        result_key: 0,
+      },
+      inputs: ["slashGPTAgent"],
+      agentId: "slashGPTFuncitons2TextAgent",
+    },
+    slashGPTAgent0: {
+      agentId: "slashGPTAgent",
+      params: {
+        debug: true,
+        manifest: {
+          prompt: "ユーザの問い合わせにある文章の専門家です。専門家として、ユーザのアイデアに対して実現可能なシナリオを800文字で書いてください。",
+        },
+      },
+      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 () => {
+  const graph = new GraphAI(graph_data, { slashGPTAgent, slashGPTFuncitons2TextAgent });
+  const result = await graph.run();
+  console.log(result);
+};
+
+const main = async () => {
+  await runAgent();
+  console.log("COMPLETE 1");
+};
+main();

package/samples/sample_gpt.ts

@@ -8,8 +8,8 @@ const config = new ChatConfig(path.resolve(__dirname));
 const slashGPTAgent: AgentFunction<{ manifest: ManifestData; prompt: string }, { answer: 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"]);
+  const prompt = context.inputs.reduce((prompt, input, index) => {
+    return prompt.replace("${" + index + "}", input["answer"]);
   }, context.params.prompt);
   session.append_user_question(prompt);
 

package/samples/sample_paper_ai.ts

@@ -7,14 +7,10 @@ import { readGraphaiData } from "~/utils/file_utils";
 import { slashGPTAgent } from "./agents/slashgpt_agent";
 import { arxivAgent, arxiv2TextAgent } from "./agents/arxiv_agent";
 
-export const parrotingAgent: AgentFunction = async (context) => {
-  return {};
-};
-
 const runAgent = async (file: string) => {
   const file_path = path.resolve(__dirname) + file;
   const graph_data = readGraphaiData(file_path);
-  const graph = new GraphAI(graph_data, { default: parrotingAgent, arxivAgent: arxivAgent, arxiv2TextAgent, slashGPTAgent });
+  const graph = new GraphAI(graph_data, { arxivAgent: arxivAgent, arxiv2TextAgent, slashGPTAgent });
   const result = await graph.run();
   console.log(result);
 };

package/src/graphai.ts

@@ -15,12 +15,11 @@ export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType; // Ag
 type NodeData = {
   inputs?: Array<string>;
   params: NodeDataParams;
-  payloadMapping?: Record<string, string>;
   retry?: number;
   timeout?: number; // msec
   agentId?: string;
   source?: boolean;
-  dispatch?: Record<string, string>; // route to node
+  outputs?: Record<string, string>; // mapping from routeId to nodeId
 };
 
 export type GraphData = {
@@ -36,7 +35,7 @@ export type TransactionLog = {
   retryCount: number;
   agentId?: string;
   params?: NodeDataParams;
-  payload?: ResultDataDictonary<ResultData>;
+  inputs?: Array<ResultData>;
   errorMessage?: string;
   result?: ResultData;
 };
@@ -45,7 +44,7 @@ export type AgentFunctionContext<ParamsType, ResultType, PreviousResultType> = {
   nodeId: string;
   retry: number;
   params: NodeDataParams<ParamsType>;
-  payload: ResultDataDictonary<PreviousResultType>;
+  inputs: Array<PreviousResultType>;
 };
 
 export type AgentFunction<ParamsType = Record<string, any>, ResultType = Record<string, any>, PreviousResultType = Record<string, any>> = (
@@ -58,11 +57,10 @@ class Node {
   public nodeId: string;
   public params: NodeDataParams; // Agent-specific parameters
   public inputs: Array<string>; // List of nodes this node needs data from.
-  public payloadMapping: Record<string, string>;
   public pendings: Set<string>; // List of nodes this node is waiting data from.
   public waitlist = new Set<string>(); // List of nodes which need data from this node.
   public state = NodeState.Waiting;
-  public agentId: string;
+  public agentId?: string;
   public result: ResultData = undefined;
   public retryLimit: number;
   public retryCount: number = 0;
@@ -70,21 +68,20 @@ class Node {
   public timeout?: number; // msec
   public error?: Error;
   public source: boolean;
-  public dispatch?: Record<string, string>; // outputId to nodeId mapping
+  public outputs?: Record<string, string>; // Mapping from routeId to nodeId
 
   private graph: GraphAI;
 
   constructor(nodeId: string, data: NodeData, graph: GraphAI) {
     this.nodeId = nodeId;
     this.inputs = data.inputs ?? [];
-    this.payloadMapping = data.payloadMapping ?? {};
     this.pendings = new Set(this.inputs);
     this.params = data.params;
-    this.agentId = data.agentId ?? "default";
+    this.agentId = data.agentId;
     this.retryLimit = data.retry ?? 0;
     this.timeout = data.timeout;
     this.source = data.source === true;
-    this.dispatch = data.dispatch;
+    this.outputs = data.outputs;
     this.graph = graph;
   }
 
@@ -112,17 +109,6 @@ class Node {
     }
   }
 
-  public payload() {
-    return this.inputs.reduce((results: ResultDataDictonary, nodeId) => {
-      if (this.payloadMapping && this.payloadMapping[nodeId]) {
-        results[this.payloadMapping[nodeId]] = this.graph.nodes[nodeId].result;
-      } else {
-        results[nodeId] = this.graph.nodes[nodeId].result;
-      }
-      return results;
-    }, {});
-  }
-
   public pushQueueIfReady() {
     if (this.pendings.size === 0 && !this.source) {
       this.graph.pushQueue(this);
@@ -136,6 +122,7 @@ class Node {
         retryCount: this.retryCount,
         state: NodeState.Injected,
         startTime: Date.now(),
+        result,
       };
       log.endTime = log.startTime;
       this.graph.appendLog(log);
@@ -156,7 +143,6 @@ class Node {
   }
 
   public async execute() {
-    const payload = this.payload();
     const log: TransactionLog = {
       nodeId: this.nodeId,
       retryCount: this.retryCount,
@@ -164,8 +150,11 @@ class Node {
       startTime: Date.now(),
       agentId: this.agentId,
       params: this.params,
-      payload,
     };
+    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;
@@ -189,7 +178,7 @@ class Node {
         nodeId: this.nodeId,
         retry: this.retryCount,
         params: this.params,
-        payload,
+        inputs: results,
       });
       if (this.transactionId !== transactionId) {
         console.log(`-- ${this.nodeId}: transactionId mismatch`);
@@ -199,10 +188,10 @@ class Node {
       log.endTime = Date.now();
       log.result = result;
 
-      const dispatch = this.dispatch;
-      if (dispatch !== undefined) {
+      const outputs = this.outputs;
+      if (outputs !== undefined) {
         Object.keys(result).forEach((outputId) => {
-          const nodeId = dispatch[outputId];
+          const nodeId = outputs[outputId];
           this.graph.injectResult(nodeId, result[outputId]);
         });
         log.state = NodeState.Dispatched;
@@ -247,10 +236,7 @@ export class GraphAI {
   private logs: Array<TransactionLog> = [];
 
   constructor(data: GraphData, callbackDictonary: AgentFunctionDictonary | AgentFunction<any, any, any>) {
-    this.callbackDictonary = typeof callbackDictonary === "function" ? { default: callbackDictonary } : callbackDictonary;
-    if (this.callbackDictonary["default"] === undefined) {
-      throw new Error("No default function");
-    }
+    this.callbackDictonary = typeof callbackDictonary === "function" ? { _default: callbackDictonary } : callbackDictonary;
     this.concurrency = data.concurrency ?? defaultConcurrency;
     this.onComplete = () => {
       console.error("-- SOMETHING IS WRONG: onComplete is called without run()");
@@ -270,11 +256,13 @@ export class GraphAI {
     });
   }
 
-  public getCallback(agentId: string) {
-    if (agentId && this.callbackDictonary[agentId]) {
+  public getCallback(_agentId?: string) {
+    const agentId = _agentId ?? "_default";
+    console.log(agentId);
+    if (this.callbackDictonary[agentId]) {
       return this.callbackDictonary[agentId];
     }
-    return this.callbackDictonary["default"];
+    throw new Error("No agent: " + agentId);
   }
 
   public asString() {
@@ -372,4 +360,10 @@ export class GraphAI {
       console.error("-- Invalid nodeId", nodeId);
     }
   }
+
+  public resultsOf(nodeIds: Array<string>) {
+    return nodeIds.map((nodeId) => {
+      return this.nodes[nodeId].result;
+    });
+  }
 }

package/tests/agents/agents.ts

@@ -2,7 +2,7 @@ import { AgentFunction } from "@/graphai";
 import { sleep } from "~/utils/utils";
 
 export const testAgent: AgentFunction<{ delay: number; fail: boolean }> = async (context) => {
-  const { nodeId, retry, params, payload } = context;
+  const { nodeId, retry, params, inputs } = context;
   console.log("executing", nodeId);
   await sleep(params.delay / (retry + 1));
 
@@ -11,10 +11,9 @@ export const testAgent: AgentFunction<{ delay: number; fail: boolean }> = async
     console.log("failed (intentional)", nodeId, retry);
     throw new Error("Intentional Failure");
   } else {
-    const result = Object.keys(payload).reduce(
-      (result, key) => {
-        result = { ...result, ...payload[key] };
-        return result;
+    const result = inputs.reduce(
+      (result: Record<string, any>, input: Record<string, any>) => {
+        return { ...result, ...input };
       },
       { [nodeId]: "output" },
     );

package/tests/graphai/test_dispatch.ts

@@ -8,8 +8,8 @@ import { testAgent } from "~/agents/agents";
 import test from "node:test";
 import assert from "node:assert";
 
-const dispatchAgent: AgentFunction<{ delay: number; fail: boolean }> = async (context) => {
-  const { nodeId, retry, params, payload } = context;
+const dispatchAgent: AgentFunction<{ delay: number; fail: boolean }, Record<string, any>, Record<string, any>> = async (context) => {
+  const { nodeId, retry, params, inputs } = context;
   console.log("executing", nodeId);
   await sleep(params.delay / (retry + 1));
 
@@ -18,10 +18,9 @@ const dispatchAgent: AgentFunction<{ delay: number; fail: boolean }> = async (co
     console.log("failed (intentional)", nodeId, retry);
     throw new Error("Intentional Failure");
   } else {
-    const result = Object.keys(payload).reduce(
-      (result, key) => {
-        result = { ...result, ...payload[key] };
-        return result;
+    const result = inputs.reduce(
+      (result: Record<string, any>, input: Record<string, any>) => {
+        return { ...result, ...input };
       },
       { [nodeId]: "dispatch" },
     );
@@ -31,7 +30,7 @@ const dispatchAgent: AgentFunction<{ delay: number; fail: boolean }> = async (co
 };
 
 test("test dispatch", async () => {
-  const result = await fileTestRunner("/graphs/test_dispatch.yml", { default: testAgent, alt: dispatchAgent });
+  const result = await fileTestRunner("/graphs/test_dispatch.yml", { test: testAgent, alt: dispatchAgent });
   assert.deepStrictEqual(result, {
     node1: { node1: "output" },
     node20: { node2: "dispatch" },

package/tests/graphai/test_http_client.ts

@@ -5,8 +5,8 @@ import test from "node:test";
 import assert from "node:assert";
 
 const httpClientAgent: AgentFunction<Record<string, string>> = async (context) => {
-  const { nodeId, retry, params, payload } = context;
-  console.log("executing", nodeId, params, payload);
+  const { nodeId, retry, params, inputs } = context;
+  console.log("executing", nodeId, params, inputs);
 
   const response = await fetch(params.url);
   const result = await response.json();
@@ -31,7 +31,7 @@ const graph_data = {
   },
 };
 
-test("test sample1", async () => {
+test("test http client", async () => {
   const result = await graphDataTestRunner("http.log", graph_data, httpClientAgent);
   assert.deepStrictEqual(result, {
     node1: { result: true, messages: ["hello"] },

package/tests/graphai/test_multiple_functions.ts

@@ -33,8 +33,8 @@ const numberTestAgent: AgentFunction<{ number: number }, { [key: string]: number
   return result;
 };
 
-test("test sample1", async () => {
-  const result = await fileTestRunner("/graphs/test_multiple_functions_1.yml", { default: testAgent1, test2: testAgent2, numberTestAgent });
+test("test multiple function", async () => {
+  const result = await fileTestRunner("/graphs/test_multiple_functions_1.yml", { test1: testAgent1, test2: testAgent2, numberTestAgent });
   assert.deepStrictEqual(result, {
     node1: { node1: "output 1" },
     node2: { node2: "output 1" },

package/tests/graphs/test_dispatch.yml

@@ -2,23 +2,28 @@ nodes:
   node1:
     params:
       delay: 500
+    agentId: test
   node2:
     params:
       delay: 100
     agentId: alt
-    dispatch:
+    outputs:
       output1: node20
   node20:
     source: true
+    agentId: test
   node3:
     params:
       delay: 500
     inputs: [node1, node20]
+    agentId: test
   node4:
     params:
       delay: 100
     inputs: [node3]
+    agentId: test
   node5:
     params:
       delay: 500
     inputs: [node20, node4]
+    agentId: test

package/tests/graphs/test_multiple_functions_1.yml

@@ -2,9 +2,11 @@ nodes:
   node1:
     params:
       delay: 500
+    agentId: test1
   node2:
     params:
       delay: 100
+    agentId: test1
   node3:
     params:
       delay: 500
@@ -14,6 +16,7 @@ nodes:
     params:
       delay: 100
     inputs: [node3]
+    agentId: test1
   node5:
     params:
       delay: 500