Merge pull request elastic#19 from Samiul-TheSoccerFan/summarizatioMm…

…odel Summarization model
joemcelroy · Feb 23, 2024 · 7072aa8 · 7072aa8
2 parents 026c1a0 + 798d3c2
commit 7072aa8
Show file tree

Hide file tree

Showing 5 changed files with 168 additions and 6 deletions.
diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
@@ -24,7 +24,12 @@ import { i18n } from '@kbn/i18n';
 
 import { ChatSidebar } from './chat_sidebar';
 import { useChat } from '../hooks/useChat';
-import { ChatForm, ChatFormFields, MessageRole } from '../types';
+import {
+  ChatForm,
+  ChatFormFields,
+  MessageRole,
+  SummarizationModelName
+} from '../types';
 
 import { MessageList } from './message_list/message_list';
 import { QuestionInput } from './question_input';
@@ -55,6 +60,7 @@ export const Chat = () => {
           api_key: data[ChatFormFields.openAIKey],
           citations: data[ChatFormFields.citations].toString(),
           elasticsearchQuery: JSON.stringify(data[ChatFormFields.elasticsearchQuery]),
+          summarization_model: data[ChatFormFields.summarizationModel] ?? SummarizationModelName.gpt3_5_turbo_1106,
         },
       }
     );

diff --git a/packages/kbn-ai-playground/components/summarization_panel/open_ai_icon.tsx b/packages/kbn-ai-playground/components/summarization_panel/open_ai_icon.tsx
@@ -0,0 +1,17 @@
+/*
+ * 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 from 'react';
+
+export const OpenAIIcon = () => (
+  <svg width="20" height="20" viewBox="0 0 24 24" fill='none' xmlns="http://www.w3.org/2000/svg">
+    <path
+      d="M22.282 9.821a5.985 5.985 0 0 0-.516-4.91 6.046 6.046 0 0 0-6.51-2.9A6.065 6.065 0 0 0 4.981 4.18a5.985 5.985 0 0 0-3.998 2.9 6.046 6.046 0 0 0 .743 7.097 5.98 5.98 0 0 0 .51 4.911 6.051 6.051 0 0 0 6.515 2.9A5.985 5.985 0 0 0 13.26 24a6.056 6.056 0 0 0 5.772-4.206 5.99 5.99 0 0 0 3.997-2.9 6.056 6.056 0 0 0-.747-7.073zM13.26 22.43a4.476 4.476 0 0 1-2.876-1.04l.141-.081 4.779-2.758a.795.795 0 0 0 .392-.681v-6.737l2.02 1.168a.071.071 0 0 1 .038.052v5.583a4.504 4.504 0 0 1-4.494 4.494zM3.6 18.304a4.47 4.47 0 0 1-.535-3.014l.142.085 4.783 2.759a.771.771 0 0 0 .78 0l5.843-3.369v2.332a.08.08 0 0 1-.033.062L9.74 19.95a4.5 4.5 0 0 1-6.14-1.646zM2.34 7.896a4.485 4.485 0 0 1 2.366-1.973V11.6a.766.766 0 0 0 .388.676l5.815 3.355-2.02 1.168a.076.076 0 0 1-.071 0l-4.83-2.786A4.504 4.504 0 0 1 2.34 7.872zm16.597 3.855l-5.833-3.387L15.119 7.2a.076.076 0 0 1 .071 0l4.83 2.791a4.494 4.494 0 0 1-.676 8.105v-5.678a.79.79 0 0 0-.407-.667zm2.01-3.023l-.141-.085-4.774-2.782a.776.776 0 0 0-.785 0L9.409 9.23V6.897a.066.066 0 0 1 .028-.061l4.83-2.787a4.5 4.5 0 0 1 6.68 4.66zm-12.64 4.135l-2.02-1.164a.08.08 0 0 1-.038-.057V6.075a4.5 4.5 0 0 1 7.375-3.453l-.142.08L8.704 5.46a.795.795 0 0 0-.393.681zm1.097-2.365l2.602-1.5 2.607 1.5v2.999l-2.597 1.5-2.607-1.5z"
+      fill='currentColor'
+    />
+  </svg>
+);
diff --git a/packages/kbn-ai-playground/components/summarization_panel/open_ai_summarization_model.tsx b/packages/kbn-ai-playground/components/summarization_panel/open_ai_summarization_model.tsx
@@ -0,0 +1,108 @@
+/*
+ * 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 from "react";
+
+import {
+  EuiButtonEmpty,
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiFormRow,
+  EuiIcon,
+  EuiSuperSelect,
+  EuiSuperSelectOption,
+  EuiToolTip
+} from "@elastic/eui";
+
+import { i18n } from "@kbn/i18n";
+import { SummarizationModelName } from "@kbn/ai-playground/types";
+
+import { OpenAIIcon } from "./open_ai_icon";
+
+
+const renderSelectOptions = (label: string) => (
+  <EuiFlexGroup alignItems="center">
+    <EuiFlexItem grow={false}>
+      <EuiIcon type={OpenAIIcon} />
+    </EuiFlexItem>
+    <EuiFlexItem>
+      {label}
+    </EuiFlexItem>
+  </EuiFlexGroup>
+);
+
+const SummarizationModel: EuiSuperSelectOption<string>[] = [
+  {
+    value: SummarizationModelName.gpt3_5,
+    inputDisplay: renderSelectOptions(SummarizationModelName.gpt3_5),
+  },
+  {
+    value: SummarizationModelName.gpt3_5_turbo_1106,
+    inputDisplay: renderSelectOptions(SummarizationModelName.gpt3_5_turbo_1106),
+  },
+  {
+
+    value: SummarizationModelName.gpt3_5_turbo_16k,
+    inputDisplay: renderSelectOptions(SummarizationModelName.gpt3_5_turbo_16k),
+  },
+  {
+    value: SummarizationModelName.gpt3_5_turbo_16k_0613,
+    inputDisplay: renderSelectOptions(SummarizationModelName.gpt3_5_turbo_16k_0613),
+  },
+  {
+    value: SummarizationModelName.gpt3_5_turbo,
+    inputDisplay: renderSelectOptions(SummarizationModelName.gpt3_5_turbo),
+  },
+];
+
+interface OpenAISummarizationModelProps {
+  openAIFlyOutOpen: () => void;
+  model: string;
+  onSelect: (key: string) => void;
+}
+
+export const OpenAISummarizationModel: React.FC<OpenAISummarizationModelProps> = ({ model = SummarizationModelName.gpt3_5_turbo_1106, onSelect, openAIFlyOutOpen }) => {
+
+  const onChange = (value: string) => {
+    onSelect(value);
+  };
+
+  return (
+    <EuiFormRow
+      label={
+        <EuiToolTip
+          content={i18n.translate('aiPlayground.sidebar.summarizationModel.help', {
+            defaultMessage:
+              'The large language model used to summarize your documents.',
+          })}
+        >
+          <>
+            <span>
+              {i18n.translate('aiPlayground.sidebar.summarizationModel.label', {
+                defaultMessage: 'Summarization Model',
+              })}
+            </span>
+            <EuiIcon type="questionInCircle" color="subdued" />
+          </>
+        </EuiToolTip>
+      }
+      labelAppend={
+        <EuiButtonEmpty flush="both" size="xs" onClick={() => openAIFlyOutOpen()}>
+          {i18n.translate('aiPlayground.sidebar.summarizationModel.editLabel', {
+            defaultMessage: 'Edit OpenAI API key',
+          })}
+        </EuiButtonEmpty>
+      }
+    >
+      <EuiSuperSelect
+        options={SummarizationModel}
+        valueOfSelected={model}
+        onChange={onChange}
+      />
+    </EuiFormRow>
+  );
+};
diff --git a/packages/kbn-ai-playground/components/summarization_panel/summarization_panel.tsx b/packages/kbn-ai-playground/components/summarization_panel/summarization_panel.tsx
@@ -5,23 +5,43 @@
  * 2.0.
  */
 
-import React from 'react';
+import React, { useState } from 'react';
 import { Controller, useFormContext } from 'react-hook-form';
+
 import { IncludeCitationsField } from './include_citations_field';
 import { InstructionsField } from './instructions_field';
-import { OpenAIKeyField } from './open_ai_key_field';
 import { ChatFormFields } from '../../types';
+import { OpenAIKeyFlyOut } from './open_ai_key_flyout';
+import { OpenAISummarizationModel } from './open_ai_summarization_model';
 
 export const SummarizationPanel: React.FC = () => {
   const { control } = useFormContext();
+  const [isOpenAIFlyOutOpen, setIsOpenAIFlyOutOpen] = useState<boolean>(false);
+
+  const onCloseOpenAIFlyOut = () => {
+    setIsOpenAIFlyOutOpen(!isOpenAIFlyOutOpen);
+  };
+  const handleOpenAIFlyOut = () => {
+    setIsOpenAIFlyOutOpen(true);
+  };
 
   return (
     <>
+      {isOpenAIFlyOutOpen && (
+        <Controller
+          name={ChatFormFields.openAIKey}
+          control={control}
+          defaultValue=""
+          render={({ field }) => <OpenAIKeyFlyOut openAPIKey={field.value} onSave={field.onChange} onClose={onCloseOpenAIFlyOut} />}
+        />
+      )}
+
       <Controller
-        name={ChatFormFields.openAIKey}
+        name={ChatFormFields.summarizationModel}
         control={control}
-        defaultValue=""
-        render={({ field }) => <OpenAIKeyField apiKey={field.value} onSave={field.onChange} />}
+        render={({ field }) => (
+          <OpenAISummarizationModel model={field.value} onSelect={field.onChange} openAIFlyOutOpen={handleOpenAIFlyOut} />
+        )}
       />
 
       <Controller
@@ -31,6 +51,7 @@ export const SummarizationPanel: React.FC = () => {
         render={({ field }) => <InstructionsField value={field.value} onChange={field.onChange} />}
       />
 
+
       <Controller
         name={ChatFormFields.citations}
         control={control}

diff --git a/packages/kbn-ai-playground/types.ts b/packages/kbn-ai-playground/types.ts
@@ -47,6 +47,7 @@ export enum ChatFormFields {
   openAIKey = 'api_key',
   indices = 'indices',
   elasticsearchQuery = 'elasticsearch_query',
+  summarizationModel = 'summarization_model'
 }
 
 export interface ChatForm {
@@ -55,6 +56,7 @@ export interface ChatForm {
   [ChatFormFields.citations]: boolean;
   [ChatFormFields.openAIKey]: string;
   [ChatFormFields.indices]: string[];
+  [ChatFormFields.summarizationModel]: string;
   [ChatFormFields.elasticsearchQuery]: QueryDslQueryContainer;
 }
 
@@ -97,3 +99,11 @@ export interface QuerySourceFields {
   bm25_query_fields: string[];
   source_fields: string[];
 }
+
+export enum SummarizationModelName {
+  gpt3_5 = 'gpt-3.5-turbo',
+  gpt3_5_turbo_1106 = 'gpt-3.5-turbo-1106',
+  gpt3_5_turbo_16k = 'gpt-3.5-turbo-16k',
+  gpt3_5_turbo_16k_0613 = 'gpt-3.5-turbo-16k-0613',
+  gpt3_5_turbo = 'gpt-3.5-turbo-instruct',
+}