Skip to content

Commit f0d4663

Browse files
stephmilovicstephmilovic
committed
[Security Assistant] Fixes conversation message not appended/title not updating (#233219)
(cherry picked from commit b81bbd0)
1 parent c35b3db commit f0d4663

File tree

11 files changed

+674
-218
lines changed

11 files changed

+674
-218
lines changed

‎x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,4 +145,34 @@ describe('use chat send', () => {
145145
});
146146
});
147147
});
148+
it('retries getConversation up to 5 times if title is empty, and stops when title is found', async () => {
149+
const promptText = 'test prompt';
150+
const getConversationMock = jest.fn();
151+
// First 3 calls return empty title, 4th returns non-empty
152+
getConversationMock
153+
.mockResolvedValueOnce({ title: '' })
154+
.mockResolvedValueOnce({ title: '' })
155+
.mockResolvedValueOnce({ title: '' })
156+
.mockResolvedValueOnce({ title: 'Final Title' });
157+
(useConversation as jest.Mock).mockReturnValue({
158+
removeLastMessage,
159+
clearConversation,
160+
getConversation: getConversationMock,
161+
createConversation: jest.fn(),
162+
});
163+
const { result } = renderHook(
164+
() =>
165+
useChatSend({
166+
...testProps,
167+
currentConversation: { ...emptyWelcomeConvo, id: 'convo-id', title: '' },
168+
}),
169+
{ wrapper: TestProviders }
170+
);
171+
await act(async () => {
172+
await result.current.handleChatSend(promptText);
173+
});
174+
// Should call getConversation 4 times (until non-empty title)
175+
expect(getConversationMock).toHaveBeenCalledTimes(4);
176+
expect(getConversationMock).toHaveBeenLastCalledWith('convo-id');
177+
});
148178
});

‎x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx‎

Lines changed: 74 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -61,12 +61,15 @@ export const useChatSend = ({
6161
const { setLastConversation } = useAssistantLastConversation({ spaceId });
6262
const [userPrompt, setUserPrompt] = useState<string | null>(null);
6363

64-
const { isLoading, sendMessage, abortStream } = useSendMessage();
64+
const { sendMessage, abortStream } = useSendMessage();
6565
const { clearConversation, createConversation, getConversation, removeLastMessage } =
6666
useConversation();
6767
const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
6868
const isSetupComplete = kbStatus?.elser_exists && kbStatus?.security_labs_exists;
6969

70+
// Local loading state that persists until the entire message flow is complete
71+
const [isLoadingChatSend, setIsLoadingChatSend] = useState(false);
72+
7073
// Handles sending latest user prompt to API
7174
const handleSendMessage = useCallback(
7275
async (promptText: string) => {
@@ -99,61 +102,80 @@ export const useChatSend = ({
99102
selectedPromptContexts,
100103
});
101104

102-
const baseReplacements: Replacements = userMessage.replacements ?? convo.replacements;
103-
104-
const selectedPromptContextsReplacements = Object.values(
105-
selectedPromptContexts
106-
).reduce<Replacements>((acc, context) => ({ ...acc, ...context.replacements }), {});
107-
108-
const replacements: Replacements = {
109-
...baseReplacements,
110-
...selectedPromptContextsReplacements,
111-
};
112-
const updatedMessages = [...convo.messages, userMessage].map((m) => ({
113-
...m,
114-
content: m.content ?? '',
115-
}));
116-
setCurrentConversation({
117-
...convo,
118-
replacements,
119-
messages: updatedMessages,
120-
});
105+
const baseReplacements: Replacements = userMessage.replacements ?? convo.replacements;
121106

122-
// Reset prompt context selection and preview before sending:
123-
setSelectedPromptContexts({});
107+
const selectedPromptContextsReplacements = Object.values(
108+
selectedPromptContexts
109+
).reduce<Replacements>((acc, context) => ({ ...acc, ...context.replacements }), {});
124110

125-
const rawResponse = await sendMessage({
126-
apiConfig,
127-
http,
128-
message: userMessage.content ?? '',
129-
conversationId: convo.id,
130-
replacements,
131-
});
111+
const replacements: Replacements = {
112+
...baseReplacements,
113+
...selectedPromptContextsReplacements,
114+
};
115+
const updatedMessages = [...convo.messages, userMessage].map((m) => ({
116+
...m,
117+
content: m.content ?? '',
118+
}));
119+
setCurrentConversation({
120+
...convo,
121+
replacements,
122+
messages: updatedMessages,
123+
});
132124

133-
assistantTelemetry?.reportAssistantMessageSent({
134-
role: userMessage.role,
135-
actionTypeId: apiConfig.actionTypeId,
136-
model: apiConfig.model,
137-
provider: apiConfig.provider,
138-
isEnabledKnowledgeBase: isSetupComplete ?? false,
139-
});
125+
// Reset prompt context selection and preview before sending:
126+
setSelectedPromptContexts({});
127+
128+
const rawResponse = await sendMessage({
129+
apiConfig,
130+
http,
131+
message: userMessage.content ?? '',
132+
conversationId: convo.id,
133+
replacements,
134+
});
135+
136+
assistantTelemetry?.reportAssistantMessageSent({
137+
role: userMessage.role,
138+
actionTypeId: apiConfig.actionTypeId,
139+
model: apiConfig.model,
140+
provider: apiConfig.provider,
141+
isEnabledKnowledgeBase: isSetupComplete ?? false,
142+
});
140143

141-
const responseMessage: ClientMessage = getMessageFromRawResponse(rawResponse);
142-
if (convo.title === '') {
143-
convo.title = (await getConversation(convo.id))?.title ?? '';
144+
const responseMessage: ClientMessage = getMessageFromRawResponse(rawResponse);
145+
if (convo.title === '') {
146+
// Retry getConversation up to 5 times if title is empty
147+
let retryCount = 0;
148+
const maxRetries = 5;
149+
while (retryCount < maxRetries) {
150+
const conversation = await getConversation(convo.id);
151+
convo.title = conversation?.title ?? '';
152+
153+
if (convo.title !== '') {
154+
break; // Title found, exit retry loop
155+
}
156+
157+
retryCount++;
158+
if (retryCount < maxRetries) {
159+
// Wait 1 second before next retry
160+
await new Promise((resolve) => setTimeout(resolve, 1000));
161+
}
162+
}
163+
}
164+
setCurrentConversation({
165+
...convo,
166+
replacements,
167+
messages: [...updatedMessages, responseMessage],
168+
});
169+
assistantTelemetry?.reportAssistantMessageSent({
170+
role: responseMessage.role,
171+
actionTypeId: apiConfig.actionTypeId,
172+
model: apiConfig.model,
173+
provider: apiConfig.provider,
174+
isEnabledKnowledgeBase: isSetupComplete ?? false,
175+
});
176+
} finally {
177+
setIsLoadingChatSend(false);
144178
}
145-
setCurrentConversation({
146-
...convo,
147-
replacements,
148-
messages: [...updatedMessages, responseMessage],
149-
});
150-
assistantTelemetry?.reportAssistantMessageSent({
151-
role: responseMessage.role,
152-
actionTypeId: apiConfig.actionTypeId,
153-
model: apiConfig.model,
154-
provider: apiConfig.provider,
155-
isEnabledKnowledgeBase: isSetupComplete ?? false,
156-
});
157179
},
158180
[
159181
assistantTelemetry,
@@ -241,7 +263,7 @@ export const useChatSend = ({
241263
handleChatSend,
242264
abortStream,
243265
handleRegenerateResponse,
244-
isLoading,
266+
isLoading: isLoadingChatSend,
245267
userPrompt,
246268
setUserPrompt,
247269
};

‎x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/conversations_schema.mock.ts‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,3 +244,36 @@ export const getEsCreateConversationSchemaMock = (
244244
namespace: 'default',
245245
...rest,
246246
});
247+
248+
export const getEsConversationSchemaMock = (
249+
rest?: Partial<EsConversationSchema>
250+
): EsConversationSchema => ({
251+
'@timestamp': '2020-04-20T15:25:31.830Z',
252+
created_at: '2020-04-20T15:25:31.830Z',
253+
title: 'title-1',
254+
updated_at: '2020-04-20T15:25:31.830Z',
255+
messages: [],
256+
id: '1',
257+
namespace: 'default',
258+
exclude_from_last_conversation_storage: false,
259+
api_config: {
260+
action_type_id: '.gen-ai',
261+
connector_id: 'c1',
262+
default_system_prompt_id: 'prompt-1',
263+
model: 'test',
264+
provider: 'Azure OpenAI',
265+
},
266+
category: 'assistant',
267+
users: [
268+
{
269+
id: '1111',
270+
name: 'elastic',
271+
},
272+
],
273+
created_by: {
274+
id: '1111',
275+
name: 'elastic',
276+
},
277+
replacements: undefined,
278+
...rest,
279+
});

0 commit comments

Comments
 (0)