graphai 0.0.4 → 0.0.6

package/samples/agents/arxiv_agent.ts

@@ -0,0 +1,46 @@
+import search from "arXiv-api-ts";
+
+import { AgentFunction } from "@/graphai";
+
+type arxivData = { id: string; title: string; summary: string };
+
+const search_arxiv_papers = async (keywords: string[], limit = 10) => {
+  const includes = keywords.map((k) => {
+    return { name: k };
+  });
+  const papers = await search({
+    searchQueryParams: [
+      {
+        include: includes,
+      },
+    ],
+    sortBy: "lastUpdatedDate",
+    sortOrder: "descending",
+    start: 0,
+    maxResults: limit,
+  });
+  return papers.entries || [];
+};
+
+export const arxivAgent: AgentFunction<{ keywords: string[]; limit: number }, arxivData[]> = async (context) => {
+  const { keywords, limit } = context.params;
+  const arxivResult = await search_arxiv_papers(keywords, limit);
+  // console.log("executing", arxivResult, context.params.keywords);
+
+  const result = arxivResult.map((r: any) => {
+    const { id, title, summary } = r;
+    return { id, title, summary };
+  });
+  return result;
+};
+
+export const arxiv2TextAgent: AgentFunction<{}, string, string[]> = async (context) => {
+  const result = (context?.payload?.inputData || [])
+    .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/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/{tests → samples}/sample_gpt.ts

@@ -1,12 +1,12 @@
 import path from "path";
-import { GraphAI, NodeExecuteContext } from "../src/graphai";
-import { ChatSession, ChatConfig } from "slashgpt";
-import { readManifestData } from "./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 = async (context: NodeExecuteContext<Record<string, string>>) => {
-  console.log("executing", context.nodeId, context.params, context.payload);
+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"]);
@@ -19,19 +19,19 @@ const testFunction = async (context: NodeExecuteContext<Record<string, string>>)
     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/sample3.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,68 +1,90 @@
-import { AssertionError } from "assert";
-
 export enum NodeState {
   Waiting,
   Executing,
   Failed,
   TimedOut,
   Completed,
+  Injected,
+  Dispatched,
 }
 type ResultData<ResultType = Record<string, any>> = ResultType | undefined;
 type ResultDataDictonary<ResultType = Record<string, any>> = Record<string, ResultData<ResultType>>;
 
-export type NodeDataParams = Record<string, any>; // App-specific parameters
+export type NodeDataParams<ParamsType = Record<string, any>> = ParamsType; // 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;
+  concurrency?: number;
+};
+
+export type TransactionLog = {
+  nodeId: string;
+  state: NodeState;
+  startTime: number;
+  endTime?: number;
+  retryCount: number;
+  agentId?: string;
+  params?: NodeDataParams;
+  payload?: ResultDataDictonary<ResultData>;
+  errorMessage?: string;
+  result?: ResultData;
 };
 
-export type NodeExecuteContext<ResultType> = {
+export type AgentFunctionContext<ParamsType, ResultType, PreviousResultType> = {
   nodeId: string;
   retry: number;
-  params: NodeDataParams;
-  payload: ResultDataDictonary<ResultType>;
+  params: NodeDataParams<ParamsType>;
+  payload: ResultDataDictonary<PreviousResultType>;
 };
 
-type NodeExecute<ResultType> = (context: NodeExecuteContext<ResultType>) => Promise<ResultData<ResultType>>;
+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<ResultType = Record<string, 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<ResultType>;
+  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 timeout?: number; // msec
+  public error?: Error;
+  public source: boolean;
+  public dispatch?: Record<string, string>; // outputId to nodeId mapping
 
-  private graph: GraphAI<ResultType>;
+  private graph: GraphAI;
 
-  constructor(nodeId: string, data: NodeData, graph: GraphAI<ResultType>) {
+  constructor(nodeId: string, data: NodeData, graph: GraphAI) {
     this.nodeId = nodeId;
     this.inputs = data.inputs ?? [];
+    this.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;
   }
 
@@ -70,101 +92,171 @@ class Node<ResultType = Record<string, any>> {
     return `${this.nodeId}: ${this.state} ${[...this.waitlist]}`;
   }
 
-  private retry(state: NodeState, result: ResultData<ResultType>) {
+  private retry(state: NodeState, error: Error) {
     if (this.retryCount < this.retryLimit) {
       this.retryCount++;
       this.execute();
     } else {
       this.state = state;
-      this.result = result;
+      this.result = undefined;
+      this.error = error;
+      this.transactionId = 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<ResultType>, nodeId) => {
-      results[nodeId] = this.graph.nodes[nodeId].result;
+    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) {
+    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(),
+      agentId: this.agentId,
+      params: this.params,
+      payload,
+    };
+    this.graph.appendLog(log);
     this.state = NodeState.Executing;
-    const transactionId = Date.now();
+    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("*** timeout", this.timeout);
-          this.retry(NodeState.TimedOut, undefined);
+          console.log(`-- ${this.nodeId}: timeout ${this.timeout}`);
+          log.errorMessage = "Timeout";
+          log.state = NodeState.TimedOut;
+          log.endTime = Date.now();
+          this.retry(NodeState.TimedOut, Error("Timeout"));
         }
       }, this.timeout);
     }
 
     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("****** transactionId mismatch (success)");
+        console.log(`-- ${this.nodeId}: transactionId mismatch`);
         return;
       }
-      this.state = NodeState.Completed;
-      this.result = result;
-      this.waitlist.forEach((nodeId) => {
-        const node = this.graph.nodes[nodeId];
-        node.removePending(this.nodeId);
-      });
+
+      log.endTime = Date.now();
+      log.result = result;
+
+      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 (e) {
+    } catch (error) {
       if (this.transactionId !== transactionId) {
-        console.log("****** transactionId mismatch (failed)");
+        console.log(`-- ${this.nodeId}: transactionId mismatch(error)`);
         return;
       }
-      this.retry(NodeState.Failed, undefined);
+      log.state = NodeState.Failed;
+      log.endTime = Date.now();
+      if (error instanceof Error) {
+        log.errorMessage = error.message;
+        this.retry(NodeState.Failed, error);
+      } else {
+        console.error(`-- ${this.nodeId}: Unexpecrted error was caught`);
+        log.errorMessage = "Unknown";
+        this.retry(NodeState.Failed, Error("Unknown"));
+      }
     }
   }
 }
 
-type GraphNodes<ResultType> = Record<string, Node<ResultType>>;
+type GraphNodes = Record<string, Node>;
 
-type NodeExecuteDictonary<ResultType> = Record<string, NodeExecute<ResultType>>;
+const defaultConcurrency = 8;
 
-export class GraphAI<ResultType = Record<string, any>> {
-  public nodes: GraphNodes<ResultType>;
-  public callbackDictonary: NodeExecuteDictonary<ResultType>;
-  private runningNodes: Set<string>;
-  private nodeQueue: Array<Node<ResultType>>;
+export class GraphAI {
+  public nodes: GraphNodes;
+  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<ResultType> | NodeExecute<ResultType>) {
+  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 ?? 2;
-    this.runningNodes = new Set<string>();
-    this.nodeQueue = [];
-    this.onComplete = () => {};
-    this.nodes = Object.keys(data.nodes).reduce((nodes: GraphNodes<ResultType>, nodeId: string) => {
-      nodes[nodeId] = new Node<ResultType>(nodeId, data.nodes[nodeId], this);
+    this.concurrency = data.concurrency ?? defaultConcurrency;
+    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;
     }, {});
 
@@ -178,9 +270,9 @@ export class GraphAI<ResultType = Record<string, any>> {
     });
   }
 
-  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"];
   }
@@ -193,7 +285,31 @@ export class GraphAI<ResultType = Record<string, any>> {
       .join("\n");
   }
 
+  public results() {
+    return Object.keys(this.nodes).reduce((results: ResultDataDictonary, nodeId) => {
+      const node = this.nodes[nodeId];
+      if (node.result !== undefined) {
+        results[nodeId] = node.result;
+      }
+      return results;
+    }, {});
+  }
+
+  public errors() {
+    return Object.keys(this.nodes).reduce((errors: Record<string, Error>, nodeId) => {
+      const node = this.nodes[nodeId];
+      if (node.error !== undefined) {
+        errors[nodeId] = node.error;
+      }
+      return errors;
+    }, {});
+  }
+
   public async run() {
+    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];
@@ -202,21 +318,24 @@ export class GraphAI<ResultType = Record<string, any>> {
 
     return new Promise((resolve, reject) => {
       this.onComplete = () => {
-        const results = Object.keys(this.nodes).reduce((results: ResultDataDictonary<ResultType>, nodeId) => {
-          results[nodeId] = this.nodes[nodeId].result;
-          return results;
-        }, {});
-        resolve(results);
+        this.isRunning = false;
+        const errors = this.errors();
+        const nodeIds = Object.keys(errors);
+        if (nodeIds.length > 0) {
+          reject(errors[nodeIds[0]]);
+        } else {
+          resolve(this.results());
+        }
       };
     });
   }
 
-  private runNode(node: Node<ResultType>) {
+  private runNode(node: Node) {
     this.runningNodes.add(node.nodeId);
     node.execute();
   }
 
-  public pushQueue(node: Node<ResultType>) {
+  public pushQueue(node: Node) {
     if (this.runningNodes.size < this.concurrency) {
       this.runNode(node);
     } else {
@@ -224,7 +343,7 @@ export class GraphAI<ResultType = Record<string, any>> {
     }
   }
 
-  public removeRunning(node: Node<ResultType>) {
+  public removeRunning(node: Node) {
     this.runningNodes.delete(node.nodeId);
     if (this.nodeQueue.length > 0) {
       const n = this.nodeQueue.shift();
@@ -236,4 +355,21 @@ export class GraphAI<ResultType = Record<string, any>> {
       this.onComplete();
     }
   }
+
+  public appendLog(log: TransactionLog) {
+    this.logs.push(log);
+  }
+
+  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;
+  }
+};