feat: Smooth status (#615)

Yonom · Jul 30, 2024 · 0fca66b · 0fca66b
1 parent f30285c
commit 0fca66b
Show file tree

Hide file tree

Showing 16 changed files with 291 additions and 136 deletions.
diff --git a/packages/react-markdown/package.json b/packages/react-markdown/package.json
@@ -40,6 +40,7 @@
     "build": "tsx scripts/build.mts"
   },
   "dependencies": {
+    "@radix-ui/react-primitive": "^2.0.0",
     "@radix-ui/react-use-callback-ref": "^1.1.0",
     "classnames": "^2.5.1",
     "lucide-react": "^0.416.0",

diff --git a/packages/react-markdown/src/primitives/MarkdownText.tsx b/packages/react-markdown/src/primitives/MarkdownText.tsx
@@ -1,8 +1,13 @@
 "use client";
 
 import { useContentPartText } from "@assistant-ui/react";
-import { useSmooth } from "@assistant-ui/react/internal";
-import type { ComponentType, FC } from "react";
+import {
+  ElementRef,
+  ElementType,
+  forwardRef,
+  type ComponentPropsWithoutRef,
+  type ComponentType,
+} from "react";
 import ReactMarkdown, { type Options } from "react-markdown";
 import { SyntaxHighlighterProps, CodeHeaderProps } from "../overrides/types";
 import { PreOverride } from "../overrides/PreOverride";
@@ -14,11 +19,19 @@ import {
 } from "../overrides/defaultComponents";
 import { useCallbackRef } from "@radix-ui/react-use-callback-ref";
 import { CodeOverride } from "../overrides/CodeOverride";
+import { useSmooth } from "@assistant-ui/react/internal";
+import { Primitive } from "@radix-ui/react-primitive";
+import classNames from "classnames";
+
+type MarkdownTextPrimitiveElement = ElementRef<typeof Primitive.div>;
+type PrimitiveDivProps = ComponentPropsWithoutRef<typeof Primitive.div>;
 
 export type MarkdownTextPrimitiveProps = Omit<
   Options,
-  "components" | "children"
+  "components" | "children" | "asChild"
 > & {
+  containerProps?: Omit<PrimitiveDivProps, "children">;
+  containerComponent?: ElementType;
   components?: NonNullable<Options["components"]> & {
     SyntaxHighlighter?: ComponentType<SyntaxHighlighterProps>;
     CodeHeader?: ComponentType<CodeHeaderProps>;
@@ -32,44 +45,64 @@ export type MarkdownTextPrimitiveProps = Omit<
   };
   smooth?: boolean;
 };
-export const MarkdownTextPrimitive: FC<MarkdownTextPrimitiveProps> = ({
-  smooth = true,
-  components: userComponents,
-  ...rest
-}) => {
-  const {
-    part: { text },
-  } = useContentPartText();
-  const smoothText = useSmooth(text, smooth); // TODO loading indicator disappears before smooth animation ends
+export const MarkdownTextPrimitive = forwardRef<
+  MarkdownTextPrimitiveElement,
+  MarkdownTextPrimitiveProps
+>(
+  (
+    {
+      components: userComponents,
+      className,
+      containerProps,
+      containerComponent: Container = "div",
+      ...rest
+    },
+    forwardedRef,
+    smooth = true,
+  ) => {
+    const {
+      part: { text },
+      status,
+    } = useSmooth(useContentPartText(), smooth);
 
-  const {
-    pre = DefaultPre,
-    code = DefaultCode,
-    SyntaxHighlighter = DefaultCodeBlockContent,
-    CodeHeader = DefaultCodeHeader,
-    by_language,
-    ...componentsRest
-  } = userComponents ?? {};
-  const components: typeof userComponents = {
-    ...componentsRest,
-    pre: PreOverride,
-    code: useCallbackRef((props) => (
-      <CodeOverride
-        components={{
-          Pre: pre,
-          Code: code,
-          SyntaxHighlighter,
-          CodeHeader,
-          by_language,
-        }}
-        {...props}
-      />
-    )),
-  };
+    const {
+      pre = DefaultPre,
+      code = DefaultCode,
+      SyntaxHighlighter = DefaultCodeBlockContent,
+      CodeHeader = DefaultCodeHeader,
+      by_language,
+      ...componentsRest
+    } = userComponents ?? {};
+    const components: typeof userComponents = {
+      ...componentsRest,
+      pre: PreOverride,
+      code: useCallbackRef((props) => (
+        <CodeOverride
+          components={{
+            Pre: pre,
+            Code: code,
+            SyntaxHighlighter,
+            CodeHeader,
+            by_language,
+          }}
+          {...props}
+        />
+      )),
+    };
 
-  return (
-    <ReactMarkdown components={components} {...rest}>
-      {smoothText}
-    </ReactMarkdown>
-  );
-};
+    return (
+      <Container
+        data-status={status.type}
+        {...containerProps}
+        className={classNames(className, containerProps?.className)}
+        ref={forwardedRef}
+      >
+        <ReactMarkdown components={components} {...rest}>
+          {text}
+        </ReactMarkdown>
+      </Container>
+    );
+  },
+);
+
+MarkdownTextPrimitive.displayName = "MarkdownTextPrimitive";
diff --git a/packages/react-markdown/src/styles/tailwindcss/markdown.css b/packages/react-markdown/src/styles/tailwindcss/markdown.css
@@ -1,107 +1,100 @@
-/* in progress indicator */
-:where(.aui-md-in-progress:empty)::after,
-:where(.aui-md-in-progress > :not(ol):not(ul):not(pre):last-child)::after,
-:where(.aui-md-in-progress > pre:last-child code)::after,
-:where(
-    .aui-md-in-progress
-      > :is(ol, ul):last-child
-      > li:last-child:not(:has(* > li))
-  )::after,
-:where(
-    .aui-md-in-progress
-      > :is(ol, ul):last-child
-      > li:last-child
-      > :is(ol, ul):last-child
-      > li:last-child:not(:has(* > li))
-  )::after,
-:where(
-    .aui-md-in-progress
-      > :is(ol, ul):last-child
-      > li:last-child
-      > :is(ol, ul):last-child
-      > li:last-child
-      > :is(ol, ul):last-child
-      > li:last-child
-  )::after {
+/* running indicator */
+:where(.aui-md-running):empty::after,
+:where(.aui-md-running) > :where(:not(ol):not(ul):not(pre)):last-child::after,
+:where(.aui-md-running) > pre:last-child code::after,
+:where(.aui-md-running)
+  > :where(:is(ol, ul):last-child)
+  > :where(li:last-child:not(:has(* > li)))::after,
+:where(.aui-md-running)
+  > :where(:is(ol, ul):last-child)
+  > :where(li:last-child)
+  > :where(:is(ol, ul):last-child)
+  > :where(li:last-child:not(:has(* > li)))::after,
+:where(.aui-md-running)
+  > :where(:is(ol, ul):last-child)
+  > :where(li:last-child)
+  > :where(:is(ol, ul):last-child)
+  > :where(li:last-child)
+  > :where(:is(ol, ul):last-child)
+  > :where(li:last-child)::after {
   @apply animate-pulse font-sans content-['\25CF'] ltr:ml-1 rtl:mr-1;
 }
-
 /* typography */
 
-:where(.aui-md-root) h1 {
+.aui-md-root h1 {
   @apply mb-8 scroll-m-20 text-4xl font-extrabold tracking-tight last:mb-0;
 }
 
-:where(.aui-md-root) h2 {
+.aui-md-root h2 {
   @apply mb-4 mt-8 scroll-m-20 text-3xl font-semibold tracking-tight first:mt-0 last:mb-0;
 }
 
-:where(.aui-md-root) h3 {
+.aui-md-root h3 {
   @apply mb-4 mt-6 scroll-m-20 text-2xl font-semibold tracking-tight first:mt-0 last:mb-0;
 }
 
-:where(.aui-md-root) h4 {
+.aui-md-root h4 {
   @apply mb-4 mt-6 scroll-m-20 text-xl font-semibold tracking-tight first:mt-0 last:mb-0;
 }
 
-:where(.aui-md-root) h5 {
+.aui-md-root h5 {
   @apply my-4 text-lg font-semibold first:mt-0 last:mb-0;
 }
 
-:where(.aui-md-root) h6 {
+.aui-md-root h6 {
   @apply my-4 font-semibold first:mt-0 last:mb-0;
 }
 
-:where(.aui-md-root) p {
+.aui-md-root p {
   @apply mb-5 mt-5 leading-7 first:mt-0 last:mb-0;
 }
 
-:where(.aui-md-root) a {
+.aui-md-root a {
   @apply text-aui-primary font-medium underline underline-offset-4;
 }
 
-:where(.aui-md-root) blockquote {
+.aui-md-root blockquote {
   @apply border-l-2 pl-6 italic;
 }
 
-:where(.aui-md-root) ul {
+.aui-md-root ul {
   @apply my-5 ml-6 list-disc [&>li]:mt-2;
 }
 
-:where(.aui-md-root) ol {
+.aui-md-root ol {
   @apply my-5 ml-6 list-decimal [&>li]:mt-2;
 }
 
-:where(.aui-md-root) hr {
+.aui-md-root hr {
   @apply my-5 border-b;
 }
 
-:where(.aui-md-root) table {
+.aui-md-root table {
   @apply my-5 w-full border-separate border-spacing-0 overflow-y-auto;
 }
 
-:where(.aui-md-root) th {
+.aui-md-root th {
   @apply bg-aui-muted px-4 py-2 text-left font-bold first:rounded-tl-lg last:rounded-tr-lg [&[align=center]]:text-center [&[align=right]]:text-right;
 }
 
-:where(.aui-md-root) td {
+.aui-md-root td {
   @apply border-b border-l px-4 py-2 text-left last:border-r [&[align=center]]:text-center [&[align=right]]:text-right;
 }
 
-:where(.aui-md-root) tr {
+.aui-md-root tr {
   @apply m-0 border-b p-0 first:border-t [&:last-child>td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg;
 }
 
-:where(.aui-md-root) sup {
+.aui-md-root sup {
   @apply [&>a]:text-xs [&>a]:no-underline;
 }
 
-:where(.aui-md-root) pre {
+.aui-md-root pre {
   @apply overflow-x-auto rounded-b-lg bg-black p-4 text-white;
 }
 
-:where(.aui-md-root) > code,
-:where(.aui-md-root) :not(:where(pre)) code {
+.aui-md-root > code,
+.aui-md-root :not(:where(pre)) code {
   @apply bg-aui-muted rounded border font-semibold;
 }
 

diff --git a/packages/react-markdown/src/ui/markdown-text.tsx b/packages/react-markdown/src/ui/markdown-text.tsx
@@ -1,11 +1,14 @@
-import { TextContentPartProps } from "@assistant-ui/react";
 import { FC, memo } from "react";
 import { CodeHeader } from "./code-header";
 import classNames from "classnames";
 import {
   MarkdownTextPrimitive,
   MarkdownTextPrimitiveProps,
 } from "../primitives/MarkdownText";
+import {
+  withSmoothContextProvider,
+  useSmoothStatus,
+} from "@assistant-ui/react/internal";
 
 export type MakeMarkdownTextProps = MarkdownTextPrimitiveProps;
 
@@ -19,23 +22,21 @@ export const makeMarkdownText = ({
     CodeHeader: userComponents?.CodeHeader ?? CodeHeader,
   };
 
-  const MarkdownTextImpl: FC<TextContentPartProps> = ({ status }) => {
+  const MarkdownTextImpl: FC = () => {
+    const status = useSmoothStatus();
     return (
-      <div
+      <MarkdownTextPrimitive
+        components={components}
+        {...rest}
         className={classNames(
           "aui-md-root",
-          status.type === "running" && "aui-md-in-progress",
+          status.type === "running" && "aui-md-running",
           className,
         )}
-      >
-        <MarkdownTextPrimitive components={components} {...rest} />
-      </div>
+      />
     );
   };
   MarkdownTextImpl.displayName = "MarkdownText";
 
-  return memo(
-    MarkdownTextImpl,
-    (prev, next) => prev.status.type === next.status.type,
-  );
+  return memo(withSmoothContextProvider(MarkdownTextImpl), () => true);
 };
diff --git a/packages/react/src/context/providers/ContentPartProvider.tsx b/packages/react/src/context/providers/ContentPartProvider.tsx
@@ -12,7 +12,7 @@ import {
   ThreadAssistantContentPart,
   ThreadMessage,
   ThreadUserContentPart,
-  ToolContentPartStatus,
+  ToolCallContentPartStatus,
 } from "../../types/AssistantTypes";
 
 type ContentPartProviderProps = PropsWithChildren<{
@@ -27,7 +27,7 @@ const toContentPartStatus = (
   message: ThreadMessage,
   partIndex: number,
   part: ThreadUserContentPart | ThreadAssistantContentPart,
-): ToolContentPartStatus => {
+): ToolCallContentPartStatus => {
   if (message.role !== "assistant") return COMPLETE_STATUS;
 
   const isLastPart = partIndex === Math.max(0, message.content.length - 1);
@@ -48,7 +48,7 @@ const toContentPartStatus = (
     return COMPLETE_STATUS;
   }
 
-  return message.status as ToolContentPartStatus;
+  return message.status as ToolCallContentPartStatus;
 };
 
 export const EMPTY_CONTENT = Object.freeze({ type: "text", text: "" });

diff --git a/packages/react/src/context/stores/ContentPart.ts b/packages/react/src/context/stores/ContentPart.ts
@@ -3,7 +3,7 @@ import type {
   ImageContentPart,
   TextContentPart,
   ToolCallContentPart,
-  ToolContentPartStatus,
+  ToolCallContentPartStatus,
   UIContentPart,
 } from "../../types/AssistantTypes";
 
@@ -23,12 +23,12 @@ export type UIContentPartState = Readonly<{
 }>;
 
 export type ToolCallContentPartState = Readonly<{
-  status: ToolContentPartStatus;
+  status: ToolCallContentPartStatus;
   part: ToolCallContentPart;
 }>;
 
 export type ContentPartState = Readonly<{
-  status: ToolContentPartStatus;
+  status: ContentPartStatus | ToolCallContentPartStatus;
   part:
     | TextContentPart
     | ImageContentPart