From beeea66a38d90ba2f1dca0bc696585d4a2d945b3 Mon Sep 17 00:00:00 2001 From: SuZhou-Joe Date: Mon, 11 Dec 2023 16:03:38 +0800 Subject: [PATCH 1/5] feat: integrate regenerate API Signed-off-by: SuZhou-Joe --- public/hooks/use_chat_actions.tsx | 8 +- public/tabs/chat/messages/message_bubble.tsx | 8 +- public/types.ts | 2 +- server/routes/chat_routes.ts | 49 +++-- server/services/chat/chat_service.ts | 11 +- .../services/chat/olly_chat_service.test.ts | 176 ++++++++++++++++++ server/services/chat/olly_chat_service.ts | 132 ++++++++----- 7 files changed, 301 insertions(+), 85 deletions(-) create mode 100644 server/services/chat/olly_chat_service.test.ts diff --git a/public/hooks/use_chat_actions.tsx b/public/hooks/use_chat_actions.tsx index e1975b46..5d5f141a 100644 --- a/public/hooks/use_chat_actions.tsx +++ b/public/hooks/use_chat_actions.tsx @@ -160,7 +160,7 @@ export const useChatActions = (): AssistantActions => { } }; - const regenerate = async () => { + const regenerate = async (interactionId: string) => { if (chatContext.sessionId) { const abortController = new AbortController(); abortControllerRef = abortController; @@ -168,7 +168,11 @@ export const useChatActions = (): AssistantActions => { try { const response = await core.services.http.put(`${ASSISTANT_API.REGENERATE}`, { - body: JSON.stringify({ sessionId: chatContext.sessionId }), + body: JSON.stringify({ + sessionId: chatContext.sessionId, + rootAgentId: chatContext.rootAgentId, + interactionId, + }), }); if (abortController.signal.aborted) { diff --git a/public/tabs/chat/messages/message_bubble.tsx b/public/tabs/chat/messages/message_bubble.tsx index 34199640..8a704e2a 100644 --- a/public/tabs/chat/messages/message_bubble.tsx +++ b/public/tabs/chat/messages/message_bubble.tsx @@ -30,7 +30,7 @@ type MessageBubbleProps = { showActionBar: boolean; showRegenerate?: boolean; shouldActionBarVisibleOnHover?: boolean; - onRegenerate?: () => void; + onRegenerate?: (interactionId: string) => void; } & ( | { message: IMessage; @@ -192,17 +192,17 @@ export const MessageBubble: React.FC = React.memo((props) => )} - {props.showRegenerate && ( + {props.showRegenerate && props.interaction?.interaction_id ? ( props.onRegenerate?.(props.interaction?.interaction_id || '')} title="regenerate message" color="text" iconType="refresh" /> - )} + ) : null} {showFeedback && ( // After feedback, only corresponding thumb icon will be kept and disabled. <> diff --git a/public/types.ts b/public/types.ts index d8f152a3..5fd027e2 100644 --- a/public/types.ts +++ b/public/types.ts @@ -16,7 +16,7 @@ export interface AssistantActions { openChatUI: (sessionId?: string) => void; executeAction: (suggestedAction: ISuggestedAction, message: IMessage) => void; abortAction: (sessionId?: string) => void; - regenerate: () => void; + regenerate: (interactionId: string) => void; } export interface AppPluginStartDependencies { diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts index 22328032..e43c7c93 100644 --- a/server/routes/chat_routes.ts +++ b/server/routes/chat_routes.ts @@ -13,7 +13,6 @@ import { } from '../../../../src/core/server'; import { ASSISTANT_API } from '../../common/constants/llm'; import { OllyChatService } from '../services/chat/olly_chat_service'; -import { IMessage, IInput } from '../../common/types/chat_saved_object_attributes'; import { AgentFrameworkStorageService } from '../services/storage/agent_framework_storage_service'; import { RoutesOptions } from '../types'; import { ChatService } from '../services/chat/chat_service'; @@ -64,6 +63,7 @@ const regenerateRoute = { body: schema.object({ sessionId: schema.string(), rootAgentId: schema.string(), + interactionId: schema.string(), }), }, }; @@ -314,42 +314,39 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) request, response ): Promise> => { - const { sessionId, rootAgentId } = request.body; + const { sessionId: sessionIdInRequestBody, rootAgentId, interactionId } = request.body; const storageService = createStorageService(context); - let messages: IMessage[] = []; const chatService = createChatService(); + let outputs: Awaited> | undefined; + + /** + * Get final answer from Agent framework + */ try { - const session = await storageService.getSession(sessionId); - messages.push(...session.messages); + outputs = await chatService.regenerate( + { sessionId: sessionIdInRequestBody, rootAgentId, interactionId }, + context + ); } catch (error) { - return response.custom({ statusCode: error.statusCode || 500, body: error.message }); + context.assistant_plugin.logger.error(error); } - const lastInputIndex = messages.findLastIndex((msg) => msg.type === 'input'); - // Find last input message - const input = messages[lastInputIndex] as IInput; - // Take the messages before last input message as memory as regenerate will exclude the last outputs - messages = messages.slice(0, lastInputIndex); - + /** + * Retrieve latest interactions from memory + */ + const sessionId = sessionIdInRequestBody; try { - const outputs = await chatService.requestLLM( - { messages, input, sessionId, rootAgentId }, - context - ); - const title = input.content.substring(0, 50); - const saveMessagesResponse = await storageService.saveMessages( - title, - sessionId, - [...messages, input, ...outputs.messages].filter( - (message) => message.content !== 'AbortError' - ) - ); + const conversation = await storageService.getSession(sessionId); + return response.ok({ - body: { ...saveMessagesResponse, title }, + body: { + ...conversation, + sessionId, + }, }); } catch (error) { - context.assistant_plugin.logger.warn(error); + context.assistant_plugin.logger.error(error); return response.custom({ statusCode: error.statusCode || 500, body: error.message }); } } diff --git a/server/services/chat/chat_service.ts b/server/services/chat/chat_service.ts index ac15adf6..25fe703f 100644 --- a/server/services/chat/chat_service.ts +++ b/server/services/chat/chat_service.ts @@ -10,8 +10,15 @@ import { LLMRequestSchema } from '../../routes/chat_routes'; export interface ChatService { requestLLM( payload: { messages: IMessage[]; input: IInput; sessionId?: string }, - context: RequestHandlerContext, - request: OpenSearchDashboardsRequest + context: RequestHandlerContext + ): Promise<{ + messages: IMessage[]; + memoryId: string; + }>; + + regenerate( + payload: { sessionId: string; interactionId: string; rootAgentId: string }, + context: RequestHandlerContext ): Promise<{ messages: IMessage[]; memoryId: string; diff --git a/server/services/chat/olly_chat_service.test.ts b/server/services/chat/olly_chat_service.test.ts new file mode 100644 index 00000000..1d5f563d --- /dev/null +++ b/server/services/chat/olly_chat_service.test.ts @@ -0,0 +1,176 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { OllyChatService } from './olly_chat_service'; +import { CoreRouteHandlerContext } from '../../../../../src/core/server/core_route_handler_context'; +import { coreMock, httpServerMock } from '../../../../../src/core/server/mocks'; +import { loggerMock } from '../../../../../src/core/server/logging/logger.mock'; + +describe('OllyChatService', () => { + const ollyChatService = new OllyChatService(); + const coreContext = new CoreRouteHandlerContext( + coreMock.createInternalStart(), + httpServerMock.createOpenSearchDashboardsRequest() + ); + const mockedTransport = coreContext.opensearch.client.asCurrentUser.transport + .request as jest.Mock; + const contextMock = { + core: coreContext, + assistant_plugin: { + logger: loggerMock.create(), + }, + }; + beforeEach(() => { + mockedTransport.mockClear(); + }); + it('requestLLM should invoke client call with correct params', async () => { + mockedTransport.mockImplementationOnce(() => { + return { + body: { + inference_results: [ + { + output: [ + { + name: 'memory_id', + result: 'foo', + }, + ], + }, + ], + }, + }; + }); + const result = await ollyChatService.requestLLM( + { + messages: [], + input: { + type: 'input', + contentType: 'text', + content: 'content', + }, + sessionId: '', + rootAgentId: 'rootAgentId', + }, + contextMock + ); + expect(mockedTransport.mock.calls).toMatchInlineSnapshot(` + Array [ + Array [ + Object { + "body": Object { + "parameters": Object { + "question": "content", + "verbose": true, + }, + }, + "method": "POST", + "path": "/_plugins/_ml/agents/rootAgentId/_execute", + }, + Object { + "maxRetries": 0, + "requestTimeout": 300000, + }, + ], + ] + `); + expect(result).toMatchInlineSnapshot(` + Object { + "memoryId": "foo", + "messages": Array [], + } + `); + }); + + it('requestLLM should throw error when transport.request throws error', async () => { + mockedTransport.mockImplementationOnce(() => { + throw new Error('error'); + }); + expect( + ollyChatService.requestLLM( + { + messages: [], + input: { + type: 'input', + contentType: 'text', + content: 'content', + }, + sessionId: '', + rootAgentId: 'rootAgentId', + }, + contextMock + ) + ).rejects.toMatchInlineSnapshot(`[Error: error]`); + }); + + it('regenerate should invoke client call with correct params', async () => { + mockedTransport.mockImplementationOnce(() => { + return { + body: { + inference_results: [ + { + output: [ + { + name: 'memory_id', + result: 'foo', + }, + ], + }, + ], + }, + }; + }); + const result = await ollyChatService.regenerate( + { + sessionId: 'sessionId', + rootAgentId: 'rootAgentId', + interactionId: 'interactionId', + }, + contextMock + ); + expect(mockedTransport.mock.calls).toMatchInlineSnapshot(` + Array [ + Array [ + Object { + "body": Object { + "parameters": Object { + "memory_id": "sessionId", + "regenerate_interaction_id": "interactionId", + "verbose": true, + }, + }, + "method": "POST", + "path": "/_plugins/_ml/agents/rootAgentId/_execute", + }, + Object { + "maxRetries": 0, + "requestTimeout": 300000, + }, + ], + ] + `); + expect(result).toMatchInlineSnapshot(` + Object { + "memoryId": "foo", + "messages": Array [], + } + `); + }); + + it('regenerate should throw error when transport.request throws error', async () => { + mockedTransport.mockImplementationOnce(() => { + throw new Error('error'); + }); + expect( + ollyChatService.regenerate( + { + sessionId: 'sessionId', + rootAgentId: 'rootAgentId', + interactionId: 'interactionId', + }, + contextMock + ) + ).rejects.toMatchInlineSnapshot(`[Error: error]`); + }); +}); diff --git a/server/services/chat/olly_chat_service.ts b/server/services/chat/olly_chat_service.ts index 13981eda..5535f241 100644 --- a/server/services/chat/olly_chat_service.ts +++ b/server/services/chat/olly_chat_service.ts @@ -14,6 +14,70 @@ const MEMORY_ID_FIELD = 'memory_id'; export class OllyChatService implements ChatService { static abortControllers: Map = new Map(); + private async requestAgentRun( + payload: { input?: IInput; sessionId?: string; rootAgentId: string; interactionId?: string }, + context: RequestHandlerContext + ) { + const { input, sessionId, interactionId, rootAgentId } = payload; + const parametersPayload: { + question?: string; + verbose?: boolean; + memory_id?: string; + regenerate_interaction_id?: string; + } = { + verbose: true, + }; + + if (input) { + parametersPayload.question = input.content; + } + + if (interactionId) { + parametersPayload.regenerate_interaction_id = interactionId; + } + + if (sessionId) { + parametersPayload.memory_id = sessionId; + } + + const opensearchClient = context.core.opensearch.client.asCurrentUser; + const agentFrameworkResponse = (await opensearchClient.transport.request( + { + method: 'POST', + path: `${ML_COMMONS_BASE_API}/agents/${rootAgentId}/_execute`, + body: { + parameters: parametersPayload, + }, + }, + { + /** + * It is time-consuming for LLM to generate final answer + * Give it a large timeout window + */ + requestTimeout: 5 * 60 * 1000, + /** + * Do not retry + */ + maxRetries: 0, + } + )) as ApiResponse<{ + inference_results: Array<{ + output: Array<{ name: string; result?: string }>; + }>; + }>; + const outputBody = agentFrameworkResponse.body.inference_results?.[0]?.output; + const memoryIdItem = outputBody?.find((item) => item.name === MEMORY_ID_FIELD); + + return { + /** + * Interactions will be stored in Agent framework, + * thus we do not need to return the latest message back. + */ + messages: [], + memoryId: memoryIdItem?.result || '', + }; + } + public async requestLLM( payload: { messages: IMessage[]; input: IInput; sessionId?: string; rootAgentId: string }, context: RequestHandlerContext @@ -21,63 +85,31 @@ export class OllyChatService implements ChatService { messages: IMessage[]; memoryId: string; }> { - const { input, sessionId, rootAgentId } = payload; - const opensearchClient = context.core.opensearch.client.asCurrentUser; - if (payload.sessionId) { OllyChatService.abortControllers.set(payload.sessionId, new AbortController()); } try { - /** - * Wait for an API to fetch root agent id. - */ - const parametersPayload: { - question: string; - verbose?: boolean; - memory_id?: string; - } = { - question: input.content, - verbose: true, - }; - if (sessionId) { - parametersPayload.memory_id = sessionId; + return await this.requestAgentRun(payload, context); + } catch (error) { + throw error; + } finally { + if (payload.sessionId) { + OllyChatService.abortControllers.delete(payload.sessionId); } - const agentFrameworkResponse = (await opensearchClient.transport.request( - { - method: 'POST', - path: `${ML_COMMONS_BASE_API}/agents/${rootAgentId}/_execute`, - body: { - parameters: parametersPayload, - }, - }, - { - /** - * It is time-consuming for LLM to generate final answer - * Give it a large timeout window - */ - requestTimeout: 5 * 60 * 1000, - /** - * Do not retry - */ - maxRetries: 0, - } - )) as ApiResponse<{ - inference_results: Array<{ - output: Array<{ name: string; result?: string }>; - }>; - }>; - const outputBody = agentFrameworkResponse.body.inference_results?.[0]?.output; - const memoryIdItem = outputBody?.find((item) => item.name === MEMORY_ID_FIELD); + } + } - return { - /** - * Interactions will be stored in Agent framework, - * thus we do not need to return the latest message back. - */ - messages: [], - memoryId: memoryIdItem?.result || '', - }; + async regenerate( + payload: { sessionId: string; interactionId: string; rootAgentId: string }, + context: RequestHandlerContext + ): Promise<{ messages: IMessage[]; memoryId: string }> { + if (payload.sessionId) { + OllyChatService.abortControllers.set(payload.sessionId, new AbortController()); + } + + try { + return await this.requestAgentRun(payload, context); } catch (error) { throw error; } finally { From 06b5adbeea62df93a433db651534e974723c6b21 Mon Sep 17 00:00:00 2001 From: SuZhou-Joe Date: Mon, 11 Dec 2023 16:26:45 +0800 Subject: [PATCH 2/5] feat: optimize code Signed-off-by: SuZhou-Joe --- server/services/chat/olly_chat_service.ts | 61 +++++++++++++---------- 1 file changed, 34 insertions(+), 27 deletions(-) diff --git a/server/services/chat/olly_chat_service.ts b/server/services/chat/olly_chat_service.ts index 5535f241..8eb9a929 100644 --- a/server/services/chat/olly_chat_service.ts +++ b/server/services/chat/olly_chat_service.ts @@ -15,38 +15,22 @@ export class OllyChatService implements ChatService { static abortControllers: Map = new Map(); private async requestAgentRun( - payload: { input?: IInput; sessionId?: string; rootAgentId: string; interactionId?: string }, - context: RequestHandlerContext - ) { - const { input, sessionId, interactionId, rootAgentId } = payload; - const parametersPayload: { + rootAgentId: string, + payload: { question?: string; verbose?: boolean; memory_id?: string; regenerate_interaction_id?: string; - } = { - verbose: true, - }; - - if (input) { - parametersPayload.question = input.content; - } - - if (interactionId) { - parametersPayload.regenerate_interaction_id = interactionId; - } - - if (sessionId) { - parametersPayload.memory_id = sessionId; - } - + }, + context: RequestHandlerContext + ) { const opensearchClient = context.core.opensearch.client.asCurrentUser; const agentFrameworkResponse = (await opensearchClient.transport.request( { method: 'POST', path: `${ML_COMMONS_BASE_API}/agents/${rootAgentId}/_execute`, body: { - parameters: parametersPayload, + parameters: payload, }, }, { @@ -85,12 +69,26 @@ export class OllyChatService implements ChatService { messages: IMessage[]; memoryId: string; }> { + const { input, sessionId, rootAgentId } = payload; if (payload.sessionId) { OllyChatService.abortControllers.set(payload.sessionId, new AbortController()); } + const parametersPayload: { + question?: string; + verbose?: boolean; + memory_id?: string; + } = { + question: input.content, + verbose: true, + }; + + if (sessionId) { + parametersPayload.memory_id = sessionId; + } + try { - return await this.requestAgentRun(payload, context); + return await this.requestAgentRun(rootAgentId, parametersPayload, context); } catch (error) { throw error; } finally { @@ -104,12 +102,21 @@ export class OllyChatService implements ChatService { payload: { sessionId: string; interactionId: string; rootAgentId: string }, context: RequestHandlerContext ): Promise<{ messages: IMessage[]; memoryId: string }> { - if (payload.sessionId) { - OllyChatService.abortControllers.set(payload.sessionId, new AbortController()); - } + const { sessionId, interactionId, rootAgentId } = payload; + const parametersPayload: { + verbose?: boolean; + memory_id?: string; + regenerate_interaction_id?: string; + } = { + memory_id: sessionId, + regenerate_interaction_id: interactionId, + verbose: true, + }; + + OllyChatService.abortControllers.set(sessionId, new AbortController()); try { - return await this.requestAgentRun(payload, context); + return await this.requestAgentRun(rootAgentId, parametersPayload, context); } catch (error) { throw error; } finally { From 8562436564466d512265ecb9e16f5f8c29958dcd Mon Sep 17 00:00:00 2001 From: SuZhou-Joe Date: Mon, 11 Dec 2023 16:39:15 +0800 Subject: [PATCH 3/5] fix: unit test error Signed-off-by: SuZhou-Joe --- public/tabs/chat/messages/message_bubble.test.tsx | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/public/tabs/chat/messages/message_bubble.test.tsx b/public/tabs/chat/messages/message_bubble.test.tsx index 4293dec4..d4652f32 100644 --- a/public/tabs/chat/messages/message_bubble.test.tsx +++ b/public/tabs/chat/messages/message_bubble.test.tsx @@ -130,6 +130,13 @@ describe('', () => { contentType: 'markdown', content: 'here are the indices in your cluster: .alert', }} + interaction={{ + input: 'foo', + response: 'bar', + conversation_id: 'foo', + interaction_id: 'bar', + create_time: new Date().toLocaleString(), + }} /> ); expect(screen.queryAllByTitle('regenerate message')).toHaveLength(1); From 07c9b046c82421d56cafd561793fed17216df6ed Mon Sep 17 00:00:00 2001 From: SuZhou-Joe Date: Tue, 12 Dec 2023 13:52:45 +0800 Subject: [PATCH 4/5] feat: optimize code Signed-off-by: SuZhou-Joe --- server/routes/chat_routes.ts | 8 +- server/services/chat/olly_chat_service.ts | 98 ++++++++++------------- 2 files changed, 46 insertions(+), 60 deletions(-) diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts index e43c7c93..5e195820 100644 --- a/server/routes/chat_routes.ts +++ b/server/routes/chat_routes.ts @@ -314,7 +314,7 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) request, response ): Promise> => { - const { sessionId: sessionIdInRequestBody, rootAgentId, interactionId } = request.body; + const { sessionId, rootAgentId, interactionId } = request.body; const storageService = createStorageService(context); const chatService = createChatService(); @@ -324,10 +324,7 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) * Get final answer from Agent framework */ try { - outputs = await chatService.regenerate( - { sessionId: sessionIdInRequestBody, rootAgentId, interactionId }, - context - ); + outputs = await chatService.regenerate({ sessionId, rootAgentId, interactionId }, context); } catch (error) { context.assistant_plugin.logger.error(error); } @@ -335,7 +332,6 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) /** * Retrieve latest interactions from memory */ - const sessionId = sessionIdInRequestBody; try { const conversation = await storageService.getSession(sessionId); diff --git a/server/services/chat/olly_chat_service.ts b/server/services/chat/olly_chat_service.ts index 8eb9a929..694270f5 100644 --- a/server/services/chat/olly_chat_service.ts +++ b/server/services/chat/olly_chat_service.ts @@ -24,42 +24,53 @@ export class OllyChatService implements ChatService { }, context: RequestHandlerContext ) { + if (payload.memory_id) { + OllyChatService.abortControllers.set(payload.memory_id, new AbortController()); + } const opensearchClient = context.core.opensearch.client.asCurrentUser; - const agentFrameworkResponse = (await opensearchClient.transport.request( - { - method: 'POST', - path: `${ML_COMMONS_BASE_API}/agents/${rootAgentId}/_execute`, - body: { - parameters: payload, + + try { + const agentFrameworkResponse = (await opensearchClient.transport.request( + { + method: 'POST', + path: `${ML_COMMONS_BASE_API}/agents/${rootAgentId}/_execute`, + body: { + parameters: payload, + }, }, - }, - { - /** - * It is time-consuming for LLM to generate final answer - * Give it a large timeout window - */ - requestTimeout: 5 * 60 * 1000, + { + /** + * It is time-consuming for LLM to generate final answer + * Give it a large timeout window + */ + requestTimeout: 5 * 60 * 1000, + /** + * Do not retry + */ + maxRetries: 0, + } + )) as ApiResponse<{ + inference_results: Array<{ + output: Array<{ name: string; result?: string }>; + }>; + }>; + const outputBody = agentFrameworkResponse.body.inference_results?.[0]?.output; + const memoryIdItem = outputBody?.find((item) => item.name === MEMORY_ID_FIELD); + return { /** - * Do not retry + * Interactions will be stored in Agent framework, + * thus we do not need to return the latest message back. */ - maxRetries: 0, + messages: [], + memoryId: memoryIdItem?.result || '', + }; + } catch (error) { + throw error; + } finally { + if (payload.memory_id) { + OllyChatService.abortControllers.delete(payload.memory_id); } - )) as ApiResponse<{ - inference_results: Array<{ - output: Array<{ name: string; result?: string }>; - }>; - }>; - const outputBody = agentFrameworkResponse.body.inference_results?.[0]?.output; - const memoryIdItem = outputBody?.find((item) => item.name === MEMORY_ID_FIELD); - - return { - /** - * Interactions will be stored in Agent framework, - * thus we do not need to return the latest message back. - */ - messages: [], - memoryId: memoryIdItem?.result || '', - }; + } } public async requestLLM( @@ -70,9 +81,6 @@ export class OllyChatService implements ChatService { memoryId: string; }> { const { input, sessionId, rootAgentId } = payload; - if (payload.sessionId) { - OllyChatService.abortControllers.set(payload.sessionId, new AbortController()); - } const parametersPayload: { question?: string; @@ -87,15 +95,7 @@ export class OllyChatService implements ChatService { parametersPayload.memory_id = sessionId; } - try { - return await this.requestAgentRun(rootAgentId, parametersPayload, context); - } catch (error) { - throw error; - } finally { - if (payload.sessionId) { - OllyChatService.abortControllers.delete(payload.sessionId); - } - } + return await this.requestAgentRun(rootAgentId, parametersPayload, context); } async regenerate( @@ -113,17 +113,7 @@ export class OllyChatService implements ChatService { verbose: true, }; - OllyChatService.abortControllers.set(sessionId, new AbortController()); - - try { - return await this.requestAgentRun(rootAgentId, parametersPayload, context); - } catch (error) { - throw error; - } finally { - if (payload.sessionId) { - OllyChatService.abortControllers.delete(payload.sessionId); - } - } + return await this.requestAgentRun(rootAgentId, parametersPayload, context); } abortAgentExecution(sessionId: string) { From ccc75b23aba6cede049406aa23c46aa5cd8bbfb2 Mon Sep 17 00:00:00 2001 From: SuZhou-Joe Date: Tue, 12 Dec 2023 16:02:20 +0800 Subject: [PATCH 5/5] feat: optimize code Signed-off-by: SuZhou-Joe --- server/services/chat/olly_chat_service.ts | 29 ++++++++++------------- server/types.ts | 2 +- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/server/services/chat/olly_chat_service.ts b/server/services/chat/olly_chat_service.ts index 694270f5..ed2cb57b 100644 --- a/server/services/chat/olly_chat_service.ts +++ b/server/services/chat/olly_chat_service.ts @@ -9,6 +9,13 @@ import { IMessage, IInput } from '../../../common/types/chat_saved_object_attrib import { ChatService } from './chat_service'; import { ML_COMMONS_BASE_API } from '../../utils/constants'; +interface AgentRunPayload { + question?: string; + verbose?: boolean; + memory_id?: string; + regenerate_interaction_id?: string; +} + const MEMORY_ID_FIELD = 'memory_id'; export class OllyChatService implements ChatService { @@ -16,12 +23,7 @@ export class OllyChatService implements ChatService { private async requestAgentRun( rootAgentId: string, - payload: { - question?: string; - verbose?: boolean; - memory_id?: string; - regenerate_interaction_id?: string; - }, + payload: AgentRunPayload, context: RequestHandlerContext ) { if (payload.memory_id) { @@ -82,11 +84,7 @@ export class OllyChatService implements ChatService { }> { const { input, sessionId, rootAgentId } = payload; - const parametersPayload: { - question?: string; - verbose?: boolean; - memory_id?: string; - } = { + const parametersPayload: Pick = { question: input.content, verbose: true, }; @@ -103,11 +101,10 @@ export class OllyChatService implements ChatService { context: RequestHandlerContext ): Promise<{ messages: IMessage[]; memoryId: string }> { const { sessionId, interactionId, rootAgentId } = payload; - const parametersPayload: { - verbose?: boolean; - memory_id?: string; - regenerate_interaction_id?: string; - } = { + const parametersPayload: Pick< + AgentRunPayload, + 'regenerate_interaction_id' | 'verbose' | 'memory_id' + > = { memory_id: sessionId, regenerate_interaction_id: interactionId, verbose: true, diff --git a/server/types.ts b/server/types.ts index 948ed5aa..3e93ce8f 100644 --- a/server/types.ts +++ b/server/types.ts @@ -4,7 +4,7 @@ */ import { IMessage, Interaction } from '../common/types/chat_saved_object_attributes'; -import { ILegacyClusterClient, Logger } from '../../../src/core/server'; +import { Logger } from '../../../src/core/server'; // eslint-disable-next-line @typescript-eslint/no-empty-interface export interface AssistantPluginSetup {}