feat: smooth text streaming (#401)

.changeset/happy-snakes-design.md

@@ -0,0 +1,7 @@
+---
+"@assistant-ui/react-markdown": patch
+"@assistant-ui/react-ui": patch
+"@assistant-ui/react": patch
+---
+
+feat: smooth text streaming

packages/react-markdown/src/MarkdownText.tsx

@@ -1,20 +1,19 @@
-import { useContentPartContext } from "@assistant-ui/react";
+import { INTERNAL, useContentPartText } from "@assistant-ui/react";
 import type { FC } from "react";
 import ReactMarkdown, { type Options } from "react-markdown";
 
-export type MarkdownTextPrimitiveProps = Omit<Options, "children">;
+const { useSmooth } = INTERNAL;
+
+export type MarkdownTextPrimitiveProps = Omit<Options, "children"> & {
+  smooth?: boolean;
+};
 
 export const MarkdownTextPrimitive: FC<MarkdownTextPrimitiveProps> = (
   options,
 ) => {
-  const { useContentPart } = useContentPartContext();
-  const text = useContentPart((c) => {
-    if (c.part.type !== "text")
-      throw new Error(
-        "This component can only be used inside text content parts.",
-      );
-
-    return c.part.text;
-  });
-  return <ReactMarkdown {...options}>{text}</ReactMarkdown>;
+  const {
+    part: { text },
+  } = useContentPartText();
+  const smoothText = useSmooth(text, options.smooth);
+  return <ReactMarkdown {...options}>{smoothText}</ReactMarkdown>;
 };

packages/react-ui/src/components/markdown-text.tsx

@@ -18,6 +18,7 @@ const MarkdownTextImpl: FC<
         (!!className ? " " + className : "")
       }
       remarkPlugins={[remarkGfm, ...(remarkPlugins ?? [])]}
+      smooth // TODO figure out the default for this
       {...rest}
     />
   );

packages/react-ui/src/components/text.tsx

@@ -11,7 +11,7 @@ export const Text: FC<TextContentPartProps> = ({ status }) => {
         "aui-text" + (status === "in_progress" ? " aui-text-in-progress" : "")
       }
     >
-      <ContentPartPrimitive.Text />
+      <ContentPartPrimitive.Text smooth />
     </p>
   );
 };

packages/react/src/internal.ts

@@ -1,3 +1,4 @@
 export { ProxyConfigProvider } from "./utils/ProxyConfigProvider";
 export { MessageRepository } from "./runtime/utils/MessageRepository";
 export { BaseAssistantRuntime } from "./runtime/core/BaseAssistantRuntime";
+export { useSmooth } from "./utils/hooks/useSmooth";

packages/react/src/primitive-hooks/contentPart/useContentPartDisplay.tsx

@@ -1,4 +1,5 @@
 import { useContentPartContext } from "../../context/react/ContentPartContext";
+import { UIContentPartState } from "../../context/stores/ContentPart";
 
 export const useContentPartDisplay = () => {
   const { useContentPart } = useContentPartContext();
@@ -9,7 +10,7 @@ export const useContentPartDisplay = () => {
         "This component can only be used inside ui content parts.",
       );
 
-    return c.part.display;
+    return c as UIContentPartState;
   });
 
   return display;

packages/react/src/primitive-hooks/contentPart/useContentPartImage.tsx

@@ -1,4 +1,5 @@
 import { useContentPartContext } from "../../context/react/ContentPartContext";
+import { ImageContentPartState } from "../../context/stores/ContentPart";
 
 export const useContentPartImage = () => {
   const { useContentPart } = useContentPartContext();
@@ -9,7 +10,7 @@ export const useContentPartImage = () => {
         "ContentPartImage can only be used inside image content parts.",
       );
 
-    return c.part.image;
+    return c as ImageContentPartState;
   });
 
   return image;

packages/react/src/primitives/contentPart/ContentPartDisplay.tsx

@@ -6,7 +6,9 @@ export type ContentPartPrimitiveDisplayProps = {};
 export const ContentPartPrimitiveDisplay: FC<
   ContentPartPrimitiveDisplayProps
 > = () => {
-  const display = useContentPartDisplay();
+  const {
+    part: { display },
+  } = useContentPartDisplay();
   return display ?? null;
 };
 

packages/react/src/primitives/contentPart/ContentPartImage.tsx

@@ -11,7 +11,9 @@ export const ContentPartPrimitiveImage = forwardRef<
   ContentPartPrimitiveImageElement,
   ContentPartPrimitiveImageProps
 >((props, forwardedRef) => {
-  const image = useContentPartImage();
+  const {
+    part: { image },
+  } = useContentPartImage();
   return <Primitive.img src={image} {...props} ref={forwardedRef} />;
 });
 

packages/react/src/primitives/contentPart/ContentPartText.tsx

@@ -1,27 +1,29 @@
 import { Primitive } from "@radix-ui/react-primitive";
 import { type ElementRef, forwardRef, ComponentPropsWithoutRef } from "react";
 import { useContentPartText } from "../../primitive-hooks/contentPart/useContentPartText";
+import { useSmooth } from "../../utils/hooks/useSmooth";
 
 type ContentPartPrimitiveTextElement = ElementRef<typeof Primitive.span>;
 type PrimitiveSpanProps = ComponentPropsWithoutRef<typeof Primitive.span>;
 
 export type ContentPartPrimitiveTextProps = Omit<
   PrimitiveSpanProps,
   "children"
->;
+> & { smooth?: boolean };
 
 export const ContentPartPrimitiveText = forwardRef<
   ContentPartPrimitiveTextElement,
   ContentPartPrimitiveTextProps
->((props, forwardedRef) => {
+>(({ smooth, ...rest }, forwardedRef) => {
   const {
-    part: { text },
     status,
+    part: { text },
   } = useContentPartText();
+  const smoothText = useSmooth(text, smooth);
 
   return (
-    <Primitive.span data-status={status} {...props} ref={forwardedRef}>
-      {text}
+    <Primitive.span data-status={status} {...rest} ref={forwardedRef}>
+      {smoothText}
     </Primitive.span>
   );
 });

packages/react/src/utils/hooks/useSmooth.tsx

@@ -0,0 +1,85 @@
+import { useEffect, useState } from "react";
+
+class TextStreamAnimator {
+  private animationFrameId: number | null = null;
+  private lastUpdateTime: number = Date.now();
+  private decayFactor: number = 0.99;
+
+  private _targetText: string = "";
+  get targetText() {
+    return this._targetText;
+  }
+  set targetText(targetText: string) {
+    this._targetText = targetText;
+    if (this.animationFrameId === null) {
+      this.animate();
+    }
+  }
+
+  constructor(
+    private setText: (callback: (prevText: string) => string) => void,
+  ) {}
+
+  stop() {
+    if (this.animationFrameId !== null) {
+      cancelAnimationFrame(this.animationFrameId);
+      this.animationFrameId = null;
+    }
+  }
+
+  private animate = () => {
+    const currentTime = Date.now();
+    const deltaTime = currentTime - this.lastUpdateTime;
+    this.lastUpdateTime = currentTime;
+
+    this.setText((currentText) => {
+      const targetText = this._targetText;
+
+      if (currentText === targetText) {
+        this.animationFrameId = null;
+        return currentText;
+      }
+
+      const remainingChars = targetText.length - currentText.length;
+      const charsToAdd = Math.max(
+        1,
+        Math.floor(
+          remainingChars * (1 - Math.pow(this.decayFactor, deltaTime)),
+        ),
+      );
+      const newText = targetText.slice(0, currentText.length + charsToAdd);
+      this.animationFrameId = requestAnimationFrame(this.animate);
+      return newText;
+    });
+  };
+}
+
+export const useSmooth = (text: string, smooth: boolean = false) => {
+  const [displayedText, setDisplayedText] = useState(text);
+  const [animatorRef] = useState<TextStreamAnimator>(
+    new TextStreamAnimator(setDisplayedText),
+  );
+
+  useEffect(() => {
+    if (!smooth) {
+      animatorRef.stop();
+      return;
+    }
+
+    if (!text.startsWith(animatorRef.targetText)) {
+      setDisplayedText(text);
+      animatorRef.stop();
+      return;
+    }
+
+    animatorRef.targetText = text;
+  }, [animatorRef, smooth, text]);
+
+  useEffect(() => {
+    return () => {
+      animatorRef.stop();
+    };
+  }, [animatorRef]);
+
+  return smooth ? displayedText : text;
+};