Merge pull request elastic#12 from yansavitski/ai-playground-ai-messa…

…ge-update Ai playground ai message update
joemcelroy · Feb 15, 2024 · f892a4b · f892a4b
2 parents 4a77158 + f2880dd
commit f892a4b
Show file tree

Hide file tree

Showing 12 changed files with 604 additions and 84 deletions.
diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React from 'react';
+import React, { useMemo } from 'react';
 
 import { Controller, useForm } from 'react-hook-form';
 import {
@@ -14,6 +14,7 @@ import {
   EuiFlexItem,
   EuiForm,
   EuiHorizontalRule,
+  EuiSpacer,
   useEuiTheme,
 } from '@elastic/eui';
 import { v4 as uuidv4 } from 'uuid';
@@ -37,14 +38,8 @@ import { IncludeCitationsField } from './include_citations_field';
 
 import { TelegramIcon } from './telegram_icon';
 import { useKibana } from '@kbn/kibana-react-plugin/public';
-
-const transformFromChatMessages = (messages: UseChatHelpers['messages']): Message[] =>
-  messages.map(({ id, content, createdAt, role }) => ({
-    id,
-    content,
-    createdAt,
-    role: role === 'assistant' ? MessageRole.assistant : MessageRole.user,
-  }));
+import { transformFromChatMessages } from '../utils/transformToMessages';
+import { SourcesPanel } from '@kbn/ai-playground/components/sources_panel/sources_panel';
 
 export const Chat = () => {
   const { euiTheme } = useEuiTheme();
@@ -85,13 +80,17 @@ export const Chat = () => {
 
     resetField(ChatFormFields.question);
   };
-  const initialMessages = [
-    {
-      id: uuidv4(),
-      role: MessageRole.system,
-      content: 'You can start chat now',
-    },
-  ];
+  const chatMessages = useMemo(
+    () => [
+      {
+        id: uuidv4(),
+        role: MessageRole.system,
+        content: 'You can start chat now',
+      },
+      ...transformFromChatMessages(messages),
+    ],
+    [messages]
+  );
 
   return (
     <EuiForm
@@ -104,19 +103,29 @@ export const Chat = () => {
           grow={2}
           css={{
             borderRight: euiTheme.border.thin,
-            padding: euiTheme.size.l,
+            paddingTop: euiTheme.size.l,
+            paddingBottom: euiTheme.size.l,
           }}
         >
-          <EuiFlexGroup direction="column">
-            <EuiFlexItem grow={1}>
-              <MessageList
-                messages={[...initialMessages, ...transformFromChatMessages(messages)]}
-              />
+          <EuiFlexGroup direction="column" className="eui-fullHeight">
+            {/*// Set scroll at the border of parent element*/}
+            <EuiFlexItem
+              grow={1}
+              className="eui-yScroll"
+              css={{ paddingLeft: euiTheme.size.l, paddingRight: euiTheme.size.l }}
+              tabIndex={0}
+            >
+              <MessageList messages={chatMessages} />
             </EuiFlexItem>
 
-            <EuiHorizontalRule margin="none" />
+            <EuiFlexItem
+              grow={false}
+              css={{ paddingLeft: euiTheme.size.l, paddingRight: euiTheme.size.l }}
+            >
+              <EuiHorizontalRule margin="none" />
+
+              <EuiSpacer size="m" />
 
-            <EuiFlexItem grow={false}>
               <Controller
                 name={ChatFormFields.question}
                 control={control}
@@ -192,6 +201,8 @@ export const Chat = () => {
               <IncludeCitationsField checked={field.value} onChange={field.onChange} />
             )}
           />
+
+          <SourcesPanel />
         </EuiFlexItem>
       </EuiFlexGroup>
     </EuiForm>

diff --git a/packages/kbn-ai-playground/components/message_list/assistant_message.tsx b/packages/kbn-ai-playground/components/message_list/assistant_message.tsx
@@ -5,20 +5,39 @@
  * 2.0.
  */
 
-import React from 'react';
+import React, { useState } from 'react';
 
 import moment from 'moment';
 
-import { EuiComment, EuiText } from '@elastic/eui';
+import {
+  EuiButtonEmpty,
+  EuiComment,
+  EuiFlexGroup,
+  EuiSpacer,
+  EuiText,
+  EuiTitle,
+} from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 
-import type { Message as MessageType } from '../../types';
+import type { AIMessage as AIMessageType } from '../../types';
 
 import { CopyActionButton } from './copy_action_button';
+import { CitationsTable } from './citations_table';
+import { RetrievalDocsFlyout } from './retrieval_docs_flyout';
 
-type AssistantMessageProps = Pick<MessageType, 'content' | 'createdAt'>;
+type AssistantMessageProps = Pick<
+  AIMessageType,
+  'content' | 'createdAt' | 'citations' | 'retrievalDocs'
+>;
+
+export const AssistantMessage: React.FC<AssistantMessageProps> = ({
+  content,
+  createdAt,
+  citations,
+  retrievalDocs,
+}) => {
+  const [isDocsFlyoutOpen, setIsDocsFlyoutOpen] = useState(false);
 
-export const AssistantMessage: React.FC<AssistantMessageProps> = ({ content, createdAt }) => {
   return (
     <EuiComment
       username={i18n.translate('aiPlayground.chat.message.assistant.username', {
@@ -52,9 +71,51 @@ export const AssistantMessage: React.FC<AssistantMessageProps> = ({ content, cre
         />
       }
     >
+      <EuiFlexGroup justifyContent="spaceBetween" alignItems="center">
+        <EuiTitle size="xs">
+          <p>
+            {i18n.translate('aiPlayground.chat.message.assistant.title', {
+              defaultMessage: 'Summary',
+            })}
+          </p>
+        </EuiTitle>
+        {!!retrievalDocs?.length && (
+          <>
+            <EuiButtonEmpty onClick={() => setIsDocsFlyoutOpen(true)}>
+              {i18n.translate('aiPlayground.chat.message.assistant.retrievalDocButton', {
+                defaultMessage:
+                  '{count} {count, plural, one {document} other {documents}} retrieved',
+                values: { count: retrievalDocs.length },
+              })}
+            </EuiButtonEmpty>
+
+            {isDocsFlyoutOpen && (
+              <RetrievalDocsFlyout
+                onClose={() => setIsDocsFlyoutOpen(false)}
+                retrievalDocs={retrievalDocs}
+              />
+            )}
+          </>
+        )}
+      </EuiFlexGroup>
+      <EuiSpacer size="m" />
       <EuiText size="s">
         <p>{content}</p>
       </EuiText>
+      {!!citations?.length && (
+        <>
+          <EuiSpacer size="l" />
+          <EuiTitle size="xs">
+            <p>
+              {i18n.translate('aiPlayground.chat.message.assistant.citations.title', {
+                defaultMessage: 'Citations',
+              })}
+            </p>
+          </EuiTitle>
+          <EuiSpacer size="xs" />
+          <CitationsTable citations={citations} />
+        </>
+      )}
     </EuiComment>
   );
 };
diff --git a/packages/kbn-ai-playground/components/message_list/citations_table.tsx b/packages/kbn-ai-playground/components/message_list/citations_table.tsx
@@ -0,0 +1,69 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React, { useState } from 'react';
+import { i18n } from '@kbn/i18n';
+import { EuiBasicTable, EuiButtonEmpty, EuiText } from '@elastic/eui';
+import { AIMessage as AIMessageType, Doc } from '@kbn/ai-playground/types';
+
+interface CitationsTableProps extends Pick<AIMessageType, 'citations'> {}
+
+export const CitationsTable: React.FC<CitationsTableProps> = ({ citations }) => {
+  const [itemIdToExpandedRowMap, setItemIdToExpandedRowMap] = useState<
+    Record<string, React.ReactNode>
+  >({});
+  const toggleDetails = (citation: Doc) => {
+    let itemIdToExpandedRowMapValues = { ...itemIdToExpandedRowMap };
+
+    if (itemIdToExpandedRowMapValues[citation.id]) {
+      delete itemIdToExpandedRowMapValues[citation.id];
+    } else {
+      itemIdToExpandedRowMapValues[citation.id] = <EuiText size="s">{citation.content}</EuiText>;
+    }
+
+    setItemIdToExpandedRowMap(itemIdToExpandedRowMapValues);
+  };
+
+  return (
+    <EuiBasicTable
+      columns={[
+        {
+          field: 'id',
+          name: i18n.translate('aiPlayground.chat.message.assistant.citations.idField', {
+            defaultMessage: 'Index Id',
+          }),
+          truncateText: true,
+        },
+        {
+          truncateText: false,
+          align: 'right',
+          isExpander: true,
+          render: (citation: Doc) => {
+            const itemIdToExpandedRowMapValues = { ...itemIdToExpandedRowMap };
+
+            return (
+              <EuiButtonEmpty
+                iconSide="right"
+                size="s"
+                onClick={() => toggleDetails(citation)}
+                iconType={itemIdToExpandedRowMapValues[citation.id] ? 'arrowDown' : 'arrowRight'}
+              >
+                {i18n.translate('aiPlayground.chat.message.assistant.citations.snippet', {
+                  defaultMessage: 'Snippet',
+                })}
+              </EuiButtonEmpty>
+            );
+          },
+        },
+      ]}
+      items={citations}
+      itemId="id"
+      itemIdToExpandedRowMap={itemIdToExpandedRowMap}
+      isExpandable
+    />
+  );
+};
diff --git a/packages/kbn-ai-playground/components/message_list/message_list.tsx b/packages/kbn-ai-playground/components/message_list/message_list.tsx
@@ -9,7 +9,7 @@ import React from 'react';
 
 import { EuiCommentList } from '@elastic/eui';
 
-import { Message, MessageRole } from '../../types';
+import { AIMessage, Message, MessageRole } from '../../types';
 
 import { AssistantMessage } from './assistant_message';
 import { SystemMessage } from './system_message';
@@ -19,24 +19,27 @@ interface MessageListProps {
   messages: Message[];
 }
 
-export const MessageList: React.FC<MessageListProps> = ({ messages }) => {
-  const mapRoleToMessageComponent = {
-    [MessageRole.system]: (message: Message) => <SystemMessage content={message.content} />,
-    [MessageRole.user]: (message: Message) => (
-      <UserMessage content={message.content} createdAt={message.createdAt} />
-    ),
-    [MessageRole.assistant]: (message: Message) => (
-      <AssistantMessage content={message.content} createdAt={message.createdAt} />
-    ),
-  };
-
-  return (
-    <EuiCommentList gutterSize="m">
-      {messages.map((message) => (
-        <React.Fragment key={message.id}>
-          {mapRoleToMessageComponent[message.role](message)}
-        </React.Fragment>
-      ))}
-    </EuiCommentList>
-  );
+const mapRoleToMessageComponent = {
+  [MessageRole.system]: (message: Message) => <SystemMessage content={message.content} />,
+  [MessageRole.user]: (message: Message) => (
+    <UserMessage content={message.content} createdAt={message.createdAt} />
+  ),
+  [MessageRole.assistant]: (message: Message) => (
+    <AssistantMessage
+      content={message.content}
+      createdAt={message.createdAt}
+      citations={(message as AIMessage).citations}
+      retrievalDocs={(message as AIMessage).retrievalDocs}
+    />
+  ),
 };
+
+export const MessageList: React.FC<MessageListProps> = ({ messages }) => (
+  <EuiCommentList gutterSize="m">
+    {messages.map((message) => (
+      <React.Fragment key={message.id}>
+        {mapRoleToMessageComponent[message.role](message)}
+      </React.Fragment>
+    ))}
+  </EuiCommentList>
+);