langgraph[minor]: Expose MessagesState annotation

langchain-ai · Aug 20, 2024 · 292c09b · 292c09b
1 parent bcefdd0
commit 292c09b
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 9 deletions.
diff --git a/langgraph/src/graph/index.ts b/langgraph/src/graph/index.ts
@@ -5,4 +5,8 @@ export {
   StateGraph,
   type CompiledStateGraph,
 } from "./state.js";
-export { MessageGraph, messagesStateReducer } from "./message.js";
+export {
+  MessageGraph,
+  messagesStateReducer,
+  MessagesState,
+} from "./message.js";
diff --git a/langgraph/src/graph/message.ts b/langgraph/src/graph/message.ts
@@ -5,6 +5,7 @@ import {
 } from "@langchain/core/messages";
 import { v4 } from "uuid";
 import { StateGraph } from "./state.js";
+import { Annotation } from "./annotation.js";
 
 type Messages =
   | Array<BaseMessage | BaseMessageLike>
@@ -76,6 +77,9 @@ export class MessageGraph extends StateGraph<
   }
 }
 
-export interface MessagesState {
-  messages: BaseMessage[];
-}
+export const MessagesState = Annotation.Root({
+  messages: Annotation<BaseMessage[]>({
+    reducer: messagesStateReducer,
+    default: () => [],
+  }),
+});
diff --git a/langgraph/src/prebuilt/react_agent_executor.ts b/langgraph/src/prebuilt/react_agent_executor.ts
@@ -40,7 +40,7 @@ export type N = typeof START | "agent" | "tools";
 export type CreateReactAgentParams = {
   llm: BaseChatModel;
   tools:
-    | ToolNode<MessagesState>
+    | ToolNode<typeof MessagesState.State>
     | (StructuredToolInterface | RunnableToolLike)[];
   messageModifier?:
     | SystemMessage

diff --git a/langgraph/src/prebuilt/tool_node.ts b/langgraph/src/prebuilt/tool_node.ts
@@ -17,7 +17,7 @@ export type ToolNodeOptions = {
 };
 
 export class ToolNode<
-  T extends BaseMessage[] | MessagesState
+  T extends BaseMessage[] | typeof MessagesState.State
 > extends RunnableCallable<T, T> {
   /**
   A node that runs the tools requested in the last AIMessage. It can be used
@@ -41,9 +41,9 @@ export class ToolNode<
   }
 
   private async run(
-    input: BaseMessage[] | MessagesState,
+    input: BaseMessage[] | typeof MessagesState.State,
     config: RunnableConfig
-  ): Promise<BaseMessage[] | MessagesState> {
+  ): Promise<BaseMessage[] | typeof MessagesState.State> {
     const message = Array.isArray(input)
       ? input[input.length - 1]
       : input.messages[input.messages.length - 1];
@@ -92,7 +92,7 @@ export class ToolNode<
 }
 
 export function toolsCondition(
-  state: BaseMessage[] | MessagesState
+  state: BaseMessage[] | typeof MessagesState.State
 ): "tools" | typeof END {
   const message = Array.isArray(state)
     ? state[state.length - 1]

diff --git a/langgraph/src/web.ts b/langgraph/src/web.ts
@@ -8,6 +8,7 @@ export {
   MessageGraph,
   messagesStateReducer,
   Annotation,
+  MessagesState,
   type StateType,
   type UpdateType,
   type CompiledGraph,