ui: update chat box (#168)

* ui: update chat box * remove log * hide back ref
pingcap · Jun 21, 2024 · dcb3512 · dcb3512
1 parent 3084d2b
commit dcb3512
Show file tree

Hide file tree

Showing 14 changed files with 313 additions and 87 deletions.
diff --git a/embeddable-javascript/src/components/Chatbot/ChatActionBar.tsx b/embeddable-javascript/src/components/Chatbot/ChatActionBar.tsx
@@ -1,18 +1,15 @@
-import * as React from 'react';
-import { styled, css } from '@mui/system';
+import { css, styled } from '@mui/system';
+import { SxProps } from '@mui/system/styleFunctionSx';
 import copy from 'copy-to-clipboard';
+import * as React from 'react';
+import { useMemo } from 'react';
 
 import { ChatActionButton } from '../Button';
-import {
-  StopFilledIcon,
-  ThumbUpIcon,
-  ThumbDownIcon,
-  CopyIcon,
-  CopiedIcon,
-  RefreshIcon,
-} from '../Icons';
+import { CopiedIcon, CopyIcon, RefreshIcon, StopFilledIcon, ThumbDownIcon, ThumbUpIcon } from '../Icons';
+import { getChatMessageAnnotations, type MyChatMessageAnnotation } from './ChatItem.tsx';
+import { useMessageFeedback } from './use-message-feedback.ts';
 
-export function ChatActionBar(props: {
+export function ChatActionBar (props: {
   handleStop: () => void;
   className?: string;
 }) {
@@ -42,14 +39,33 @@ export function ChatActionBar(props: {
   );
 }
 
-export function ChatItemActionBar(props: {
+export function ChatItemActionBar (props: {
+  sessionId?: string | null;
+  annotations?: MyChatMessageAnnotation[];
   handleReload?: () => void;
   className?: string;
   content?: string;
 }) {
-  const { className, handleReload, content } = props;
+  const { className, handleReload, content, sessionId, annotations } = props;
   const [copied, setCopied] = React.useState(false);
 
+  const annotation = useMemo(() => {
+    return getChatMessageAnnotations(annotations);
+  }, []);
+
+  const {
+    feedbackData,
+    feedback,
+  } = useMessageFeedback(Number(sessionId), annotation.messageId, annotation.state === 'FINISHED');
+
+  const like = () => {
+    feedback('like', {}, 'From tidb.ai ChatBot');
+  }
+
+  const dislike = () => {
+    feedback('dislike', {}, 'From tidb.ai ChatBot');
+  }
+
   React.useEffect(() => {
     if (copied) {
       const timeout = setTimeout(() => {
@@ -69,7 +85,7 @@ export function ChatItemActionBar(props: {
             <ActionButton
               onClick={handleReload}
               Icon={RefreshIcon}
-              name='Regenerate'
+              name="Regenerate"
               className={className}
             >
               Regenerate
@@ -82,7 +98,7 @@ export function ChatItemActionBar(props: {
               setCopied(true);
             }}
             Icon={copied ? CopiedIcon : CopyIcon}
-            name='Copy'
+            name="Copy"
             className={className}
           >
             {copied ? 'Copied!' : 'Copy'}
@@ -94,15 +110,27 @@ export function ChatItemActionBar(props: {
         >
           <ActionButton
             Icon={ThumbUpIcon}
-            name='GoodAnswer'
+            name="GoodAnswer"
             className={className}
+            disabled={!!feedbackData}
+            sx={feedbackData?.action === 'like' ? ({
+              backgroundColor: 'rgb(34 197 94 / 0.05)',
+              color: 'rgb(34 197 94)'
+            }) : !!feedbackData ? {opacity: 0.5, cursor: 'not-allowed'} : undefined}
+            onClick={like}
           >
             Good Answer
           </ActionButton>
           <ActionButton
             Icon={ThumbDownIcon}
-            name='BadAnswer'
+            name="BadAnswer"
             className={className}
+            disabled={!!feedbackData}
+            sx={feedbackData?.action === 'dislike' ? ({
+              backgroundColor: 'rgb(239 68 68 / 0.05)',
+              color: 'rgb(239 68 68)',
+            }) : !!feedbackData ? {opacity: 0.5, cursor: 'not-allowed'} : undefined}
+            onClick={dislike}
           >
             Bad Answer
           </ActionButton>
@@ -112,17 +140,20 @@ export function ChatItemActionBar(props: {
   );
 }
 
-function ActionButton(props: {
+function ActionButton (props: {
   name: string;
   className?: string;
   onClick?: () => void;
   children: React.ReactNode;
   Icon: React.ElementType;
+  disabled?: boolean;
+  sx?: SxProps;
 }) {
-  const { name, className, onClick, children, Icon } = props;
+  const { name, className, onClick, children, Icon, disabled, sx } = props;
 
   return (
     <ChatActionButton
+      disabled={disabled}
       onClick={onClick}
       className={className + `-${name}`}
       sx={{
@@ -131,6 +162,7 @@ function ActionButton(props: {
         fontWeight: 'normal',
         fontSize: '0.75rem',
         color: 'mutedForeground',
+        ...sx,
       }}
     >
       <Icon
@@ -171,5 +203,5 @@ const StyledChatActionBar = styled('div')(({ theme }) =>
         gap: '0.5rem',
       },
     },
-  })
+  }),
 );
diff --git a/embeddable-javascript/src/components/Chatbot/ChatContainer.tsx b/embeddable-javascript/src/components/Chatbot/ChatContainer.tsx
@@ -30,7 +30,7 @@ export default function ChatContainer(props: {
     securityMode,
   } = props;
 
-  const { session, setSession } = useLocalCreateRAGSessionId();
+  const { session, setSession, sessionId, setSessionId } = useLocalCreateRAGSessionId();
   // const contextRef = React.useRef<string[]>([]);
 
   const cfg = React.useContext(CfgContext);
@@ -59,6 +59,9 @@ export default function ChatContainer(props: {
       if (session !== response.headers.get('X-CreateRag-Session')) {
         setSession(response.headers.get('X-CreateRag-Session') ?? undefined);
       }
+      if (sessionId !== response.headers.get('X-CreateRag-Session-Id')) {
+        setSessionId(response.headers.get('X-CreateRag-Session-Id') ?? undefined);
+      }
     },
   });
 
@@ -128,6 +131,8 @@ export default function ChatContainer(props: {
                   Action={
                     showAction ? (
                       <ChatItemActionBar
+                        sessionId={sessionId}
+                        annotations={m.annotations as any}
                         handleReload={idx === 0 ? reload : undefined}
                         content={m.content}
                         className={

diff --git a/embeddable-javascript/src/components/Chatbot/ChatItem.tsx b/embeddable-javascript/src/components/Chatbot/ChatItem.tsx
@@ -12,6 +12,7 @@ export type AppChatStreamSource = { title: string, uri: string };
 export const enum AppChatStreamState {
   CONNECTING = 'CONNECTING', // only client side
   CREATING = 'CREATING',
+  KG_RETRIEVING = 'KG_RETRIEVING',
   SEARCHING = 'SEARCHING',
   RERANKING = 'RERANKING',
   GENERATING = 'GENERATING',
@@ -20,6 +21,8 @@ export const enum AppChatStreamState {
 }
 
 export type MyChatMessageAnnotation = {
+  ts: number;
+  messageId: number;
   traceURL?: string,
   context?: AppChatStreamSource[],
   state?: AppChatStreamState,
@@ -96,6 +99,9 @@ export function ChatItemLoading({ annotations }: { annotations: MyChatMessageAnn
     case 'CREATING':
       text = 'Preparing to ask...';
       break;
+    case 'KG_RETRIEVING':
+      text = 'Retrieving knowledge...';
+      break;
     case 'SEARCHING':
       text = 'Gathering resources...';
       break;
@@ -348,7 +354,7 @@ export const StyledExampleQuestionWrapper = styled('div')(() =>
 //   })
 // );
 
-function getChatMessageAnnotations (annotations: MyChatMessageAnnotation[] | undefined) {
+export function getChatMessageAnnotations (annotations: MyChatMessageAnnotation[] | undefined) {
   return ((annotations ?? []) as MyChatMessageAnnotation[])
-    .reduce((annotation, next) => Object.assign(annotation, next), {});
+    .reduce((annotation, next) => Object.assign(annotation, next), {} as MyChatMessageAnnotation);
 }
diff --git a/embeddable-javascript/src/components/Chatbot/Markdown.tsx b/embeddable-javascript/src/components/Chatbot/Markdown.tsx
@@ -1,3 +1,4 @@
+// import { styled } from '@mui/system';
 import * as React from 'react';
 import ReactMarkdown, { Options } from 'react-markdown';
 import remarkGfm from 'remark-gfm';
@@ -9,14 +10,26 @@ export const MemoizedReactMarkdown: React.FC<Options> = React.memo(
   ReactMarkdown,
   (prevProps, nextProps) =>
     prevProps.children === nextProps.children &&
-    prevProps.className === nextProps.className
+    prevProps.className === nextProps.className,
 );
 
+// const SROnlySection = styled('section')`
+//   position: absolute;
+//   width: 1px;
+//   height: 1px;
+//   padding: 0;
+//   margin: -1px;
+//   overflow: hidden;
+//   clip: rect(0, 0, 0, 0);
+//   white-space: nowrap;
+//   border-width: 0;
+// `;
+
 export const MDMessage = (props: { children: string }) => {
   return (
     <>
       <MemoizedReactMarkdown
-        className='markdown-body'
+        className="markdown-body"
         remarkPlugins={[remarkGfm, remarkMath]}
         components={{
           // code(props) {
@@ -34,13 +47,26 @@ export const MDMessage = (props: { children: string }) => {
           //     </code>
           //   );
           // },
-          a(props) {
+          // section ({ ...props }) {
+          //   // eslint-disable-next-line react-hooks/rules-of-hooks
+          //
+          //   if (!(props as any)['data-footnotes']) return <section {...props} />;
+          //   return (
+          //     <SROnlySection {...props}>
+          //       {props.children}
+          //     </SROnlySection>
+          //   );
+          // },
+          a (props) {
+            if (props.href?.startsWith('#')) {
+              return <a {...props} >{props.children}</a>
+            }
             return (
-              <a {...props} target='_blank' rel='noopener noreferrer'>
+              <a {...props} target="_blank" rel="noopener noreferrer">
                 {props.children}
               </a>
             );
-          }
+          },
         }}
       >
         {props.children}

diff --git a/embeddable-javascript/src/components/Chatbot/use-message-feedback.ts b/embeddable-javascript/src/components/Chatbot/use-message-feedback.ts
@@ -0,0 +1,68 @@
+'use client';
+
+import { useContext, useState } from 'react';
+import useSWR, { mutate } from 'swr';
+import { CfgContext } from '../../context';
+import { fetcher } from '../../lib/fetch';
+
+export interface UseMessageFeedbackReturns {
+  feedbackData: Feedback | undefined;
+  disabled: boolean;
+
+  // source?: ContentSource;
+  // sourceLoading: boolean;
+
+  feedback (action: 'like' | 'dislike', details: Record<string, 'like' | 'dislike'>, comment: string): Promise<void>;
+}
+
+export type ContentSource = {
+  query: string
+  markdownSources: {
+    kgRelationshipUrls: string[]
+    restUrls: string[]
+  }
+  kgSources: Record<string, any>
+}
+
+export function useMessageFeedback (chatId: number, messageId: number, enabled: boolean): UseMessageFeedbackReturns {
+  const { baseUrl } = useContext(CfgContext);
+
+  const { data: feedback, isLoading, isValidating } = useSWR(enabled ? ['get', `${baseUrl}/api/v1/chats/${chatId}/messages/${messageId}/feedback`] : undefined, fetcher<Feedback>);
+  const [acting, setActing] = useState(false);
+  const disabled = isValidating || isLoading || acting || !enabled;
+
+  // const contentData = useSWR((enabled && !disabled) ? ['get', `/api/v1/chats/${chatId}/messages/${messageId}/content-sources`] : undefined, fetcher<ContentSource>, { keepPreviousData: true, revalidateIfStale: false, revalidateOnReconnect: false });
+
+  return {
+    feedbackData: feedback,
+    disabled,
+    feedback: (action, detail, comment) => {
+      setActing(true);
+      return addFeedback(baseUrl, chatId, messageId, { action, knowledge_graph_detail: detail, comment }).finally(() => setActing(false));
+    },
+    // source: contentData.data,
+    // sourceLoading: contentData.isLoading || contentData.isValidating,
+  };
+}
+
+async function addFeedback (baseUrl: string, chatId: number, messageId: number, data: any) {
+  mutate(['get', `${baseUrl}/api/v1/chats/${chatId}/messages/${messageId}/feedback`], () => data, false);
+  await fetch(`${baseUrl}/api/v1/chats/${chatId}/messages/${messageId}/feedback`, {
+    method: 'post',
+    body: JSON.stringify(data),
+  });
+}
+
+type Feedback = {
+  action: 'dislike' | 'like';
+  chat_id: number;
+  comment: string;
+  created_at: Date;
+  created_by: string;
+  id: number;
+  knowledge_graph_detail: Record<string, 'like' | 'dislike'>;
+  knowledge_graph_report_error: string | null;
+  knowledge_graph_reported_at: Date | null;
+  message_id: number;
+  trace_id: string;
+}