feat(nx-dev): add "new chat" button to AI page

nx-dev/feature-ai/src/lib/error-message.tsx

@@ -1,9 +1,10 @@
+import { type JSX, memo } from 'react';
 import {
-  XCircleIcon,
   ExclamationTriangleIcon,
+  XCircleIcon,
 } from '@heroicons/react/24/outline';
 
-export function ErrorMessage({ error }: { error: any }): JSX.Element {
+function ErrorMessage({ error }: { error: any }): JSX.Element {
   try {
     if (error.message) {
       error = JSON.parse(error.message);
@@ -57,3 +58,6 @@ export function ErrorMessage({ error }: { error: any }): JSX.Element {
     );
   }
 }
+
+const MemoErrorMessage = memo(ErrorMessage);
+export { MemoErrorMessage as ErrorMessage };

nx-dev/feature-ai/src/lib/feed-container.tsx

@@ -1,11 +1,20 @@
 import { sendCustomEvent } from '@nx/nx-dev/feature-analytics';
-import { RefObject, useEffect, useRef, useState } from 'react';
+import {
+  type FormEvent,
+  type JSX,
+  RefObject,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import { ErrorMessage } from './error-message';
 import { Feed } from './feed/feed';
 import { LoadingState } from './loading-state';
 import { Prompt } from './prompt';
 import { getQueryFromUid, storeQueryForUid } from '@nx/nx-dev/util-ai';
 import { Message, useChat } from 'ai/react';
+import { cx } from '@nx/nx-dev/ui-primitives';
 
 const assistantWelcome: Message = {
   id: 'first-custom-message',
@@ -18,34 +27,63 @@ export function FeedContainer(): JSX.Element {
   const [error, setError] = useState<Error | null>(null);
   const [startedReply, setStartedReply] = useState(false);
 
-  const feedContainer: RefObject<HTMLDivElement> | undefined = useRef(null);
-  const { messages, input, handleInputChange, handleSubmit, isLoading } =
-    useChat({
-      api: '/api/query-ai-handler',
-      onError: (error) => {
-        setError(error);
-      },
-      onResponse: (_response) => {
-        setStartedReply(true);
-        sendCustomEvent('ai_query', 'ai', 'query', undefined, {
-          query: input,
-        });
-        setError(null);
-      },
-      onFinish: (response: Message) => {
-        setStartedReply(false);
-        storeQueryForUid(response.id, input);
-      },
-    });
+  const wrapperRef: RefObject<HTMLDivElement> | undefined = useRef(null);
+
+  const {
+    messages,
+    setMessages,
+    input,
+    handleInputChange,
+    handleSubmit: _handleSubmit,
+    stop: handleStopGenerating,
+    reload: handleRegenerate,
+    isLoading,
+  } = useChat({
+    api: '/api/query-ai-handler',
+    onError: (error) => {
+      setError(error);
+    },
+    onResponse: (_response) => {
+      setStartedReply(true);
+      sendCustomEvent('ai_query', 'ai', 'query', undefined, {
+        query: input,
+      });
+      setError(null);
+    },
+    onFinish: (response: Message) => {
+      setStartedReply(false);
+      storeQueryForUid(response.id, input);
+    },
+  });
 
+  const hasReply = useMemo(() => messages.length > 0, [messages]);
+
+  // Scroll the wrapper to the bottom if the user is at the bottom.
+  // Otherwise, if the user manually scrolled away from the bottom, don't scroll.
+  const userHasScrolled = useRef(false);
+  const handleScroll = () => {
+    userHasScrolled.current = true;
+  };
   useEffect(() => {
-    if (feedContainer.current) {
-      const elements =
-        feedContainer.current.getElementsByClassName('feed-item');
-      elements[elements.length - 1].scrollIntoView({ behavior: 'smooth' });
+    if (wrapperRef?.current) {
+      const el = wrapperRef.current;
+      if (!userHasScrolled.current) {
+        el.scrollTo(0, el.scrollHeight);
+      }
     }
   }, [messages, isLoading]);
 
+  const handleSubmit = (event: FormEvent<HTMLFormElement>) => {
+    userHasScrolled.current = false;
+    _handleSubmit(event);
+  };
+
+  const handleNewChat = () => {
+    setMessages([]);
+    setError(null);
+    setStartedReply(false);
+  };
+
   const handleFeedback = (statement: 'good' | 'bad', chatItemUid: string) => {
     const query = getQueryFromUid(chatItemUid);
     sendCustomEvent('ai_feedback', 'ai', statement, undefined, {
@@ -57,6 +95,8 @@ export function FeedContainer(): JSX.Element {
     <>
       {/*WRAPPER*/}
       <div
+        ref={wrapperRef}
+        onScroll={handleScroll}
         id="wrapper"
         data-testid="wrapper"
         className="relative flex flex-grow flex-col items-stretch justify-start overflow-y-scroll"
@@ -68,28 +108,31 @@ export function FeedContainer(): JSX.Element {
           >
             <div className="relative min-w-0 flex-auto">
               {/*MAIN CONTENT*/}
-              <div
-                ref={feedContainer}
-                data-document="main"
-                className="relative"
-              >
+              <div data-document="main" className="relative pb-24">
                 <Feed
                   activity={!!messages.length ? messages : [assistantWelcome]}
-                  handleFeedback={(statement, chatItemUid) =>
-                    handleFeedback(statement, chatItemUid)
-                  }
+                  onFeedback={handleFeedback}
                 />
 
                 {/* Change this message if it's loading but it's writing as well  */}
                 {isLoading && !startedReply && <LoadingState />}
                 {error && <ErrorMessage error={error} />}
 
-                <div className="sticky bottom-0 left-0 right-0 w-full pt-6 pb-4 bg-gradient-to-t from-white via-white dark:from-slate-900 dark:via-slate-900">
+                <div
+                  className={cx(
+                    'fixed bottom-0 left0 right-0 w-full py-4 px-4 lg:py-6 lg:px-0',
+                    'bg-gradient-to-t from-white via-white/75 dark:from-slate-900 dark:via-slate-900/75'
+                  )}
+                >
                   <Prompt
-                    handleSubmit={handleSubmit}
-                    handleInputChange={handleInputChange}
+                    onSubmit={handleSubmit}
+                    onInputChange={handleInputChange}
+                    onNewChat={handleNewChat}
+                    onStopGenerating={handleStopGenerating}
+                    onRegenerate={handleRegenerate}
                     input={input}
-                    isDisabled={isLoading}
+                    isGenerating={isLoading}
+                    showNewChatCta={!isLoading && hasReply}
                   />
                 </div>
               </div>

nx-dev/feature-ai/src/lib/feed/feed-answer.tsx

@@ -66,7 +66,7 @@ export function FeedAnswer({
           <ReactMarkdown children={content} />
         </div>
         {!isFirst && (
-          <div className="group text-xs flex-1 md:flex md:justify-end gap-4 md:items-center text-slate-400 hover:text-slate-500 transition">
+          <div className="group text-md flex-1 md:flex md:justify-end gap-4 md:items-center text-slate-400 hover:text-slate-500 transition">
             {feedbackStatement ? (
               <p className="italic group-hover:flex">
                 {feedbackStatement === 'good'
@@ -89,7 +89,7 @@ export function FeedAnswer({
                 title="Bad"
               >
                 <span className="sr-only">Bad answer</span>
-                <HandThumbDownIcon className="h-5 w-5" aria-hidden="true" />
+                <HandThumbDownIcon className="h-6 w-6" aria-hidden="true" />
               </button>
               <button
                 className={cx(
@@ -101,7 +101,7 @@ export function FeedAnswer({
                 title="Good"
               >
                 <span className="sr-only">Good answer</span>
-                <HandThumbUpIcon className="h-5 w-5" aria-hidden="true" />
+                <HandThumbUpIcon className="h-6 w-6" aria-hidden="true" />
               </button>
             </div>
           </div>

nx-dev/feature-ai/src/lib/feed/feed.tsx

@@ -4,10 +4,10 @@ import { Message } from 'ai/react';
 
 export function Feed({
   activity,
-  handleFeedback,
+  onFeedback,
 }: {
   activity: Message[];
-  handleFeedback: (statement: 'bad' | 'good', chatItemUid: string) => void;
+  onFeedback: (statement: 'bad' | 'good', chatItemUid: string) => void;
 }) {
   return (
     <div className="flow-root my-12">
@@ -21,7 +21,7 @@ export function Feed({
               <FeedAnswer
                 content={activityItem.content}
                 feedbackButtonCallback={(statement) =>
-                  handleFeedback(statement, activityItem.id)
+                  onFeedback(statement, activityItem.id)
                 }
                 isFirst={activityItemIdx === 0}
               />

nx-dev/feature-ai/src/lib/prompt.tsx

@@ -1,40 +1,114 @@
-import { ChangeEvent, FormEvent, useEffect, useRef } from 'react';
-import { PaperAirplaneIcon } from '@heroicons/react/24/outline';
+import { ChangeEvent, FormEvent, useEffect, useRef, useState } from 'react';
+import {
+  ArrowPathIcon,
+  PaperAirplaneIcon,
+  PlusIcon,
+  StopIcon,
+} from '@heroicons/react/24/outline';
 import { Button } from '@nx/nx-dev/ui-common';
 import Textarea from 'react-textarea-autosize';
 import { ChatRequestOptions } from 'ai';
+import { cx } from '@nx/nx-dev/ui-primitives';
 
 export function Prompt({
-  isDisabled,
-  handleSubmit,
-  handleInputChange,
+  isGenerating,
+  showNewChatCta,
+  onSubmit,
+  onInputChange,
+  onNewChat,
+  onStopGenerating,
+  onRegenerate,
   input,
 }: {
-  isDisabled: boolean;
-  handleSubmit: (
+  isGenerating: boolean;
+  showNewChatCta: boolean;
+  onSubmit: (
     e: FormEvent<HTMLFormElement>,
     chatRequestOptions?: ChatRequestOptions | undefined
   ) => void;
-  handleInputChange: (
+  onInputChange: (
     e: ChangeEvent<HTMLTextAreaElement> | ChangeEvent<HTMLInputElement>
   ) => void;
+  onNewChat: () => void;
+  onStopGenerating: () => void;
+  onRegenerate: () => void;
   input: string;
 }) {
   const formRef = useRef<HTMLFormElement>(null);
   const inputRef = useRef<HTMLTextAreaElement>(null);
+  const [isStopped, setStopped] = useState(false);
 
   useEffect(() => {
-    if (inputRef.current) {
-      inputRef.current.focus();
-    }
-  }, []);
+    if (!isGenerating) inputRef.current?.focus();
+  }, [isGenerating]);
+  const handleSubmit = (event: FormEvent<HTMLFormElement>) => {
+    setStopped(false);
+    if (inputRef.current?.value.trim()) onSubmit(event);
+    else event.preventDefault();
+  };
+
+  const handleNewChat = () => {
+    onNewChat();
+    inputRef.current?.focus();
+  };
+
+  const handleStopGenerating = () => {
+    setStopped(true);
+    onStopGenerating();
+    inputRef.current?.focus();
+  };
+
+  const handleOnRegenerate = () => {
+    setStopped(false);
+    onRegenerate();
+  };
 
   return (
     <form
       ref={formRef}
       onSubmit={handleSubmit}
-      className="relative flex gap-2 max-w-2xl mx-auto py-0 px-2 shadow-lg rounded-md border border-slate-300 bg-white dark:border-slate-900 dark:bg-slate-700"
+      className="relative flex gap-2 max-w-3xl mx-auto py-0 px-2 shadow-lg rounded-md border border-slate-300 bg-white dark:border-slate-900 dark:bg-slate-700"
     >
+      <div
+        className={cx(
+          'absolute -top-full left-1/2 mt-1 -translate-x-1/2',
+          'flex gap-4'
+        )}
+      >
+        {isGenerating && (
+          <Button
+            variant="secondary"
+            size="small"
+            className={cx('bg-white dark:bg-slate-900')}
+            onClick={handleStopGenerating}
+          >
+            <StopIcon aria-hidden="true" className="h-5 w-5" />
+            <span className="text-base">Stop generating</span>
+          </Button>
+        )}
+        {showNewChatCta && (
+          <Button
+            variant="secondary"
+            size="small"
+            className={cx('bg-white dark:bg-slate-900')}
+            onClick={handleNewChat}
+          >
+            <PlusIcon aria-hidden="true" className="h-5 w-5" />
+            <span className="text-base">New Chat</span>
+          </Button>
+        )}
+        {isStopped && (
+          <Button
+            variant="secondary"
+            size="small"
+            className={cx('bg-white dark:bg-slate-900')}
+            onClick={handleOnRegenerate}
+          >
+            <ArrowPathIcon aria-hidden="true" className="h-5 w-5" />
+            <span className="text-base">Regenerated</span>
+          </Button>
+        )}
+      </div>
       <div className="overflow-y-auto w-full h-full max-h-[300px]">
         <Textarea
           onKeyDown={(event) => {
@@ -49,10 +123,10 @@ export function Prompt({
           }}
           ref={inputRef}
           value={input}
-          onChange={handleInputChange}
+          onChange={onInputChange}
           id="query-prompt"
           name="query"
-          disabled={isDisabled}
+          disabled={isGenerating}
           className="block w-full p-0 resize-none bg-transparent text-sm placeholder-slate-500 pl-2 py-[1.3rem] focus-within:outline-none focus:placeholder-slate-400 dark:focus:placeholder-slate-300 dark:text-white focus:outline-none focus:ring-0 border-none disabled:cursor-not-allowed"
           placeholder="How does caching work?"
           rows={1}
@@ -63,7 +137,7 @@ export function Prompt({
           variant="primary"
           size="small"
           type="submit"
-          disabled={isDisabled}
+          disabled={isGenerating}
           className="self-end w-12 h-12 disabled:cursor-not-allowed"
         >
           <div hidden className="sr-only">