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 index c1c6f8e6a67aa..0182c4061866b 100644 --- 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 index 6a11466f9faa0..a9d2142fc3963 100644 --- 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 index e656291c6adce..593dc922e0b74 100644 --- 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') {