Merge pull request #143 from get-convex/ian/ui-stream-chunks

Stream UIMessage chunks

Author: ianmacartney

CHANGELOG.md

@@ -1,5 +1,20 @@
 # Changelog
 
+## 0.2.8 alpha
+
+- Adds useUIMessages and useUIStreamingMessages for UIMessage-first client
+  usage. These have more metadata around reasoning streaming status, and
+  provide support for metadata and custom data parts.
+- Uses the UIMessageChunk stream when streaming text, which fixes many bugs
+  which stemmed from the AI SDK's onChunk callback.
+- Enables having streaming hooks skip streamIds so you can stream via HTTP
+  for some clients and have others get the deltas.
+- Compresses the deltas saved to the DB - combining text deltas & reasoning
+  deltas that come from the same part `id`.
+- `optimisticallySendMessage` will set fields for both MessageDoc and
+  UIMessage, so it should transparently work for either hook strategy, though
+  you may see more fields than you expect.
+
 ## 0.2.7
 
 - Updates types to match the latest `ai` package types - note: you should update

example/convex/_generated/api.d.ts

@@ -2559,6 +2559,7 @@ export declare const components: {
         "internal",
         {
           agentName?: string;
+          format?: "UIMessageChunk" | "TextStreamPart";
           model?: string;
           order: number;
           provider?: string;
@@ -2623,6 +2624,7 @@ export declare const components: {
         },
         Array<{
           agentName?: string;
+          format?: "UIMessageChunk" | "TextStreamPart";
           model?: string;
           order: number;
           provider?: string;

example/convex/chat/streaming.ts

@@ -4,6 +4,7 @@ import {
   createThread,
   listMessages,
   syncStreams,
+  toUIMessages,
   vStreamArgs,
 } from "@convex-dev/agent";
 import { components, internal } from "../_generated/api";
@@ -111,6 +112,7 @@ export const listThreadMessages = query({
     return {
       ...paginated,
       streams,
+      page: toUIMessages(paginated.page),
 
       // ... you can return other metadata here too.
       // note: this function will be called with various permutations of delta

example/package-lock.json

@@ -3,9 +3,8 @@ import { Toaster } from "../components/ui/toaster";
 import { api } from "../../convex/_generated/api";
 import {
   optimisticallySendMessage,
-  toUIMessages,
   useSmoothText,
-  useThreadMessages,
+  useUIMessages,
   type UIMessage,
 } from "@convex-dev/agent/react";
 import { useCallback, useEffect, useRef, useState } from "react";
@@ -68,11 +67,27 @@ export default function ChatStreaming() {
 }
 
 function Story({ threadId, reset }: { threadId: string; reset: () => void }) {
-  const messages = useThreadMessages(
+  const { results: messages } = useUIMessages(
     api.chat.streaming.listThreadMessages,
     { threadId },
     { initialNumItems: 10, stream: true },
   );
+
+  // If you don't want to use UIMessages, you can use this hook:
+  // (note: you'll need to return MessageDoc from listThreadMessages)
+  // const { results } = useThreadMessages(
+  //   api.chat.streaming.listThreadMessages,
+  //   { threadId },
+  //   { initialNumItems: 10, stream: true },
+  // );
+  // const messages = toUIMessages(results ?? []);
+
+  // If you only want to show streaming messages, you can use this hook:
+  // const messages =
+  //   useStreamingUIMessages(api.chat.streaming.listThreadMessages, {
+  //     threadId,
+  //     paginationOpts: { numItems: 10, cursor: null },
+  //   }) ?? [];
   const sendMessage = useMutation(
     api.chat.streaming.initiateAsyncStreaming,
   ).withOptimisticUpdate(
@@ -90,22 +105,22 @@ function Story({ threadId, reset }: { threadId: string; reset: () => void }) {
 
   useEffect(() => {
     scrollToBottom();
-  }, [messages.results]);
+  }, [messages]);
 
   function onSendClicked() {
     if (prompt.trim() === "") return;
     void sendMessage({ threadId, prompt }).catch(() => setPrompt(prompt));
-    setPrompt("");
+    setPrompt("Continue the story...");
   }
 
   return (
     <>
       <div className="h-full flex flex-col max-w-4xl mx-auto w-full">
         {/* Messages area - scrollable */}
         <div className="flex-1 overflow-y-auto p-6">
-          {messages.results?.length > 0 ? (
+          {messages.length > 0 ? (
             <div className="flex flex-col gap-4 whitespace-pre">
-              {toUIMessages(messages.results ?? []).map((m) => (
+              {messages.map((m) => (
                 <Message key={m.key} message={m} />
               ))}
               <div ref={messagesEndRef} />
@@ -132,17 +147,17 @@ function Story({ threadId, reset }: { threadId: string; reset: () => void }) {
               onChange={(e) => setPrompt(e.target.value)}
               className="flex-1 px-4 py-2 rounded-lg border border-gray-300 focus:outline-none focus:ring-2 focus:ring-blue-400 bg-gray-50"
               placeholder={
-                messages.results?.length > 0
+                messages.length > 0
                   ? "Continue the story..."
                   : "Tell me a story..."
               }
             />
-            {messages.results.find((m) => m.streaming) ? (
+            {messages.find((m) => m.status === "streaming") ? (
               <button
                 className="px-4 py-2 rounded-lg bg-red-100 text-red-700 hover:bg-red-200 transition font-medium self-end"
                 onClick={() => {
                   const order =
-                    messages.results.find((m) => m.streaming)?.order ?? 0;
+                    messages.find((m) => m.status === "streaming")?.order ?? 0;
                   void abortStreamByOrder({ threadId, order });
                 }}
                 type="button"
@@ -158,7 +173,7 @@ function Story({ threadId, reset }: { threadId: string; reset: () => void }) {
                 Send
               </button>
             )}
-            {messages.results?.length > 0 && (
+            {messages.length > 0 && (
               <button
                 className="px-4 py-2 rounded-lg bg-red-100 text-red-700 hover:bg-red-200 transition font-medium self-end"
                 onClick={() => {
@@ -187,6 +202,15 @@ function Message({ message }: { message: UIMessage }) {
     // wouldn't start streaming until the second chunk was received.
     startStreaming: message.status === "streaming",
   });
+  const [reasoningText] = useSmoothText(
+    message.parts
+      .filter((p) => p.type === "reasoning")
+      .map((p) => p.text)
+      .join("\n") ?? "",
+    {
+      startStreaming: message.status === "streaming",
+    },
+  );
   return (
     <div className={cn("flex", isUser ? "justify-end" : "justify-start")}>
       <div
@@ -199,6 +223,9 @@ function Message({ message }: { message: UIMessage }) {
           },
         )}
       >
+        {reasoningText && (
+          <div className="text-xs text-gray-500">💭{reasoningText}</div>
+        )}
         {visibleText || "..."}
       </div>
     </div>

example/ui/chat/ChatStreaming.tsx

@@ -7,7 +7,7 @@
     "email": "support@convex.dev",
     "url": "https://github.com/get-convex/agent/issues"
   },
-  "version": "0.2.7",
+  "version": "0.2.8-alpha.0",
   "license": "Apache-2.0",
   "keywords": [
     "convex",

package-lock.json

@@ -5,7 +5,7 @@ import reactRefresh from "eslint-plugin-react-refresh";
 import tseslint from "typescript-eslint";
 
 export default tseslint.config(
-  { ignores: ["dist"] },
+  { ignores: ["dist", "bin/agent-playground.js"] },
   {
     extends: [js.configs.recommended, ...tseslint.configs.recommended],
     files: ["**/*.{ts,tsx}"],

package.json

@@ -1,20 +1,31 @@
-import type {
-  UIMessage as AIUIMessage,
-  DeepPartial,
-  ReasoningUIPart,
-  SourceDocumentUIPart,
-  SourceUrlUIPart,
-  StepStartUIPart,
-  TextUIPart,
-  ToolUIPart,
-  UIDataTypes,
-  UITools,
+import {
+  convertToModelMessages,
+  type UIMessage as AIUIMessage,
+  type DeepPartial,
+  type ReasoningUIPart,
+  type SourceDocumentUIPart,
+  type SourceUrlUIPart,
+  type StepStartUIPart,
+  type TextUIPart,
+  type ToolUIPart,
+  type UIDataTypes,
+  type UITools,
 } from "ai";
-import { extractText, type MessageDoc } from "../client/index.js";
-import { deserializeMessage, toUIFilePart } from "../mapping.js";
-import type { MessageStatus, SourcePart, vSource } from "../validators.js";
-import { sorted } from "../shared.js";
 import type { Infer } from "convex/values";
+import {
+  deserializeMessage,
+  fromModelMessage,
+  toUIFilePart,
+} from "./mapping.js";
+import { extractReasoning, extractText, isTool, sorted } from "./shared.js";
+import type {
+  MessageDoc,
+  MessageStatus,
+  ProviderOptions,
+  SourcePart,
+  vSource,
+} from "./validators.js";
+import { omit, pick } from "convex-helpers";
 
 export type UIMessage<
   METADATA = unknown,
@@ -30,11 +41,117 @@ export type UIMessage<
   _creationTime: number;
 };
 
+/**
+ * Converts a list of UIMessages to MessageDocs, along with extra metadata that
+ * may be available to associate with the MessageDocs.
+ * @param messages - The UIMessages to convert to MessageDocs.
+ * @param meta - The metadata to add to the MessageDocs.
+ * @returns
+ */
+export function fromUIMessages<METADATA = unknown>(
+  messages: UIMessage<METADATA>[],
+  meta: {
+    threadId: string;
+    userId?: string;
+    model?: string;
+    provider?: string;
+    providerOptions?: ProviderOptions;
+    metadata?: METADATA;
+  },
+): (MessageDoc & { streaming: boolean; metadata?: METADATA })[] {
+  return messages.flatMap((uiMessage) => {
+    const stepOrder = uiMessage.stepOrder;
+    const commonFields = {
+      ...pick(meta, [
+        "threadId",
+        "userId",
+        "model",
+        "provider",
+        "providerOptions",
+        "metadata",
+      ]),
+      ...omit(uiMessage, ["parts", "role", "key", "text"]),
+      status: uiMessage.status === "streaming" ? "pending" : "success",
+      streaming: uiMessage.status === "streaming",
+      // to override
+      _id: uiMessage.id,
+      tool: false,
+    } satisfies MessageDoc & { streaming: boolean; metadata?: METADATA };
+    const modelMessages = convertToModelMessages([uiMessage]);
+    return modelMessages
+      .map((modelMessage, i) => {
+        if (modelMessage.content.length === 0) {
+          return undefined;
+        }
+        const message = fromModelMessage(modelMessage);
+        const tool = isTool(message);
+        const doc: MessageDoc & { streaming: boolean; metadata?: METADATA } = {
+          ...commonFields,
+          _id: uiMessage.id + `-${i}`,
+          stepOrder: stepOrder + i,
+          message,
+          tool,
+          text: extractText(message),
+          reasoning: extractReasoning(message),
+          finishReason: tool ? "tool-calls" : "stop",
+          sources: fromSourceParts(uiMessage.parts),
+        };
+        if (Array.isArray(modelMessage.content)) {
+          const providerOptions = modelMessage.content.find(
+            (c) => c.providerOptions,
+          )?.providerOptions;
+          if (providerOptions) {
+            // convertToModelMessages changes providerMetadata to providerOptions
+            doc.providerMetadata = providerOptions;
+            doc.providerOptions ??= providerOptions;
+          }
+        }
+        return doc;
+      })
+      .filter((d) => d !== undefined);
+  });
+}
+
+function fromSourceParts(parts: UIMessage["parts"]): Infer<typeof vSource>[] {
+  return parts
+    .map((part) => {
+      if (part.type === "source-url") {
+        return {
+          type: "source",
+          sourceType: "url",
+          url: part.url,
+          id: part.sourceId,
+          providerMetadata: part.providerMetadata,
+          title: part.title,
+        } satisfies Infer<typeof vSource>;
+      }
+      if (part.type === "source-document") {
+        return {
+          type: "source",
+          sourceType: "document",
+          mediaType: part.mediaType,
+          id: part.sourceId,
+          providerMetadata: part.providerMetadata,
+          title: part.title,
+        } satisfies Infer<typeof vSource>;
+      }
+      return undefined;
+    })
+    .filter((p) => p !== undefined);
+}
+
 type ExtraFields<METADATA = unknown> = {
   streaming?: boolean;
   metadata?: METADATA;
 };
 
+/**
+ * Converts a list of MessageDocs to UIMessages.
+ * This is somewhat lossy, as many fields are not supported by UIMessages, e.g.
+ * the model, provider, userId, etc.
+ * The UIMessage type is the augmented type that includes more fields such as
+ * key, order, stepOrder, status, agentName, text, etc.
+ */
 export function toUIMessages<
   METADATA = unknown,
   DATA_PARTS extends UIDataTypes = UIDataTypes,