From 798d3c20de1b376c629a09319472b4372ad9295d Mon Sep 17 00:00:00 2001
From: Samiul Monir <samiul.monir@elastic.co>
Date: Fri, 23 Feb 2024 11:21:52 -0500
Subject: [PATCH] Summarization model types support

---
 .../kbn-ai-playground/components/chat.tsx     |  9 +++++--
 .../open_ai_summarization_model.tsx           | 26 ++++++-------------
 packages/kbn-ai-playground/types.ts           |  8 ++++++
 3 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
index 7459bbc6afb35..cb0280a695ce1 100644
--- 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,7 +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],
+          summarization_model: data[ChatFormFields.summarizationModel] ?? SummarizationModelName.gpt3_5_turbo_1106,
         },
       }
     );
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
index ee2e7680ac66c..6d78222cbdd0c 100644
--- 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
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React, { useState } from "react";
+import React from "react";
 
 import {
   EuiButtonEmpty,
@@ -19,9 +19,11 @@ import {
 } 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}>
@@ -33,16 +35,7 @@ const renderSelectOptions = (label: string) => (
   </EuiFlexGroup>
 );
 
-
-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',
-}
-
-const Summarization_Model: EuiSuperSelectOption<string>[] = [
+const SummarizationModel: EuiSuperSelectOption<string>[] = [
   {
     value: SummarizationModelName.gpt3_5,
     inputDisplay: renderSelectOptions(SummarizationModelName.gpt3_5),
@@ -72,12 +65,9 @@ interface OpenAISummarizationModelProps {
   onSelect: (key: string) => void;
 }
 
-export const OpenAISummarizationModel: React.FC<OpenAISummarizationModelProps> = ({ model, onSelect, openAIFlyOutOpen }) => {
-
-  const [selectedModel, setSelectedModel] = useState(model ?? SummarizationModelName.gpt3_5_turbo_1106);
+export const OpenAISummarizationModel: React.FC<OpenAISummarizationModelProps> = ({ model = SummarizationModelName.gpt3_5_turbo_1106, onSelect, openAIFlyOutOpen }) => {
 
   const onChange = (value: string) => {
-    setSelectedModel(value);
     onSelect(value);
   };
 
@@ -109,9 +99,9 @@ export const OpenAISummarizationModel: React.FC<OpenAISummarizationModelProps> =
       }
     >
       <EuiSuperSelect
-        options={Summarization_Model}
-        valueOfSelected={selectedModel}
-        onChange={(value) => onChange(value)}
+        options={SummarizationModel}
+        valueOfSelected={model}
+        onChange={onChange}
       />
     </EuiFormRow>
   );
diff --git a/packages/kbn-ai-playground/types.ts b/packages/kbn-ai-playground/types.ts
index 2f42778a54a14..7af27058fb75b 100644
--- a/packages/kbn-ai-playground/types.ts
+++ b/packages/kbn-ai-playground/types.ts
@@ -99,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',
+}