graphai 0.0.5 → 0.0.7

package/samples/agents/slashgpt_agent.ts

@@ -0,0 +1,21 @@
+import path from "path";
+import { AgentFunction } from "@/graphai";
+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 ?? {});
+
+  const prompt = [context.params?.prompt, context.payload.inputData].join("\n\n");
+  session.append_user_question(prompt);
+
+  await session.call_loop(() => {});
+  const message = session.history.last_message();
+  if (message === undefined) {
+    throw new Error("No message in the history");
+  }
+  const result = { answer: message.content };
+  return result;
+};

package/samples/express.ts

@@ -0,0 +1,47 @@
+//  npx ts-node samples/express.ts
+import { GraphAI, AgentFunction } from "@/graphai";
+
+import express from "express";
+
+const app = express();
+
+const graphAISample = async (req: express.Request, res: express.Response) => {
+  const graph_data = {
+    nodes: {
+      node1: {
+        params: {},
+      },
+    },
+    concurrency: 8,
+  };
+  const testFunction: AgentFunction<Record<string, string>> = async (context) => {
+    console.log("hello");
+    return {};
+  };
+  const graph = new GraphAI(graph_data, testFunction);
+  const response = await graph.run();
+  res.json({ result: response });
+  res.end();
+};
+
+const hello = async (req: express.Request, res: express.Response) => {
+  const { params, query } = req;
+  res.json({
+    result: [
+      {
+        message: "hello",
+        params,
+        query,
+      },
+    ],
+  });
+  res.end();
+};
+
+app.use(express.json());
+app.get("/", hello);
+app.get("/mock", graphAISample);
+
+const server = app.listen(8080, () => {
+  console.log("Running Server");
+});

package/samples/graphs/arxiv.yml

@@ -0,0 +1,29 @@
+nodes:
+  searchArxiv:
+    params:
+      keywords:
+        - llm
+        - gpt
+      limit: 10
+    functionName: arxivAgent
+  arxiv2TextAgent:
+    inputs: [searchArxiv]
+    functionName: arxiv2TextAgent
+    payloadMapping:
+      searchArxiv: inputData
+  slashGPTAgent:
+    inputs: [arxiv2TextAgent]
+    payloadMapping:
+      arxiv2TextAgent: inputData
+    functionName: slashGPTAgent
+    params:
+      prompt: |
+        与えられたそれぞれの論文の要点をまとめ、以下の項目で日本語で出力せよ。それぞれの項目は最大でも180文字以内に要約せよ。
+        ```
+        論文名:タイトルの日本語訳
+        キーワード:この論文のキーワード
+        課題:この論文が解決する課題
+        手法:この論文が提案する手法,
+        結果:提案手法によって得られた結果
+        ```
+      

package/samples/sample_gpt.ts

@@ -1,11 +1,11 @@
 import path from "path";
-import { GraphAI, NodeExecute } from "../src/graphai";
-import { ChatSession, ChatConfig } from "slashgpt";
-import { readManifestData } from "../tests/file_utils";
+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 testFunction: NodeExecute<Record<string, string>> = async (context) => {
+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) => {
@@ -19,19 +19,19 @@ const testFunction: NodeExecute<Record<string, string>> = async (context) => {
     throw new Error("No message in the history");
   }
   const result = { answer: message.content };
-  console.log(result);
   return result;
 };
 
-const test = async (file: string) => {
+const runAgent = async (file: string) => {
   const file_path = path.resolve(__dirname) + file;
-  const graph_data = readManifestData(file_path);
-  const graph = new GraphAI(graph_data, testFunction);
-  await graph.run();
+  const graph_data = readGraphaiData(file_path);
+  const graph = new GraphAI(graph_data, slashGPTAgent);
+  const result = await graph.run();
+  console.log(result);
 };
 
 const main = async () => {
-  await test("/graphs/slash_gpt.yml");
+  await runAgent("/graphs/slash_gpt.yml");
   console.log("COMPLETE 1");
 };
 main();

package/samples/sample_paper_ai.ts

@@ -0,0 +1,26 @@
+import path from "path";
+import search from "arXiv-api-ts";
+
+import { GraphAI, AgentFunction } from "@/graphai";
+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 result = await graph.run();
+  console.log(result);
+};
+
+const main = async () => {
+  await runAgent("/graphs/arxiv.yml");
+  console.log("COMPLETE 1");
+};
+main();

package/src/graphai.ts

@@ -1,81 +1,90 @@
-import { AssertionError } from "assert";
-
 export enum NodeState {
-  Waiting,
-  Executing,
-  Failed,
-  TimedOut,
-  Completed,
+  Waiting = "waiting",
+  Executing = "executing",
+  Failed = "failed",
+  TimedOut = "timed-out",
+  Completed = "completed",
+  Injected = "injected",
+  Dispatched = "dispatched",
 }
 type ResultData<ResultType = Record<string, any>> = ResultType | undefined;
 type ResultDataDictonary<ResultType = Record<string, any>> = Record<string, ResultData<ResultType>>;
 
-export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType; // App-specific parameters
+export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType; // Agent-specific parameters
 
 type NodeData = {
-  inputs: undefined | Array<string>;
+  inputs?: Array<string>;
   params: NodeDataParams;
-  retry: undefined | number;
-  timeout: undefined | number; // msec
-  functionName: undefined | string;
+  payloadMapping?: Record<string, string>;
+  retry?: number;
+  timeout?: number; // msec
+  agentId?: string;
+  source?: boolean;
+  dispatch?: Record<string, string>; // route to node
 };
 
-type GraphData = {
+export type GraphData = {
   nodes: Record<string, NodeData>;
-  concurrency: number;
-};
-
-type NodeExecuteContext<ResultType, ParamsType> = {
-  nodeId: string;
-  retry: number;
-  params: NodeDataParams<ParamsType>;
-  payload: ResultDataDictonary<ResultType>;
+  concurrency?: number;
 };
 
 export type TransactionLog = {
   nodeId: string;
   state: NodeState;
-  startTime: undefined | number;
-  endTime: undefined | number;
+  startTime: number;
+  endTime?: number;
   retryCount: number;
-  error: undefined | Error;
+  agentId?: string;
+  params?: NodeDataParams;
+  payload?: ResultDataDictonary<ResultData>;
+  errorMessage?: string;
   result?: ResultData;
 };
 
-export type NodeExecute<ResultType = Record<string, any>, ParamsType = Record<string, any>> = (
-  context: NodeExecuteContext<ResultType, ParamsType>,
+export type AgentFunctionContext<ParamsType, ResultType, PreviousResultType> = {
+  nodeId: string;
+  retry: number;
+  params: NodeDataParams<ParamsType>;
+  payload: ResultDataDictonary<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>>;
+
 class Node {
   public nodeId: string;
-  public params: NodeDataParams; // App-specific parameters
+  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: Set<string>; // List of nodes which need data from this node.
-  public state: NodeState;
-  public functionName: string;
-  public result: ResultData;
+  public waitlist = new Set<string>(); // List of nodes which need data from this node.
+  public state = NodeState.Waiting;
+  public agentId: string;
+  public result: ResultData = undefined;
   public retryLimit: number;
-  public retryCount: number;
+  public retryCount: number = 0;
   public transactionId: undefined | number; // To reject callbacks from timed-out transactions
-  public timeout: number; // msec
-  public error: undefined | Error;
+  public timeout?: number; // msec
+  public error?: Error;
+  public source: boolean;
+  public dispatch?: Record<string, string>; // outputId to nodeId mapping
 
   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.waitlist = new Set<string>();
-    this.state = NodeState.Waiting;
-    this.functionName = data.functionName ?? "default";
-    this.result = undefined;
+    this.agentId = data.agentId ?? "default";
     this.retryLimit = data.retry ?? 0;
-    this.retryCount = 0;
-    this.timeout = data.timeout ?? 0;
-
+    this.timeout = data.timeout;
+    this.source = data.source === true;
+    this.dispatch = data.dispatch;
     this.graph = graph;
   }
 
@@ -91,48 +100,82 @@ class Node {
       this.state = state;
       this.result = undefined;
       this.error = error;
-      this.transactionId = 0; // This is necessary for timeout case
+      this.transactionId = undefined; // This is necessary for timeout case
       this.graph.removeRunning(this);
     }
   }
 
   public removePending(nodeId: string) {
     this.pendings.delete(nodeId);
-    this.pushQueueIfReady();
+    if (this.graph.isRunning) {
+      this.pushQueueIfReady();
+    }
   }
 
   public payload() {
     return this.inputs.reduce((results: ResultDataDictonary, nodeId) => {
-      results[nodeId] = this.graph.nodes[nodeId].result;
+      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) {
+    if (this.pendings.size === 0 && !this.source) {
       this.graph.pushQueue(this);
     }
   }
 
+  public injectResult(result: ResultData) {
+    if (this.source) {
+      const log: TransactionLog = {
+        nodeId: this.nodeId,
+        retryCount: this.retryCount,
+        state: NodeState.Injected,
+        startTime: Date.now(),
+      };
+      log.endTime = log.startTime;
+      this.graph.appendLog(log);
+      this.setResult(result, NodeState.Injected);
+    } else {
+      console.error("- injectResult called on non-source node.", this.nodeId);
+    }
+  }
+
+  private setResult(result: ResultData, state: NodeState) {
+    this.state = state;
+    this.result = result;
+    this.waitlist.forEach((nodeId) => {
+      const node = this.graph.nodes[nodeId];
+      // Todo: Avoid running before Run()
+      node.removePending(this.nodeId);
+    });
+  }
+
   public async execute() {
+    const payload = this.payload();
     const log: TransactionLog = {
       nodeId: this.nodeId,
       retryCount: this.retryCount,
       state: NodeState.Executing,
       startTime: Date.now(),
-      endTime: undefined,
-      error: undefined,
+      agentId: this.agentId,
+      params: this.params,
+      payload,
     };
     this.graph.appendLog(log);
     this.state = NodeState.Executing;
     const transactionId = log.startTime;
     this.transactionId = transactionId;
 
-    if (this.timeout > 0) {
+    if (this.timeout && this.timeout > 0) {
       setTimeout(() => {
         if (this.state === NodeState.Executing && this.transactionId === transactionId) {
           console.log(`-- ${this.nodeId}: timeout ${this.timeout}`);
-          log.error = Error("Timeout");
+          log.errorMessage = "Timeout";
           log.state = NodeState.TimedOut;
           log.endTime = Date.now();
           this.retry(NodeState.TimedOut, Error("Timeout"));
@@ -141,26 +184,34 @@ class Node {
     }
 
     try {
-      const callback = this.graph.getCallback(this.functionName);
+      const callback = this.graph.getCallback(this.agentId);
       const result = await callback({
         nodeId: this.nodeId,
         retry: this.retryCount,
         params: this.params,
-        payload: this.payload(),
+        payload,
       });
       if (this.transactionId !== transactionId) {
         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) => {
-        const node = this.graph.nodes[nodeId];
-        node.removePending(this.nodeId);
-      });
+
+      const dispatch = this.dispatch;
+      if (dispatch !== undefined) {
+        Object.keys(result).forEach((outputId) => {
+          const nodeId = dispatch[outputId];
+          this.graph.injectResult(nodeId, result[outputId]);
+        });
+        log.state = NodeState.Dispatched;
+        this.state = NodeState.Dispatched;
+        this.graph.removeRunning(this);
+        return;
+      }
+      log.state = NodeState.Completed;
+      this.setResult(result, NodeState.Completed);
       this.graph.removeRunning(this);
     } catch (error) {
       if (this.transactionId !== transactionId) {
@@ -170,11 +221,11 @@ class Node {
       log.state = NodeState.Failed;
       log.endTime = Date.now();
       if (error instanceof Error) {
-        log.error = error;
+        log.errorMessage = error.message;
         this.retry(NodeState.Failed, error);
       } else {
         console.error(`-- ${this.nodeId}: Unexpecrted error was caught`);
-        log.error = Error("Unknown");
+        log.errorMessage = "Unknown";
         this.retry(NodeState.Failed, Error("Unknown"));
       }
     }
@@ -183,28 +234,27 @@ class Node {
 
 type GraphNodes = Record<string, Node>;
 
-type NodeExecuteDictonary = Record<string, NodeExecute>;
-
 const defaultConcurrency = 8;
 
 export class GraphAI {
   public nodes: GraphNodes;
-  public callbackDictonary: NodeExecuteDictonary;
-  private runningNodes: Set<string>;
-  private nodeQueue: Array<Node>;
+  public callbackDictonary: AgentFunctionDictonary;
+  public isRunning = false;
+  private runningNodes = new Set<string>();
+  private nodeQueue: Array<Node> = [];
   private onComplete: () => void;
   private concurrency: number;
   private logs: Array<TransactionLog> = [];
 
-  constructor(data: GraphData, callbackDictonary: NodeExecuteDictonary | NodeExecute) {
+  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.concurrency = data.concurrency ?? defaultConcurrency;
-    this.runningNodes = new Set<string>();
-    this.nodeQueue = [];
-    this.onComplete = () => {};
+    this.onComplete = () => {
+      console.error("-- SOMETHING IS WRONG: onComplete is called without run()");
+    };
     this.nodes = Object.keys(data.nodes).reduce((nodes: GraphNodes, nodeId: string) => {
       nodes[nodeId] = new Node(nodeId, data.nodes[nodeId], this);
       return nodes;
@@ -220,9 +270,9 @@ export class GraphAI {
     });
   }
 
-  public getCallback(functionName: string) {
-    if (functionName && this.callbackDictonary[functionName]) {
-      return this.callbackDictonary[functionName];
+  public getCallback(agentId: string) {
+    if (agentId && this.callbackDictonary[agentId]) {
+      return this.callbackDictonary[agentId];
     }
     return this.callbackDictonary["default"];
   }
@@ -256,6 +306,10 @@ export class GraphAI {
   }
 
   public async run() {
+    if (this.isRunning) {
+      console.error("-- Already Running");
+    }
+    this.isRunning = true;
     // Nodes without pending data should run immediately.
     Object.keys(this.nodes).forEach((nodeId) => {
       const node = this.nodes[nodeId];
@@ -264,6 +318,7 @@ export class GraphAI {
 
     return new Promise((resolve, reject) => {
       this.onComplete = () => {
+        this.isRunning = false;
         const errors = this.errors();
         const nodeIds = Object.keys(errors);
         if (nodeIds.length > 0) {
@@ -308,4 +363,13 @@ export class GraphAI {
   public transactionLogs() {
     return this.logs;
   }
+
+  public injectResult(nodeId: string, result: ResultData) {
+    const node = this.nodes[nodeId];
+    if (node) {
+      node.injectResult(result);
+    } else {
+      console.error("-- Invalid nodeId", nodeId);
+    }
+  }
 }

package/tests/agents/agents.ts

@@ -0,0 +1,24 @@
+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;
+  console.log("executing", nodeId);
+  await sleep(params.delay / (retry + 1));
+
+  if (params.fail && retry < 2) {
+    const result = { [nodeId]: "failed" };
+    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;
+      },
+      { [nodeId]: "output" },
+    );
+    console.log("completing", nodeId);
+    return result;
+  }
+};

package/tests/graphai/test_dispatch.ts

@@ -0,0 +1,42 @@
+import { AgentFunction } from "@/graphai";
+import { fileTestRunner } from "~/utils/runner";
+
+import { sleep } from "~/utils/utils";
+
+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;
+  console.log("executing", nodeId);
+  await sleep(params.delay / (retry + 1));
+
+  if (params.fail && retry < 2) {
+    const result = { [nodeId]: "failed" };
+    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;
+      },
+      { [nodeId]: "dispatch" },
+    );
+    console.log("completing", nodeId);
+    return { output1: result };
+  }
+};
+
+test("test dispatch", async () => {
+  const result = await fileTestRunner("/graphs/test_dispatch.yml", { default: testAgent, alt: dispatchAgent });
+  assert.deepStrictEqual(result, {
+    node1: { node1: "output" },
+    node20: { node2: "dispatch" },
+    node3: { node3: "output", node1: "output", node2: "dispatch" },
+    node4: { node4: "output", node3: "output", node1: "output", node2: "dispatch" },
+    node5: { node5: "output", node4: "output", node3: "output", node1: "output", node2: "dispatch" },
+  });
+});

package/tests/graphai/test_http_client.ts

@@ -0,0 +1,40 @@
+import { GraphAI, AgentFunction } from "@/graphai";
+import { graphDataTestRunner } from "~/utils/runner";
+
+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 response = await fetch(params.url);
+  const result = await response.json();
+
+  console.log("completing", nodeId, result);
+  return result;
+};
+
+const graph_data = {
+  nodes: {
+    node1: {
+      params: {
+        url: "http://127.0.0.1:8080/llm.json",
+      },
+    },
+    node2: {
+      params: {
+        url: "http://127.0.0.1:8080/llm2.json",
+      },
+      inputs: ["node1"],
+    },
+  },
+};
+
+test("test sample1", async () => {
+  const result = await graphDataTestRunner("http.log", graph_data, httpClientAgent);
+  assert.deepStrictEqual(result, {
+    node1: { result: true, messages: ["hello"] },
+    node2: { result: true, messages: ["hello2"] },
+  });
+});