fix tests

elastic · Jun 10, 2024 · 3eafc96 · 3eafc96
1 parent 7f0957f
commit 3eafc96
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.test.tsx b/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}',
         }
       );
     });

diff --git a/x-pack/packages/kbn-langchain/server/language_models/simple_chat_model.test.ts b/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);
     });

diff --git a/x-pack/packages/kbn-langchain/server/language_models/simple_chat_model.ts b/x-pack/packages/kbn-langchain/server/language_models/simple_chat_model.ts
@@ -149,7 +149,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') {