@@ -14,9 +14,24 @@ import { mockOllyChatService } from '../services/chat/olly_chat_service.mock';
14
14
import { loggerMock } from '../../../../src/core/server/logging/logger.mock' ;
15
15
import { registerChatRoutes } from './chat_routes' ;
16
16
import { ASSISTANT_API } from '../../common/constants/llm' ;
17
+ import { getOpenSearchClientTransport } from '../utils/get_opensearch_client_transport' ;
17
18
18
- const mockedLogger = loggerMock . create ( ) ;
19
+ jest . mock ( '../utils/get_opensearch_client_transport' ) ;
20
+
21
+ beforeEach ( ( ) => {
22
+ ( getOpenSearchClientTransport as jest . Mock ) . mockImplementation ( ( { dataSourceId } ) => {
23
+ if ( dataSourceId ) {
24
+ return 'dataSource-client' ;
25
+ } else {
26
+ return 'client' ;
27
+ }
28
+ } ) ;
29
+ } ) ;
30
+ afterEach ( ( ) => {
31
+ ( getOpenSearchClientTransport as jest . Mock ) . mockClear ( ) ;
32
+ } ) ;
19
33
34
+ const mockedLogger = loggerMock . create ( ) ;
20
35
const router = new Router (
21
36
'' ,
22
37
mockedLogger ,
@@ -30,38 +45,90 @@ registerChatRoutes(router, {
30
45
messageParsers : [ ] ,
31
46
} ) ;
32
47
33
- const triggerDeleteConversation = ( conversationId : string ) =>
48
+ const triggerDeleteConversation = ( conversationId : string , dataSourceId ?: string ) =>
34
49
triggerHandler ( router , {
35
50
method : 'delete' ,
36
51
path : `${ ASSISTANT_API . CONVERSATION } /{conversationId}` ,
37
- req : httpServerMock . createRawRequest ( { params : { conversationId } } ) ,
52
+ req : httpServerMock . createRawRequest ( {
53
+ params : { conversationId } ,
54
+ ...( dataSourceId
55
+ ? {
56
+ query : {
57
+ dataSourceId,
58
+ } ,
59
+ }
60
+ : { } ) ,
61
+ } ) ,
38
62
} ) ;
39
63
const triggerUpdateConversation = (
40
64
params : { conversationId : string } ,
41
- payload : { title : string }
65
+ payload : { title : string } ,
66
+ dataSourceId ?: string
42
67
) =>
43
68
triggerHandler ( router , {
44
69
method : 'put' ,
45
70
path : `${ ASSISTANT_API . CONVERSATION } /{conversationId}` ,
46
- req : httpServerMock . createRawRequest ( { params, payload } ) ,
71
+ req : httpServerMock . createRawRequest ( {
72
+ params,
73
+ payload,
74
+ ...( dataSourceId
75
+ ? {
76
+ query : {
77
+ dataSourceId,
78
+ } ,
79
+ }
80
+ : { } ) ,
81
+ } ) ,
47
82
} ) ;
48
- const triggerGetTrace = ( interactionId : string ) =>
83
+ const triggerGetTrace = ( interactionId : string , dataSourceId ?: string ) =>
49
84
triggerHandler ( router , {
50
85
method : 'get' ,
51
86
path : `${ ASSISTANT_API . TRACE } /{interactionId}` ,
52
- req : httpServerMock . createRawRequest ( { params : { interactionId } } ) ,
87
+ req : httpServerMock . createRawRequest ( {
88
+ params : { interactionId } ,
89
+ ...( dataSourceId
90
+ ? {
91
+ query : {
92
+ dataSourceId,
93
+ } ,
94
+ }
95
+ : { } ) ,
96
+ } ) ,
53
97
} ) ;
54
- const triggerAbortAgentExecution = ( conversationId : string ) =>
98
+ const triggerAbortAgentExecution = ( conversationId : string , dataSourceId ?: string ) =>
55
99
triggerHandler ( router , {
56
100
method : 'post' ,
57
101
path : ASSISTANT_API . ABORT_AGENT_EXECUTION ,
58
- req : httpServerMock . createRawRequest ( { payload : { conversationId } } ) ,
102
+ req : httpServerMock . createRawRequest ( {
103
+ payload : { conversationId } ,
104
+ ...( dataSourceId
105
+ ? {
106
+ query : {
107
+ dataSourceId,
108
+ } ,
109
+ }
110
+ : { } ) ,
111
+ } ) ,
59
112
} ) ;
60
- const triggerFeedback = ( params : { interactionId : string } , payload : { satisfaction : boolean } ) =>
113
+ const triggerFeedback = (
114
+ params : { interactionId : string } ,
115
+ payload : { satisfaction : boolean } ,
116
+ dataSourceId ?: string
117
+ ) =>
61
118
triggerHandler ( router , {
62
119
method : 'put' ,
63
120
path : `${ ASSISTANT_API . FEEDBACK } /{interactionId}` ,
64
- req : httpServerMock . createRawRequest ( { params, payload } ) ,
121
+ req : httpServerMock . createRawRequest ( {
122
+ params,
123
+ payload,
124
+ ...( dataSourceId
125
+ ? {
126
+ query : {
127
+ dataSourceId,
128
+ } ,
129
+ }
130
+ : { } ) ,
131
+ } ) ,
65
132
} ) ;
66
133
67
134
describe ( 'chat routes' , ( ) => {
@@ -80,6 +147,22 @@ describe('chat routes', () => {
80
147
81
148
expect ( mockAgentFrameworkStorageService . deleteConversation ) . not . toHaveBeenCalled ( ) ;
82
149
const result = ( await triggerDeleteConversation ( 'foo' ) ) as ResponseObject ;
150
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'client' ) ;
151
+ expect ( mockAgentFrameworkStorageService . deleteConversation ) . toHaveBeenCalledWith ( 'foo' ) ;
152
+ expect ( result . source ) . toMatchInlineSnapshot ( `
153
+ Object {
154
+ "success": true,
155
+ }
156
+ ` ) ;
157
+ } ) ;
158
+
159
+ it ( 'should call delete conversation with passed data source id and get data source transport' , async ( ) => {
160
+ mockAgentFrameworkStorageService . deleteConversation . mockResolvedValueOnce ( {
161
+ success : true ,
162
+ } ) ;
163
+ expect ( mockAgentFrameworkStorageService . deleteConversation ) . not . toHaveBeenCalled ( ) ;
164
+ const result = ( await triggerDeleteConversation ( 'foo' , 'data_source_id' ) ) as ResponseObject ;
165
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'dataSource-client' ) ;
83
166
expect ( mockAgentFrameworkStorageService . deleteConversation ) . toHaveBeenCalledWith ( 'foo' ) ;
84
167
expect ( result . source ) . toMatchInlineSnapshot ( `
85
168
Object {
@@ -109,6 +192,7 @@ describe('chat routes', () => {
109
192
{ conversationId : 'foo' } ,
110
193
{ title : 'new-title' }
111
194
) ) as ResponseObject ;
195
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'client' ) ;
112
196
expect ( mockAgentFrameworkStorageService . updateConversation ) . toHaveBeenCalledWith (
113
197
'foo' ,
114
198
'new-title'
@@ -120,6 +204,29 @@ describe('chat routes', () => {
120
204
` ) ;
121
205
} ) ;
122
206
207
+ it ( 'should call update conversation with passed data source id and title then get data source transport' , async ( ) => {
208
+ mockAgentFrameworkStorageService . updateConversation . mockResolvedValueOnce ( {
209
+ success : true ,
210
+ } ) ;
211
+
212
+ expect ( mockAgentFrameworkStorageService . updateConversation ) . not . toHaveBeenCalled ( ) ;
213
+ const result = ( await triggerUpdateConversation (
214
+ { conversationId : 'foo' } ,
215
+ { title : 'new-title' } ,
216
+ 'data_source_id'
217
+ ) ) as ResponseObject ;
218
+ expect ( mockAgentFrameworkStorageService . updateConversation ) . toHaveBeenCalledWith (
219
+ 'foo' ,
220
+ 'new-title'
221
+ ) ;
222
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'dataSource-client' ) ;
223
+ expect ( result . source ) . toMatchInlineSnapshot ( `
224
+ Object {
225
+ "success": true,
226
+ }
227
+ ` ) ;
228
+ } ) ;
229
+
123
230
it ( 'should log error and return 500 error when failed to update conversation' , async ( ) => {
124
231
mockAgentFrameworkStorageService . updateConversation . mockRejectedValueOnce ( new Error ( ) ) ;
125
232
@@ -149,6 +256,27 @@ describe('chat routes', () => {
149
256
150
257
expect ( mockAgentFrameworkStorageService . getTraces ) . not . toHaveBeenCalled ( ) ;
151
258
const result = ( await triggerGetTrace ( 'interaction-1' ) ) as ResponseObject ;
259
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'client' ) ;
260
+ expect ( mockAgentFrameworkStorageService . getTraces ) . toHaveBeenCalledWith ( 'interaction-1' ) ;
261
+ expect ( result . source ) . toEqual ( getTraceResultMock ) ;
262
+ } ) ;
263
+
264
+ it ( 'should call get traces with passed data source id and get data source transport' , async ( ) => {
265
+ const getTraceResultMock = [
266
+ {
267
+ interactionId : 'interaction-1' ,
268
+ createTime : '' ,
269
+ input : 'foo' ,
270
+ output : 'bar' ,
271
+ origin : '' ,
272
+ traceNumber : 0 ,
273
+ } ,
274
+ ] ;
275
+ mockAgentFrameworkStorageService . getTraces . mockResolvedValueOnce ( getTraceResultMock ) ;
276
+
277
+ expect ( mockAgentFrameworkStorageService . getTraces ) . not . toHaveBeenCalled ( ) ;
278
+ const result = ( await triggerGetTrace ( 'interaction-1' , 'data_source_id' ) ) as ResponseObject ;
279
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'dataSource-client' ) ;
152
280
expect ( mockAgentFrameworkStorageService . getTraces ) . toHaveBeenCalledWith ( 'interaction-1' ) ;
153
281
expect ( result . source ) . toEqual ( getTraceResultMock ) ;
154
282
} ) ;
@@ -169,6 +297,16 @@ describe('chat routes', () => {
169
297
170
298
await triggerAbortAgentExecution ( 'foo' ) ;
171
299
expect ( mockOllyChatService . abortAgentExecution ) . toHaveBeenCalledWith ( 'foo' ) ;
300
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'client' ) ;
301
+ expect ( mockedLogger . info ) . toHaveBeenCalledWith ( 'Abort agent execution: foo' ) ;
302
+ } ) ;
303
+
304
+ it ( 'should call get abort agent with passed data source id and get data source transport ' , async ( ) => {
305
+ expect ( mockOllyChatService . abortAgentExecution ) . not . toHaveBeenCalled ( ) ;
306
+
307
+ await triggerAbortAgentExecution ( 'foo' , 'data_source_id' ) ;
308
+ expect ( mockOllyChatService . abortAgentExecution ) . toHaveBeenCalledWith ( 'foo' ) ;
309
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'dataSource-client' ) ;
172
310
expect ( mockedLogger . info ) . toHaveBeenCalledWith ( 'Abort agent execution: foo' ) ;
173
311
} ) ;
174
312
@@ -200,6 +338,31 @@ describe('chat routes', () => {
200
338
{ interactionId : 'foo' } ,
201
339
{ satisfaction : true }
202
340
) ) as ResponseObject ;
341
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'client' ) ;
342
+ expect ( mockAgentFrameworkStorageService . updateInteraction ) . toHaveBeenCalledWith ( 'foo' , {
343
+ feedback : {
344
+ satisfaction : true ,
345
+ } ,
346
+ } ) ;
347
+ expect ( result . source ) . toMatchInlineSnapshot ( `
348
+ Object {
349
+ "success": true,
350
+ }
351
+ ` ) ;
352
+ } ) ;
353
+
354
+ it ( 'should call update interaction with passed data source id and get data source transport' , async ( ) => {
355
+ mockAgentFrameworkStorageService . updateConversation . mockResolvedValueOnce ( {
356
+ success : true ,
357
+ } ) ;
358
+
359
+ expect ( mockAgentFrameworkStorageService . updateConversation ) . not . toHaveBeenCalled ( ) ;
360
+ const result = ( await triggerFeedback (
361
+ { interactionId : 'foo' } ,
362
+ { satisfaction : true } ,
363
+ 'data_source_id'
364
+ ) ) as ResponseObject ;
365
+ expect ( getOpenSearchClientTransport . mock . results [ 0 ] . value ) . toBe ( 'dataSource-client' ) ;
203
366
expect ( mockAgentFrameworkStorageService . updateInteraction ) . toHaveBeenCalledWith ( 'foo' , {
204
367
feedback : {
205
368
satisfaction : true ,
0 commit comments