graphai 0.0.2 → 0.0.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -1,6 +1,8 @@
1
1
  # GraphAI
2
2
 
3
- GraphAI is a TypeScript library, which makes it easy to create AI applications that need to call asynchrouns AI API calls multiple times. You just need to describe dependencies among those API calls in a single JSON/YAML file, create a GraphAI object with it, and run it.
3
+ 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.
4
+
5
+ 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.
4
6
 
5
7
  Here is an example:
6
8
 
@@ -18,3 +20,25 @@ nodes:
18
20
  inputs: [taskA, taskB]
19
21
  ```
20
22
 
23
+ ``` TypeScript
24
+ const nodeExecute = async (context: NodeExecuteContext) => {
25
+ const {
26
+ nodeId, // taskA, taskB or taskC
27
+ params, // app-specific/task-specific parameters specified in the graph definition file
28
+ payload // for taskC, { taskA: resultA, taskB: resultB }
29
+ } = context;
30
+ // App-specific code (such as calling OpenAI's chat.completions API)
31
+ ...
32
+ return result;
33
+ }
34
+
35
+ ...
36
+ const file = fs.readFileSync(pathToYamlFile, "utf8");
37
+ const graphdata = YAML.parse(file);
38
+ const graph = new GraphAI(graph_data, nodeExecute);
39
+ const results = await graph.run();
40
+ return results["taskC"];
41
+
42
+ ```
43
+
44
+
@@ -0,0 +1,58 @@
1
+ export declare enum NodeState {
2
+ Waiting = 0,
3
+ Executing = 1,
4
+ Failed = 2,
5
+ TimedOut = 3,
6
+ Completed = 4
7
+ }
8
+ type ResultData = Record<string, any>;
9
+ export type NodeDataParams = Record<string, any>;
10
+ type NodeData = {
11
+ inputs: undefined | Array<string>;
12
+ params: NodeDataParams;
13
+ retry: undefined | number;
14
+ timeout: undefined | number;
15
+ };
16
+ type GraphData = {
17
+ nodes: Record<string, NodeData>;
18
+ };
19
+ export type NodeExecuteContext = {
20
+ nodeId: string;
21
+ retry: number;
22
+ params: NodeDataParams;
23
+ payload: ResultData;
24
+ };
25
+ type NodeExecute = (context: NodeExecuteContext) => Promise<ResultData>;
26
+ declare class Node {
27
+ nodeId: string;
28
+ params: NodeDataParams;
29
+ inputs: Array<string>;
30
+ pendings: Set<string>;
31
+ waitlist: Set<string>;
32
+ state: NodeState;
33
+ result: ResultData;
34
+ retryLimit: number;
35
+ retryCount: number;
36
+ transactionId: undefined | number;
37
+ timeout: number;
38
+ constructor(nodeId: string, data: NodeData);
39
+ asString(): string;
40
+ private retry;
41
+ removePending(nodeId: string, graph: GraphAI): void;
42
+ payload(graph: GraphAI): ResultData;
43
+ executeIfReady(graph: GraphAI): void;
44
+ private execute;
45
+ }
46
+ type GraphNodes = Record<string, Node>;
47
+ export declare class GraphAI {
48
+ nodes: GraphNodes;
49
+ callback: NodeExecute;
50
+ private runningNodes;
51
+ private onComplete;
52
+ constructor(data: GraphData, callback: NodeExecute);
53
+ asString(): string;
54
+ run(): Promise<unknown>;
55
+ addRunning(node: Node): void;
56
+ removeRunning(node: Node): void;
57
+ }
58
+ export {};
package/lib/graphai.js ADDED
@@ -0,0 +1,141 @@
1
+ "use strict";
2
+ Object.defineProperty(exports, "__esModule", { value: true });
3
+ exports.GraphAI = exports.NodeState = void 0;
4
+ var NodeState;
5
+ (function (NodeState) {
6
+ NodeState[NodeState["Waiting"] = 0] = "Waiting";
7
+ NodeState[NodeState["Executing"] = 1] = "Executing";
8
+ NodeState[NodeState["Failed"] = 2] = "Failed";
9
+ NodeState[NodeState["TimedOut"] = 3] = "TimedOut";
10
+ NodeState[NodeState["Completed"] = 4] = "Completed";
11
+ })(NodeState || (exports.NodeState = NodeState = {}));
12
+ class Node {
13
+ constructor(nodeId, data) {
14
+ this.nodeId = nodeId;
15
+ this.inputs = data.inputs ?? [];
16
+ this.pendings = new Set(this.inputs);
17
+ this.params = data.params;
18
+ this.waitlist = new Set();
19
+ this.state = NodeState.Waiting;
20
+ this.result = {};
21
+ this.retryLimit = data.retry ?? 0;
22
+ this.retryCount = 0;
23
+ this.timeout = data.timeout ?? 0;
24
+ }
25
+ asString() {
26
+ return `${this.nodeId}: ${this.state} ${[...this.waitlist]}`;
27
+ }
28
+ retry(graph, state, result) {
29
+ if (this.retryCount < this.retryLimit) {
30
+ this.retryCount++;
31
+ this.execute(graph);
32
+ }
33
+ else {
34
+ this.state = state;
35
+ this.result = result;
36
+ graph.removeRunning(this);
37
+ }
38
+ }
39
+ removePending(nodeId, graph) {
40
+ this.pendings.delete(nodeId);
41
+ this.executeIfReady(graph);
42
+ }
43
+ payload(graph) {
44
+ return this.inputs.reduce((results, nodeId) => {
45
+ results[nodeId] = graph.nodes[nodeId].result;
46
+ return results;
47
+ }, {});
48
+ }
49
+ executeIfReady(graph) {
50
+ if (this.pendings.size == 0) {
51
+ graph.addRunning(this);
52
+ this.execute(graph);
53
+ }
54
+ }
55
+ async execute(graph) {
56
+ this.state = NodeState.Executing;
57
+ const transactionId = Date.now();
58
+ this.transactionId = transactionId;
59
+ if (this.timeout > 0) {
60
+ setTimeout(() => {
61
+ if (this.state == NodeState.Executing && this.transactionId == transactionId) {
62
+ console.log("*** timeout", this.timeout);
63
+ this.retry(graph, NodeState.TimedOut, {});
64
+ }
65
+ }, this.timeout);
66
+ }
67
+ try {
68
+ const result = await graph.callback({ nodeId: this.nodeId, retry: this.retryCount, params: this.params, payload: this.payload(graph) });
69
+ if (this.transactionId !== transactionId) {
70
+ console.log("****** tid mismatch (success)");
71
+ return;
72
+ }
73
+ this.state = NodeState.Completed;
74
+ this.result = result;
75
+ this.waitlist.forEach((nodeId) => {
76
+ const node = graph.nodes[nodeId];
77
+ node.removePending(this.nodeId, graph);
78
+ });
79
+ graph.removeRunning(this);
80
+ }
81
+ catch (e) {
82
+ if (this.transactionId !== transactionId) {
83
+ console.log("****** tid mismatch (failed)");
84
+ return;
85
+ }
86
+ this.retry(graph, NodeState.Failed, {});
87
+ }
88
+ }
89
+ }
90
+ class GraphAI {
91
+ constructor(data, callback) {
92
+ this.callback = callback;
93
+ this.runningNodes = new Set();
94
+ this.onComplete = () => { };
95
+ this.nodes = Object.keys(data.nodes).reduce((nodes, nodeId) => {
96
+ nodes[nodeId] = new Node(nodeId, data.nodes[nodeId]);
97
+ return nodes;
98
+ }, {});
99
+ // Generate the waitlist for each node
100
+ Object.keys(this.nodes).forEach((nodeId) => {
101
+ const node = this.nodes[nodeId];
102
+ node.pendings.forEach((pending) => {
103
+ const node2 = this.nodes[pending];
104
+ node2.waitlist.add(nodeId);
105
+ });
106
+ });
107
+ }
108
+ asString() {
109
+ return Object.keys(this.nodes)
110
+ .map((nodeId) => {
111
+ return this.nodes[nodeId].asString();
112
+ })
113
+ .join("\n");
114
+ }
115
+ async run() {
116
+ // Nodes without pending data should run immediately.
117
+ Object.keys(this.nodes).forEach((nodeId) => {
118
+ const node = this.nodes[nodeId];
119
+ node.executeIfReady(this);
120
+ });
121
+ return new Promise((resolve, reject) => {
122
+ this.onComplete = () => {
123
+ const results = Object.keys(this.nodes).reduce((results, nodeId) => {
124
+ results[nodeId] = this.nodes[nodeId].result;
125
+ return results;
126
+ }, {});
127
+ resolve(results);
128
+ };
129
+ });
130
+ }
131
+ addRunning(node) {
132
+ this.runningNodes.add(node.nodeId);
133
+ }
134
+ removeRunning(node) {
135
+ this.runningNodes.delete(node.nodeId);
136
+ if (this.runningNodes.size == 0) {
137
+ this.onComplete();
138
+ }
139
+ }
140
+ }
141
+ exports.GraphAI = GraphAI;
package/lib/index.d.ts CHANGED
@@ -1,2 +1,2 @@
1
- import { GraphAI } from "./flow";
1
+ import { GraphAI } from "./graphai";
2
2
  export { GraphAI };
package/lib/index.js CHANGED
@@ -1,5 +1,5 @@
1
1
  "use strict";
2
2
  Object.defineProperty(exports, "__esModule", { value: true });
3
3
  exports.GraphAI = void 0;
4
- const flow_1 = require("./flow");
5
- Object.defineProperty(exports, "GraphAI", { enumerable: true, get: function () { return flow_1.GraphAI; } });
4
+ const graphai_1 = require("./graphai");
5
+ Object.defineProperty(exports, "GraphAI", { enumerable: true, get: function () { return graphai_1.GraphAI; } });
package/package.json CHANGED
@@ -1,13 +1,13 @@
1
1
  {
2
2
  "name": "graphai",
3
- "version": "0.0.2",
3
+ "version": "0.0.3",
4
4
  "description": "Asynchronous data flow execution engine to make it simple to build LLM apps.",
5
5
  "main": "lib/index.js",
6
6
  "scripts": {
7
7
  "build": "tsc",
8
8
  "eslint": "eslint --fix --ext src/**/*.{ts} tests/**/*.ts",
9
9
  "format": "prettier --write '{src,tests}/**/*.ts'",
10
- "test": "echo \"Error: no test specified\" && exit 1"
10
+ "test": "node --test --require ts-node/register ./tests/sample_flow.ts"
11
11
  },
12
12
  "repository": {
13
13
  "type": "git",
@@ -26,6 +26,7 @@
26
26
  "eslint": "^7.32.0 || ^8.2.0",
27
27
  "eslint-plugin-import": "^2.25.2",
28
28
  "prettier": "^3.0.3",
29
+ "slashgpt": "^0.0.6",
29
30
  "ts-node": "^10.9.1",
30
31
  "typescript": "^5.2.2"
31
32
  },
package/src/graphai.ts ADDED
@@ -0,0 +1,193 @@
1
+ import { AssertionError } from "assert";
2
+
3
+ export enum NodeState {
4
+ Waiting,
5
+ Executing,
6
+ Failed,
7
+ TimedOut,
8
+ Completed,
9
+ }
10
+ type ResultData = Record<string, any>;
11
+ export type NodeDataParams = Record<string, any>; // App-specific parameters
12
+
13
+ type NodeData = {
14
+ inputs: undefined | Array<string>;
15
+ params: NodeDataParams;
16
+ retry: undefined | number;
17
+ timeout: undefined | number; // msec
18
+ };
19
+
20
+ type GraphData = {
21
+ nodes: Record<string, NodeData>;
22
+ };
23
+
24
+ export type NodeExecuteContext = {
25
+ nodeId: string;
26
+ retry: number;
27
+ params: NodeDataParams;
28
+ payload: ResultData;
29
+ };
30
+
31
+ type NodeExecute = (context: NodeExecuteContext) => Promise<ResultData>;
32
+
33
+ class Node {
34
+ public nodeId: string;
35
+ public params: NodeDataParams; // App-specific parameters
36
+ public inputs: Array<string>; // List of nodes this node needs data from.
37
+ public pendings: Set<string>; // List of nodes this node is waiting data from.
38
+ public waitlist: Set<string>; // List of nodes which need data from this node.
39
+ public state: NodeState;
40
+ public result: ResultData;
41
+ public retryLimit: number;
42
+ public retryCount: number;
43
+ public transactionId: undefined | number; // To reject callbacks from timed-out transactions
44
+ public timeout: number; // msec
45
+
46
+ constructor(nodeId: string, data: NodeData) {
47
+ this.nodeId = nodeId;
48
+ this.inputs = data.inputs ?? [];
49
+ this.pendings = new Set(this.inputs);
50
+ this.params = data.params;
51
+ this.waitlist = new Set<string>();
52
+ this.state = NodeState.Waiting;
53
+ this.result = {};
54
+ this.retryLimit = data.retry ?? 0;
55
+ this.retryCount = 0;
56
+ this.timeout = data.timeout ?? 0;
57
+ }
58
+
59
+ public asString() {
60
+ return `${this.nodeId}: ${this.state} ${[...this.waitlist]}`;
61
+ }
62
+
63
+ private retry(graph: GraphAI, state: NodeState, result: ResultData) {
64
+ if (this.retryCount < this.retryLimit) {
65
+ this.retryCount++;
66
+ this.execute(graph);
67
+ } else {
68
+ this.state = state;
69
+ this.result = result;
70
+ graph.removeRunning(this);
71
+ }
72
+ }
73
+
74
+ public removePending(nodeId: string, graph: GraphAI) {
75
+ this.pendings.delete(nodeId);
76
+ this.executeIfReady(graph);
77
+ }
78
+
79
+ public payload(graph: GraphAI) {
80
+ return this.inputs.reduce((results: ResultData, nodeId) => {
81
+ results[nodeId] = graph.nodes[nodeId].result;
82
+ return results;
83
+ }, {});
84
+ }
85
+
86
+ public executeIfReady(graph: GraphAI) {
87
+ if (this.pendings.size == 0) {
88
+ graph.addRunning(this);
89
+ this.execute(graph);
90
+ }
91
+ }
92
+
93
+ private async execute(graph: GraphAI) {
94
+ this.state = NodeState.Executing;
95
+ const transactionId = Date.now();
96
+ this.transactionId = transactionId;
97
+
98
+ if (this.timeout > 0) {
99
+ setTimeout(() => {
100
+ if (this.state == NodeState.Executing && this.transactionId == transactionId) {
101
+ console.log("*** timeout", this.timeout);
102
+ this.retry(graph, NodeState.TimedOut, {});
103
+ }
104
+ }, this.timeout);
105
+ }
106
+
107
+ try {
108
+ const result = await graph.callback({ nodeId: this.nodeId, retry: this.retryCount, params: this.params, payload: this.payload(graph) });
109
+ if (this.transactionId !== transactionId) {
110
+ console.log("****** tid mismatch (success)");
111
+ return;
112
+ }
113
+ this.state = NodeState.Completed;
114
+ this.result = result;
115
+ this.waitlist.forEach((nodeId) => {
116
+ const node = graph.nodes[nodeId];
117
+ node.removePending(this.nodeId, graph);
118
+ });
119
+ graph.removeRunning(this);
120
+ } catch (e) {
121
+ if (this.transactionId !== transactionId) {
122
+ console.log("****** tid mismatch (failed)");
123
+ return;
124
+ }
125
+ this.retry(graph, NodeState.Failed, {});
126
+ }
127
+ }
128
+ }
129
+
130
+ type GraphNodes = Record<string, Node>;
131
+
132
+ export class GraphAI {
133
+ public nodes: GraphNodes;
134
+ public callback: NodeExecute;
135
+ private runningNodes: Set<string>;
136
+ private onComplete: () => void;
137
+
138
+ constructor(data: GraphData, callback: NodeExecute) {
139
+ this.callback = callback;
140
+ this.runningNodes = new Set<string>();
141
+ this.onComplete = () => {};
142
+ this.nodes = Object.keys(data.nodes).reduce((nodes: GraphNodes, nodeId: string) => {
143
+ nodes[nodeId] = new Node(nodeId, data.nodes[nodeId]);
144
+ return nodes;
145
+ }, {});
146
+
147
+ // Generate the waitlist for each node
148
+ Object.keys(this.nodes).forEach((nodeId) => {
149
+ const node = this.nodes[nodeId];
150
+ node.pendings.forEach((pending) => {
151
+ const node2 = this.nodes[pending];
152
+ node2.waitlist.add(nodeId);
153
+ });
154
+ });
155
+ }
156
+
157
+ public asString() {
158
+ return Object.keys(this.nodes)
159
+ .map((nodeId) => {
160
+ return this.nodes[nodeId].asString();
161
+ })
162
+ .join("\n");
163
+ }
164
+
165
+ public async run() {
166
+ // Nodes without pending data should run immediately.
167
+ Object.keys(this.nodes).forEach((nodeId) => {
168
+ const node = this.nodes[nodeId];
169
+ node.executeIfReady(this);
170
+ });
171
+
172
+ return new Promise((resolve, reject) => {
173
+ this.onComplete = () => {
174
+ const results = Object.keys(this.nodes).reduce((results: ResultData, nodeId) => {
175
+ results[nodeId] = this.nodes[nodeId].result;
176
+ return results;
177
+ }, {});
178
+ resolve(results);
179
+ };
180
+ });
181
+ }
182
+
183
+ public addRunning(node: Node) {
184
+ this.runningNodes.add(node.nodeId);
185
+ }
186
+
187
+ public removeRunning(node: Node) {
188
+ this.runningNodes.delete(node.nodeId);
189
+ if (this.runningNodes.size == 0) {
190
+ this.onComplete();
191
+ }
192
+ }
193
+ }
package/src/index.ts CHANGED
@@ -1,3 +1,3 @@
1
- import { GraphAI } from "./flow";
1
+ import { GraphAI } from "./graphai";
2
2
 
3
- export { GraphAI };
3
+ export { GraphAI };
@@ -12,6 +12,12 @@ nodes:
12
12
  inputs: [node1, node2]
13
13
  retry: 2
14
14
  node4:
15
+ timeout: 200
16
+ retry: 2
15
17
  params:
16
- delay: 100
18
+ delay: 300
17
19
  inputs: [node3]
20
+ node5:
21
+ params:
22
+ delay: 100
23
+ inputs: [node4]
@@ -0,0 +1,12 @@
1
+ nodes:
2
+ node1:
3
+ params:
4
+ prompt: Come up with ten business ideas for AI startup
5
+ node2:
6
+ inputs: [node1]
7
+ params:
8
+ prompt: Please evaluate following business ideas. ${node1}
9
+ node3:
10
+ inputs: [node1, node2]
11
+ params:
12
+ prompt: Please pick the winner of this business idea contest. ${node1} ${node2}
@@ -1,39 +1,57 @@
1
1
  import path from "path";
2
- import { GraphAI, FlowCommand } from "../src/flow";
3
-
2
+ import { GraphAI, NodeExecuteContext } from "../src/graphai";
4
3
  import { readManifestData } from "../src/file_utils";
4
+ import { sleep } from "./utils";
5
+
6
+ import test from "node:test";
7
+ import assert from "node:assert";
8
+
9
+ const testFunction = async (context: NodeExecuteContext) => {
10
+ const { nodeId, retry, params } = context;
11
+ console.log("executing", nodeId, params);
12
+ await sleep(params.delay / (retry + 1));
5
13
 
6
- const test = async (file: string) => {
14
+ if (params.fail && retry < 2) {
15
+ const result = { [nodeId]: "failed" };
16
+ console.log("failed", nodeId, result, retry);
17
+ throw new Error("Intentional Failure");
18
+ } else {
19
+ const result = { [nodeId]: "output" };
20
+ console.log("completing", nodeId, result);
21
+ return result;
22
+ }
23
+ };
24
+
25
+ const runTest = async (file: string) => {
7
26
  const file_path = path.resolve(__dirname) + file;
8
27
  const graph_data = readManifestData(file_path);
9
- return new Promise((resolve, reject) => {
10
- const graph = new GraphAI(graph_data, async (params) => {
11
- if (params.cmd == FlowCommand.Execute) {
12
- const node = params.node;
13
- console.log("executing", node, params.params, params.payload)
14
- setTimeout(() => {
15
- if (params.params.fail && params.retry < 2) {
16
- const result = { [node]:"failed" };
17
- console.log("failed", node, result, params.retry)
18
- graph.reportError(node, result);
19
- } else {
20
- const result = { [node]:"output" };
21
- console.log("completing", node, result)
22
- graph.feed(node, result)
23
- }
24
- }, params.params.delay);
25
- } else if (params.cmd == FlowCommand.OnComplete) {
26
- resolve(graph);
27
- }
28
- });
29
- graph.run();
28
+
29
+ const graph = new GraphAI(graph_data, testFunction);
30
+
31
+ const results = await graph.run();
32
+ console.log(results);
33
+ return results;
34
+ };
35
+
36
+ test("test sample1", async () => {
37
+ const result = await runTest("/graphs/sample1.yml");
38
+ assert.deepStrictEqual(result, {
39
+ node1: { node1: "output" },
40
+ node2: { node2: "output" },
41
+ node3: { node3: "output" },
42
+ node4: { node4: "output" },
43
+ node5: { node5: "output" },
30
44
  });
31
- }
45
+ });
32
46
 
33
- const main = async () => {
34
- await test("/graphs/sample1.yml");
35
- console.log("COMPLETE 1");
36
- await test("/graphs/sample2.yml");
47
+ test("test sample1", async () => {
48
+ const result = await runTest("/graphs/sample2.yml");
49
+ assert.deepStrictEqual(result, {
50
+ node1: { node1: "output" },
51
+ node2: { node2: "output" },
52
+ node3: { node3: "output" },
53
+ node4: { node4: "output" },
54
+ node5: { node5: "output" },
55
+ });
37
56
  console.log("COMPLETE 2");
38
- };
39
- main();
57
+ });
@@ -0,0 +1,37 @@
1
+ import path from "path";
2
+ import { GraphAI, NodeExecuteContext } from "../src/graphai";
3
+ import { ChatSession, ChatConfig } from "slashgpt";
4
+ import { readManifestData } from "../src/file_utils";
5
+
6
+ const config = new ChatConfig(path.resolve(__dirname));
7
+
8
+ const testFunction = async (context: NodeExecuteContext) => {
9
+ console.log("executing", context.nodeId, context.params, context.payload);
10
+ const session = new ChatSession(config, context.params.manifest ?? {});
11
+ const prompt = Object.keys(context.payload).reduce((prompt, key) => {
12
+ return prompt.replace("${" + key + "}", context.payload[key]["answer"]);
13
+ }, context.params.prompt);
14
+ session.append_user_question(prompt);
15
+
16
+ await session.call_loop(() => {});
17
+ const message = session.history.last_message();
18
+ if (message === undefined) {
19
+ throw new Error("No message in the history");
20
+ }
21
+ const result = { answer: message.content };
22
+ console.log(result);
23
+ return result;
24
+ };
25
+
26
+ const test = async (file: string) => {
27
+ const file_path = path.resolve(__dirname) + file;
28
+ const graph_data = readManifestData(file_path);
29
+ const graph = new GraphAI(graph_data, testFunction);
30
+ await graph.run();
31
+ };
32
+
33
+ const main = async () => {
34
+ await test("/graphs/sample3.yml");
35
+ console.log("COMPLETE 1");
36
+ };
37
+ main();
package/tests/utils.ts ADDED
@@ -0,0 +1,3 @@
1
+ export const sleep = async (milliseconds: number) => {
2
+ return await new Promise((resolve) => setTimeout(resolve, milliseconds));
3
+ };
package/src/flow.ts DELETED
@@ -1,152 +0,0 @@
1
- export enum NodeState {
2
- Waiting,
3
- Executing,
4
- Failed,
5
- Completed
6
- }
7
-
8
- type NodeData = {
9
- inputs: undefined | Array<string>;
10
- params: any; // Application specific parameters
11
- retry: undefined | number;
12
- }
13
-
14
- type FlowData = {
15
- nodes: Record<string, NodeData>;
16
- };
17
-
18
- export enum FlowCommand {
19
- Log,
20
- Execute,
21
- OnComplete
22
- }
23
-
24
- type FlowCallback = (params: Record<string, any>) => void;
25
-
26
- class Node {
27
- public key: string;
28
- public inputs: Array<string>;
29
- public pendings: Set<string>;
30
- public params: any;
31
- public waitlist: Set<string>;
32
- public state: NodeState;
33
- public result: Record<string, any>;
34
- public retryLimit: number;
35
- public retryCount: number;
36
- constructor(key: string, data: NodeData) {
37
- this.key = key;
38
- this.inputs = data.inputs ?? [];
39
- this.pendings = new Set(this.inputs);
40
- this.params = data.params;
41
- this.waitlist = new Set<string>();
42
- this.state = NodeState.Waiting;
43
- this.result = {};
44
- this.retryLimit = data.retry ?? 0;
45
- this.retryCount = 0;
46
- }
47
-
48
- public asString() {
49
- return `${this.key}: ${this.state} ${[...this.waitlist]}`
50
- }
51
-
52
- public complete(result: Record<string, any>, nodes: Record<string, Node>, graph: GraphAI) {
53
- this.state = NodeState.Completed;
54
- this.result = result;
55
- this.waitlist.forEach(key => {
56
- const node = nodes[key];
57
- node.removePending(this.key, graph);
58
- });
59
- graph.remove(this);
60
- }
61
-
62
- public reportError(result: Record<string, any>, nodes: Record<string, Node>, graph: GraphAI) {
63
- this.state = NodeState.Failed;
64
- this.result = result;
65
- if (this.retryCount < this.retryLimit) {
66
- this.retryCount++;
67
- this.state = NodeState.Executing;
68
- graph.callback({cmd: FlowCommand.Execute, node: this.key, params: this.params, retry: this.retryCount, payload: this.payload(graph) });
69
- } else {
70
- graph.remove(this);
71
- }
72
- }
73
-
74
- public removePending(key: string, graph: GraphAI) {
75
- this.pendings.delete(key);
76
- this.executeIfReady(graph);
77
- }
78
-
79
- public payload(graph: GraphAI) {
80
- const foo: Record<string, any> = {};
81
- return this.inputs.reduce((payload, key) => {
82
- payload[key] = graph.nodes[key].result;
83
- return payload;
84
- }, foo);
85
- }
86
-
87
- public executeIfReady(graph: GraphAI) {
88
- if (this.pendings.size == 0) {
89
- this.state = NodeState.Executing;
90
- graph.add(this);
91
- graph.callback({cmd: FlowCommand.Execute, node: this.key, params: this.params, retry: 0, payload: this.payload(graph) });
92
- }
93
- }
94
- }
95
-
96
- export class GraphAI {
97
- public nodes: Record<string, Node>
98
- public callback: FlowCallback;
99
- private runningNodes: Set<string>;
100
-
101
- constructor(data: FlowData, callback: FlowCallback) {
102
- this.callback = callback;
103
- this.runningNodes = new Set<string>();
104
- const foo: Record<string, Node> = {}; // HACK: Work around
105
- this.nodes = Object.keys(data.nodes).reduce((nodes, key) => {
106
- nodes[key] = new Node(key, data.nodes[key]);
107
- return nodes;
108
- }, foo);
109
-
110
- // Generate the waitlist for each node
111
- Object.keys(this.nodes).forEach(key => {
112
- const node = this.nodes[key];
113
- node.pendings.forEach(pending => {
114
- const node2 = this.nodes[pending]
115
- node2.waitlist.add(key);
116
- });
117
- });
118
- }
119
-
120
- public asString() {
121
- return Object.keys(this.nodes).map((key) => { return this.nodes[key].asString() }).join('\n');
122
- }
123
-
124
- public run() {
125
- // Nodes without pending data should run immediately.
126
- Object.keys(this.nodes).forEach(key => {
127
- const node = this.nodes[key];
128
- node.executeIfReady(this);
129
- });
130
- }
131
-
132
- public feed(key: string, result: Record<string, any>) {
133
- const node = this.nodes[key];
134
- node.complete(result, this.nodes, this);
135
- }
136
-
137
- public reportError(key: string, result: Record<string, any>) {
138
- const node = this.nodes[key];
139
- node.reportError(result, this.nodes, this);
140
- }
141
-
142
- public add(node: Node) {
143
- this.runningNodes.add(node.key);
144
- }
145
-
146
- public remove(node: Node) {
147
- this.runningNodes.delete(node.key);
148
- if (this.runningNodes.size == 0) {
149
- this.callback({cmd: FlowCommand.OnComplete});
150
- }
151
- }
152
- }