fix: type adjusted state for langgraph

Author: ranst91

typescript-sdk/integrations/langgraph/src/agent.ts

@@ -24,6 +24,8 @@ import {
   RunMetadata,
   PredictStateTool,
   LangGraphReasoning,
+  StateEnrichment,
+  LangGraphTool,
 } from "./types";
 import {
   AbstractAgent,
@@ -181,7 +183,7 @@ export class LangGraphAgent extends AbstractAgent {
   }
 
   async prepareRegenerateStream(input: RegenerateInput, streamMode: StreamMode | StreamMode[]) {
-    const { threadId, messageCheckpoint, tools } = input;
+    const { threadId, messageCheckpoint } = input;
 
     const timeTravelCheckpoint = await this.getCheckpointByMessage(
       messageCheckpoint!.id!,
@@ -262,7 +264,7 @@ export class LangGraphAgent extends AbstractAgent {
       ...agentState,
       values: {
         ...stateValuesDiff,
-        messages: [...agentStateMessages, ...stateValuesDiff.messages],
+        messages: [...agentStateMessages, ...(stateValuesDiff.messages ?? [])],
       },
     };
     let stateValues = threadState.values;
@@ -968,7 +970,7 @@ export class LangGraphAgent extends AbstractAgent {
     }
   }
 
-  langGraphDefaultMergeState(state: State, messages: LangGraphMessage[], input: RunAgentExtendedInput): State {
+  langGraphDefaultMergeState(state: State, messages: LangGraphMessage[], input: RunAgentExtendedInput): State<StateEnrichment> {
     if (messages.length > 0 && "role" in messages[0] && messages[0].role === "system") {
       // remove system message
       messages = messages.slice(1);
@@ -980,7 +982,7 @@ export class LangGraphAgent extends AbstractAgent {
 
     const newMessages = messages.filter((message) => !existingMessageIds.has(message.id));
 
-    const langGraphTools = [...(state.tools ?? []), ...(input.tools ?? [])].map((tool) => {
+    const langGraphTools: LangGraphTool[] = [...(state.tools ?? []), ...(input.tools ?? [])].map((tool) => {
       if (tool.type) {
         return tool;
       }

typescript-sdk/integrations/langgraph/src/types.ts

@@ -1,5 +1,6 @@
-import { AssistantGraph, Message } from "@langchain/langgraph-sdk";
+import { AssistantGraph, Message as LangGraphMessage, } from "@langchain/langgraph-sdk";
 import { MessageType } from "@langchain/core/messages";
+import { RunAgentInput } from "@ag-ui/core";
 
 export enum LangGraphEventTypes {
   OnChainStart = "on_chain_start",
@@ -14,7 +15,26 @@ export enum LangGraphEventTypes {
   OnInterrupt = "on_interrupt",
 }
 
-export type State = Record<string, any>;
+export type LangGraphTool = {
+  type: "function";
+  function: {
+    name: string;
+    description: string;
+    parameters: any;
+  },
+}
+
+export type State<TDefinedState = Record<string, any>> = {
+  [k in keyof TDefinedState]: TDefinedState[k] | null;
+} & Record<string, any>;
+export interface StateEnrichment {
+  messages: LangGraphMessage[];
+  tools: LangGraphTool[];
+  'ag-ui': {
+    tools: LangGraphTool[];
+    context: RunAgentInput['context']
+  }
+}
 
 export type SchemaKeys = {
   input: string[] | null;
@@ -54,7 +74,7 @@ export interface ToolCall {
 }
 
 type BaseLangGraphPlatformMessage = Omit<
-  Message,
+  LangGraphMessage,
   | "isResultMessage"
   | "isTextMessage"
   | "isImageMessage"