[8.15] [Search] [Playground] Gemini search playground + Robustness fo…

…r Question Rewriting (elastic#187559) (elastic#187779) # Backport This will backport the following commits from `main` to `8.15`: - [[Search] [Playground] Gemini search playground + Robustness for Question Rewriting (elastic#187559)](elastic#187559)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Joe McElroy <[email protected]>
kibanamachine · Jul 8, 2024 · d11287d · d11287d
1 parent 12cbac9
commit d11287d
Show file tree

Hide file tree

Showing 14 changed files with 190 additions and 52 deletions.
diff --git a/x-pack/plugins/search_playground/common/models.ts b/x-pack/plugins/search_playground/common/models.ts
@@ -27,15 +27,27 @@ export const MODELS: ModelProvider[] = [
     provider: LLMs.openai,
   },
   {
-    name: 'Claude 3 Haiku',
+    name: 'Anthropic Claude 3 Haiku',
     model: 'anthropic.claude-3-haiku-20240307-v1:0',
     promptTokenLimit: 200000,
     provider: LLMs.bedrock,
   },
   {
-    name: 'Claude 3 Sonnet',
-    model: 'anthropic.claude-3-haiku-20240307-v1:0',
+    name: 'Anthropic Claude 3.5 Sonnet',
+    model: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
     promptTokenLimit: 200000,
     provider: LLMs.bedrock,
   },
+  {
+    name: 'Google Gemini 1.5 Pro',
+    model: 'gemini-1.5-pro-001',
+    promptTokenLimit: 2097152,
+    provider: LLMs.gemini,
+  },
+  {
+    name: 'Google Gemini 1.5 Flash',
+    model: 'gemini-1.5-flash-001',
+    promptTokenLimit: 2097152,
+    provider: LLMs.gemini,
+  },
 ];
diff --git a/x-pack/plugins/search_playground/common/prompt.ts b/x-pack/plugins/search_playground/common/prompt.ts
@@ -45,10 +45,23 @@ const AnthropicPrompt = (systemInstructions: string) => {
   `;
 };
 
+const GeminiPrompt = (systemInstructions: string) => {
+  return `
+  Instructions:
+  ${systemInstructions}
+
+  Context:
+  {context}
+
+  Question: {question}
+  Answer:
+  `;
+};
+
 interface PromptTemplateOptions {
   citations?: boolean;
   context?: boolean;
-  type?: 'openai' | 'mistral' | 'anthropic';
+  type?: 'openai' | 'mistral' | 'anthropic' | 'gemini';
 }
 
 export const Prompt = (instructions: string, options: PromptTemplateOptions): string => {
@@ -73,5 +86,20 @@ export const Prompt = (instructions: string, options: PromptTemplateOptions): st
     openai: OpenAIPrompt,
     mistral: MistralPrompt,
     anthropic: AnthropicPrompt,
+    gemini: GeminiPrompt,
+  }[options.type || 'openai'](systemInstructions);
+};
+
+interface QuestionRewritePromptOptions {
+  type: 'openai' | 'mistral' | 'anthropic' | 'gemini';
+}
+
+export const QuestionRewritePrompt = (options: QuestionRewritePromptOptions): string => {
+  const systemInstructions = `Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question. Rewrite the question in the question language. Keep the answer to a single sentence. Do not use quotes.`;
+  return {
+    openai: OpenAIPrompt,
+    mistral: MistralPrompt,
+    anthropic: AnthropicPrompt,
+    gemini: GeminiPrompt,
   }[options.type || 'openai'](systemInstructions);
 };
diff --git a/x-pack/plugins/search_playground/common/types.ts b/x-pack/plugins/search_playground/common/types.ts
@@ -40,6 +40,7 @@ export enum LLMs {
   openai = 'openai',
   openai_azure = 'openai_azure',
   bedrock = 'bedrock',
+  gemini = 'gemini',
 }
 
 export interface ChatRequestData {

diff --git a/x-pack/plugins/search_playground/public/hooks/use_llms_models.test.ts b/x-pack/plugins/search_playground/public/hooks/use_llms_models.test.ts
@@ -88,8 +88,8 @@ describe('useLLMsModels Hook', () => {
         connectorType: LLMs.bedrock,
         disabled: false,
         icon: expect.any(Function),
-        id: 'connectorId2Claude 3 Haiku',
-        name: 'Claude 3 Haiku',
+        id: 'connectorId2Anthropic Claude 3 Haiku',
+        name: 'Anthropic Claude 3 Haiku',
         showConnectorName: false,
         value: 'anthropic.claude-3-haiku-20240307-v1:0',
         promptTokenLimit: 200000,
@@ -100,10 +100,10 @@ describe('useLLMsModels Hook', () => {
         connectorType: LLMs.bedrock,
         disabled: false,
         icon: expect.any(Function),
-        id: 'connectorId2Claude 3 Sonnet',
-        name: 'Claude 3 Sonnet',
+        id: 'connectorId2Anthropic Claude 3.5 Sonnet',
+        name: 'Anthropic Claude 3.5 Sonnet',
         showConnectorName: false,
-        value: 'anthropic.claude-3-haiku-20240307-v1:0',
+        value: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
         promptTokenLimit: 200000,
       },
     ]);

diff --git a/x-pack/plugins/search_playground/public/hooks/use_llms_models.ts b/x-pack/plugins/search_playground/public/hooks/use_llms_models.ts
@@ -6,7 +6,7 @@
  */
 
 import { i18n } from '@kbn/i18n';
-import { BedrockLogo, OpenAILogo } from '@kbn/stack-connectors-plugin/public/common';
+import { BedrockLogo, OpenAILogo, GeminiLogo } from '@kbn/stack-connectors-plugin/public/common';
 import { ComponentType, useMemo } from 'react';
 import { LLMs } from '../../common/types';
 import { LLMModel } from '../types';
@@ -52,6 +52,15 @@ const mapLlmToModels: Record<
         promptTokenLimit: model.promptTokenLimit,
       })),
   },
+  [LLMs.gemini]: {
+    icon: GeminiLogo,
+    getModels: () =>
+      MODELS.filter(({ provider }) => provider === LLMs.gemini).map((model) => ({
+        label: model.name,
+        value: model.model,
+        promptTokenLimit: model.promptTokenLimit,
+      })),
+  },
 };
 
 export const useLLMsModels = (): LLMModel[] => {

diff --git a/x-pack/plugins/search_playground/public/hooks/use_load_connectors.ts b/x-pack/plugins/search_playground/public/hooks/use_load_connectors.ts
@@ -16,6 +16,7 @@ import {
   OPENAI_CONNECTOR_ID,
   OpenAiProviderType,
   BEDROCK_CONNECTOR_ID,
+  GEMINI_CONNECTOR_ID,
 } from '@kbn/stack-connectors-plugin/public/common';
 import { UserConfiguredActionConnector } from '@kbn/triggers-actions-ui-plugin/public/types';
 import { useKibana } from './use_kibana';
@@ -73,6 +74,17 @@ const connectorTypeToLLM: Array<{
       type: LLMs.bedrock,
     }),
   },
+  {
+    actionId: GEMINI_CONNECTOR_ID,
+    match: (connector) => connector.actionTypeId === GEMINI_CONNECTOR_ID,
+    transform: (connector) => ({
+      ...connector,
+      title: i18n.translate('xpack.searchPlayground.geminiConnectorTitle', {
+        defaultMessage: 'Gemini',
+      }),
+      type: LLMs.gemini,
+    }),
+  },
 ];
 
 type PlaygroundConnector = ActionConnector & { title: string; type: LLMs };

diff --git a/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts b/x-pack/plugins/search_playground/server/lib/conversational_chain.test.ts
@@ -97,6 +97,7 @@ describe('conversational chain', () => {
         inputTokensLimit: modelLimit,
       },
       prompt: 'you are a QA bot {question} {chat_history} {context}',
+      questionRewritePrompt: 'rewrite question {question} using {chat_history}"',
     });
 
     const stream = await conversationalChain.stream(aiClient, chat);
@@ -442,7 +443,7 @@ describe('conversational chain', () => {
           type: 'retrieved_docs',
         },
       ],
-      // Even with body_content of 1000, the token count should be below the model limit of 100
+      // Even with body_content of 1000, the token count should be below or equal to model limit of 100
       expectedTokens: [
         { type: 'context_token_count', count: 70 },
         { type: 'prompt_token_count', count: 97 },

diff --git a/x-pack/plugins/search_playground/server/lib/conversational_chain.ts b/x-pack/plugins/search_playground/server/lib/conversational_chain.ts
@@ -37,6 +37,7 @@ interface RAGOptions {
 interface ConversationalChainOptions {
   model: BaseLanguageModel;
   prompt: string;
+  questionRewritePrompt: string;
   rag?: RAGOptions;
 }
 
@@ -46,16 +47,6 @@ interface ContextInputs {
   question: string;
 }
 
-const CONDENSE_QUESTION_TEMPLATE = `Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language. Be verbose in your answer.
-
-Chat History:
-{chat_history}
-
-Follow Up Input: {question}
-Standalone question:`;
-
-const condenseQuestionPrompt = PromptTemplate.fromTemplate(CONDENSE_QUESTION_TEMPLATE);
-
 const formatVercelMessages = (chatHistory: VercelChatMessage[]) => {
   const formattedDialogueTurns = chatHistory.map((message) => {
     if (message.role === 'user') {
@@ -160,11 +151,21 @@ class ConversationalChainFn {
       retrievalChain = retriever.pipe(buildContext);
     }
 
-    let standaloneQuestionChain: Runnable = RunnableLambda.from((input) => input.question);
+    let standaloneQuestionChain: Runnable = RunnableLambda.from((input) => {
+      return input.question;
+    });
 
     if (previousMessages.length > 0) {
+      const questionRewritePromptTemplate = PromptTemplate.fromTemplate(
+        this.options.questionRewritePrompt
+      );
       standaloneQuestionChain = RunnableSequence.from([
-        condenseQuestionPrompt,
+        {
+          context: () => '',
+          chat_history: (input) => input.chat_history,
+          question: (input) => input.question,
+        },
+        questionRewritePromptTemplate,
         this.options.model,
         new StringOutputParser(),
       ]).withConfig({

diff --git a/x-pack/plugins/search_playground/server/lib/get_chat_params.test.ts b/x-pack/plugins/search_playground/server/lib/get_chat_params.test.ts
@@ -6,12 +6,13 @@
  */
 
 import { getChatParams } from './get_chat_params';
-import { ActionsClientChatOpenAI, ActionsClientLlm } from '@kbn/langchain/server';
+import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel } from '@kbn/langchain/server';
 import {
   OPENAI_CONNECTOR_ID,
   BEDROCK_CONNECTOR_ID,
+  GEMINI_CONNECTOR_ID,
 } from '@kbn/stack-connectors-plugin/public/common';
-import { Prompt } from '../../common/prompt';
+import { Prompt, QuestionRewritePrompt } from '../../common/prompt';
 import { KibanaRequest, Logger } from '@kbn/core/server';
 import { PluginStartContract as ActionsPluginStartContract } from '@kbn/actions-plugin/server';
 
@@ -20,12 +21,13 @@ jest.mock('@kbn/langchain/server', () => {
   return {
     ...original,
     ActionsClientChatOpenAI: jest.fn(),
-    ActionsClientLlm: jest.fn(),
+    ActionsClientSimpleChatModel: jest.fn(),
   };
 });
 
 jest.mock('../../common/prompt', () => ({
   Prompt: jest.fn((instructions) => instructions),
+  QuestionRewritePrompt: jest.fn((instructions) => instructions),
 }));
 
 jest.mock('uuid', () => ({
@@ -64,10 +66,45 @@ describe('getChatParams', () => {
       context: true,
       type: 'openai',
     });
+    expect(QuestionRewritePrompt).toHaveBeenCalledWith({
+      type: 'openai',
+    });
     expect(ActionsClientChatOpenAI).toHaveBeenCalledWith(expect.anything());
     expect(result.chatPrompt).toContain('Hello, world!');
   });
 
+  it('returns the correct chat model and prompt for Gemeni', async () => {
+    mockActionsClient.get.mockResolvedValue({ id: '1', actionTypeId: GEMINI_CONNECTOR_ID });
+
+    const result = await getChatParams(
+      {
+        connectorId: '1',
+        model: 'gemini-1.5-pro',
+        prompt: 'Hello, world!',
+        citations: true,
+      },
+      { actions, request, logger }
+    );
+    expect(Prompt).toHaveBeenCalledWith('Hello, world!', {
+      citations: true,
+      context: true,
+      type: 'gemini',
+    });
+    expect(QuestionRewritePrompt).toHaveBeenCalledWith({
+      type: 'gemini',
+    });
+    expect(ActionsClientSimpleChatModel).toHaveBeenCalledWith({
+      temperature: 0,
+      llmType: 'gemini',
+      logger: expect.anything(),
+      model: 'gemini-1.5-pro',
+      connectorId: '1',
+      actionsClient: expect.anything(),
+      streaming: true,
+    });
+    expect(result.chatPrompt).toContain('Hello, world!');
+  });
+
   it('returns the correct chat model and prompt for BEDROCK_CONNECTOR_ID', async () => {
     mockActionsClient.get.mockResolvedValue({ id: '2', actionTypeId: BEDROCK_CONNECTOR_ID });
 
@@ -86,14 +123,17 @@ describe('getChatParams', () => {
       context: true,
       type: 'anthropic',
     });
-    expect(ActionsClientLlm).toHaveBeenCalledWith({
+    expect(QuestionRewritePrompt).toHaveBeenCalledWith({
+      type: 'anthropic',
+    });
+    expect(ActionsClientSimpleChatModel).toHaveBeenCalledWith({
       temperature: 0,
       llmType: 'bedrock',
-      traceId: 'test-uuid',
       logger: expect.anything(),
       model: 'custom-model',
       connectorId: '2',
       actionsClient: expect.anything(),
+      streaming: true,
     });
     expect(result.chatPrompt).toContain('How does it work?');
   });