@ai.ntellect/core 0.6.7 → 0.6.9

package/create-llm-to-select-multiple-graph copy.ts

@@ -1,5 +1,5 @@
 import { openai } from "@ai-sdk/openai";
-import { streamObject } from "ai";
+import { generateObject } from "ai";
 import { parseEther } from "ethers";
 import { z } from "zod";
 import { GraphEngine } from "./graph/engine";
@@ -10,8 +10,7 @@ import { stringifyZodSchema } from "./utils/stringifiy-zod-schema";
 // ---- I. Create a simple workflow ----
 // 1. Define the context
 type MainContext = {
-  messages: string;
-  reasoning?: boolean;
+  results?: Record<string, CombinedContext>[];
   actions?: Action[];
 };
 
@@ -25,9 +24,7 @@ type ContextA = {
   type: string;
 };
 
-type ContextB = {
-  messages: string;
-};
+type ContextB = {};
 
 // // 2. Define the graph
 
@@ -52,17 +49,18 @@ const graphDefinitionA: GraphDefinition<ContextA> = {
             "Examples networks: ethereum, arbitrum, base. IMPORTANT: You must respect the network name."
           ),
       }),
-      execute: async (params, state) => {
-        console.log("💰 Prepare transaction", { params, state });
-        const networkConfig = networkConfigs[params.chain.toLowerCase()];
+      execute: async (state) => {
+        console.log("💰 Prepare transaction", { state });
+        const networkConfig =
+          networkConfigs[(state.chain as any).toLowerCase()];
 
         if (!networkConfig) {
-          throw new Error(`Network ${params.chain} not found`);
+          throw new Error(`Network ${state.chain} not found`);
         }
 
         return {
-          to: params.to,
-          value: parseEther(params.value).toString(),
+          to: state.to,
+          value: parseEther(state.value).toString(),
           chain: {
             id: networkConfig.id || 0,
             rpc: networkConfig.rpc,
@@ -99,10 +97,7 @@ const graphDefinitionB: GraphDefinition<ContextB> = {
   },
 };
 // 3. Define the initial state
-const initialState: SharedState<MainContext> = {
-  messages: "",
-  reasoning: true,
-};
+const initialState: SharedState<MainContext> = {};
 type CombinedContext = ContextA | ContextB;
 
 const graphMaps: GraphDefinition<CombinedContext>[] = [
@@ -123,8 +118,8 @@ const graphMaps: GraphDefinition<CombinedContext>[] = [
       reasoning: {
         name: "reasoning",
         description: "Reasoning",
-        execute: async (params, state) => {
-          console.log("Reasoning", { params, state });
+        execute: async (state) => {
+          console.log("Reasoning", { state });
           const startNodesSchema = graphMaps.map(
             (graphDefinition) =>
               graphDefinition.nodes[graphDefinition.entryNode]
@@ -137,7 +132,7 @@ const graphMaps: GraphDefinition<CombinedContext>[] = [
         workflows:${stringifyZodSchema(startNodesSchema)}`;
           console.log(system);
 
-          const llm = await streamObject({
+          const llm = await generateObject({
             model: openai("gpt-4o"),
             prompt,
             schema: z.object({
@@ -158,12 +153,6 @@ const graphMaps: GraphDefinition<CombinedContext>[] = [
           });
           console.dir(llm.object, { depth: null });
 
-          for await (const chunk of llm.partialObjectStream) {
-            if (chunk.response) {
-              // console.log(chunk.response);
-            }
-          }
-
           const actions = (await llm.object).actions;
           console.log({ actions });
 
@@ -173,15 +162,12 @@ const graphMaps: GraphDefinition<CombinedContext>[] = [
             actions: selectedWorkflows,
           };
         },
-        condition: (state) => {
-          return state.reasoning === true;
-        },
         relationships: [{ name: "actions" }],
       },
       actions: {
         name: "actions",
         description: "Actions",
-        execute: async (parameters, state) => {
+        execute: async (state) => {
           if (!state.actions) {
             throw new Error("No actions found");
           }
@@ -209,27 +195,25 @@ const graphMaps: GraphDefinition<CombinedContext>[] = [
               selectedGraphs,
               startNodes,
               initialStates,
+              state.actions,
               (graph) => {
                 console.log(`Graph ${graph.name} updated`, graph.getState());
               },
-              (error, nodeName, state) => {
+              (
+                error: Error,
+                nodeName: string,
+                state: SharedState<CombinedContext>
+              ) => {
                 console.error(`Erreur dans ${nodeName}`, error, state);
               }
             );
 
-          console.log({ results });
-
           return {
-            messages: "Hello, world!",
-            results,
+            actions: results,
           };
         },
         condition: (state) => {
-          if (
-            state.reasoning === true &&
-            state.actions &&
-            state.actions.length > 0
-          ) {
+          if (state.actions && state.actions.length > 0) {
             return true;
           }
           return false;
@@ -239,5 +223,15 @@ const graphMaps: GraphDefinition<CombinedContext>[] = [
   };
 
   const reasoningGraph = new GraphEngine(reasoningGraphDefinition);
-  await reasoningGraph.execute(initialState, "reasoning");
+  await reasoningGraph.execute(
+    initialState,
+    "reasoning",
+    (graph) => {
+      console.log("Main graph updated", graph.name);
+      console.dir(graph.getState(), { depth: null });
+    },
+    (error, nodeName, state) => {
+      console.error(`Erreur dans ${nodeName}`, error, state);
+    }
+  );
 })();

package/dist/create-llm-to-select-multiple-graph copy.js

@@ -8,13 +8,6 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
         step((generator = generator.apply(thisArg, _arguments || [])).next());
     });
 };
-var __asyncValues = (this && this.__asyncValues) || function (o) {
-    if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
-    var m = o[Symbol.asyncIterator], i;
-    return m ? m.call(o) : (o = typeof __values === "function" ? __values(o) : o[Symbol.iterator](), i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i);
-    function verb(n) { i[n] = o[n] && function (v) { return new Promise(function (resolve, reject) { v = o[n](v), settle(resolve, reject, v.done, v.value); }); }; }
-    function settle(resolve, reject, d, v) { Promise.resolve(v).then(function(v) { resolve({ value: v, done: d }); }, reject); }
-};
 Object.defineProperty(exports, "__esModule", { value: true });
 const openai_1 = require("@ai-sdk/openai");
 const ai_1 = require("ai");
@@ -43,15 +36,15 @@ const graphDefinitionA = {
                     .string()
                     .describe("Examples networks: ethereum, arbitrum, base. IMPORTANT: You must respect the network name."),
             }),
-            execute: (params, state) => __awaiter(void 0, void 0, void 0, function* () {
-                console.log("💰 Prepare transaction", { params, state });
-                const networkConfig = index_copy_1.networkConfigs[params.chain.toLowerCase()];
+            execute: (state) => __awaiter(void 0, void 0, void 0, function* () {
+                console.log("💰 Prepare transaction", { state });
+                const networkConfig = index_copy_1.networkConfigs[state.chain.toLowerCase()];
                 if (!networkConfig) {
-                    throw new Error(`Network ${params.chain} not found`);
+                    throw new Error(`Network ${state.chain} not found`);
                 }
                 return {
-                    to: params.to,
-                    value: (0, ethers_1.parseEther)(params.value).toString(),
+                    to: state.to,
+                    value: (0, ethers_1.parseEther)(state.value).toString(),
                     chain: {
                         id: networkConfig.id || 0,
                         rpc: networkConfig.rpc,
@@ -86,10 +79,7 @@ const graphDefinitionB = {
     },
 };
 // 3. Define the initial state
-const initialState = {
-    messages: "",
-    reasoning: true,
-};
+const initialState = {};
 const graphMaps = [
     graphDefinitionA,
     graphDefinitionB,
@@ -106,16 +96,15 @@ const graphMaps = [
             reasoning: {
                 name: "reasoning",
                 description: "Reasoning",
-                execute: (params, state) => __awaiter(void 0, void 0, void 0, function* () {
-                    var _a, e_1, _b, _c;
-                    console.log("Reasoning", { params, state });
+                execute: (state) => __awaiter(void 0, void 0, void 0, function* () {
+                    console.log("Reasoning", { state });
                     const startNodesSchema = graphMaps.map((graphDefinition) => graphDefinition.nodes[graphDefinition.entryNode]);
                     // Pass the start nodes into stringifyZodSchema
                     const system = `You are a wallet assistant. Don't reuse the same workflow multiple times.
         Here the available workflows and parameters:
         workflows:${(0, stringifiy_zod_schema_1.stringifyZodSchema)(startNodesSchema)}`;
                     console.log(system);
-                    const llm = yield (0, ai_1.streamObject)({
+                    const llm = yield (0, ai_1.generateObject)({
                         model: (0, openai_1.openai)("gpt-4o"),
                         prompt,
                         schema: zod_1.z.object({
@@ -131,23 +120,6 @@ const graphMaps = [
                         system,
                     });
                     console.dir(llm.object, { depth: null });
-                    try {
-                        for (var _d = true, _e = __asyncValues(llm.partialObjectStream), _f; _f = yield _e.next(), _a = _f.done, !_a; _d = true) {
-                            _c = _f.value;
-                            _d = false;
-                            const chunk = _c;
-                            if (chunk.response) {
-                                // console.log(chunk.response);
-                            }
-                        }
-                    }
-                    catch (e_1_1) { e_1 = { error: e_1_1 }; }
-                    finally {
-                        try {
-                            if (!_d && !_a && (_b = _e.return)) yield _b.call(_e);
-                        }
-                        finally { if (e_1) throw e_1.error; }
-                    }
                     const actions = (yield llm.object).actions;
                     console.log({ actions });
                     const selectedWorkflows = (yield llm.object).actions;
@@ -156,15 +128,12 @@ const graphMaps = [
                         actions: selectedWorkflows,
                     };
                 }),
-                condition: (state) => {
-                    return state.reasoning === true;
-                },
                 relationships: [{ name: "actions" }],
             },
             actions: {
                 name: "actions",
                 description: "Actions",
-                execute: (parameters, state) => __awaiter(void 0, void 0, void 0, function* () {
+                execute: (state) => __awaiter(void 0, void 0, void 0, function* () {
                     if (!state.actions) {
                         throw new Error("No actions found");
                     }
@@ -174,21 +143,17 @@ const graphMaps = [
                     };
                     const { initialStates, graphs: selectedGraphs, startNodes, } = (0, setup_graphs_1.setupGraphsWithActions)(state.actions, baseStateMapping, graphMaps);
                     // Execute graphs with dynamically determined initial states
-                    const results = yield engine_1.GraphEngine.executeGraphsInSequence(selectedGraphs, startNodes, initialStates, (graph) => {
+                    const results = yield engine_1.GraphEngine.executeGraphsInSequence(selectedGraphs, startNodes, initialStates, state.actions, (graph) => {
                         console.log(`Graph ${graph.name} updated`, graph.getState());
                     }, (error, nodeName, state) => {
                         console.error(`Erreur dans ${nodeName}`, error, state);
                     });
-                    console.log({ results });
                     return {
-                        messages: "Hello, world!",
-                        results,
+                        actions: results,
                     };
                 }),
                 condition: (state) => {
-                    if (state.reasoning === true &&
-                        state.actions &&
-                        state.actions.length > 0) {
+                    if (state.actions && state.actions.length > 0) {
                         return true;
                     }
                     return false;
@@ -197,5 +162,10 @@ const graphMaps = [
         },
     };
     const reasoningGraph = new engine_1.GraphEngine(reasoningGraphDefinition);
-    yield reasoningGraph.execute(initialState, "reasoning");
+    yield reasoningGraph.execute(initialState, "reasoning", (graph) => {
+        console.log("Main graph updated", graph.name);
+        console.dir(graph.getState(), { depth: null });
+    }, (error, nodeName, state) => {
+        console.error(`Erreur dans ${nodeName}`, error, state);
+    });
 }))();

package/dist/graph/engine.js

@@ -44,7 +44,7 @@ class GraphEngine {
             throw new Error("Cycle détecté dans le workflow");
         }
         if (options === null || options === void 0 ? void 0 : options.initialState) {
-            this.setState(options.initialState);
+            // this.setState(options.initialState);
         }
     }
     /**
@@ -189,7 +189,7 @@ class GraphEngine {
      */
     execute(state, startNode, onStream, onError) {
         return __awaiter(this, void 0, void 0, function* () {
-            var _a, _b;
+            var _a;
             try {
                 // Valide l'état initial via le schéma global (si défini)
                 if (this.schema) {
@@ -209,7 +209,7 @@ class GraphEngine {
                     this.executedNodes.add(currentNodeName);
                     const currentNode = this.nodes.get(currentNodeName);
                     if (!currentNode) {
-                        throw new Error(`Nœud ${currentNodeName} introuvable.`);
+                        throw new Error(`Node ${currentNodeName} introuvable.`);
                     }
                     // Vérification de condition (si présente)
                     if (currentNode.condition &&
@@ -224,8 +224,7 @@ class GraphEngine {
                                 node: currentNodeName,
                             });
                         }
-                        const params = (_a = currentNode.schema) === null || _a === void 0 ? void 0 : _a.parse(this.currentState);
-                        const newState = yield currentNode.execute(params || {}, this.currentState);
+                        const newState = yield currentNode.execute(this.currentState);
                         if (newState) {
                             this.setState(newState);
                             if (onStream)
@@ -268,7 +267,7 @@ class GraphEngine {
                     }
                     else {
                         // Cas normal : un seul chemin
-                        currentNodeName = ((_b = relationsNodes[0]) === null || _b === void 0 ? void 0 : _b.name) || "";
+                        currentNodeName = ((_a = relationsNodes[0]) === null || _a === void 0 ? void 0 : _a.name) || "";
                     }
                 }
                 return this.getState();
@@ -568,17 +567,26 @@ class GraphEngine {
      * @param onError Callback d'erreur
      * @returns Tableau des états finaux de chaque graphe
      */
-    static executeGraphsInSequence(graphs, startNodes, initialStates, onStream, onError) {
+    static executeGraphsInSequence(graphs, startNodes, initialStates, actions, // Pass actions here
+    onStream, onError) {
         return __awaiter(this, void 0, void 0, function* () {
+            // Return updated actions directly
             const finalStates = [];
             for (let i = 0; i < graphs.length; i++) {
                 const graph = graphs[i];
                 const startNode = startNodes[i];
                 const initialState = initialStates[i];
                 const result = yield graph.execute(initialState, startNode, onStream, onError);
-                finalStates.push(result);
+                finalStates.push({
+                    name: graph.name,
+                    result,
+                });
             }
-            return finalStates;
+            // Map results to actions
+            return actions.map((action) => {
+                const result = finalStates.find((state) => state.name === action.name);
+                return Object.assign(Object.assign({}, action), { result: result ? result.result : null });
+            });
         });
     }
     /**

package/graph/engine.ts

@@ -1,5 +1,5 @@
 import { Persistence, RealTimeNotifier } from "@/interfaces";
-import { GraphDefinition, Node, SharedState } from "@/types";
+import { Action, GraphDefinition, Node, SharedState } from "@/types";
 import EventEmitter from "events";
 import { z } from "zod";
 
@@ -68,7 +68,7 @@ export class GraphEngine<T> {
     }
 
     if (options?.initialState) {
-      this.setState(options.initialState);
+      // this.setState(options.initialState);
     }
   }
 
@@ -112,7 +112,6 @@ export class GraphEngine<T> {
   setNotifier(notifier: RealTimeNotifier): void {
     this.notifier = notifier;
   }
-
   /**
    * Charge un workflow à partir d'une définition.
    * @private
@@ -260,7 +259,7 @@ export class GraphEngine<T> {
         this.executedNodes.add(currentNodeName);
         const currentNode = this.nodes.get(currentNodeName);
         if (!currentNode) {
-          throw new Error(`Nœud ${currentNodeName} introuvable.`);
+          throw new Error(`Node ${currentNodeName} introuvable.`);
         }
 
         // Vérification de condition (si présente)
@@ -280,12 +279,7 @@ export class GraphEngine<T> {
             });
           }
 
-          const params = currentNode.schema?.parse(this.currentState);
-          const newState = await currentNode.execute(
-            params || {},
-            this.currentState
-          );
-
+          const newState = await currentNode.execute(this.currentState);
           if (newState) {
             this.setState(newState);
             if (onStream) onStream(this);
@@ -702,10 +696,12 @@ export class GraphEngine<T> {
     graphs: GraphEngine<U>[],
     startNodes: string[],
     initialStates: SharedState<U>[],
+    actions: Action[], // Pass actions here
     onStream?: (graph: GraphEngine<U>) => void,
     onError?: (error: Error, nodeName: string, state: SharedState<U>) => void
-  ): Promise<SharedState<U>[]> {
-    const finalStates: SharedState<U>[] = [];
+  ): Promise<Action[]> {
+    // Return updated actions directly
+    const finalStates: { name: string; result: SharedState<U> }[] = [];
 
     for (let i = 0; i < graphs.length; i++) {
       const graph = graphs[i];
@@ -717,10 +713,20 @@ export class GraphEngine<T> {
         onStream,
         onError
       );
-      finalStates.push(result);
+      finalStates.push({
+        name: graph.name,
+        result,
+      });
     }
 
-    return finalStates;
+    // Map results to actions
+    return actions.map((action) => {
+      const result = finalStates.find((state) => state.name === action.name);
+      return {
+        ...action,
+        result: result ? result.result : null,
+      };
+    });
   }
 
   /**

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ai.ntellect/core",
-  "version": "0.6.7",
+  "version": "0.6.9",
   "description": "",
   "main": "dist/index.js",
   "scripts": {

package/test/graph/controller.test.ts

@@ -19,12 +19,10 @@ describe("Controller", () => {
     nodes: {
       start: {
         name: "start",
-        execute: async (_params: any, state: any) => ({
-          context: {
-            ...state.context,
-            status: "completed",
-            count: state.context.count + 1,
-          },
+        execute: async (state: any) => ({
+          ...state,
+          status: "completed",
+          count: state.count + 1,
         }),
         relationships: [],
       },
@@ -38,12 +36,10 @@ describe("Controller", () => {
     nodes: {
       first: {
         name: "first",
-        execute: async (_params: any, state: any) => ({
-          context: {
-            ...state.context,
-            status: "step1",
-            count: state.context.count + 2,
-          },
+        execute: async (state: any) => ({
+          ...state,
+          status: "step1",
+          count: state.count + 2,
         }),
         relationships: [],
       },

package/test/graph/engine.test.ts

@@ -33,9 +33,9 @@ describe("Graph", () => {
       start: {
         name: "start",
         description: "Starting node",
-        execute: async (_params: any, state: SharedState<TestState>) => {
+        execute: async (state: SharedState<TestState>) => {
           return graph.updateState({
-            ...state.context,
+            ...state,
             status: "started",
             step: 1,
           });
@@ -45,22 +45,22 @@ describe("Graph", () => {
       process: {
         name: "process",
         description: "Processing node",
-        execute: async (_params: any, state: SharedState<TestState>) => {
+        execute: async (state: SharedState<TestState>) => {
           return graph.updateState({
-            ...state.context,
+            ...state,
             status: "processing",
             step: 2,
           });
         },
-        condition: (state) => state.context.step === 1,
+        condition: (state) => state.step === 1,
         relationships: [{ name: "end" }],
       },
       end: {
         name: "end",
         description: "End node",
-        execute: async (_params: any, state: SharedState<TestState>) => {
+        execute: async (state: SharedState<TestState>) => {
           return graph.updateState({
-            ...state.context,
+            ...state,
             status: "completed",
             step: 3,
           });
@@ -82,10 +82,8 @@ describe("Graph", () => {
      */
     it("should execute the complete workflow sequence", async () => {
       const initialState: SharedState<TestState> = {
-        context: {
-          status: "init",
-          step: 0,
-        },
+        status: "init",
+        step: 0,
       };
 
       // Initialiser le graph avec l'état initial
@@ -98,7 +96,7 @@ describe("Graph", () => {
       await graph.execute(initialState, "start");
       const result = graph.getState();
 
-      expect(result.context).to.deep.equal({
+      expect(result).to.deep.equal({
         status: "completed",
         step: 3,
       });
@@ -110,10 +108,8 @@ describe("Graph", () => {
      */
     it("should respect conditions in workflow", async () => {
       const initialState: SharedState<TestState> = {
-        context: {
-          status: "init",
-          step: 2,
-        },
+        status: "init",
+        step: 2,
       };
 
       // Initialiser le graph avec l'état initial
@@ -125,7 +121,7 @@ describe("Graph", () => {
       await graph.execute(initialState, "process");
       const result = graph.getState();
 
-      expect(result.context).to.deep.equal({
+      expect(result).to.deep.equal({
         status: "init",
         step: 2,
       });
@@ -137,9 +133,9 @@ describe("Graph", () => {
       const newNode = {
         name: "new-node",
         description: "A new test node",
-        execute: async (_params: any, state: SharedState<TestState>) => {
+        execute: async (state: SharedState<TestState>) => {
           return graph.updateState({
-            ...state.context,
+            ...state,
             status: "new",
             step: 4,
           });
@@ -163,9 +159,9 @@ describe("Graph", () => {
           "new-step": {
             name: "new-step",
             description: "New step node",
-            execute: async (_params: any, state: SharedState<TestState>) => {
+            execute: async (state: SharedState<TestState>) => {
               return graph.updateState({
-                ...state.context,
+                ...state,
                 status: "new-step",
                 step: 4,
               });
@@ -184,24 +180,20 @@ describe("Graph", () => {
   describe("State Management", () => {
     it("should properly update and retrieve state", async () => {
       const newState: SharedState<TestState> = {
-        context: {
-          status: "test",
-          step: 5,
-        },
+        status: "test",
+        step: 5,
       };
 
       graph.setState(newState);
       const retrievedState = graph.getState();
 
-      expect(retrievedState.context).to.deep.equal(newState.context);
+      expect(retrievedState).to.deep.equal(newState);
     });
 
     it("should merge states correctly when updating partially", () => {
       const initialState: SharedState<TestState> = {
-        context: {
-          status: "initial",
-          step: 1,
-        },
+        status: "initial",
+        step: 1,
       };
 
       graph.setState(initialState);
@@ -212,7 +204,7 @@ describe("Graph", () => {
 
       const updatedState = graph.updateState(partialUpdate);
 
-      expect(updatedState.context).to.deep.equal({
+      expect(updatedState).to.deep.equal({
         status: "updated",
         step: 1,
       });
@@ -233,7 +225,7 @@ describe("Graph", () => {
       let errorCaught = false;
       try {
         await graph.execute(
-          { context: { status: "test", step: 1 } },
+          { status: "test", step: 1 },
           "error-node",
           undefined,
           (error) => {
@@ -254,10 +246,8 @@ describe("Graph", () => {
       const strictGraph = new GraphEngine(testDefinition, {
         schema: TestSchema,
         initialState: {
-          context: {
-            status: "init",
-            step: 0,
-          },
+          status: "init",
+          step: 0,
         },
       });
 
@@ -295,7 +285,7 @@ describe("Graph", () => {
 
       const parallelNodes = ["node1", "node2", "node3"].map((name) => ({
         name,
-        execute: async (_params: any, state: SharedState<TestState>) => {
+        execute: async (state: SharedState<TestState>) => {
           executionOrder.push(name);
           await new Promise((resolve) => setTimeout(resolve, 10));
           return state;
@@ -307,7 +297,7 @@ describe("Graph", () => {
       });
 
       await graph.executeParallel(
-        { context: { status: "test", step: 1 } },
+        { status: "test", step: 1 },
         ["node1", "node2", "node3"],
         2
       );
@@ -326,9 +316,9 @@ describe("Graph", () => {
     it("should emit and handle events correctly", async () => {
       const eventNode = {
         name: "event-node",
-        execute: async (_params: any, state: SharedState<TestState>) => {
+        execute: async (state: SharedState<TestState>) => {
           return graph.updateState({
-            ...state.context,
+            ...state,
             status: "event-triggered",
             step: 10,
           });
@@ -347,8 +337,7 @@ describe("Graph", () => {
       await new Promise((resolve) => setTimeout(resolve, 50));
 
       const state = graph.getState();
-      expect(state.context.status).to.equal("event-triggered");
-      expect(state.context.status).to.equal("event-triggered");
+      expect(state.status).to.equal("event-triggered");
     });
   });
 
@@ -365,9 +354,9 @@ describe("Graph", () => {
         nodes: {
           "sub-start": {
             name: "sub-start",
-            execute: async (_params: any, state: SharedState<TestState>) => {
+            execute: async (state: SharedState<TestState>) => {
               return graph.updateState({
-                ...state.context,
+                ...state,
                 status: "sub-completed",
                 step: 100,
               });
@@ -382,17 +371,15 @@ describe("Graph", () => {
       graph.addSubGraph(subGraph, "sub-start", "sub-workflow");
 
       const initialState: SharedState<TestState> = {
-        context: {
-          status: "init",
-          step: 0,
-        },
+        status: "init",
+        step: 0,
       };
 
       await graph.execute(initialState, "sub-workflow");
       const state = graph.getState();
 
-      expect(state.context.status).to.equal("sub-completed");
-      expect(state.context.step).to.equal(100);
+      expect(state.status).to.equal("sub-completed");
+      expect(state.step).to.equal(100);
     });
   });
 
@@ -463,14 +450,14 @@ describe("Graph", () => {
       const mockPersistence: Persistence<TestState> = {
         saveState: async (graphName, state, currentNode) => {
           expect(graphName).to.equal("simple-workflow");
-          expect(state.context).to.exist;
+          expect(state).to.exist;
           expect(currentNode).to.exist;
         },
         loadState: async () => null,
       };
 
       graph.setPersistence(mockPersistence);
-      await graph.execute({ context: { status: "init", step: 0 } }, "start");
+      await graph.execute({ status: "init", step: 0 }, "start");
     });
   });
 
@@ -492,7 +479,7 @@ describe("Graph", () => {
       };
 
       graph.setNotifier(mockNotifier);
-      await graph.execute({ context: { status: "init", step: 0 } }, "start");
+      await graph.execute({ status: "init", step: 0 }, "start");
 
       expect(notifications).to.have.length.greaterThan(0);
       expect(notifications[0].event).to.equal("nodeExecutionStarted");

package/types/index.ts

@@ -209,7 +209,7 @@ export type SharedState<T> = T;
 export type Node<T, P = any> = {
   name: string;
   description?: string;
-  execute: (params: P, state: SharedState<T>) => Promise<SharedState<T> | void>;
+  execute: (state: SharedState<T>) => Promise<SharedState<T> | void>;
   condition?: (state: SharedState<T>) => boolean;
   relationships?: NodeRelationship[];
   schema?: z.ZodSchema<P>;