@assistant-ui/react-langgraph 0.5.5 → 0.5.7

package/src/useLangGraphMessages.ts

@@ -1,6 +1,16 @@
-import { useState, useCallback, useRef } from "react";
+import { useState, useCallback, useRef, useMemo } from "react";
 import { v4 as uuidv4 } from "uuid";
 import { LangGraphMessageAccumulator } from "./LangGraphMessageAccumulator";
+import {
+  EventType,
+  LangChainMessageTupleEvent,
+  LangGraphKnownEventTypes,
+  LangChainMessageChunk,
+  OnCustomEventCallback,
+  OnErrorEventCallback,
+  OnInfoEventCallback,
+  OnMetadataEventCallback,
+} from "./types";
 
 export type LangGraphCommand = {
   resume: string;
@@ -12,15 +22,10 @@ export type LangGraphSendMessageConfig = {
 };
 
 export type LangGraphMessagesEvent<TMessage> = {
-  event:
-    | "messages"
-    | "messages/partial"
-    | "messages/complete"
-    | "metadata"
-    | "updates"
-    | string;
+  event: EventType;
   data: TMessage[] | any;
 };
+
 export type LangGraphStreamCallback<TMessage> = (
   messages: TMessage[],
   config: LangGraphSendMessageConfig & { abortSignal: AbortSignal },
@@ -40,12 +45,35 @@ const DEFAULT_APPEND_MESSAGE = <TMessage>(
   curr: TMessage,
 ) => curr;
 
+const isLangChainMessageChunk = (
+  value: unknown,
+): value is LangChainMessageChunk => {
+  if (!value || typeof value !== "object") return false;
+  const chunk = value as any;
+  return (
+    "type" in chunk &&
+    chunk.type === "AIMessageChunk" &&
+    (chunk.content === undefined ||
+      typeof chunk.content === "string" ||
+      Array.isArray(chunk.content)) &&
+    (chunk.tool_call_chunks === undefined ||
+      Array.isArray(chunk.tool_call_chunks))
+  );
+};
+
 export const useLangGraphMessages = <TMessage extends { id?: string }>({
   stream,
   appendMessage = DEFAULT_APPEND_MESSAGE,
+  eventHandlers,
 }: {
   stream: LangGraphStreamCallback<TMessage>;
   appendMessage?: (prev: TMessage | undefined, curr: TMessage) => TMessage;
+  eventHandlers?: {
+    onMetadata?: OnMetadataEventCallback;
+    onInfo?: OnInfoEventCallback;
+    onError?: OnErrorEventCallback;
+    onCustomEvent?: OnCustomEventCallback;
+  };
 }) => {
   const [interrupt, setInterrupt] = useState<
     LangGraphInterruptState | undefined
@@ -53,36 +81,87 @@ export const useLangGraphMessages = <TMessage extends { id?: string }>({
   const [messages, setMessages] = useState<TMessage[]>([]);
   const abortControllerRef = useRef<AbortController | null>(null);
 
+  const { onMetadata, onInfo, onError, onCustomEvent } = useMemo(
+    () => eventHandlers ?? {},
+    [eventHandlers],
+  );
+
   const sendMessage = useCallback(
     async (newMessages: TMessage[], config: LangGraphSendMessageConfig) => {
       // ensure all messages have an ID
-      newMessages = newMessages.map((m) => (m.id ? m : { ...m, id: uuidv4() }));
+      const newMessagesWithId = newMessages.map((m) =>
+        m.id ? m : { ...m, id: uuidv4() },
+      );
 
       const accumulator = new LangGraphMessageAccumulator({
         initialMessages: messages,
         appendMessage,
       });
-      setMessages(accumulator.addMessages(newMessages));
+      setMessages(accumulator.addMessages(newMessagesWithId));
 
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
-      const response = await stream(newMessages, {
+      const response = await stream(newMessagesWithId, {
         ...config,
         abortSignal: abortController.signal,
       });
 
       for await (const chunk of response) {
-        if (
-          chunk.event === "messages/partial" ||
-          chunk.event === "messages/complete"
-        ) {
-          setMessages(accumulator.addMessages(chunk.data));
-        } else if (chunk.event === "updates") {
-          setInterrupt(chunk.data.__interrupt__?.[0]);
+        switch (chunk.event) {
+          case LangGraphKnownEventTypes.MessagesPartial:
+          case LangGraphKnownEventTypes.MessagesComplete:
+            setMessages(accumulator.addMessages(chunk.data));
+            break;
+          case LangGraphKnownEventTypes.Updates:
+            setInterrupt(chunk.data.__interrupt__?.[0]);
+            break;
+          case LangGraphKnownEventTypes.Messages: {
+            const [messageChunk] = (chunk as LangChainMessageTupleEvent).data;
+            if (!isLangChainMessageChunk(messageChunk)) {
+              console.warn(
+                "Received invalid message chunk format:",
+                messageChunk,
+              );
+              break;
+            }
+            const updatedMessages = accumulator.addMessages([
+              messageChunk as unknown as TMessage,
+            ]);
+            setMessages(updatedMessages);
+            break;
+          }
+          case LangGraphKnownEventTypes.Metadata:
+            onMetadata?.(chunk.data);
+            break;
+          case LangGraphKnownEventTypes.Info:
+            onInfo?.(chunk.data);
+            break;
+          case LangGraphKnownEventTypes.Error:
+            onError?.(chunk.data);
+            break;
+          default:
+            if (onCustomEvent) {
+              onCustomEvent(chunk.event, chunk.data);
+            } else {
+              console.warn(
+                "Unhandled event received:",
+                chunk.event,
+                chunk.data,
+              );
+            }
+            break;
         }
       }
     },
-    [messages, stream, appendMessage],
+    [
+      messages,
+      appendMessage,
+      stream,
+      onMetadata,
+      onInfo,
+      onError,
+      onCustomEvent,
+    ],
   );
 
   const cancel = useCallback(() => {

package/src/useLangGraphRuntime.test.tsx

@@ -0,0 +1,276 @@
+import { describe, it, expect, vi } from "vitest";
+import { act, renderHook, waitFor } from "@testing-library/react";
+import {
+  AssistantRuntime,
+  AssistantRuntimeProvider,
+} from "@assistant-ui/react";
+import { useLangGraphRuntime, useLangGraphSend } from "./useLangGraphRuntime";
+import { mockStreamCallbackFactory } from "./testUtils";
+import React, { ReactNode } from "react";
+
+const metadataEvent = {
+  event: "metadata",
+  data: {
+    thread_id: "123",
+    run_attempt: 1,
+  },
+};
+
+const infoEvent = {
+  event: "info",
+  data: {
+    message: "Processing request",
+  },
+};
+
+const errorEvent = {
+  event: "error",
+  data: {
+    message: "Something went wrong",
+  },
+};
+
+const customEvent = {
+  event: "custom",
+  data: {
+    type: "test",
+    value: "custom data",
+  },
+};
+
+describe("useLangGraphRuntime", () => {
+  const wrapperFactory =
+    (runtime: AssistantRuntime) =>
+    // eslint-disable-next-line react/display-name
+    ({ children }: { children: ReactNode }) => (
+      <AssistantRuntimeProvider runtime={runtime}>
+        {children}
+      </AssistantRuntimeProvider>
+    );
+
+  it("should handle metadata events", async () => {
+    const onMetadata = vi.fn();
+
+    const streamMock = vi
+      .fn()
+      .mockImplementation(() => mockStreamCallbackFactory([metadataEvent])());
+
+    const { result: runtimeResult } = renderHook(
+      () =>
+        useLangGraphRuntime({
+          stream: streamMock,
+          eventHandlers: {
+            onMetadata,
+          },
+        }),
+      {},
+    );
+
+    const wrapper = wrapperFactory(runtimeResult.current);
+
+    const { result: sendResult } = renderHook(() => useLangGraphSend(), {
+      wrapper,
+    });
+
+    act(() => {
+      sendResult.current(
+        [
+          {
+            type: "human",
+            content: "Hello, world!",
+          },
+        ],
+        {},
+      );
+    });
+
+    await waitFor(() => {
+      expect(streamMock).toHaveBeenCalled();
+      expect(onMetadata).toHaveBeenCalledWith(metadataEvent.data);
+    });
+  });
+
+  it("should handle info events", async () => {
+    const onInfo = vi.fn();
+
+    const streamMock = vi
+      .fn()
+      .mockImplementation(() => mockStreamCallbackFactory([infoEvent])());
+
+    const { result: runtimeResult } = renderHook(
+      () =>
+        useLangGraphRuntime({
+          stream: streamMock,
+          eventHandlers: {
+            onInfo,
+          },
+        }),
+      {},
+    );
+
+    const wrapper = ({ children }: { children: ReactNode }) => (
+      <AssistantRuntimeProvider runtime={runtimeResult.current}>
+        {children}
+      </AssistantRuntimeProvider>
+    );
+
+    const { result: sendResult } = renderHook(() => useLangGraphSend(), {
+      wrapper,
+    });
+
+    act(() => {
+      sendResult.current(
+        [
+          {
+            type: "human",
+            content: "Hello, world!",
+          },
+        ],
+        {},
+      );
+    });
+
+    await waitFor(() => {
+      expect(streamMock).toHaveBeenCalled();
+      expect(onInfo).toHaveBeenCalledWith(infoEvent.data);
+    });
+  });
+
+  it("should handle error events", async () => {
+    const onError = vi.fn();
+
+    const streamMock = vi
+      .fn()
+      .mockImplementation(() => mockStreamCallbackFactory([errorEvent])());
+
+    const { result: runtimeResult } = renderHook(
+      () =>
+        useLangGraphRuntime({
+          stream: streamMock,
+          eventHandlers: {
+            onError,
+          },
+        }),
+      {},
+    );
+
+    const wrapper = ({ children }: { children: ReactNode }) => (
+      <AssistantRuntimeProvider runtime={runtimeResult.current}>
+        {children}
+      </AssistantRuntimeProvider>
+    );
+
+    const { result: sendResult } = renderHook(() => useLangGraphSend(), {
+      wrapper,
+    });
+
+    act(() => {
+      sendResult.current(
+        [
+          {
+            type: "human",
+            content: "Hello, world!",
+          },
+        ],
+        {},
+      );
+    });
+
+    await waitFor(() => {
+      expect(streamMock).toHaveBeenCalled();
+      expect(onError).toHaveBeenCalledWith(errorEvent.data);
+    });
+  });
+
+  it("should handle custom events", async () => {
+    const onCustomEvent = vi.fn();
+
+    const streamMock = vi
+      .fn()
+      .mockImplementation(() => mockStreamCallbackFactory([customEvent])());
+
+    const { result: runtimeResult } = renderHook(
+      () =>
+        useLangGraphRuntime({
+          stream: streamMock,
+          eventHandlers: {
+            onCustomEvent,
+          },
+        }),
+      {},
+    );
+
+    const wrapper = wrapperFactory(runtimeResult.current);
+
+    const { result: sendResult } = renderHook(() => useLangGraphSend(), {
+      wrapper,
+    });
+
+    act(() => {
+      sendResult.current(
+        [
+          {
+            type: "human",
+            content: "Hello, world!",
+          },
+        ],
+        {},
+      );
+    });
+
+    await waitFor(() => {
+      expect(streamMock).toHaveBeenCalled();
+      expect(onCustomEvent).toHaveBeenCalledWith(
+        customEvent.event,
+        customEvent.data,
+      );
+    });
+  });
+
+  it("should work without any provided callbacks", async () => {
+    const streamMock = vi
+      .fn()
+      .mockImplementation(() =>
+        mockStreamCallbackFactory([
+          metadataEvent,
+          infoEvent,
+          errorEvent,
+          customEvent,
+        ])(),
+      );
+
+    const { result: runtimeResult } = renderHook(
+      () =>
+        useLangGraphRuntime({
+          stream: streamMock,
+          eventHandlers: {},
+        }),
+      {},
+    );
+
+    const wrapper = wrapperFactory(runtimeResult.current);
+
+    const { result: sendResult } = renderHook(() => useLangGraphSend(), {
+      wrapper,
+    });
+
+    act(() => {
+      sendResult.current(
+        [
+          {
+            type: "human",
+            content: "Hello, world!",
+          },
+        ],
+        {},
+      );
+    });
+
+    await waitFor(() => {
+      expect(streamMock).toHaveBeenCalled();
+    });
+
+    // Should not throw any errors even when events are processed without handlers
+    expect(runtimeResult.current).toBeDefined();
+  });
+});

package/src/useLangGraphRuntime.ts

@@ -1,5 +1,12 @@
 import { useEffect, useRef, useState } from "react";
-import { LangChainMessage, LangChainToolCall } from "./types";
+import {
+  LangChainMessage,
+  LangChainToolCall,
+  OnCustomEventCallback,
+  OnErrorEventCallback,
+  OnInfoEventCallback,
+  OnMetadataEventCallback,
+} from "./types";
 import {
   useExternalMessageConverter,
   useExternalStoreRuntime,
@@ -115,6 +122,7 @@ export const useLangGraphRuntime = ({
   threadId,
   onSwitchToNewThread,
   onSwitchToThread,
+  eventHandlers,
 }: {
   /**
    * @deprecated For thread management use `useCloudThreadListRuntime` instead. This option will be removed in a future version.
@@ -138,6 +146,29 @@ export const useLangGraphRuntime = ({
         feedback?: FeedbackAdapter;
       }
     | undefined;
+  /**
+   * Event handlers for various LangGraph stream events
+   */
+  eventHandlers?:
+    | {
+        /**
+         * Called when metadata is received from the LangGraph stream
+         */
+        onMetadata?: OnMetadataEventCallback;
+        /**
+         * Called when informational messages are received from the LangGraph stream
+         */
+        onInfo?: OnInfoEventCallback;
+        /**
+         * Called when errors occur during LangGraph stream processing
+         */
+        onError?: OnErrorEventCallback;
+        /**
+         * Called when custom events are received from the LangGraph stream
+         */
+        onCustomEvent?: OnCustomEventCallback;
+      }
+    | undefined;
 }) => {
   const {
     interrupt,
@@ -149,6 +180,7 @@ export const useLangGraphRuntime = ({
   } = useLangGraphMessages({
     appendMessage: appendLangChainChunk,
     stream,
+    ...(eventHandlers && { eventHandlers }),
   });
 
   const [isRunning, setIsRunning] = useState(false);