Gemini Connector Assistant Integration

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.test.tsx

@@ -134,7 +134,7 @@ describe('API tests', () => {
       );
     });
 
-    it('calls the non-stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is true', async () => {
+    it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is true', async () => {
       const testProps: FetchConnectorExecuteAction = {
         ...fetchConnectorArgs,
         apiConfig: apiConfig.gemini,
@@ -145,13 +145,13 @@ describe('API tests', () => {
       expect(mockHttp.fetch).toHaveBeenCalledWith(
         '/internal/elastic_assistant/actions/connector/foo/_execute',
         {
-          ...staticDefaults,
-          body: '{"message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
+          ...streamingDefaults,
+          body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
         }
       );
     });
 
-    it('calls the non-stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
+    it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
       const testProps: FetchConnectorExecuteAction = {
         ...fetchConnectorArgs,
         apiConfig: apiConfig.gemini,
@@ -164,8 +164,8 @@ describe('API tests', () => {
       expect(mockHttp.fetch).toHaveBeenCalledWith(
         '/internal/elastic_assistant/actions/connector/foo/_execute',
         {
-          ...staticDefaults,
-          body: '{"message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
+          ...streamingDefaults,
+          body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
         }
       );
     });

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx

@@ -64,13 +64,7 @@ export const fetchConnectorExecuteAction = async ({
   traceOptions,
 }: FetchConnectorExecuteAction): Promise<FetchConnectorExecuteResponse> => {
   // TODO add streaming support for gemini with langchain on
-  const isStream =
-    assistantStreamingEnabled &&
-    (apiConfig.actionTypeId === '.gen-ai' ||
-      apiConfig.actionTypeId === '.bedrock' ||
-      // TODO add streaming support for gemini with langchain on
-      // tracked here: https://github.com/elastic/security-team/issues/7363
-      (apiConfig.actionTypeId === '.gemini' && !isEnabledRAGAlerts && !isEnabledKnowledgeBase));
+  const isStream = assistantStreamingEnabled;
 
   const optionalRequestParams = getOptionalRequestParams({
     isEnabledRAGAlerts,

x-pack/packages/kbn-elastic-assistant/impl/connectorland/use_load_connectors/index.tsx

@@ -30,6 +30,7 @@ export interface Props {
 const actionTypeKey = {
   bedrock: '.bedrock',
   openai: '.gen-ai',
+  gemini: '.gemini',
 };
 
 export const useLoadConnectors = ({
@@ -44,7 +45,9 @@ export const useLoadConnectors = ({
         (acc: AIConnector[], connector) => [
           ...acc,
           ...(!connector.isMissingSecrets &&
-          [actionTypeKey.bedrock, actionTypeKey.openai].includes(connector.actionTypeId)
+          [actionTypeKey.bedrock, actionTypeKey.openai, actionTypeKey.gemini].includes(
+            connector.actionTypeId
+          )
             ? [
                 {
                   ...connector,

x-pack/packages/kbn-langchain/server/index.ts

@@ -9,10 +9,12 @@ import { ActionsClientChatOpenAI } from './language_models/chat_openai';
 import { ActionsClientLlm } from './language_models/llm';
 import { ActionsClientSimpleChatModel } from './language_models/simple_chat_model';
 import { parseBedrockStream } from './utils/bedrock';
+import { parseGeminiResponse } from './utils/gemini';
 import { getDefaultArguments } from './language_models/constants';
 
 export {
   parseBedrockStream,
+  parseGeminiResponse,
   getDefaultArguments,
   ActionsClientChatOpenAI,
   ActionsClientLlm,

x-pack/packages/kbn-langchain/server/language_models/constants.ts

@@ -11,6 +11,10 @@ export const getDefaultArguments = (llmType?: string, temperature?: number, stop
         temperature: temperature ?? DEFAULT_BEDROCK_TEMPERATURE,
         stopSequences: stop ?? DEFAULT_BEDROCK_STOP_SEQUENCES,
       }
+    : llmType === 'gemini'
+    ? {
+        temperature: temperature ?? DEFAULT_GEMINI_TEMPERATURE,
+      }
     : { n: 1, stop: stop ?? null, temperature: temperature ?? DEFAULT_OPEN_AI_TEMPERATURE };
 
 export const DEFAULT_OPEN_AI_TEMPERATURE = 0.2;
@@ -19,4 +23,5 @@ export const DEFAULT_OPEN_AI_TEMPERATURE = 0.2;
 export const DEFAULT_OPEN_AI_MODEL = 'gpt-4';
 const DEFAULT_BEDROCK_TEMPERATURE = 0;
 const DEFAULT_BEDROCK_STOP_SEQUENCES = ['\n\nHuman:', '\nObservation:'];
+const DEFAULT_GEMINI_TEMPERATURE = 0;
 export const DEFAULT_TIMEOUT = 180000;

x-pack/packages/kbn-langchain/server/language_models/simple_chat_model.test.ts

@@ -14,6 +14,7 @@ import { mockActionResponse } from './mocks';
 import { BaseMessage } from '@langchain/core/messages';
 import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
 import { parseBedrockStream } from '../utils/bedrock';
+import { parseGeminiStream } from '../utils/gemini';
 
 const connectorId = 'mock-connector-id';
 
@@ -94,6 +95,7 @@ const defaultArgs = {
   streaming: false,
 };
 jest.mock('../utils/bedrock');
+jest.mock('../utils/gemini');
 
 describe('ActionsClientSimpleChatModel', () => {
   beforeEach(() => {
@@ -216,6 +218,7 @@ describe('ActionsClientSimpleChatModel', () => {
   describe('_call streaming: true', () => {
     beforeEach(() => {
       (parseBedrockStream as jest.Mock).mockResolvedValue(mockActionResponse.message);
+      (parseGeminiStream as jest.Mock).mockResolvedValue(mockActionResponse.message);
     });
     it('returns the expected content when _call is invoked with streaming and llmType is Bedrock', async () => {
       const actionsClientSimpleChatModel = new ActionsClientSimpleChatModel({
@@ -238,7 +241,7 @@ describe('ActionsClientSimpleChatModel', () => {
     it('returns the expected content when _call is invoked with streaming and llmType is Gemini', async () => {
       const actionsClientSimpleChatModel = new ActionsClientSimpleChatModel({
         ...defaultArgs,
-        actions: mockActions,
+        actions: mockStreamActions,
         llmType: 'gemini',
         streaming: true,
       });
@@ -248,8 +251,8 @@ describe('ActionsClientSimpleChatModel', () => {
         callOptions,
         callRunManager
       );
-      const subAction = mockExecute.mock.calls[0][0].params.subAction;
-      expect(subAction).toEqual('invokeAI');
+      const subAction = mockStreamExecute.mock.calls[0][0].params.subAction;
+      expect(subAction).toEqual('invokeStream');
 
       expect(result).toEqual(mockActionResponse.message);
     });

x-pack/packages/kbn-langchain/server/language_models/simple_chat_model.ts

@@ -17,6 +17,7 @@ import { KibanaRequest } from '@kbn/core-http-server';
 import { v4 as uuidv4 } from 'uuid';
 import { get } from 'lodash/fp';
 import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
+import { parseGeminiStream } from '../utils/gemini';
 import { parseBedrockStream } from '../utils/bedrock';
 import { getDefaultArguments } from './constants';
 
@@ -75,8 +76,7 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel {
     this.llmType = llmType ?? 'ActionsClientSimpleChatModel';
     this.model = model;
     this.temperature = temperature;
-    // only enable streaming for bedrock
-    this.streaming = streaming && llmType === 'bedrock';
+    this.streaming = streaming;
   }
 
   _llmType() {
@@ -154,7 +154,6 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel {
       return content; // per the contact of _call, return a string
     }
 
-    // Bedrock streaming
     const readable = get('data', actionResult) as Readable;
 
     if (typeof readable?.read !== 'function') {
@@ -182,13 +181,9 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel {
         }
       }
     };
+    const streamParser = this.llmType === 'bedrock' ? parseBedrockStream : parseGeminiStream;
 
-    const parsed = await parseBedrockStream(
-      readable,
-      this.#logger,
-      this.#signal,
-      handleLLMNewToken
-    );
+    const parsed = await streamParser(readable, this.#logger, this.#signal, handleLLMNewToken);
 
     return parsed; // per the contact of _call, return a string
   }

x-pack/packages/kbn-langchain/server/utils/bedrock.ts

@@ -6,17 +6,10 @@
  */
 
 import { finished } from 'stream/promises';
-import { Readable } from 'stream';
 import { Logger } from '@kbn/core/server';
 import { EventStreamCodec } from '@smithy/eventstream-codec';
 import { fromUtf8, toUtf8 } from '@smithy/util-utf8';
-
-type StreamParser = (
-  responseStream: Readable,
-  logger: Logger,
-  abortSignal?: AbortSignal,
-  tokenHandler?: (token: string) => void
-) => Promise<string>;
+import { StreamParser } from './types';
 
 export const parseBedrockStream: StreamParser = async (
   responseStream,

x-pack/packages/kbn-langchain/server/utils/gemini.test.ts

@@ -0,0 +1,89 @@
+/*
+ * 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 { Readable } from 'stream';
+import { parseGeminiStream, parseGeminiResponse } from './gemini';
+import { loggerMock } from '@kbn/logging-mocks';
+
+describe('parseGeminiStream', () => {
+  const mockLogger = loggerMock.create();
+  let mockStream: Readable;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockStream = new Readable({
+      read() {},
+    });
+  });
+
+  it('should parse the stream correctly', async () => {
+    const data =
+      'data: {"candidates":[{"content":{"role":"system","parts":[{"text":"Hello"}]},"finishReason":"stop","safetyRatings":[{"category":"safe","probability":"low"}]}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":10,"totalTokenCount":20}}\n';
+    mockStream.push(data);
+    mockStream.push(null);
+
+    const result = await parseGeminiStream(mockStream, mockLogger);
+    expect(result).toBe('Hello');
+  });
+
+  it('should handle abort signal correctly', async () => {
+    const abortSignal = new AbortController().signal;
+    setTimeout(() => {
+      abortSignal.dispatchEvent(new Event('abort'));
+    }, 100);
+
+    const result = parseGeminiStream(mockStream, mockLogger, abortSignal);
+
+    await expect(result).resolves.toBe('');
+    expect(mockLogger.info).toHaveBeenCalledWith('Bedrock stream parsing was aborted.');
+  });
+
+  it('should call tokenHandler with correct tokens', async () => {
+    const data =
+      'data: {"candidates":[{"content":{"role":"system","parts":[{"text":"Hello world"}]},"finishReason":"stop","safetyRatings":[{"category":"safe","probability":"low"}]}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":10,"totalTokenCount":20}}\n';
+    mockStream.push(data);
+    mockStream.push(null);
+
+    const tokenHandler = jest.fn();
+    await parseGeminiStream(mockStream, mockLogger, undefined, tokenHandler);
+
+    expect(tokenHandler).toHaveBeenCalledWith('Hello ');
+    expect(tokenHandler).toHaveBeenCalledWith('world ');
+  });
+
+  it('should handle stream error correctly', async () => {
+    const error = new Error('Stream error');
+    const resultPromise = parseGeminiStream(mockStream, mockLogger);
+
+    mockStream.emit('error', error);
+
+    await expect(resultPromise).rejects.toThrow('Stream error');
+  });
+});
+
+describe('parseGeminiResponse', () => {
+  it('should parse response correctly', () => {
+    const response =
+      'data: {"candidates":[{"content":{"role":"system","parts":[{"text":"Hello"}]},"finishReason":"stop","safetyRatings":[{"category":"safe","probability":"low"}]}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":10,"totalTokenCount":20}}\n';
+    const result = parseGeminiResponse(response);
+    expect(result).toBe('Hello');
+  });
+
+  it('should ignore lines that do not start with data: ', () => {
+    const response =
+      'invalid line\ndata: {"candidates":[{"content":{"role":"system","parts":[{"text":"Hello"}]},"finishReason":"stop","safetyRatings":[{"category":"safe","probability":"low"}]}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":10,"totalTokenCount":20}}\n';
+    const result = parseGeminiResponse(response);
+    expect(result).toBe('Hello');
+  });
+
+  it('should ignore lines that end with [DONE]', () => {
+    const response =
+      'data: {"candidates":[{"content":{"role":"system","parts":[{"text":"Hello"}]},"finishReason":"stop","safetyRatings":[{"category":"safe","probability":"low"}]}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":10,"totalTokenCount":20}}\ndata: [DONE]';
+    const result = parseGeminiResponse(response);
+    expect(result).toBe('Hello');
+  });
+});

x-pack/packages/kbn-langchain/server/utils/gemini.ts

@@ -0,0 +1,80 @@
+/*
+ * 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 { StreamParser } from './types';
+
+export const parseGeminiStream: StreamParser = async (
+  stream,
+  logger,
+  abortSignal,
+  tokenHandler
+) => {
+  let responseBody = '';
+  stream.on('data', (chunk) => {
+    const decoded = chunk.toString();
+    const parsed = parseGeminiResponse(decoded);
+    if (tokenHandler) {
+      const splitByQuotes = parsed.split(`"`);
+      splitByQuotes.forEach((chunkk, index) => {
+        // add quote back on except for last chunk
+        const splitBySpace = `${chunkk}${index === splitByQuotes.length - 1 ? '' : '"'}`.split(` `);
+
+        for (const char of splitBySpace) {
+          tokenHandler(`${char} `);
+        }
+      });
+    }
+    responseBody += parsed;
+  });
+  return new Promise((resolve, reject) => {
+    stream.on('end', () => {
+      resolve(responseBody);
+    });
+    stream.on('error', (err) => {
+      reject(err);
+    });
+    if (abortSignal) {
+      abortSignal.addEventListener('abort', () => {
+        logger.info('Bedrock stream parsing was aborted.');
+        stream.destroy();
+        resolve(responseBody);
+      });
+    }
+  });
+};
+
+/** Parse Gemini stream response body */
+export const parseGeminiResponse = (responseBody: string) => {
+  return responseBody
+    .split('\n')
+    .filter((line) => line.startsWith('data: ') && !line.endsWith('[DONE]'))
+    .map((line) => JSON.parse(line.replace('data: ', '')))
+    .filter(
+      (
+        line
+      ): line is {
+        candidates: Array<{
+          content: { role: string; parts: Array<{ text: string }> };
+          finishReason: string;
+          safetyRatings: Array<{ category: string; probability: string }>;
+        }>;
+        usageMetadata: {
+          promptTokenCount: number;
+          candidatesTokenCount: number;
+          totalTokenCount: number;
+        };
+      } => 'candidates' in line
+    )
+    .reduce((prev, line) => {
+      if (line.candidates[0] && line.candidates[0].content) {
+        const parts = line.candidates[0].content.parts;
+        const text = parts.map((part) => part.text).join('');
+        return prev + text;
+      }
+      return prev;
+    }, '');
+};