createReactAgent does not support stateModifier with custom stateSchema types

[luizzappa]

Issue Description

The createReactAgent function does not accept a stateSchema different from MessagesAnnotation.

Current Behavior

When trying to use a custom stateSchema, the linter throws an error as shown below:

import { ChatOpenAI } from "@langchain/openai";
import { Annotation} from "@langchain/langgraph";
import { MessagesAnnotation } from "@langchain/langgraph";
import { createReactAgent } from '@langchain/langgraph/prebuilt';

const StateAnnotation = Annotation.Root({
  ...MessagesAnnotation.spec,
  // user provided
  lastName: Annotation<string>,
  // updated by the tool
  userInfo: Annotation<Record<string, any>>,
  });

const stateModifier = (state: typeof StateAnnotation.State) => {
  const userInfo = state.userInfo;
  if (userInfo == null) {
    return state.messages;
  }
  const systemMessage = `User name is ${userInfo.name}. User lives in ${userInfo.location}`;
  return [{
    role: "system",
    content: systemMessage,
  }, ...state.messages];
};

const model = new ChatOpenAI({
  model: "gpt-4o",
});

const agent = createReactAgent({
  llm: model,
  stateSchema: StateAnnotation,
  stateModifier: stateModifier,
})

Additional Information

Repository Version: "@langchain/langgraph": "^0.2.38",

[jacoblee93]

Thanks for flagging, will have a look