feat(ai-sdk): attachments support (#764)

Yonom · Sep 7, 2024 · 7809584 · 7809584
1 parent 9e00772
commit 7809584
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 26 deletions.
diff --git a/.changeset/olive-olives-learn.md b/.changeset/olive-olives-learn.md
@@ -1,5 +1,7 @@
 ---
 "@assistant-ui/react": patch
+"@assistant-ui/react-playground": patch
+"@assistant-ui/react-ai-sdk": patch
 ---
 
 feat: add composer attachments state
diff --git a/.changeset/purple-knives-grow.md b/.changeset/purple-knives-grow.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react-ai-sdk": patch
+---
+
+feat: allow image content types
diff --git a/.changeset/yellow-cameras-count.md b/.changeset/yellow-cameras-count.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react-ai-sdk": patch
+---
+
+feat: AI SDK attachments support
diff --git a/packages/react-ai-sdk/src/ui/use-assistant/useVercelUseAssistantRuntime.tsx b/packages/react-ai-sdk/src/ui/use-assistant/useVercelUseAssistantRuntime.tsx
@@ -3,6 +3,7 @@ import { useExternalStoreRuntime } from "@assistant-ui/react";
 import { useCachedChunkedMessages } from "../utils/useCachedChunkedMessages";
 import { convertMessage } from "../utils/convertMessage";
 import { useInputSync } from "../utils/useInputSync";
+import { toCreateMessage } from "../utils/toCreateMessage";
 
 export const useVercelUseAssistantRuntime = (
   assistantHelpers: ReturnType<typeof useAssistant>,
@@ -13,15 +14,7 @@ export const useVercelUseAssistantRuntime = (
     messages,
     onCancel: async () => assistantHelpers.stop(),
     onNew: async (message) => {
-      if (message.content.length !== 1 || message.content[0]?.type !== "text")
-        throw new Error(
-          "VercelUseAssistantRuntime only supports text content.",
-        );
-
-      await assistantHelpers.append({
-        role: message.role,
-        content: message.content[0].text,
-      });
+      await assistantHelpers.append(await toCreateMessage(message));
     },
     onNewThread: () => {
       assistantHelpers.messages = [];

diff --git a/packages/react-ai-sdk/src/ui/use-chat/useVercelUseChatRuntime.tsx b/packages/react-ai-sdk/src/ui/use-chat/useVercelUseChatRuntime.tsx
@@ -4,42 +4,29 @@ import { convertMessage } from "../utils/convertMessage";
 import { useExternalStoreRuntime } from "@assistant-ui/react";
 import { useInputSync } from "../utils/useInputSync";
 import { sliceMessagesUntil } from "../utils/sliceMessagesUntil";
+import { toCreateMessage } from "../utils/toCreateMessage";
 
 export const useVercelUseChatRuntime = (
   chatHelpers: ReturnType<typeof useChat>,
 ) => {
   const messages = useCachedChunkedMessages(chatHelpers.messages);
+
   const runtime = useExternalStoreRuntime({
     isRunning: chatHelpers.isLoading,
     messages,
     setMessages: (messages) => chatHelpers.setMessages(messages.flat()),
     onCancel: async () => chatHelpers.stop(),
     onNew: async (message) => {
-      if (message.content.length !== 1 || message.content[0]?.type !== "text")
-        throw new Error(
-          "Only text content is supported by VercelUseChatRuntime. Use the Edge runtime for image support.",
-        );
-      await chatHelpers.append({
-        role: message.role,
-        content: message.content[0].text,
-      });
+      await chatHelpers.append(await toCreateMessage(message));
     },
     onEdit: async (message) => {
-      if (message.content.length !== 1 || message.content[0]?.type !== "text")
-        throw new Error(
-          "Only text content is supported by VercelUseChatRuntime. Use the Edge runtime for image support.",
-        );
-
       const newMessages = sliceMessagesUntil(
         chatHelpers.messages,
         message.parentId,
       );
       chatHelpers.setMessages(newMessages);
 
-      await chatHelpers.append({
-        role: message.role,
-        content: message.content[0].text,
-      });
+      await chatHelpers.append(await toCreateMessage(message));
     },
     onReload: async (parentId: string | null) => {
       const newMessages = sliceMessagesUntil(chatHelpers.messages, parentId);

diff --git a/packages/react-ai-sdk/src/ui/utils/toCreateMessage.ts b/packages/react-ai-sdk/src/ui/utils/toCreateMessage.ts
@@ -0,0 +1,44 @@
+import { AppendMessage } from "@assistant-ui/react";
+import { CreateMessage } from "ai";
+
+export const toCreateMessage = async (
+  message: AppendMessage,
+): Promise<CreateMessage> => {
+  const content = message.content
+    .filter((part) => part.type === "text")
+    .map((t) => t.text)
+    .join("\n\n");
+
+  const images = message.content
+    .filter((part) => part.type === "image")
+    .map((part) => ({ url: part.image }));
+
+  return {
+    role: message.role,
+    content,
+    experimental_attachments: [
+      ...images,
+      ...(await Promise.all(
+        message.attachments.map(async (m) => {
+          if (m.file == null)
+            throw new Error("Attachment did not contain a file");
+          return {
+            contentType: m.file.type,
+            name: m.file.name,
+            url: await getFileDataURL(m.file),
+          };
+        }),
+      )),
+    ],
+  };
+};
+
+const getFileDataURL = (file: File) =>
+  new Promise<string>((resolve, reject) => {
+    const reader = new FileReader();
+
+    reader.onload = () => resolve(reader.result as string);
+    reader.onerror = (error) => reject(error);
+
+    reader.readAsDataURL(file);
+  });