implement streaming into ai chat (#1599)

Co-authored-by: michaeljguarino <[email protected]>
pluralsh · Nov 21, 2024 · a33c448 · a33c448
1 parent 6a410d7
commit a33c448
Show file tree

Hide file tree

Showing 8 changed files with 258 additions and 98 deletions.
diff --git a/assets/src/components/ai/AIPanel.tsx b/assets/src/components/ai/AIPanel.tsx
@@ -6,109 +6,122 @@ import {
   IconFrame,
 } from '@pluralsh/design-system'
 import { useTheme } from 'styled-components'
-import { ReactNode } from 'react'
+import { forwardRef, Ref, ReactNode } from 'react'
 import { AIPanelOverlay } from './AIPanelOverlay.tsx'
 
-export default function AIPanel({
-  open,
-  onClose,
-  showCloseIcon = false,
-  showClosePanel = false,
-  header,
-  subheader,
-  footer,
-  children,
-  ...props
-}: {
-  open: boolean
-  onClose: () => void
-  showCloseIcon?: boolean
-  showClosePanel?: boolean
-  header: string
-  subheader: string
-  footer?: ReactNode
-  children: ReactNode
-} & CardProps) {
-  const theme = useTheme()
+const AIPanel = forwardRef(
+  (
+    {
+      open,
+      onClose,
+      showCloseIcon = false,
+      showClosePanel = false,
+      header,
+      subheader,
+      footer,
+      children,
+      ...props
+    }: {
+      open: boolean
+      onClose: () => void
+      CloseIcon?: boolean
+      showClosePanel?: boolean
+      showCloseIcon?: boolean
+      header: string
+      subheader: string
+      footer?: ReactNode
+      children: ReactNode
+    } & CardProps,
+    ref: Ref<HTMLDivElement>
+  ) => {
+    const theme = useTheme()
 
-  return (
-    <AIPanelOverlay
-      open={open}
-      onClose={onClose}
-    >
-      <Card
-        fillLevel={1}
-        css={{
-          border: theme.borders.input,
-          display: 'flex',
-          flexDirection: 'column',
-          flexGrow: 1,
-          maxHeight: 720,
-          minHeight: 0,
-          overflow: 'hidden',
-        }}
-        {...props}
+    return (
+      <AIPanelOverlay
+        open={open}
+        onClose={onClose}
       >
-        <div
+        <Card
+          fillLevel={1}
           css={{
-            alignItems: 'center',
-            backgroundColor: theme.colors['fill-two'],
-            borderBottom: theme.borders.input,
+            border: theme.borders.input,
             display: 'flex',
-            gap: theme.spacing.small,
-            padding: theme.spacing.large,
+            flexDirection: 'column',
+            flexGrow: 1,
+            maxHeight: 720,
+            minHeight: 0,
+            overflow: 'hidden',
           }}
+          {...props}
         >
-          <div css={{ flexGrow: 1 }}>
-            <div css={{ ...theme.partials.text.subtitle2 }}>{header}</div>
-            <div
-              css={{
-                ...theme.partials.text.body2,
-                color: theme.colors['text-light'],
-              }}
-            >
-              {subheader}
-            </div>
-          </div>
-          {showCloseIcon && (
-            <IconFrame
-              clickable
-              icon={<CloseIcon />}
-              onClick={onClose}
-              tooltip="Close"
-            />
-          )}
-        </div>
-        <div
-          css={{ flexGrow: 1, overflow: 'auto', padding: theme.spacing.medium }}
-        >
-          {children}
-        </div>
-        {showClosePanel && (
           <div
             css={{
               alignItems: 'center',
               backgroundColor: theme.colors['fill-two'],
-              borderTop: theme.borders.input,
+              borderBottom: theme.borders.input,
               display: 'flex',
               gap: theme.spacing.small,
               padding: theme.spacing.large,
-              '> *': {
-                flexGrow: 1,
-              },
             }}
           >
-            <Button
-              onClick={onClose}
-              secondary={!!footer}
-              floating={!!footer}
-            >
-              Got it, thanks!
-            </Button>
-            {footer && footer}
+            <div css={{ flexGrow: 1 }}>
+              <div css={{ ...theme.partials.text.subtitle2 }}>{header}</div>
+              <div
+                css={{
+                  ...theme.partials.text.body2,
+                  color: theme.colors['text-light'],
+                }}
+              >
+                {subheader}
+              </div>
+            </div>
+            {showCloseIcon && (
+              <IconFrame
+                clickable
+                icon={<CloseIcon />}
+                onClick={onClose}
+                tooltip="Close"
+              />
+            )}
+          </div>
+          <div
+            css={{
+              flexGrow: 1,
+              overflow: 'auto',
+              padding: theme.spacing.medium,
+            }}
+            ref={ref}
+          >
+            {children}
           </div>
-        )}
-      </Card>
-    </AIPanelOverlay>
-  )
-}
+          {showClosePanel && (
+            <div
+              css={{
+                alignItems: 'center',
+                backgroundColor: theme.colors['fill-two'],
+                borderTop: theme.borders.input,
+                display: 'flex',
+                gap: theme.spacing.small,
+                padding: theme.spacing.large,
+                '> *': {
+                  flexGrow: 1,
+                },
+              }}
+            >
+              <Button
+                onClick={onClose}
+                secondary={!!footer}
+                floating={!!footer}
+              >
+                Got it, thanks!
+              </Button>
+              {footer && footer}
+            </div>
+          )}
+        </Card>
+      </AIPanelOverlay>
+    )
+  }
+)
+
+export default AIPanel
diff --git a/assets/src/components/ai/chatbot/AISuggestFix.tsx b/assets/src/components/ai/chatbot/AISuggestFix.tsx
@@ -1,9 +1,18 @@
 import { Markdown } from '@pluralsh/design-system'
-import { ReactNode, useCallback, useState } from 'react'
 import {
+  Dispatch,
+  ReactNode,
+  SetStateAction,
+  useCallback,
+  useRef,
+  useState,
+} from 'react'
+import {
+  AiDelta,
   AiInsightFragment,
   AiRole,
   ChatMessage,
+  useAiChatStreamSubscription,
   useAiSuggestedFixLazyQuery,
 } from '../../../generated/graphql.ts'
 import { GqlError } from '../../utils/Alert.tsx'
@@ -23,9 +32,58 @@ function fixMessage(fix: string): ChatMessage {
   }
 }
 
+function Loading({
+  insightId,
+  scrollToBottom,
+  setStreaming,
+}: {
+  insightId: string
+  scrollToBottom: () => void
+  setStreaming: Dispatch<SetStateAction<boolean>>
+}): ReactNode {
+  const [streamedMessage, setStreamedMessage] = useState<AiDelta[]>([])
+  useAiChatStreamSubscription({
+    variables: { insightId },
+    onData: ({ data: { data } }) => {
+      setStreaming(true)
+      if ((data?.aiStream?.seq ?? 1) % 120 === 0) scrollToBottom()
+      setStreamedMessage((streamedMessage) => [
+        ...streamedMessage,
+        {
+          seq: data?.aiStream?.seq ?? 0,
+          content: data?.aiStream?.content ?? '',
+        },
+      ])
+    },
+  })
+
+  if (!streamedMessage.length) {
+    return <LoadingIndicator />
+  }
+
+  return (
+    <Markdown
+      text={streamedMessage
+        .sort((a, b) => a.seq - b.seq)
+        .map((delta) => delta.content)
+        .join('')}
+    />
+  )
+}
+
 function AISuggestFix({ insight }: AISuggestFixProps): ReactNode {
+  const ref = useRef<HTMLDivElement>(null)
+  const [streaming, setStreaming] = useState<boolean>(false)
+  const scrollToBottom = useCallback(() => {
+    ref.current?.scrollTo({
+      top: ref.current.scrollHeight,
+      behavior: 'smooth',
+    })
+  }, [ref])
+
   const [getSuggestion, { loading, data, error }] = useAiSuggestedFixLazyQuery({
     variables: { insightID: insight?.id ?? '' },
+    onCompleted: () => streaming && scrollToBottom(),
     fetchPolicy: 'network-only',
   })
 
@@ -47,6 +105,7 @@ function AISuggestFix({ insight }: AISuggestFixProps): ReactNode {
     >
       <AISuggestFixButton onClick={showPanel} />
       <AIPanel
+        ref={ref}
         open={open}
         onClose={() => setOpen(false)}
         showCloseIcon
@@ -65,7 +124,13 @@ function AISuggestFix({ insight }: AISuggestFixProps): ReactNode {
         }
       >
         {data?.aiSuggestedFix && <Markdown text={data?.aiSuggestedFix} />}
-        {loading && !data && <LoadingIndicator />}
+        {loading && !data && (
+          <Loading
+            insightId={insight.id}
+            scrollToBottom={scrollToBottom}
+            setStreaming={setStreaming}
+          />
+        )}
         {!loading && error && <GqlError error={error} />}
       </AIPanel>
     </div>