feat: change langgraph callback to accept an array for tool call canc…

…ellation support (#818)
Yonom · Sep 11, 2024 · 184d836 · 184d836
1 parent 3b2a502
commit 184d836
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 23 deletions.
diff --git a/.changeset/thick-kings-buy.md b/.changeset/thick-kings-buy.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react-langgraph": patch
+---
+
+feat: allow multiple message sends to support pending tool call cancellations
diff --git a/packages/react-langgraph/src/index.ts b/packages/react-langgraph/src/index.ts
@@ -8,8 +8,3 @@ export type {
   LangChainToolCall,
   LangChainToolCallChunk,
 } from "./types";
-
-/**
- * @deprecated Use `useLangGraphRuntime` instead. This will be removed in 0.1.0.
- */
-export { useLangGraphRuntime as useLangChainLangGraphRuntime } from "./useLangGraphRuntime";
diff --git a/packages/react-langgraph/src/useLangGraphMessages.ts b/packages/react-langgraph/src/useLangGraphMessages.ts
@@ -3,7 +3,7 @@ import { useState, useCallback } from "react";
 export const useLangGraphMessages = <TMessage>({
   stream,
 }: {
-  stream: (message: TMessage) => Promise<
+  stream: (messages: TMessage[]) => Promise<
     AsyncGenerator<{
       event: string;
       data: any;
@@ -13,12 +13,12 @@ export const useLangGraphMessages = <TMessage>({
   const [messages, setMessages] = useState<TMessage[]>([]);
 
   const sendMessage = useCallback(
-    async (message: TMessage) => {
-      if (message !== null) {
-        setMessages((currentMessages) => [...currentMessages, message]);
+    async (messages: TMessage[]) => {
+      if (messages.length > 0) {
+        setMessages((currentMessages) => [...currentMessages, ...messages]);
       }
 
-      const response = await stream(message);
+      const response = await stream(messages);
 
       const completeMessages: TMessage[] = [];
       let partialMessages: Map<string, TMessage> = new Map();

diff --git a/packages/react-langgraph/src/useLangGraphRuntime.ts b/packages/react-langgraph/src/useLangGraphRuntime.ts
@@ -13,7 +13,7 @@ export const useLangGraphRuntime = ({
   stream,
 }: {
   threadId?: string | undefined;
-  stream: (message: LangChainMessage) => Promise<
+  stream: (messages: LangChainMessage[]) => Promise<
     AsyncGenerator<{
       event: string;
       data: any;
@@ -25,10 +25,10 @@ export const useLangGraphRuntime = ({
   });
 
   const [isRunning, setIsRunning] = useState(false);
-  const handleSendMessage = async (message: LangChainMessage) => {
+  const handleSendMessage = async (messages: LangChainMessage[]) => {
     try {
       setIsRunning(true);
-      await sendMessage(message);
+      await sendMessage(messages);
     } catch (error) {
       console.error("Error streaming messages:", error);
     } finally {
@@ -49,18 +49,22 @@ export const useLangGraphRuntime = ({
     onNew: (msg) => {
       if (msg.content.length !== 1 || msg.content[0]?.type !== "text")
         throw new Error("Only text messages are supported");
-      return handleSendMessage({
-        type: "human",
-        content: msg.content[0].text,
-      });
+      return handleSendMessage([
+        {
+          type: "human",
+          content: msg.content[0].text,
+        },
+      ]);
     },
     onAddToolResult: async ({ toolCallId, toolName, result }) => {
-      await handleSendMessage({
-        type: "tool",
-        name: toolName,
-        tool_call_id: toolCallId,
-        content: JSON.stringify(result),
-      });
+      await handleSendMessage([
+        {
+          type: "tool",
+          name: toolName,
+          tool_call_id: toolCallId,
+          content: JSON.stringify(result),
+        },
+      ]);
     },
   });
 };