@assistant-ui/react-langgraph 0.5.4 → 0.5.6

package/src/useLangGraphMessages.ts

@@ -1,6 +1,12 @@
 import { useState, useCallback, useRef } from "react";
 import { v4 as uuidv4 } from "uuid";
 import { LangGraphMessageAccumulator } from "./LangGraphMessageAccumulator";
+import {
+  EventType,
+  LangChainMessageTupleEvent,
+  LangGraphKnownEventTypes,
+  LangChainMessageChunk,
+} from "./types";
 
 export type LangGraphCommand = {
   resume: string;
@@ -12,15 +18,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,6 +41,22 @@ 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,
@@ -56,29 +73,52 @@ export const useLangGraphMessages = <TMessage extends { id?: string }>({
   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:
+            // currently this is a no-op
+            break;
+          default:
+            console.warn(`The event type ${chunk.event} is not supported.`);
         }
       }
     },

package/src/useLangGraphRuntime.ts

@@ -232,6 +232,7 @@ export const useLangGraphRuntime = ({
                   name: t.name,
                   tool_call_id: t.id,
                   content: JSON.stringify({ cancelled: true }),
+                  status: "error",
                 }) satisfies LangChainMessage & { type: "tool" },
             )
           : [];
@@ -249,7 +250,13 @@ export const useLangGraphRuntime = ({
         },
       );
     },
-    onAddToolResult: async ({ toolCallId, toolName, result }) => {
+    onAddToolResult: async ({
+      toolCallId,
+      toolName,
+      result,
+      isError,
+      artifact,
+    }) => {
       // TODO parallel human in the loop calls
       await handleSendMessage(
         [
@@ -258,6 +265,8 @@ export const useLangGraphRuntime = ({
             name: toolName,
             tool_call_id: toolCallId,
             content: JSON.stringify(result),
+            artifact,
+            status: isError ? "error" : "success",
           },
         ],
         // TODO reuse runconfig here!