From 820e34e41d54b587ce1591ffced7646a5386648b Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Wed, 23 Oct 2024 22:56:37 -0600 Subject: [PATCH] [8.x] [Security assistant] Conversation pagination patch MIN (#197305) (#197558) # Backport This will backport the following commits from `main` to `8.x`: - [[Security assistant] Conversation pagination patch MIN (#197305)](https://github.com/elastic/kibana/pull/197305) ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport) --- ..._fetch_current_user_conversations.test.tsx | 2 +- .../use_fetch_current_user_conversations.ts | 8 +- .../impl/assistant/helpers.test.ts | 2 +- .../scripts/create_conversations.js | 9 + .../scripts/create_conversations_script.ts | 165 ++++++++++++++++++ .../__mocks__/conversations_schema.mock.ts | 5 +- .../server/ai_assistant_data_clients/find.ts | 90 ++++++++-- .../server/ai_assistant_data_clients/index.ts | 6 + .../routes/user_conversations/create_route.ts | 17 -- .../routes/user_conversations/find_route.ts | 15 +- .../e2e/ai_assistant/conversations.cy.ts | 12 +- 11 files changed, 278 insertions(+), 53 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/scripts/create_conversations.js create mode 100644 x-pack/plugins/elastic_assistant/scripts/create_conversations_script.ts diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.test.tsx index 652764212e996..f10c7d07a35d6 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.test.tsx @@ -52,7 +52,7 @@ describe('useFetchCurrentUserConversations', () => { method: 'GET', query: { page: 1, - perPage: 100, + per_page: 99, }, version: '2023-10-31', signal: undefined, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.ts index 68612e3e22397..9006ca387e086 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/conversations/use_fetch_current_user_conversations.ts @@ -15,7 +15,7 @@ import { Conversation } from '../../../assistant_context/types'; export interface FetchConversationsResponse { page: number; - perPage: number; + per_page: number; total: number; data: Conversation[]; } @@ -40,13 +40,13 @@ export interface UseFetchCurrentUserConversationsParams { */ const query = { page: 1, - perPage: 100, + per_page: 99, }; export const CONVERSATIONS_QUERY_KEYS = [ ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_FIND, query.page, - query.perPage, + query.per_page, API_VERSIONS.public.v1, ]; @@ -69,7 +69,7 @@ export const useFetchCurrentUserConversations = ({ { select: (data) => onFetch(data), keepPreviousData: true, - initialData: { page: 1, perPage: 100, total: 0, data: [] }, + initialData: { ...query, total: 0, data: [] }, refetchOnWindowFocus, enabled: isAssistantEnabled, } diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.test.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.test.ts index 5559e273f06b5..26609dea82164 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.test.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/helpers.test.ts @@ -110,7 +110,7 @@ describe('helpers', () => { }; const conversationsData = { page: 1, - perPage: 10, + per_page: 10, total: 2, data: Object.values(baseConversations).map((c) => c), }; diff --git a/x-pack/plugins/elastic_assistant/scripts/create_conversations.js b/x-pack/plugins/elastic_assistant/scripts/create_conversations.js new file mode 100644 index 0000000000000..b08f1419af0c9 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/scripts/create_conversations.js @@ -0,0 +1,9 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +require('../../../../src/setup_node_env'); +require('./create_conversations_script').create(); diff --git a/x-pack/plugins/elastic_assistant/scripts/create_conversations_script.ts b/x-pack/plugins/elastic_assistant/scripts/create_conversations_script.ts new file mode 100644 index 0000000000000..2fd388e299f6f --- /dev/null +++ b/x-pack/plugins/elastic_assistant/scripts/create_conversations_script.ts @@ -0,0 +1,165 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { randomBytes } from 'node:crypto'; +import yargs from 'yargs/yargs'; +import { ToolingLog } from '@kbn/tooling-log'; +import axios from 'axios'; +import { + API_VERSIONS, + ConversationCreateProps, + ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL, +} from '@kbn/elastic-assistant-common'; +import pLimit from 'p-limit'; +import { getCreateConversationSchemaMock } from '../server/__mocks__/conversations_schema.mock'; + +/** + * Developer script for creating conversations. + * node x-pack/plugins/elastic_assistant/scripts/create_conversations + */ +export const create = async () => { + const logger = new ToolingLog({ + level: 'info', + writeTo: process.stdout, + }); + const argv = yargs(process.argv.slice(2)) + .option('count', { + type: 'number', + description: 'Number of conversations to create', + default: 100, + }) + .option('kibana', { + type: 'string', + description: 'Kibana url including auth', + default: `http://elastic:changeme@localhost:5601`, + }) + .parse(); + const kibanaUrl = removeTrailingSlash(argv.kibana); + const count = Number(argv.count); + logger.info(`Kibana URL: ${kibanaUrl}`); + const connectorsApiUrl = `${kibanaUrl}/api/actions/connectors`; + const conversationsCreateUrl = `${kibanaUrl}${ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL}`; + + try { + logger.info(`Fetching available connectors...`); + const { data: connectors } = await axios.get(connectorsApiUrl, { + headers: requestHeaders, + }); + const aiConnectors = connectors.filter( + ({ connector_type_id: connectorTypeId }: { connector_type_id: string }) => + AllowedActionTypeIds.includes(connectorTypeId) + ); + if (aiConnectors.length === 0) { + throw new Error('No AI connectors found, create an AI connector to use this script'); + } + + logger.info(`Creating ${count} conversations...`); + if (count > 999) { + logger.info(`This may take a couple of minutes...`); + } + + const promises = Array.from({ length: count }, (_, i) => + limit(() => + retryRequest( + () => + axios.post( + conversationsCreateUrl, + getCreateConversationSchemaMock({ + ...getMockConversationContent(), + apiConfig: { + actionTypeId: aiConnectors[0].connector_type_id, + connectorId: aiConnectors[0].id, + }, + }), + { headers: requestHeaders } + ), + 3, // Retry up to 3 times + 1000 // Delay of 1 second between retries + ) + ) + ); + + const results = await Promise.allSettled(promises); + + const successfulResults = results.filter((result) => result.status === 'fulfilled'); + const errorResults = results.filter( + (result) => result.status === 'rejected' + ) as PromiseRejectedResult[]; + const conversationsCreated = successfulResults.length; + + if (count > conversationsCreated) { + const errorExample = + errorResults.length > 0 ? errorResults[0]?.reason?.message ?? 'unknown' : 'unknown'; + throw new Error( + `Failed to create all conversations. Expected count: ${count}, Created count: ${conversationsCreated}. Reason: ${errorExample}` + ); + } + logger.info(`Successfully created ${successfulResults.length} conversations.`); + } catch (e) { + logger.error(e); + } +}; +// Set the concurrency limit (e.g., 50 requests at a time) +const limit = pLimit(50); + +// Retry helper function +const retryRequest = async ( + fn: () => Promise, + retries: number = 3, + delay: number = 1000 +): Promise => { + try { + return await fn(); + } catch (e) { + if (retries > 0) { + await new Promise((res) => setTimeout(res, delay)); + return retryRequest(fn, retries - 1, delay); + } + throw e; // If retries are exhausted, throw the error + } +}; + +const getMockConversationContent = (): Partial => ({ + title: `A ${randomBytes(4).toString('hex')} title`, + isDefault: false, + messages: [ + { + content: 'Hello robot', + role: 'user', + timestamp: '2019-12-13T16:40:33.400Z', + traceData: { + traceId: '1', + transactionId: '2', + }, + }, + { + content: 'Hello human', + role: 'assistant', + timestamp: '2019-12-13T16:41:33.400Z', + traceData: { + traceId: '3', + transactionId: '4', + }, + }, + ], +}); + +export const AllowedActionTypeIds = ['.bedrock', '.gen-ai', '.gemini']; + +const requestHeaders = { + 'kbn-xsrf': 'xxx', + 'Content-Type': 'application/json', + 'elastic-api-version': API_VERSIONS.public.v1, +}; + +function removeTrailingSlash(url: string) { + if (url.endsWith('/')) { + return url.slice(0, -1); + } else { + return url; + } +} diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/conversations_schema.mock.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/conversations_schema.mock.ts index 7594839bd21b4..278dfc9fe829b 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/conversations_schema.mock.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/conversations_schema.mock.ts @@ -60,7 +60,9 @@ export const getConversationSearchEsMock = () => { return searchResponse; }; -export const getCreateConversationSchemaMock = (): ConversationCreateProps => ({ +export const getCreateConversationSchemaMock = ( + rest?: Partial +): ConversationCreateProps => ({ title: 'Welcome', apiConfig: { actionTypeId: '.gen-ai', @@ -82,6 +84,7 @@ export const getCreateConversationSchemaMock = (): ConversationCreateProps => ({ }, ], category: 'assistant', + ...rest, }); export const getUpdateConversationSchemaMock = ( diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/find.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/find.ts index 101354a6802b7..9f37e45250a9c 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/find.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/find.ts @@ -9,6 +9,7 @@ import { AggregationsAggregationContainer, MappingRuntimeFields, Sort, + SearchResponse, } from '@elastic/elasticsearch/lib/api/types'; import { ElasticsearchClient, Logger } from '@kbn/core/server'; @@ -27,6 +28,10 @@ interface FindOptions { runtimeMappings?: MappingRuntimeFields | undefined; logger: Logger; aggs?: Record; + mSearch?: { + filter: string; + perPage: number; + }; } export interface FindResponse { @@ -47,6 +52,7 @@ export const findDocuments = async ({ sortOrder, logger, aggs, + mSearch, }: FindOptions): Promise> => { const query = getQueryFilter({ filter }); let sort: Sort | undefined; @@ -61,28 +67,78 @@ export const findDocuments = async ({ }; } try { - const response = await esClient.search({ - body: { - query, - track_total_hits: true, - sort, - }, - _source: true, - from: (page - 1) * perPage, + if (mSearch == null) { + const response = await esClient.search({ + body: { + query, + track_total_hits: true, + sort, + }, + _source: true, + from: (page - 1) * perPage, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + size: perPage, + aggs, + }); + + return { + data: response, + page, + perPage, + total: + (typeof response.hits.total === 'number' + ? response.hits.total + : response.hits.total?.value) ?? 0, + }; + } + const mSearchQueryBody = { + body: [ + { index }, + { + query, + size: perPage, + aggs, + seq_no_primary_term: true, + from: (page - 1) * perPage, + sort, + _source: true, + }, + { index }, + { + query: getQueryFilter({ filter: mSearch.filter }), + size: mSearch.perPage, + aggs, + seq_no_primary_term: true, + from: (page - 1) * mSearch.perPage, + sort, + _source: true, + }, + ], ignore_unavailable: true, index, - seq_no_primary_term: true, - size: perPage, - aggs, + }; + const response = await esClient.msearch>(mSearchQueryBody); + let responseStats: Omit, 'hits'> = { + took: 0, + _shards: { total: 0, successful: 0, skipped: 0, failed: 0 }, + timed_out: false, + }; + // flatten the results of the combined find queries into a single array of hits: + const results = response.responses.flatMap((res) => { + const mResponse = res as SearchResponse; + const { hits, ...responseBody } = mResponse; + // assign whatever the last stats are, they are only used for type + responseStats = { ...responseStats, ...responseBody }; + return hits?.hits ?? []; }); + return { - data: response, + data: { ...responseStats, hits: { hits: results } }, page, - perPage, - total: - (typeof response.hits.total === 'number' - ? response.hits.total // This format is to be removed in 8.0 - : response.hits.total?.value) ?? 0, + perPage: perPage + mSearch.perPage, + total: results.length, }; } catch (err) { logger.error(`Error fetching documents: ${err}`); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.ts index cc74e1f03d3d9..706e4444488d9 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.ts @@ -100,6 +100,7 @@ export class AIAssistantDataClient { filter, fields, aggs, + mSearch, }: { perPage: number; page: number; @@ -108,6 +109,10 @@ export class AIAssistantDataClient { filter?: string; fields?: string[]; aggs?: Record; + mSearch?: { + filter: string; + perPage: number; + }; }): Promise>> => { const esClient = await this.options.elasticsearchClientPromise; return findDocuments({ @@ -121,6 +126,7 @@ export class AIAssistantDataClient { sortOrder: sortOrder as estypes.SortOrder, logger: this.options.logger, aggs, + mSearch, }); }; } diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts index f997121f95aa8..ea056f81e7c6c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts @@ -54,23 +54,6 @@ export const createConversationRoute = (router: ElasticAssistantPluginRouter): v } const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); - const currentUser = ctx.elasticAssistant.getCurrentUser(); - const userFilter = currentUser?.username - ? `name: "${currentUser?.username}"` - : `id: "${currentUser?.profile_uid}"`; - const result = await dataClient?.findDocuments({ - perPage: 100, - page: 1, - filter: `users:{ ${userFilter} } AND title:${request.body.title}`, - fields: ['title'], - }); - if (result?.data != null && result.total > 0) { - return assistantResponse.error({ - statusCode: 409, - body: `conversation title: "${request.body.title}" already exists`, - }); - } - const createdConversation = await dataClient?.createConversation({ conversation: request.body, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts index 6a2c3afc41374..e7ce80039beb0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts @@ -61,13 +61,24 @@ export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) const userFilter = currentUser?.username ? `name: "${currentUser?.username}"` : `id: "${currentUser?.profile_uid}"`; + + const MAX_CONVERSATION_TOTAL = query.per_page; + // TODO remove once we have pagination https://github.com/elastic/kibana/issues/192714 + // do a separate search for default conversations and non-default conversations to ensure defaults always get included + // MUST MATCH THE LENGTH OF BASE_SECURITY_CONVERSATIONS from 'x-pack/plugins/security_solution/public/assistant/content/conversations/index.tsx' + const MAX_DEFAULT_CONVERSATION_TOTAL = 7; + const nonDefaultSize = MAX_CONVERSATION_TOTAL - MAX_DEFAULT_CONVERSATION_TOTAL; const result = await dataClient?.findDocuments({ - perPage: query.per_page, + perPage: nonDefaultSize, page: query.page, sortField: query.sort_field, sortOrder: query.sort_order, - filter: `users:{ ${userFilter} }${additionalFilter}`, + filter: `users:{ ${userFilter} }${additionalFilter} and not is_default: true`, fields: query.fields, + mSearch: { + filter: `users:{ ${userFilter} }${additionalFilter} and is_default: true`, + perPage: MAX_DEFAULT_CONVERSATION_TOTAL, + }, }); if (result) { diff --git a/x-pack/test/security_solution_cypress/cypress/e2e/ai_assistant/conversations.cy.ts b/x-pack/test/security_solution_cypress/cypress/e2e/ai_assistant/conversations.cy.ts index c91ee7de475e3..4d87cce1fdaa8 100644 --- a/x-pack/test/security_solution_cypress/cypress/e2e/ai_assistant/conversations.cy.ts +++ b/x-pack/test/security_solution_cypress/cypress/e2e/ai_assistant/conversations.cy.ts @@ -23,7 +23,6 @@ import { typeAndSendMessage, assertErrorResponse, selectRule, - assertErrorToastShown, updateConversationTitle, } from '../../tasks/assistant'; import { deleteConversations } from '../../tasks/api_calls/assistant'; @@ -146,7 +145,7 @@ describe('AI Assistant Conversations', { tags: ['@ess', '@serverless'] }, () => assertConnectorSelected(bedrockConnectorAPIPayload.name); assertMessageSent('goodbye'); }); - it('Correctly titles new conversations, and only allows one conversation called "New chat" at a time', () => { + it('Correctly creates and titles new conversations, and allows title updates', () => { visitGetStartedPage(); openAssistant(); createNewChat(); @@ -155,14 +154,7 @@ describe('AI Assistant Conversations', { tags: ['@ess', '@serverless'] }, () => typeAndSendMessage('hello'); assertMessageSent('hello'); assertConversationTitle('Unexpected API Error: - Connection error.'); - updateConversationTitle('New chat'); - selectConversation('Welcome'); - createNewChat(); - assertErrorToastShown('Error creating conversation with title New chat'); - selectConversation('New chat'); - updateConversationTitle('My other chat'); - createNewChat(); - assertNewConversation(false, 'New chat'); + updateConversationTitle('Something else'); }); }); });