From 9f302e4e8d1c095181885533f2eb734377fb383c Mon Sep 17 00:00:00 2001
From: Tony Holdstock-Brown <tonyhb@gmail.com>
Date: Thu, 19 Dec 2024 12:12:45 -0800
Subject: [PATCH] Fix network -> agent tate in run (#27)

* Fix network -> agent tate in run

For some reason, we were not using the actual run state when calling
agent.run in networks.  This is via a large refactor.

* Changeset
---
 .changeset/clean-jeans-cheer.md | 5 +++++
 examples/swebench/inngest.ts    | 7 ++++++-
 src/agent.ts                    | 7 +++----
 src/state.ts                    | 4 ++++
 4 files changed, 18 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/clean-jeans-cheer.md

diff --git a/.changeset/clean-jeans-cheer.md b/.changeset/clean-jeans-cheer.md
new file mode 100644
index 0000000..5a99094
--- /dev/null
+++ b/.changeset/clean-jeans-cheer.md
@@ -0,0 +1,5 @@
+---
+"@inngest/agent-kit": patch
+---
+
+Network state flow to agents
diff --git a/examples/swebench/inngest.ts b/examples/swebench/inngest.ts
index 18c959f..5ae5642 100644
--- a/examples/swebench/inngest.ts
+++ b/examples/swebench/inngest.ts
@@ -2,6 +2,7 @@ import { execSync } from "child_process";
 import fs from "fs";
 import { EventSchemas, Inngest } from "inngest";
 import { z } from "zod";
+import { State } from "@inngest/agent-kit";
 import { codeWritingNetwork } from "./networks/codeWritingNetwork";
 
 export const inngest = new Inngest({
@@ -49,8 +50,12 @@ export const fn = inngest.createFunction(
       execSync(`cd ${dir} && git reset --hard FETCH_HEAD`);
     });
 
+    // Create new state and store the repo in KV for access via tools.
+    const state = new State();
+    state.kv.set("repo", event.data.repo);
+
     await codeWritingNetwork.run(event.data.problem_statement, {
-      state: { repo: event.data.repo },
+      state,
     });
   }
 );
diff --git a/src/agent.ts b/src/agent.ts
index 7a554fe..e1b0a7e 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -1,6 +1,5 @@
 import { type AiAdapter } from "inngest";
 import { createAgenticModelFromAiAdapter, type AgenticModel } from "./model";
-import { type Network } from "./network";
 import { NetworkRun } from "./networkRun";
 import {
   InferenceResult,
@@ -119,8 +118,8 @@ export class Agent {
     const p = createAgenticModelFromAiAdapter(rawModel);
 
     // input state always overrides the network state.
-    const s = state || network?.defaultState?.clone();
-    const run = network && new NetworkRun(network, s || new State());
+    const s = state || network?.state?.clone() || new State();
+    const run = network && new NetworkRun(network, s);
 
     let history = s ? s.format() : [];
     let prompt = await this.agentPrompt(input, run);
@@ -361,7 +360,7 @@ export namespace Agent {
 
   export interface RunOptions {
     model?: AiAdapter.Any;
-    network?: Network;
+    network?: NetworkRun;
     /**
      * State allows you to pass custom state into a single agent run call.  This should only
      * be provided if you are running agents outside of a network.  Networks automatically
diff --git a/src/state.ts b/src/state.ts
index 441510f..daebc45 100644
--- a/src/state.ts
+++ b/src/state.ts
@@ -68,6 +68,7 @@ export class State {
     get: <T = any>(key: string) => T | undefined;
     delete: (key: string) => boolean;
     has: (key: string) => boolean;
+    all: () => Record<string, unknown>;
   };
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -96,6 +97,9 @@ export class State {
       has: (key: string) => {
         return this._kv.has(key);
       },
+      all: () => {
+        return Object.fromEntries(this._kv);
+      },
     };
   }