feat(graph): passthrough input types to invoke/stream

libs/langgraph/src/graph/graph.ts

@@ -441,37 +441,39 @@ export class Graph<
 export class CompiledGraph<
   N extends string,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  RunInput = any,
+  State = any,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  RunOutput = any,
+  Update = any,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   ConfigurableFieldType extends Record<string, any> = Record<string, any>
 > extends Pregel<
-  Record<N | typeof START, PregelNode<RunInput, RunOutput>>,
+  Record<N | typeof START, PregelNode<State, Update>>,
   Record<N | typeof START | typeof END | string, BaseChannel>,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  ConfigurableFieldType & Record<string, any>
+  ConfigurableFieldType & Record<string, any>,
+  Update,
+  State
 > {
   declare NodeType: N;
 
-  declare RunInput: RunInput;
+  declare RunInput: State;
 
-  declare RunOutput: RunOutput;
+  declare RunOutput: Update;
 
-  builder: Graph<N, RunInput, RunOutput>;
+  builder: Graph<N, State, Update>;
 
   constructor({
     builder,
     ...rest
-  }: { builder: Graph<N, RunInput, RunOutput> } & PregelParams<
-    Record<N | typeof START, PregelNode<RunInput, RunOutput>>,
+  }: { builder: Graph<N, State, Update> } & PregelParams<
+    Record<N | typeof START, PregelNode<State, Update>>,
     Record<N | typeof START | typeof END | string, BaseChannel>
   >) {
     super(rest);
     this.builder = builder;
   }
 
-  attachNode(key: N, node: NodeSpec<RunInput, RunOutput>): void {
+  attachNode(key: N, node: NodeSpec<State, Update>): void {
     this.channels[key] = new EphemeralValue();
     this.nodes[key] = new PregelNode({
       channels: [],
@@ -503,7 +505,7 @@ export class CompiledGraph<
   attachBranch(
     start: N | typeof START,
     name: string,
-    branch: Branch<RunInput, N>
+    branch: Branch<State, N>
   ) {
     // add hidden start node
     if (start === START && this.nodes[START]) {
@@ -588,7 +590,7 @@ export class CompiledGraph<
 
     for (const [key, nodeSpec] of Object.entries(this.builder.nodes) as [
       N,
-      NodeSpec<RunInput, RunOutput>
+      NodeSpec<State, Update>
     ][]) {
       const displayKey = _escapeMermaidKeywords(key);
       const node = nodeSpec.runnable;
@@ -769,7 +771,7 @@ export class CompiledGraph<
 
     for (const [key, nodeSpec] of Object.entries(this.builder.nodes) as [
       N,
-      NodeSpec<RunInput, RunOutput>
+      NodeSpec<State, Update>
     ][]) {
       const displayKey = _escapeMermaidKeywords(key);
       const node = nodeSpec.runnable;

libs/langgraph/src/prebuilt/react_agent_executor.ts

@@ -20,14 +20,9 @@ import {
 } from "@langchain/core/language_models/base";
 import { ChatPromptTemplate } from "@langchain/core/prompts";
 import { All, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
-import {
-  END,
-  messagesStateReducer,
-  START,
-  StateGraph,
-} from "../graph/index.js";
+import { END, START, StateGraph } from "../graph/index.js";
 import { MessagesAnnotation } from "../graph/messages_annotation.js";
-import { CompiledStateGraph, StateGraphArgs } from "../graph/state.js";
+import { CompiledStateGraph } from "../graph/state.js";
 import { ToolNode } from "./tool_node.js";
 
 export interface AgentState {
@@ -107,11 +102,12 @@ export type CreateReactAgentParams = {
  * // Returns the messages in the state at each step of execution
  * ```
  */
+
 export function createReactAgent(
   params: CreateReactAgentParams
 ): CompiledStateGraph<
-  AgentState,
-  Partial<AgentState>,
+  (typeof MessagesAnnotation)["State"],
+  (typeof MessagesAnnotation)["Update"],
   typeof START | "agent" | "tools"
 > {
   const {
@@ -122,12 +118,6 @@ export function createReactAgent(
     interruptBefore,
     interruptAfter,
   } = params;
-  const schema: StateGraphArgs<AgentState>["channels"] = {
-    messages: {
-      value: messagesStateReducer,
-      default: () => [],
-    },
-  };
 
   let toolClasses: (StructuredToolInterface | DynamicTool | RunnableToolLike)[];
   if (!Array.isArray(tools)) {
@@ -160,9 +150,7 @@ export function createReactAgent(
     return { messages: [await modelRunnable.invoke(messages, config)] };
   };
 
-  const workflow = new StateGraph<AgentState>({
-    channels: schema,
-  })
+  const workflow = new StateGraph(MessagesAnnotation)
     .addNode("agent", callModel)
     .addNode("tools", new ToolNode<AgentState>(toolClasses))
     .addEdge(START, "agent")

libs/langgraph/src/pregel/index.ts

@@ -194,11 +194,13 @@ export class Pregel<
     Nn extends StrRecord<string, PregelNode>,
     Cc extends StrRecord<string, BaseChannel | ManagedValueSpec>,
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    ConfigurableFieldType extends Record<string, any> = StrRecord<string, any>
+    ConfigurableFieldType extends Record<string, any> = StrRecord<string, any>,
+    InputType = PregelInputType,
+    OutputType = PregelOutputType
   >
   extends Runnable<
-    PregelInputType,
-    PregelOutputType,
+    InputType | null,
+    OutputType,
     PregelOptions<Nn, Cc, ConfigurableFieldType>
   >
   implements
@@ -901,7 +903,7 @@ export class Pregel<
    * @param options.debug Whether to print debug information during execution.
    */
   override async stream(
-    input: PregelInputType,
+    input: InputType | null,
     options?: Partial<PregelOptions<Nn, Cc, ConfigurableFieldType>>
   ): Promise<IterableReadableStream<PregelOutputType>> {
     return super.stream(input, options);
@@ -1221,9 +1223,9 @@ export class Pregel<
    * @param options.debug Whether to print debug information during execution.
    */
   override async invoke(
-    input: PregelInputType,
+    input: InputType | null,
     options?: Partial<PregelOptions<Nn, Cc, ConfigurableFieldType>>
-  ): Promise<PregelOutputType> {
+  ): Promise<OutputType> {
     const streamMode = options?.streamMode ?? "values";
     const config = {
       ...options,
@@ -1238,6 +1240,6 @@ export class Pregel<
     if (streamMode === "values") {
       return chunks[chunks.length - 1];
     }
-    return chunks;
+    return chunks as OutputType;
   }
 }

libs/langgraph/src/tests/graph.test.ts

@@ -23,9 +23,13 @@ describe("State", () => {
   it("should allow reducers with different argument types", async () => {
     const StateAnnotation = Annotation.Root({
       val: Annotation<number>,
-      testval: Annotation<string[], string>({
-        reducer: (left, right) =>
-          right ? left.concat([right.toString()]) : left,
+      testval: Annotation<string[], string | string[]>({
+        reducer: (left, right) => {
+          if (typeof right === "string") {
+            return right ? left.concat([right.toString()]) : left;
+          }
+          return right.length ? left.concat(right) : left;
+        },
       }),
     });
     const stateGraph = new StateGraph(StateAnnotation);
@@ -41,6 +45,7 @@ describe("State", () => {
       .addEdge(START, "testnode")
       .addEdge("testnode", END)
       .compile();
+
     expect(await graph.invoke({ testval: ["hello"] })).toEqual({
       testval: ["hello", "hi!"],
       val: 3,
@@ -51,12 +56,16 @@ describe("State", () => {
     const stateGraph = new StateGraph<
       unknown,
       { testval: string[] },
-      { testval: string }
+      { testval: string | string[] }
     >({
       channels: {
         testval: {
-          reducer: (left: string[], right?: string) =>
-            right ? left.concat([right.toString()]) : left,
+          reducer: (left, right) => {
+            if (typeof right === "string") {
+              return right ? left.concat([right.toString()]) : left;
+            }
+            return right.length ? left.concat(right) : left;
+          },
         },
       },
     });

libs/langgraph/src/tests/prebuilt.int.test.ts

@@ -63,7 +63,9 @@ describe("createReactAgent", () => {
     expect(response.messages.length > 1).toBe(true);
     const lastMessage = response.messages[response.messages.length - 1];
     expect(lastMessage._getType()).toBe("ai");
-    expect(lastMessage.content.toLowerCase()).toContain("not too cold");
+    expect((lastMessage.content as string).toLowerCase()).toContain(
+      "not too cold"
+    );
   });
 
   it("can stream a tool call with a checkpointer", async () => {

libs/langgraph/src/tests/pregel.test.ts

@@ -3541,6 +3541,7 @@ export function runPregelTests(
           hello: "there",
           bye: "world",
           messages: ["hello"],
+          // @ts-expect-error This should emit a TS error
           now: 345, // ignored because not in input schema
         })
       ).toEqual({
@@ -3553,6 +3554,7 @@ export function runPregelTests(
             hello: "there",
             bye: "world",
             messages: ["hello"],
+            // @ts-expect-error This should emit a TS error
             now: 345, // ignored because not in input schema
           })
         )
@@ -3712,6 +3714,7 @@ export function runPregelTests(
       };
       const res = await app.invoke(
         {
+          // @ts-expect-error Messages is not in schema
           messages: ["initial input"],
         },
         config

libs/langgraph/src/tests/tracing.int.test.ts

@@ -470,10 +470,16 @@ Only add steps to the plan that still NEED to be done. Do not return previously
     state: PlanExecuteState
   ): Promise<Partial<PlanExecuteState>> {
     const task = state.input;
-    const agentResponse = await agentExecutor.invoke({ input: task });
-    return {
-      pastSteps: [task, agentResponse.agentOutcome.returnValues.output],
-    };
+    const agentResponse = await agentExecutor.invoke({
+      input: task ?? undefined,
+    });
+
+    const outcome = agentResponse.agentOutcome;
+    if (!outcome || !("returnValues" in outcome)) {
+      throw new Error("Agent did not return a valid outcome.");
+    }
+
+    return { pastSteps: [task, outcome.returnValues.output] };
   }
 
   async function planStep(