From 2e8b823b1e84f953e854762edc28b3fedf9ac532 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 1 May 2024 10:14:34 -0600 Subject: [PATCH 01/20] Adds feature flag and cleaned up 'feature capabilities' touch points --- .../impl/capabilities/index.ts | 1 + .../capabilities/get_capabilities_route.gen.ts | 1 + .../get_capabilities_route.schema.yaml | 3 +++ .../use_fetch_anonymization_fields.test.tsx | 5 ++--- .../capabilities/__mocks__/use_capabilities.tsx | 17 +++++++++++++++++ .../api/capabilities/use_capabilities.test.tsx | 10 ++-------- ...se_fetch_current_user_conversations.test.tsx | 5 ++--- .../settings/assistant_settings.test.tsx | 2 +- .../assistant/settings/assistant_settings.tsx | 2 +- .../settings/assistant_settings_management.tsx | 2 +- .../impl/assistant_context/index.tsx | 14 +++++--------- .../routes/evaluate/post_evaluate.test.ts | 5 +++-- .../server/services/app_context.test.ts | 10 +++++----- .../common/experimental_features.ts | 5 +++++ .../plugins/security_solution/server/plugin.ts | 1 + 15 files changed, 50 insertions(+), 33 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/__mocks__/use_capabilities.tsx diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index 86c3b26f200e6..409e1f1836b20 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -14,6 +14,7 @@ export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]: * Default features available to the elastic assistant */ export const defaultAssistantFeatures = Object.freeze({ + assistantKnowledgeBaseByDefault: false, assistantModelEvaluation: false, attackDiscoveryEnabled: false, }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts index 6e90724097fc8..52c249753afcd 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts @@ -18,6 +18,7 @@ import { z } from 'zod'; export type GetCapabilitiesResponse = z.infer; export const GetCapabilitiesResponse = z.object({ + assistantKnowledgeBaseByDefault: z.boolean(), assistantModelEvaluation: z.boolean(), attackDiscoveryEnabled: z.boolean(), }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml index 8d28a6b95d55a..6ddc4f6f81408 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml @@ -19,11 +19,14 @@ paths: schema: type: object properties: + assistantKnowledgeBaseByDefault: + type: boolean assistantModelEvaluation: type: boolean attackDiscoveryEnabled: type: boolean required: + - assistantKnowledgeBaseByDefault - assistantModelEvaluation - attackDiscoveryEnabled '400': diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/anonymization_fields/use_fetch_anonymization_fields.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/anonymization_fields/use_fetch_anonymization_fields.test.tsx index aabdcb6909ccb..eefa9f3593f61 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/anonymization_fields/use_fetch_anonymization_fields.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/anonymization_fields/use_fetch_anonymization_fields.test.tsx @@ -13,11 +13,10 @@ import React from 'react'; import { useFetchAnonymizationFields } from './use_fetch_anonymization_fields'; import { HttpSetup } from '@kbn/core-http-browser'; import { useAssistantContext } from '../../../assistant_context'; - -const statusResponse = { assistantModelEvaluation: true, assistantStreamingEnabled: false }; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; const http = { - fetch: jest.fn().mockResolvedValue(statusResponse), + fetch: jest.fn().mockResolvedValue(defaultAssistantFeatures), } as unknown as HttpSetup; jest.mock('../../../assistant_context'); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/__mocks__/use_capabilities.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/__mocks__/use_capabilities.tsx new file mode 100644 index 0000000000000..d6dd239d4b6b8 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/__mocks__/use_capabilities.tsx @@ -0,0 +1,17 @@ +/* + * 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 { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; + +export const useCapabilities = jest.fn().mockReturnValue({ + isLoading: false, + isError: false, + data: defaultAssistantFeatures, + error: null, + isFetching: false, + isSuccess: true, +}); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx index df1f399dadcc3..6101782ae43b1 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx @@ -11,16 +11,10 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import type { ReactNode } from 'react'; import React from 'react'; import { useCapabilities, UseCapabilitiesParams } from './use_capabilities'; -import { API_VERSIONS } from '@kbn/elastic-assistant-common'; - -const statusResponse = { - assistantModelEvaluation: true, - assistantStreamingEnabled: false, - attackDiscoveryEnabled: false, -}; +import { API_VERSIONS, defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; const http = { - get: jest.fn().mockResolvedValue(statusResponse), + get: jest.fn().mockResolvedValue(defaultAssistantFeatures), }; const toasts = { addError: jest.fn(), 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 86b39f5ea43be..22023179ad5de 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 @@ -14,11 +14,10 @@ import { UseFetchCurrentUserConversationsParams, useFetchCurrentUserConversations, } from './use_fetch_current_user_conversations'; - -const statusResponse = { assistantModelEvaluation: true, assistantStreamingEnabled: false }; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; const http = { - fetch: jest.fn().mockResolvedValue(statusResponse), + fetch: jest.fn().mockResolvedValue(defaultAssistantFeatures), }; const onFetch = jest.fn(); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx index 4f745f51a11b7..8278cb1559535 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.test.tsx @@ -38,7 +38,7 @@ const mockContext = { basePromptContexts: MOCK_QUICK_PROMPTS, setSelectedSettingsTab, http: {}, - modelEvaluatorEnabled: true, + assistantFeatures: { assistantModelEvaluation: true }, selectedSettingsTab: 'CONVERSATIONS_TAB', assistantAvailability: { isAssistantEnabled: true, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx index b6bd193aa8534..f83e6c0d72ee6 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx @@ -86,7 +86,7 @@ export const AssistantSettings: React.FC = React.memo( }) => { const { actionTypeRegistry, - modelEvaluatorEnabled, + assistantFeatures: { assistantModelEvaluation: modelEvaluatorEnabled }, http, toasts, selectedSettingsTab, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx index b0bbd18cb0184..f453c19d34320 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx @@ -66,7 +66,7 @@ export const AssistantSettingsManagement: React.FC = React.memo( }) => { const { actionTypeRegistry, - modelEvaluatorEnabled, + assistantFeatures: { assistantModelEvaluation: modelEvaluatorEnabled }, http, selectedSettingsTab, setSelectedSettingsTab, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx index 7726f54f4637a..6a74dab81ac47 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx @@ -13,7 +13,7 @@ import type { IToasts } from '@kbn/core-notifications-browser'; import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public'; import { useLocalStorage, useSessionStorage } from 'react-use'; import type { DocLinksStart } from '@kbn/core-doc-links-browser'; -import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; +import { AssistantFeatures, defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { updatePromptContexts } from './helpers'; import type { PromptContext, @@ -97,6 +97,7 @@ export interface UseAssistantContext { actionTypeRegistry: ActionTypeRegistryContract; alertsIndexPattern: string | undefined; assistantAvailability: AssistantAvailability; + assistantFeatures: AssistantFeatures; assistantStreamingEnabled: boolean; assistantTelemetry?: AssistantTelemetry; augmentMessageCodeBlocks: ( @@ -127,7 +128,6 @@ export interface UseAssistantContext { knowledgeBase: KnowledgeBaseConfig; getLastConversationId: (conversationTitle?: string) => string; promptContexts: Record; - modelEvaluatorEnabled: boolean; nameSpace: string; registerPromptContext: RegisterPromptContext; selectedSettingsTab: SettingsTabs; @@ -276,17 +276,15 @@ export const AssistantProvider: React.FC = ({ ); // Fetch assistant capabilities - const { data: capabilities } = useCapabilities({ http, toasts }); - const { assistantModelEvaluation: modelEvaluatorEnabled, attackDiscoveryEnabled } = - capabilities ?? defaultAssistantFeatures; + const { data: assistantFeatures } = useCapabilities({ http, toasts }); const value = useMemo( () => ({ actionTypeRegistry, alertsIndexPattern, assistantAvailability, + assistantFeatures: assistantFeatures ?? defaultAssistantFeatures, assistantTelemetry, - attackDiscoveryEnabled, augmentMessageCodeBlocks, allQuickPrompts: localStorageQuickPrompts ?? [], allSystemPrompts: localStorageSystemPrompts ?? [], @@ -298,7 +296,6 @@ export const AssistantProvider: React.FC = ({ getComments, http, knowledgeBase: { ...DEFAULT_KNOWLEDGE_BASE_SETTINGS, ...localStorageKnowledgeBase }, - modelEvaluatorEnabled, promptContexts, nameSpace, registerPromptContext, @@ -325,8 +322,8 @@ export const AssistantProvider: React.FC = ({ actionTypeRegistry, alertsIndexPattern, assistantAvailability, + assistantFeatures, assistantTelemetry, - attackDiscoveryEnabled, augmentMessageCodeBlocks, localStorageQuickPrompts, localStorageSystemPrompts, @@ -338,7 +335,6 @@ export const AssistantProvider: React.FC = ({ getComments, http, localStorageKnowledgeBase, - modelEvaluatorEnabled, promptContexts, nameSpace, registerPromptContext, diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts index d0cca8f4d4e09..b3fb27fa835c1 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts @@ -9,7 +9,8 @@ import { postEvaluateRoute } from './post_evaluate'; import { serverMock } from '../../__mocks__/server'; import { requestContextMock } from '../../__mocks__/request_context'; import { getPostEvaluateRequest } from '../../__mocks__/request'; -import type { +import { + defaultAssistantFeatures, PostEvaluateRequestBodyInput, PostEvaluateRequestQueryInput, } from '@kbn/elastic-assistant-common'; @@ -45,8 +46,8 @@ describe('Post Evaluate Route', () => { describe('Capabilities', () => { it('returns a 404 if evaluate feature is not registered', async () => { context.elasticAssistant.getRegisteredFeatures.mockReturnValueOnce({ + ...defaultAssistantFeatures, assistantModelEvaluation: false, - attackDiscoveryEnabled: false, }); const response = await server.inject( diff --git a/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts index aac4e53dfae22..42a24714876b8 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts @@ -103,8 +103,8 @@ describe('AppContextService', () => { it('should register and get features for a single plugin', () => { const pluginName = 'pluginName'; const features: AssistantFeatures = { + ...defaultAssistantFeatures, assistantModelEvaluation: true, - attackDiscoveryEnabled: false, }; appContextService.start(mockAppContext); @@ -118,13 +118,13 @@ describe('AppContextService', () => { it('should register and get features for multiple plugins', () => { const pluginOne = 'plugin1'; const featuresOne: AssistantFeatures = { + ...defaultAssistantFeatures, assistantModelEvaluation: true, - attackDiscoveryEnabled: false, }; const pluginTwo = 'plugin2'; const featuresTwo: AssistantFeatures = { + ...defaultAssistantFeatures, assistantModelEvaluation: false, - attackDiscoveryEnabled: false, }; appContextService.start(mockAppContext); @@ -138,12 +138,12 @@ describe('AppContextService', () => { it('should update features if registered again', () => { const pluginName = 'pluginName'; const featuresOne: AssistantFeatures = { + ...defaultAssistantFeatures, assistantModelEvaluation: true, - attackDiscoveryEnabled: false, }; const featuresTwo: AssistantFeatures = { + ...defaultAssistantFeatures, assistantModelEvaluation: false, - attackDiscoveryEnabled: false, }; appContextService.start(mockAppContext); diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index 3ad99b95905bb..8ccd8998316d7 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -143,6 +143,11 @@ export const allowedExperimentalValues = Object.freeze({ */ assistantModelEvaluation: false, + /** + * Enables the Assistant Knowledge Base by default, introduced in `8.15.0`. + */ + assistantKnowledgeBaseByDefault: false, + /** * Enables the new user details flyout displayed on the Alerts table. */ diff --git a/x-pack/plugins/security_solution/server/plugin.ts b/x-pack/plugins/security_solution/server/plugin.ts index 998e02aba56ad..1ff04a70ea50f 100644 --- a/x-pack/plugins/security_solution/server/plugin.ts +++ b/x-pack/plugins/security_solution/server/plugin.ts @@ -559,6 +559,7 @@ export class Plugin implements ISecuritySolutionPlugin { // Assistant Tool and Feature Registration plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools()); plugins.elasticAssistant.registerFeatures(APP_UI_ID, { + assistantKnowledgeBaseByDefault: config.experimentalFeatures.assistantKnowledgeBaseByDefault, assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation, attackDiscoveryEnabled: config.experimentalFeatures.attackDiscoveryEnabled, }); From 4e4e1e72a2ce913b0fc41059b6d3eecc22963a40 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Mon, 6 May 2024 16:58:34 -0600 Subject: [PATCH 02/20] Use _inference API to automatically setup ELSER, and install ingest pipeline and esql docs via KB data client --- .../impl/assistant/index.tsx | 4 + .../install_knowledge_base_button.tsx | 55 +++++++++ .../knowledge_base_settings.tsx | 21 +++- .../server/__mocks__/request_context.ts | 15 ++- .../conversations/index.test.ts | 2 +- .../ai_assistant_data_clients/index.test.ts | 2 +- .../server/ai_assistant_data_clients/index.ts | 4 +- .../field_maps_configuration.ts | 45 +++++++ .../knowledge_base/index.ts | 114 ++++++++++++++++++ .../knowledge_base/ingest_pipeline.ts | 36 ++++++ .../server/ai_assistant_service/index.test.ts | 10 +- .../server/ai_assistant_service/index.ts | 101 ++++++++++++++-- .../elasticsearch_store.ts | 5 - .../elastic_assistant/server/plugin.ts | 1 - .../server/routes/evaluate/post_evaluate.ts | 2 +- .../delete_knowledge_base.test.ts | 2 +- .../knowledge_base/delete_knowledge_base.ts | 7 +- .../get_knowledge_base_status.test.ts | 2 +- .../get_knowledge_base_status.ts | 9 +- .../post_knowledge_base.test.ts | 2 +- .../knowledge_base/post_knowledge_base.ts | 43 ++++++- .../routes/post_actions_connector_execute.ts | 2 +- .../server/routes/register_routes.ts | 16 ++- .../server/routes/request_context_factory.ts | 12 ++ .../plugins/elastic_assistant/server/types.ts | 20 +-- 25 files changed, 461 insertions(+), 71 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx index 07f3598101709..1f9712a664980 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx @@ -90,6 +90,7 @@ import { clearPresentationData } from '../connectorland/connector_setup/helpers' import { getGenAiConfig } from '../connectorland/helpers'; import { AssistantAnimatedIcon } from './assistant_animated_icon'; import { useFetchAnonymizationFields } from './api/anonymization_fields/use_fetch_anonymization_fields'; +import { InstallKnowledgeBaseButton } from '../knowledge_base/install_knowledge_base_button'; export interface Props { conversationTitle?: string; @@ -825,6 +826,9 @@ const AssistantComponent: React.FC = ({ isFlyoutMode /> + + + diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx new file mode 100644 index 0000000000000..c58061a733210 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx @@ -0,0 +1,55 @@ +/* + * 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 React, { useCallback } from 'react'; +import { EuiButton } from '@elastic/eui'; +import { i18n } from '@kbn/i18n'; + +import { useAssistantContext } from '../..'; +import { useSetupKnowledgeBase } from './use_setup_knowledge_base'; + +const ESQL_RESOURCE = 'esql'; + +/** + * Self-contained component that renders a button to install the knowledge base. + * + * Only renders if `assistantKnowledgeBaseByDefault` feature flag is enabled. + */ +export const InstallKnowledgeBaseButton: React.FC = React.memo(() => { + const { + assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault }, + http, + } = useAssistantContext(); + + // const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); + const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http }); + + const onInstallKnowledgeBase = useCallback(() => { + setupKB(ESQL_RESOURCE); + }, [setupKB]); + + if (!enableKnowledgeBaseByDefault) { + return null; + } + + return ( + + {i18n.translate('xpack.elasticAssistant.knowledgeBase.installKnowledgeBaseButton', { + defaultMessage: 'Install Knowledge Base', + })} + + ); +}); + +InstallKnowledgeBaseButton.displayName = 'InstallKnowledgeBaseButton'; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx index 95bf754874966..6c0e2a8f5e244 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx @@ -47,7 +47,10 @@ interface Props { */ export const KnowledgeBaseSettings: React.FC = React.memo( ({ knowledgeBase, setUpdatedKnowledgeBaseSettings }) => { - const { http } = useAssistantContext(); + const { + assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault }, + http, + } = useAssistantContext(); const { data: kbStatus, isLoading, @@ -67,7 +70,9 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const isESQLAvailable = knowledgeBase.isEnabledKnowledgeBase && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled; // Prevent enabling if elser doesn't exist, but always allow to disable - const isSwitchDisabled = !kbStatus?.elser_exists && !knowledgeBase.isEnabledKnowledgeBase; + const isSwitchDisabled = enableKnowledgeBaseByDefault + ? false + : !kbStatus?.elser_exists && !knowledgeBase.isEnabledKnowledgeBase; // Calculated health state for EuiHealth component const elserHealth = isElserEnabled ? 'success' : 'subdued'; @@ -84,12 +89,18 @@ export const KnowledgeBaseSettings: React.FC = React.memo( isEnabledKnowledgeBase: event.target.checked, }); - // If enabling and ELSER exists, try to set up automatically - if (event.target.checked && kbStatus?.elser_exists) { + // If enabling and ELSER exists or automatic KB setup FF is enabled, try to set up automatically + if (event.target.checked && (enableKnowledgeBaseByDefault || kbStatus?.elser_exists)) { setupKB(ESQL_RESOURCE); } }, - [kbStatus?.elser_exists, knowledgeBase, setUpdatedKnowledgeBaseSettings, setupKB] + [ + enableKnowledgeBaseByDefault, + kbStatus?.elser_exists, + knowledgeBase, + setUpdatedKnowledgeBaseSettings, + setupKB, + ] ); const isEnabledKnowledgeBaseSwitch = useMemo(() => { diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts index 78982d0437650..6c5d9b2bc4d57 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -17,6 +17,8 @@ import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/s import { conversationsDataClientMock, dataClientMock } from './data_clients.mock'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; +import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; export const createMockClients = () => { const core = coreMock.createRequestHandlerContext(); @@ -27,11 +29,12 @@ export const createMockClients = () => { clusterClient: core.elasticsearch.client, elasticAssistant: { actions: actionsClientMock.create(), - getRegisteredFeatures: jest.fn(), + getRegisteredFeatures: jest.fn(() => defaultAssistantFeatures), getRegisteredTools: jest.fn(), logger: loggingSystemMock.createLogger(), telemetry: coreMock.createSetup().analytics, getAIAssistantConversationsDataClient: conversationsDataClientMock.create(), + getAIAssistantKnowledgeBaseDataClient: dataClientMock.create(), getAIAssistantPromptsDataClient: dataClientMock.create(), getAIAssistantAnonymizationFieldsDataClient: dataClientMock.create(), getSpaceId: jest.fn(), @@ -85,7 +88,7 @@ const createElasticAssistantRequestContextMock = ( ): jest.Mocked => { return { actions: clients.elasticAssistant.actions as unknown as ActionsPluginStart, - getRegisteredFeatures: jest.fn(), + getRegisteredFeatures: jest.fn((pluginName: string) => defaultAssistantFeatures), getRegisteredTools: jest.fn(), logger: clients.elasticAssistant.logger, @@ -106,6 +109,14 @@ const createElasticAssistantRequestContextMock = ( () => clients.elasticAssistant.getAIAssistantPromptsDataClient ) as unknown as jest.MockInstance, [], unknown> & (() => Promise), + getAIAssistantKnowledgeBaseDataClient: jest.fn( + () => clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient + ) as unknown as jest.MockInstance< + Promise, + [boolean], + unknown + > & + ((initializeKnowledgeBase: boolean) => Promise), getCurrentUser: jest.fn(), getServerBasePath: jest.fn(), getSpaceId: jest.fn(), diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts index 550d523db4c6f..38426dd06a36f 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts @@ -33,7 +33,7 @@ describe('AIAssistantConversationsDataClient', () => { logger, elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '', + indexPatternsResourceName: '', currentUser: mockUser1, kibanaVersion: '8.8.0', }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts index fa27331f6c6c5..4838c70882f19 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts @@ -30,7 +30,7 @@ describe('AIAssistantDataClient', () => { logger, elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '.kibana-elastic-ai-assistant-conversations', + indexPatternsResourceName: '.kibana-elastic-ai-assistant-conversations', currentUser: mockUser1, kibanaVersion: '8.8.0', }; 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 75f2b166f1468..5bc84e05c8087 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 @@ -21,7 +21,7 @@ export interface AIAssistantDataClientParams { kibanaVersion: string; spaceId: string; logger: Logger; - indexPatternsResorceName: string; + indexPatternsResourceName: string; currentUser: AuthenticatedUser | null; } @@ -38,7 +38,7 @@ export class AIAssistantDataClient { constructor(public readonly options: AIAssistantDataClientParams) { this.indexTemplateAndPattern = getIndexTemplateAndPattern( - this.options.indexPatternsResorceName, + this.options.indexPatternsResourceName, this.options.spaceId ?? DEFAULT_NAMESPACE_STRING ); this.currentUser = this.options.currentUser; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts new file mode 100644 index 0000000000000..e334427f9d9f5 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -0,0 +1,45 @@ +/* + * 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 { FieldMap } from '@kbn/data-stream-adapter'; + +export const knowledgeBaseFieldMap: FieldMap = { + '@timestamp': { + type: 'date', + array: false, + required: false, + }, + metadata: { + type: 'object', + array: false, + required: false, + }, + 'metadata.kbResource': { + type: 'keyword', + array: false, + required: false, + }, + 'metadata.required': { + type: 'boolean', + array: false, + required: false, + }, + 'metadata.source': { + type: 'keyword', + array: false, + required: false, + }, + vector: { + type: 'object', + array: false, + required: false, + }, + 'vector.tokens': { + type: 'rank_features', + array: false, + required: false, + }, +} as const; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts new file mode 100644 index 0000000000000..64882b7371e29 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -0,0 +1,114 @@ +/* + * 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 { + MlTrainedModelDeploymentNodesStats, + MlTrainedModelStats, +} from '@elastic/elasticsearch/lib/api/types'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; +import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; +import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; +export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { + private isInstallingElser: boolean = false; + + constructor(public readonly options: AIAssistantDataClientParams) { + super(options); + } + + /** + * Checks if the provided model is installed (deployed and allocated) in Elasticsearch + * + * @param modelId ID of the model to check + * @returns Promise indicating whether the model is installed + */ + private isModelInstalled = async (modelId: string): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + + try { + const getResponse = await esClient.ml.getTrainedModelsStats({ + model_id: modelId, + }); + + this.options.logger.debug(`modelId: ${modelId}`); + + // For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986 + const isReadyESS = (stats: MlTrainedModelStats) => + stats.deployment_stats?.state === 'started' && + stats.deployment_stats?.allocation_status.state === 'fully_allocated'; + + const isReadyServerless = (stats: MlTrainedModelStats) => + (stats.deployment_stats?.nodes as unknown as MlTrainedModelDeploymentNodesStats[]).some( + (node) => node.routing_state.routing_state === 'started' + ); + + return getResponse.trained_model_stats.some( + (stats) => isReadyESS(stats) || isReadyServerless(stats) + ); + } catch (e) { + // Returns 404 if it doesn't exist + return false; + } + }; + + /** + * Downloads and deploys ELSER (if not already) by means of the _inference API, then loads ES|QL docs + * + * @param options + * @param options.elserId ID of the recommended ELSER model + * @returns Promise + */ + public setupKnowledgeBase = async ({ + elserId, + esClient, + esStore, + }: { + elserId: string; + esClient: ElasticsearchClient; + esStore: ElasticsearchStore; + }): Promise => { + if (this.isInstallingElser) { + return; + } + + this.isInstallingElser = true; + const isInstalled = await this.isModelInstalled(elserId); + + if (isInstalled) { + this.options.logger.debug(`ELSER model '${elserId}' is already installed`); + this.options.logger.debug(`Loading KB docs!`); + const loadedKnowledgeBase = await loadESQL(esStore, this.options.logger); + this.options.logger.debug(`${loadedKnowledgeBase}`); + this.isInstallingElser = false; + + return; + } + + try { + // Temporarily use esClient for current user until `kibana_system` user has `inference_admin` role + // See https://github.com/elastic/elasticsearch/pull/108262 + // const esClient = await this.options.elasticsearchClientPromise; + const elserResponse = await esClient.inference.putModel({ + inference_id: 'elser_model_2', + task_type: 'sparse_embedding', + model_config: { + service: 'elser', + service_settings: { + model_id: elserId, + num_allocations: 1, + num_threads: 1, + }, + task_settings: {}, + }, + }); + + this.options.logger.debug(`elser response:\n: ${JSON.stringify(elserResponse, null, 2)}`); + } catch (e) { + this.options.logger.error(`Error setting up ELSER model: ${e.message}`); + } + }; +} diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts new file mode 100644 index 0000000000000..170fa0342f9d9 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/ingest_pipeline.ts @@ -0,0 +1,36 @@ +/* + * 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 { IngestPutPipelineRequest } from '@elastic/elasticsearch/lib/api/types'; + +export const knowledgeBaseIngestPipeline = ({ + id, + modelId, +}: { + id: string; + modelId: string; +}): IngestPutPipelineRequest => ({ + id, + description: 'Embedding pipeline for Elastic AI Assistant ELSER Knowledge Base', + processors: [ + { + inference: { + model_id: modelId, + target_field: 'vector', + field_map: { + text: 'text_field', + }, + inference_config: { + // @ts-expect-error + text_expansion: { + results_field: 'tokens', + }, + }, + }, + }, + ], +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts index a96088b413dcf..00babc22809d4 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts @@ -192,7 +192,7 @@ describe('AI Assistant Service', () => { logger, elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '.kibana-elastic-ai-assistant-conversations', + indexPatternsResourceName: '.kibana-elastic-ai-assistant-conversations', currentUser: mockUser1, kibanaVersion: '8.8.0', }); @@ -234,7 +234,7 @@ describe('AI Assistant Service', () => { logger, elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '.kibana-elastic-ai-assistant-conversations', + indexPatternsResourceName: '.kibana-elastic-ai-assistant-conversations', currentUser: mockUser1, kibanaVersion: '8.8.0', }); @@ -298,7 +298,7 @@ describe('AI Assistant Service', () => { expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({ elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '.kibana-elastic-ai-assistant-conversations', + indexPatternsResourceName: '.kibana-elastic-ai-assistant-conversations', currentUser: mockUser1, kibanaVersion: '8.8.0', logger, @@ -358,7 +358,7 @@ describe('AI Assistant Service', () => { expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({ elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '.kibana-elastic-ai-assistant-conversations', + indexPatternsResourceName: '.kibana-elastic-ai-assistant-conversations', currentUser: mockUser1, kibanaVersion: '8.8.0', logger, @@ -431,7 +431,7 @@ describe('AI Assistant Service', () => { expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({ elasticsearchClientPromise: Promise.resolve(clusterClient), spaceId: 'default', - indexPatternsResorceName: '.kibana-elastic-ai-assistant-conversations', + indexPatternsResourceName: '.kibana-elastic-ai-assistant-conversations', currentUser: mockUser1, kibanaVersion: '8.8.0', logger, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 9cc3efb03e195..af2d4b78a807a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -25,6 +25,9 @@ import { conversationsFieldMap } from '../ai_assistant_data_clients/conversation import { assistantPromptsFieldMap } from '../ai_assistant_data_clients/prompts/field_maps_configuration'; import { assistantAnonymizationFieldsFieldMap } from '../ai_assistant_data_clients/anonymization_fields/field_maps_configuration'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; +import { knowledgeBaseFieldMap } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; +import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; +import { knowledgeBaseIngestPipeline } from '../ai_assistant_data_clients/knowledge_base/ingest_pipeline'; const TOTAL_FIELDS_LIMIT = 2500; @@ -47,7 +50,7 @@ export interface CreateAIAssistantClientParams { } export type CreateDataStream = (params: { - resource: 'conversations' | 'prompts' | 'anonymizationFields'; + resource: 'anonymizationFields' | 'conversations' | 'knowledgeBase' | 'prompts'; fieldMap: FieldMap; kibanaVersion: string; spaceId?: string; @@ -55,8 +58,11 @@ export type CreateDataStream = (params: { export class AIAssistantService { private initialized: boolean; + // Temporary 'feature flag' to determine if we should initialize the knowledge base, toggled when accessing data client + private initializeKnowledgeBase: boolean = false; private isInitializing: boolean = false; private conversationsDataStream: DataStreamSpacesAdapter; + private knowledgeBaseDataStream: DataStreamSpacesAdapter; private promptsDataStream: DataStreamSpacesAdapter; private anonymizationFieldsDataStream: DataStreamSpacesAdapter; private resourceInitializationHelper: ResourceInstallationHelper; @@ -69,6 +75,11 @@ export class AIAssistantService { kibanaVersion: options.kibanaVersion, fieldMap: conversationsFieldMap, }); + this.knowledgeBaseDataStream = this.createDataStream({ + resource: 'knowledgeBase', + kibanaVersion: options.kibanaVersion, + fieldMap: knowledgeBaseFieldMap, + }); this.promptsDataStream = this.createDataStream({ resource: 'prompts', kibanaVersion: options.kibanaVersion, @@ -124,6 +135,39 @@ export class AIAssistantService { pluginStop$: this.options.pluginStop$, }); + if (this.initializeKnowledgeBase) { + await this.knowledgeBaseDataStream.install({ + esClient, + logger: this.options.logger, + pluginStop$: this.options.pluginStop$, + }); + + // TODO: Add generic ingest pipeline support to `kbn-data-stream-adapter` package? + // TODO: Inject `elserId` + let pipelineExists = false; + try { + const response = await esClient.ingest.getPipeline({ + id: this.resourceNames.pipelines.knowledgeBase, + }); + pipelineExists = Object.keys(response).length > 0; + } catch (e) { + // The GET /_ingest/pipeline/{pipelineId} API returns an empty object w/ 404 Not Found. + pipelineExists = false; + } + if (!pipelineExists) { + this.options.logger.info('Installing ingest pipeline'); + const response = await esClient.ingest.putPipeline( + knowledgeBaseIngestPipeline({ + id: this.resourceNames.pipelines.knowledgeBase, + modelId: '.elser_model_2', + }) + ); + this.options.logger.info(`Installed ingest pipeline: ${response.acknowledged}`); + } else { + this.options.logger.info('Ingest pipeline already exists'); + } + } + await this.promptsDataStream.install({ esClient, logger: this.options.logger, @@ -149,30 +193,30 @@ export class AIAssistantService { private readonly resourceNames: AssistantResourceNames = { componentTemplate: { conversations: getResourceName('component-template-conversations'), + knowledgeBase: getResourceName('component-template-knowledge-base'), prompts: getResourceName('component-template-prompts'), anonymizationFields: getResourceName('component-template-anonymization-fields'), - kb: getResourceName('component-template-kb'), }, aliases: { conversations: getResourceName('conversations'), + knowledgeBase: getResourceName('knowledge-base'), prompts: getResourceName('prompts'), anonymizationFields: getResourceName('anonymization-fields'), - kb: getResourceName('kb'), }, indexPatterns: { conversations: getResourceName('conversations*'), + knowledgeBase: getResourceName('knowledge-base*'), prompts: getResourceName('prompts*'), anonymizationFields: getResourceName('anonymization-fields*'), - kb: getResourceName('kb*'), }, indexTemplate: { conversations: getResourceName('index-template-conversations'), + knowledgeBase: getResourceName('index-template-knowledge-base'), prompts: getResourceName('index-template-prompts'), anonymizationFields: getResourceName('index-template-anonymization-fields'), - kb: getResourceName('index-template-kb'), }, pipelines: { - kb: getResourceName('kb-ingest-pipeline'), + knowledgeBase: getResourceName('ingest-pipeline-knowledge-base'), }, }; @@ -182,7 +226,7 @@ export class AIAssistantService { opts.spaceId ); - // If space evel resources initialization failed, retry + // If space level resources initialization failed, retry if (!initialized && error) { let initRetryPromise: Promise | undefined; @@ -236,7 +280,33 @@ export class AIAssistantService { elasticsearchClientPromise: this.options.elasticsearchClientPromise, spaceId: opts.spaceId, kibanaVersion: this.options.kibanaVersion, - indexPatternsResorceName: this.resourceNames.aliases.conversations, + indexPatternsResourceName: this.resourceNames.aliases.conversations, + currentUser: opts.currentUser, + }); + } + + public async createAIAssistantKnowledgeBaseDataClient( + opts: CreateAIAssistantClientParams & { initializeKnowledgeBase: boolean } + ): Promise { + // Note: Due to plugin lifecycle and feature flag registration timing, we need to pass in the feature flag here + // Remove this param and initialization when the `assistantKnowledgeBaseByDefault` feature flag is removed + if (opts.initializeKnowledgeBase) { + this.initializeKnowledgeBase = true; + await this.initializeResources(); + } + + const res = await this.checkResourcesInstallation(opts); + + if (res === null) { + return null; + } + + return new AIAssistantKnowledgeBaseDataClient({ + logger: this.options.logger, + elasticsearchClientPromise: this.options.elasticsearchClientPromise, + spaceId: opts.spaceId, + kibanaVersion: this.options.kibanaVersion, + indexPatternsResourceName: this.resourceNames.aliases.conversations, currentUser: opts.currentUser, }); } @@ -255,7 +325,7 @@ export class AIAssistantService { elasticsearchClientPromise: this.options.elasticsearchClientPromise, spaceId: opts.spaceId, kibanaVersion: this.options.kibanaVersion, - indexPatternsResorceName: this.resourceNames.aliases.prompts, + indexPatternsResourceName: this.resourceNames.aliases.prompts, currentUser: opts.currentUser, }); } @@ -274,7 +344,7 @@ export class AIAssistantService { elasticsearchClientPromise: this.options.elasticsearchClientPromise, spaceId: opts.spaceId, kibanaVersion: this.options.kibanaVersion, - indexPatternsResorceName: this.resourceNames.aliases.anonymizationFields, + indexPatternsResourceName: this.resourceNames.aliases.anonymizationFields, currentUser: opts.currentUser, }); } @@ -308,6 +378,15 @@ export class AIAssistantService { await this.conversationsDataStream.installSpace(spaceId); } + if (this.initializeKnowledgeBase) { + const knowledgeBaseIndexName = await this.knowledgeBaseDataStream.getInstalledSpaceName( + spaceId + ); + if (!knowledgeBaseIndexName) { + await this.knowledgeBaseDataStream.installSpace(spaceId); + } + } + const promptsIndexName = await this.promptsDataStream.getInstalledSpaceName(spaceId); if (!promptsIndexName) { await this.promptsDataStream.installSpace(spaceId); @@ -334,7 +413,7 @@ export class AIAssistantService { elasticsearchClientPromise: this.options.elasticsearchClientPromise, spaceId, kibanaVersion: this.options.kibanaVersion, - indexPatternsResorceName: this.resourceNames.aliases.anonymizationFields, + indexPatternsResourceName: this.resourceNames.aliases.anonymizationFields, currentUser: null, }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts index ae6540de5e271..d6987eae089e0 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts @@ -109,10 +109,6 @@ export class ElasticsearchStore extends VectorStore { if (!pipelineExists) { await this.createPipeline(); } - const indexExists = await this.indexExists(); - if (!indexExists) { - await this.createIndex(); - } const operations = documents.flatMap(({ pageContent, metadata }) => [ { index: { _index: this.index, _id: uuid.v4() } }, @@ -400,7 +396,6 @@ export class ElasticsearchStore extends VectorStore { }); this.logger.debug(`modelId: ${modelId}`); - this.logger.debug(`getResponse: ${JSON.stringify(getResponse, null, 2)}`); // For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986 const isReadyESS = (stats: MlTrainedModelStats) => diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 53b05857beb4b..63d4c362044e1 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -76,7 +76,6 @@ export class ElasticAssistantPlugin ); events.forEach((eventConfig) => core.analytics.registerEventType(eventConfig)); - // this.assistantService registerKBTask registerRoutes(router, this.logger, plugins); return { actions: plugins.actions, diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index d1bf9dfa26ab1..fc919d10e345c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -136,7 +136,7 @@ export const postEvaluateRoute = ( const esClient = (await context.core).elasticsearch.client.asCurrentUser; // Default ELSER model - const elserId = await getElser(request, (await context.core).savedObjects.getClient()); + const elserId = await getElser(); // Skeleton request from route to pass to the agents // params will be passed to the actions executor diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.test.ts index 0443bd1b3eedd..ad130cddc5560 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.test.ts @@ -36,7 +36,7 @@ describe('Delete Knowledge Base Route', () => { }); test('returns 500 if error is thrown when deleting resources', async () => { - context.core.elasticsearch.client.asCurrentUser.indices.delete.mockRejectedValue( + context.core.elasticsearch.client.asInternalUser.indices.delete.mockRejectedValue( new Error('Test error') ); const response = await server.inject( diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts index 6886b56f7ef32..17b01fdb5f9aa 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts @@ -32,9 +32,7 @@ export const deleteKnowledgeBaseRoute = ( access: 'internal', path: KNOWLEDGE_BASE, options: { - // Note: Relying on current user privileges to scope an esClient. - // Add `access:kbnElasticAssistant` to limit API access to only users with assistant privileges - tags: [], + tags: ['access:elasticAssistant'], }, }) .addVersion( @@ -58,8 +56,7 @@ export const deleteKnowledgeBaseRoute = ( ? decodeURIComponent(request.params.resource) : undefined; - // Get a scoped esClient for deleting the Knowledge Base index, pipeline, and documents - const esClient = (await context.core).elasticsearch.client.asCurrentUser; + const esClient = (await context.core).elasticsearch.client.asInternalUser; const esStore = new ElasticsearchStore( esClient, KNOWLEDGE_BASE_INDEX_PATTERN, diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts index f37f2a3ab2882..e6811ac7cf3d2 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts @@ -37,7 +37,7 @@ describe('Get Knowledge Base Status Route', () => { }); test('returns 500 if error is thrown in checking kb status', async () => { - context.core.elasticsearch.client.asCurrentUser.indices.exists.mockRejectedValue( + context.core.elasticsearch.client.asInternalUser.indices.exists.mockRejectedValue( new Error('Test error') ); const response = await server.inject( diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index c7d83e8c24401..cb13249423406 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -36,9 +36,7 @@ export const getKnowledgeBaseStatusRoute = ( access: 'internal', path: KNOWLEDGE_BASE, options: { - // Note: Relying on current user privileges to scope an esClient. - // Add `access:kbnElasticAssistant` to limit API access to only users with assistant privileges - tags: [], + tags: ['access:elasticAssistant'], }, }) .addVersion( @@ -57,9 +55,8 @@ export const getKnowledgeBaseStatusRoute = ( const telemetry = assistantContext.telemetry; try { - // Get a scoped esClient for finding the status of the Knowledge Base index, pipeline, and documents - const esClient = (await context.core).elasticsearch.client.asCurrentUser; - const elserId = await getElser(request, (await context.core).savedObjects.getClient()); + const esClient = (await context.core).elasticsearch.client.asInternalUser; + const elserId = await getElser(); const kbResource = getKbResource(request); const esStore = new ElasticsearchStore( esClient, diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts index ceb5f1b3879f6..b9ab3bc4c6866 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts @@ -38,7 +38,7 @@ describe('Post Knowledge Base Route', () => { }); test('returns 500 if error is thrown when creating resources', async () => { - context.core.elasticsearch.client.asCurrentUser.indices.exists.mockRejectedValue( + context.core.elasticsearch.client.asInternalUser.indices.exists.mockRejectedValue( new Error('Test error') ); const response = await server.inject( diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 17c2011fbc0f5..05c342f2deef8 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -21,10 +21,12 @@ import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elas import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants'; import { getKbResource } from './get_kb_resource'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; /** * Load Knowledge Base index, pipeline, and resources (collection of documents) * @param router + * @param getElser */ export const postKnowledgeBaseRoute = ( router: ElasticAssistantPluginRouter, @@ -35,9 +37,7 @@ export const postKnowledgeBaseRoute = ( access: 'internal', path: KNOWLEDGE_BASE, options: { - // Note: Relying on current user privileges to scope an esClient. - // Add `access:kbnElasticAssistant` to limit API access to only users with assistant privileges - tags: [], + tags: ['access:elasticAssistant'], }, }) .addVersion( @@ -58,12 +58,43 @@ export const postKnowledgeBaseRoute = ( const assistantContext = await context.elasticAssistant; const logger = assistantContext.logger; const telemetry = assistantContext.telemetry; + const elserId = await getElser(); + + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const enableKnowledgeBaseByDefault = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; + + if (enableKnowledgeBaseByDefault) { + // Go ahead and do that now, right here's the place :) + const knowledgeBaseDataClient = + await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); + // Temporarily get esClient for current user until `kibana_system` user has `inference_admin` role + // See https://github.com/elastic/elasticsearch/pull/108262 + const esClient = (await context.core).elasticsearch.client.asCurrentUser; + + // + const esStore = new ElasticsearchStore( + esClient, + `.kibana-elastic-ai-assistant-knowledge-base-default`, + logger, + telemetry, + elserId, + getKbResource(request) + ); + // + + await knowledgeBaseDataClient?.setupKnowledgeBase({ elserId, esClient, esStore }); + + return response.ok({ body: { success: true } }); + } try { const core = await context.core; - // Get a scoped esClient for creating the Knowledge Base index, pipeline, and documents - const esClient = core.elasticsearch.client.asCurrentUser; - const elserId = await getElser(request, core.savedObjects.getClient()); + const esClient = core.elasticsearch.client.asInternalUser; const kbResource = getKbResource(request); const esStore = new ElasticsearchStore( esClient, diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index e579996eb13d1..61326f157bb27 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -315,7 +315,7 @@ export const postActionsConnectorExecuteRoute = ( []) as unknown as Array> ); - const elserId = await getElser(request, (await context.core).savedObjects.getClient()); + const elserId = await getElser(); const anonymizationFieldsRes = await anonymizationFieldsDataClient?.findDocuments({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts index 325f4a84ab8c7..5ff629d6a0b45 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts @@ -56,12 +56,16 @@ export const registerRoutes = ( // Knowledge Base deleteKnowledgeBaseRoute(router); - const getElserId: GetElser = once( - async (request: KibanaRequest, savedObjectsClient: SavedObjectsClientContract) => { - return (await plugins.ml.trainedModelsProvider(request, savedObjectsClient).getELSER()) - .model_id; - } - ); + const getElserId: GetElser = once(async () => { + return ( + ( + await plugins.ml + // Force check to happen as internal user + .trainedModelsProvider({} as KibanaRequest, {} as SavedObjectsClientContract) + .getELSER() + ).model_id + ); + }); getKnowledgeBaseStatusRoute(router, getElserId); postKnowledgeBaseRoute(router, getElserId); diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index 82e21a8cd8690..0a0864882df16 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -81,6 +81,18 @@ export class RequestContextFactory implements IRequestContextFactory { telemetry: core.analytics, + // Note: Due to plugin lifecycle and feature flag registration timing, we need to pass in the feature flag here + // Remove `initializeKnowledgeBase` once 'assistantKnowledgeBaseByDefault' feature flag is removed + getAIAssistantKnowledgeBaseDataClient: memoize((initializeKnowledgeBase = false) => { + const currentUser = getCurrentUser(); + return this.assistantService.createAIAssistantKnowledgeBaseDataClient({ + spaceId: getSpaceId(), + logger: this.logger, + currentUser, + initializeKnowledgeBase, + }); + }), + getAIAssistantPromptsDataClient: memoize(() => { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantPromptsDataClient({ diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 063573be37c3e..c9787c0b6213e 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -17,7 +17,6 @@ import type { IRouter, KibanaRequest, Logger, - SavedObjectsClientContract, } from '@kbn/core/server'; import { type MlPluginSetup } from '@kbn/ml-plugin/server'; import { DynamicStructuredTool, Tool } from '@langchain/core/tools'; @@ -42,6 +41,7 @@ import { import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; import { AIAssistantDataClient } from './ai_assistant_data_clients'; +import { AIAssistantKnowledgeBaseDataClient } from './ai_assistant_data_clients/knowledge_base'; export const PLUGIN_ID = 'elasticAssistant' as const; @@ -110,6 +110,9 @@ export interface ElasticAssistantApiRequestHandlerContext { getSpaceId: () => string; getCurrentUser: () => AuthenticatedUser | null; getAIAssistantConversationsDataClient: () => Promise; + getAIAssistantKnowledgeBaseDataClient: ( + initializeKnowledgeBase: boolean + ) => Promise; getAIAssistantPromptsDataClient: () => Promise; getAIAssistantAnonymizationFieldsDataClient: () => Promise; telemetry: AnalyticsServiceSetup; @@ -129,10 +132,7 @@ export type ElasticAssistantPluginCoreSetupDependencies = CoreSetup< ElasticAssistantPluginStart >; -export type GetElser = ( - request: KibanaRequest, - savedObjectsClient: SavedObjectsClientContract -) => Promise | never; +export type GetElser = () => Promise | never; export interface InitAssistantResult { assistantResourcesInstalled: boolean; @@ -144,30 +144,30 @@ export interface InitAssistantResult { export interface AssistantResourceNames { componentTemplate: { conversations: string; + knowledgeBase: string; prompts: string; anonymizationFields: string; - kb: string; }; indexTemplate: { conversations: string; + knowledgeBase: string; prompts: string; anonymizationFields: string; - kb: string; }; aliases: { conversations: string; + knowledgeBase: string; prompts: string; anonymizationFields: string; - kb: string; }; indexPatterns: { conversations: string; + knowledgeBase: string; prompts: string; anonymizationFields: string; - kb: string; }; pipelines: { - kb: string; + knowledgeBase: string; }; } From 316e8c91c7cff1fc4ad1c0ed9ada3ef02e5bc9c3 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 7 May 2024 15:01:23 -0600 Subject: [PATCH 03/20] Fix settings tests --- .../impl/knowledge_base/knowledge_base_settings.test.tsx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx index 20ab3aab4a26f..56f6796ac16fa 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx @@ -13,9 +13,11 @@ import { KnowledgeBaseSettings } from './knowledge_base_settings'; import { TestProviders } from '../mock/test_providers/test_providers'; import { useKnowledgeBaseStatus } from './use_knowledge_base_status'; import { mockSystemPrompts } from '../mock/system_prompt'; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; const mockUseAssistantContext = { allSystemPrompts: mockSystemPrompts, + assistantFeatures: jest.fn(() => defaultAssistantFeatures), conversations: {}, http: { basePath: { From 71f4dd09184c6152421e6534348b91ea0299f472 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 8 May 2024 16:21:36 -0600 Subject: [PATCH 04/20] Switch to using installElasticModel() --- .../knowledge_base/index.ts | 59 +++++--- .../server/ai_assistant_service/helpers.ts | 107 +++++++++++++ .../server/ai_assistant_service/index.test.ts | 141 ++++-------------- .../server/ai_assistant_service/index.ts | 45 +++--- .../elastic_assistant/server/plugin.ts | 5 +- .../knowledge_base/post_knowledge_base.ts | 2 +- .../server/routes/register_routes.ts | 21 +-- 7 files changed, 201 insertions(+), 179 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 64882b7371e29..bf8db77f16e72 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -9,14 +9,23 @@ import { MlTrainedModelDeploymentNodesStats, MlTrainedModelStats, } from '@elastic/elasticsearch/lib/api/types'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { AuthenticatedUser } from '@kbn/core-security-common'; +import type { MlPluginSetup } from '@kbn/ml-plugin/server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; +import { GetElser } from '../../types'; + +interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { + ml: MlPluginSetup; + getElserId: GetElser; +} export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { private isInstallingElser: boolean = false; - constructor(public readonly options: AIAssistantDataClientParams) { + constructor(public readonly options: KnowledgeBaseDataClientParams) { super(options); } @@ -56,26 +65,33 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }; /** - * Downloads and deploys ELSER (if not already) by means of the _inference API, then loads ES|QL docs + * Downloads and deploys ELSER (if not already), then loads ES|QL docs * * @param options - * @param options.elserId ID of the recommended ELSER model + * @param options.esStore ElasticsearchStore for loading ES|QL docs * @returns Promise */ public setupKnowledgeBase = async ({ - elserId, - esClient, esStore, }: { - elserId: string; - esClient: ElasticsearchClient; esStore: ElasticsearchStore; }): Promise => { if (this.isInstallingElser) { return; } + // TODO: Before automatically installing ELSER in the background, we should perform the following deployment resource checks + // Note: ESS only, as Serverless can always auto-install if `productTier === complete` + // 1. Deployment has ML Nodes with adequate free memory + // We can just auto-install, yay! + // 2. Deployment doesn't have adequate ML resources, and ML Autoscaling is disabled (or unavailable due to cluster health). + // Refer the user to the docs for further details + // 3. Deployment doesn't have adequate ML resources, but have ML Autoscaling enabled and scale limits are are NOT WITHIN the required resources. + // Again, refer the user to the docs + // 4. Deployment doesn't have adequate ML resources, but have ML Autoscaling enabled and scale limits ARE WITHIN the required resources. + // In this instance we could auto-install, but may have it behind a user action since deployment costs would change... this.isInstallingElser = true; + const elserId = await this.options.getElserId(); const isInstalled = await this.isModelInstalled(elserId); if (isInstalled) { @@ -89,26 +105,23 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } try { - // Temporarily use esClient for current user until `kibana_system` user has `inference_admin` role - // See https://github.com/elastic/elasticsearch/pull/108262 + const elserResponse = await this.options.ml + .trainedModelsProvider({} as KibanaRequest, {} as SavedObjectsClientContract) + .installElasticModel(elserId); + // const esClient = await this.options.elasticsearchClientPromise; - const elserResponse = await esClient.inference.putModel({ - inference_id: 'elser_model_2', - task_type: 'sparse_embedding', - model_config: { - service: 'elser', - service_settings: { - model_id: elserId, - num_allocations: 1, - num_threads: 1, - }, - task_settings: {}, - }, - }); this.options.logger.debug(`elser response:\n: ${JSON.stringify(elserResponse, null, 2)}`); } catch (e) { this.options.logger.error(`Error setting up ELSER model: ${e.message}`); } }; + + public addKnowledgeBaseResource = async ({ + document, + authenticatedUser, + }: { + document: Document; + authenticatedUser: AuthenticatedUser; + }): Promise => {}; } diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts new file mode 100644 index 0000000000000..732339b18ded3 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts @@ -0,0 +1,107 @@ +/* + * 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 type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { once } from 'lodash/fp'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; +import type { MlPluginSetup } from '@kbn/ml-plugin/server'; +import { knowledgeBaseIngestPipeline } from '../ai_assistant_data_clients/knowledge_base/ingest_pipeline'; +import { GetElser } from '../types'; + +/** + * Creates a function that returns the ELSER model ID + * + * @param ml + */ +export const createGetElserId = (ml: MlPluginSetup): GetElser => { + return once(async () => { + return ( + ( + await ml + // Force check to happen as internal user + .trainedModelsProvider({} as KibanaRequest, {} as SavedObjectsClientContract) + .getELSER() + ).model_id + ); + }); +}; + +interface PipelineExistsParams { + esClient: ElasticsearchClient; + id: string; +} + +/** + * Checks if the provided ingest pipeline exists in Elasticsearch + * + * @param params params + * @param params.esClient Elasticsearch client with privileges to check for ingest pipelines + * @param params.id ID of the ingest pipeline to check + * + * @returns Promise indicating whether the pipeline exists + */ +export const pipelineExists = async ({ esClient, id }: PipelineExistsParams): Promise => { + try { + const response = await esClient.ingest.getPipeline({ + id, + }); + return Object.keys(response).length > 0; + } catch (e) { + // The GET /_ingest/pipeline/{pipelineId} API returns an empty object w/ 404 Not Found. + return false; + } +}; + +interface CreatePipelineParams { + esClient: ElasticsearchClient; + id: string; + modelId: string; +} + +/** + * Create ingest pipeline for ELSER in Elasticsearch + * + * @param params params + * @param params.esClient Elasticsearch client with privileges to check for ingest pipelines + * @param params.id ID of the ingest pipeline + * @param params.modelId ID of the ELSER model + * + * @returns Promise indicating whether the pipeline was created + */ +export const createPipeline = async ({ + esClient, + id, + modelId, +}: CreatePipelineParams): Promise => { + const response = await esClient.ingest.putPipeline( + knowledgeBaseIngestPipeline({ + id, + modelId, + }) + ); + + return response.acknowledged; +}; + +interface DeletePipelineParams { + esClient: ElasticsearchClient; + id: string; +} + +/** + * Delete ingest pipeline for ELSER in Elasticsearch + * + * @returns Promise indicating whether the pipeline was created + */ +export const deletePipeline = async ({ esClient, id }: DeletePipelineParams): Promise => { + const response = await esClient.ingest.deletePipeline({ + id, + }); + + return response.acknowledged; +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts index 00babc22809d4..dbdc01dcf9e57 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts @@ -14,8 +14,10 @@ import { AuthenticatedUser } from '@kbn/security-plugin/server'; import { DEFAULT_NAMESPACE_STRING } from '@kbn/core-saved-objects-utils-server'; import { conversationsDataClientMock } from '../__mocks__/data_clients.mock'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; -import { AIAssistantService } from '.'; +import { AIAssistantService, AIAssistantServiceOpts } from '.'; import { retryUntil } from './create_resource_installation_helper.test'; +import { mlPluginMock } from '@kbn/ml-plugin/public/mocks'; +import type { MlPluginSetup } from '@kbn/ml-plugin/server'; jest.mock('../ai_assistant_data_clients/conversations', () => ({ AIAssistantConversationsDataClient: jest.fn(), @@ -95,6 +97,7 @@ const mockUser1 = { describe('AI Assistant Service', () => { let pluginStop$: Subject; + let assistantServiceOpts: AIAssistantServiceOpts; beforeEach(() => { jest.resetAllMocks(); @@ -107,6 +110,14 @@ describe('AI Assistant Service', () => { ); clusterClient.indices.getAlias.mockImplementation(async () => GetAliasResponse); clusterClient.indices.getDataStream.mockImplementation(async () => GetDataStreamResponse); + assistantServiceOpts = { + logger, + elasticsearchClientPromise: Promise.resolve(clusterClient), + pluginStop$, + kibanaVersion: '8.8.0', + ml: mlPluginMock.createSetupContract() as unknown as MlPluginSetup, // Missing SharedServices mock + taskManager: taskManagerMock.createSetup(), + }; }); afterEach(() => { @@ -116,13 +127,7 @@ describe('AI Assistant Service', () => { describe('AIAssistantService()', () => { test('should correctly initialize common resources', async () => { - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -140,13 +145,7 @@ describe('AI Assistant Service', () => { test('should log error and set initialized to false if creating/updating common component template throws error', async () => { clusterClient.cluster.putComponentTemplate.mockRejectedValueOnce(new Error('fail')); - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil('error log called', async () => logger.error.mock.calls.length > 0); @@ -169,13 +168,7 @@ describe('AI Assistant Service', () => { }); test('should create new AIAssistantConversationsDataClient', async () => { - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -201,13 +194,7 @@ describe('AI Assistant Service', () => { test('should retry initializing common resources if common resource initialization failed', async () => { clusterClient.cluster.putComponentTemplate.mockRejectedValueOnce(new Error('fail')); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil('error log called', async () => logger.error.mock.calls.length > 0); @@ -259,13 +246,7 @@ describe('AI Assistant Service', () => { return { acknowledged: true }; }); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil('error log called', async () => logger.error.mock.calls.length > 0); @@ -336,13 +317,7 @@ describe('AI Assistant Service', () => { mappings: {}, }, })); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -397,13 +372,7 @@ describe('AI Assistant Service', () => { return SimulateTemplateResponse; }); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -475,13 +444,7 @@ describe('AI Assistant Service', () => { }, })); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -523,13 +486,7 @@ describe('AI Assistant Service', () => { throw new Error(`fail ${++failCount}`); }); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil('error log called', async () => logger.error.mock.calls.length > 0, 1); @@ -574,13 +531,7 @@ describe('AI Assistant Service', () => { test('should return null if retrying common resources initialization fails again with same error', async () => { clusterClient.cluster.putComponentTemplate.mockRejectedValue(new Error('fail')); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil('error log called', async () => logger.error.mock.calls.length > 0); @@ -633,13 +584,7 @@ describe('AI Assistant Service', () => { })); clusterClient.indices.putIndexTemplate.mockRejectedValue(new Error('fail index template')); - assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -677,13 +622,7 @@ describe('AI Assistant Service', () => { .mockRejectedValueOnce(new EsErrors.ConnectionError('foo')) .mockRejectedValueOnce(new EsErrors.TimeoutError('timeout')) .mockResolvedValue({ acknowledged: true }); - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -697,13 +636,7 @@ describe('AI Assistant Service', () => { .mockRejectedValueOnce(new EsErrors.ConnectionError('foo')) .mockRejectedValueOnce(new EsErrors.TimeoutError('timeout')) .mockResolvedValue({ acknowledged: true }); - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -724,13 +657,7 @@ describe('AI Assistant Service', () => { .mockRejectedValueOnce(new EsErrors.ConnectionError('foo')) .mockRejectedValueOnce(new EsErrors.TimeoutError('timeout')) .mockResolvedValue({ acknowledged: true }); - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -750,13 +677,7 @@ describe('AI Assistant Service', () => { .mockRejectedValueOnce(new EsErrors.ConnectionError('foo')) .mockRejectedValueOnce(new EsErrors.TimeoutError('timeout')) .mockResolvedValue({ acknowledged: true }); - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', @@ -783,13 +704,7 @@ describe('AI Assistant Service', () => { hits: { hits: [], total: { value: 0 } }, }); - const assistantService = new AIAssistantService({ - logger, - elasticsearchClientPromise: Promise.resolve(clusterClient), - pluginStop$, - kibanaVersion: '8.8.0', - taskManager: taskManagerMock.createSetup(), - }); + const assistantService = new AIAssistantService(assistantServiceOpts); await retryUntil( 'AI Assistant service initialized', diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index af2d4b78a807a..fb6f0faf33c82 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -9,10 +9,11 @@ import { DataStreamSpacesAdapter, FieldMap } from '@kbn/data-stream-adapter'; import { DEFAULT_NAMESPACE_STRING } from '@kbn/core-saved-objects-utils-server'; import type { Logger, ElasticsearchClient } from '@kbn/core/server'; import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server'; +import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import { AuthenticatedUser } from '@kbn/security-plugin/server'; import { Subject } from 'rxjs'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; -import { AssistantResourceNames } from '../types'; +import { AssistantResourceNames, GetElser } from '../types'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; import { InitializationPromise, @@ -27,7 +28,7 @@ import { assistantAnonymizationFieldsFieldMap } from '../ai_assistant_data_clien import { AIAssistantDataClient } from '../ai_assistant_data_clients'; import { knowledgeBaseFieldMap } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; -import { knowledgeBaseIngestPipeline } from '../ai_assistant_data_clients/knowledge_base/ingest_pipeline'; +import { createGetElserId, createPipeline, pipelineExists } from './helpers'; const TOTAL_FIELDS_LIMIT = 2500; @@ -35,10 +36,11 @@ function getResourceName(resource: string) { return `.kibana-elastic-ai-assistant-${resource}`; } -interface AIAssistantServiceOpts { +export interface AIAssistantServiceOpts { logger: Logger; kibanaVersion: string; elasticsearchClientPromise: Promise; + ml: MlPluginSetup; taskManager: TaskManagerSetupContract; pluginStop$: Subject; } @@ -61,6 +63,7 @@ export class AIAssistantService { // Temporary 'feature flag' to determine if we should initialize the knowledge base, toggled when accessing data client private initializeKnowledgeBase: boolean = false; private isInitializing: boolean = false; + private getElserId: GetElser; private conversationsDataStream: DataStreamSpacesAdapter; private knowledgeBaseDataStream: DataStreamSpacesAdapter; private promptsDataStream: DataStreamSpacesAdapter; @@ -70,6 +73,7 @@ export class AIAssistantService { constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; + this.getElserId = createGetElserId(options.ml); this.conversationsDataStream = this.createDataStream({ resource: 'conversations', kibanaVersion: options.kibanaVersion, @@ -142,29 +146,22 @@ export class AIAssistantService { pluginStop$: this.options.pluginStop$, }); - // TODO: Add generic ingest pipeline support to `kbn-data-stream-adapter` package? - // TODO: Inject `elserId` - let pipelineExists = false; - try { - const response = await esClient.ingest.getPipeline({ + // TODO: Pipeline creation is temporary as we'll be moving to semantic_text field once available in ES + const pipelineCreated = await pipelineExists({ + esClient, + id: this.resourceNames.pipelines.knowledgeBase, + }); + if (!pipelineCreated) { + this.options.logger.debug('Installing ingest pipeline'); + const response = await createPipeline({ + esClient, id: this.resourceNames.pipelines.knowledgeBase, + modelId: await this.getElserId(), }); - pipelineExists = Object.keys(response).length > 0; - } catch (e) { - // The GET /_ingest/pipeline/{pipelineId} API returns an empty object w/ 404 Not Found. - pipelineExists = false; - } - if (!pipelineExists) { - this.options.logger.info('Installing ingest pipeline'); - const response = await esClient.ingest.putPipeline( - knowledgeBaseIngestPipeline({ - id: this.resourceNames.pipelines.knowledgeBase, - modelId: '.elser_model_2', - }) - ); - this.options.logger.info(`Installed ingest pipeline: ${response.acknowledged}`); + + this.options.logger.debug(`Installed ingest pipeline: ${response}`); } else { - this.options.logger.info('Ingest pipeline already exists'); + this.options.logger.debug('Ingest pipeline already exists'); } } @@ -308,6 +305,8 @@ export class AIAssistantService { kibanaVersion: this.options.kibanaVersion, indexPatternsResourceName: this.resourceNames.aliases.conversations, currentUser: opts.currentUser, + ml: this.options.ml, + getElserId: this.getElserId, }); } diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 63d4c362044e1..7a092928ce99a 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -24,6 +24,7 @@ import { RequestContextFactory } from './routes/request_context_factory'; import { PLUGIN_ID } from '../common/constants'; import { registerRoutes } from './routes/register_routes'; import { appContextService } from './services/app_context'; +import { createGetElserId } from './ai_assistant_service/helpers'; export class ElasticAssistantPlugin implements @@ -53,6 +54,7 @@ export class ElasticAssistantPlugin this.assistantService = new AIAssistantService({ logger: this.logger.get('service'), + ml: plugins.ml, taskManager: plugins.taskManager, kibanaVersion: this.kibanaVersion, elasticsearchClientPromise: core @@ -76,7 +78,8 @@ export class ElasticAssistantPlugin ); events.forEach((eventConfig) => core.analytics.registerEventType(eventConfig)); - registerRoutes(router, this.logger, plugins); + const getElserId = createGetElserId(plugins.ml); + registerRoutes(router, this.logger, getElserId); return { actions: plugins.actions, getRegisteredFeatures: (pluginName: string) => { diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 05c342f2deef8..ffcfd63059bba 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -87,7 +87,7 @@ export const postKnowledgeBaseRoute = ( ); // - await knowledgeBaseDataClient?.setupKnowledgeBase({ elserId, esClient, esStore }); + await knowledgeBaseDataClient?.setupKnowledgeBase({ esStore }); return response.ok({ body: { success: true } }); } diff --git a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts index 5ff629d6a0b45..fc0e30f4a925c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts @@ -5,15 +5,10 @@ * 2.0. */ -import type { KibanaRequest, Logger, SavedObjectsClientContract } from '@kbn/core/server'; -import { once } from 'lodash/fp'; +import type { Logger } from '@kbn/core/server'; import { postAttackDiscoveryRoute } from './attack_discovery/post_attack_discovery'; -import { - ElasticAssistantPluginRouter, - ElasticAssistantPluginSetupDependencies, - GetElser, -} from '../types'; +import { ElasticAssistantPluginRouter, GetElser } from '../types'; import { createConversationRoute } from './user_conversations/create_route'; import { deleteConversationRoute } from './user_conversations/delete_route'; import { readConversationRoute } from './user_conversations/read_route'; @@ -36,7 +31,7 @@ import { findAnonymizationFieldsRoute } from './anonymization_fields/find_route' export const registerRoutes = ( router: ElasticAssistantPluginRouter, logger: Logger, - plugins: ElasticAssistantPluginSetupDependencies + getElserId: GetElser ) => { // Capabilities getCapabilitiesRoute(router); @@ -56,16 +51,6 @@ export const registerRoutes = ( // Knowledge Base deleteKnowledgeBaseRoute(router); - const getElserId: GetElser = once(async () => { - return ( - ( - await plugins.ml - // Force check to happen as internal user - .trainedModelsProvider({} as KibanaRequest, {} as SavedObjectsClientContract) - .getELSER() - ).model_id - ); - }); getKnowledgeBaseStatusRoute(router, getElserId); postKnowledgeBaseRoute(router, getElserId); From 49a9c5aa833d252f3c5572651f9fe421c6a383fa Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 9 May 2024 08:43:59 -0600 Subject: [PATCH 05/20] Resouce check note update --- .../knowledge_base/index.ts | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index bf8db77f16e72..ce920da4cbf83 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -79,17 +79,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { if (this.isInstallingElser) { return; } - // TODO: Before automatically installing ELSER in the background, we should perform the following deployment resource checks + // TODO: Before automatically installing ELSER in the background, we should perform deployment resource checks // Note: ESS only, as Serverless can always auto-install if `productTier === complete` - // 1. Deployment has ML Nodes with adequate free memory - // We can just auto-install, yay! - // 2. Deployment doesn't have adequate ML resources, and ML Autoscaling is disabled (or unavailable due to cluster health). - // Refer the user to the docs for further details - // 3. Deployment doesn't have adequate ML resources, but have ML Autoscaling enabled and scale limits are are NOT WITHIN the required resources. - // Again, refer the user to the docs - // 4. Deployment doesn't have adequate ML resources, but have ML Autoscaling enabled and scale limits ARE WITHIN the required resources. - // In this instance we could auto-install, but may have it behind a user action since deployment costs would change... - + // See ml-team issue for providing 'dry run' flag to perform these checks: https://github.com/elastic/ml-team/issues/1208 this.isInstallingElser = true; const elserId = await this.options.getElserId(); const isInstalled = await this.isModelInstalled(elserId); From 512806a2fe6a6a5cd214a71bcf98124b902f5094 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Thu, 9 May 2024 14:51:27 +0000 Subject: [PATCH 06/20] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/elastic_assistant/tsconfig.json | 2 ++ 1 file changed, 2 insertions(+) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 20146c45df5fa..4671ab75e544d 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -45,6 +45,8 @@ "@kbn/ml-response-stream", "@kbn/data-plugin", "@kbn/i18n", + "@kbn/core-security-common", + "@kbn/core-saved-objects-api-server", ], "exclude": [ "target/**/*", From dbc96e454c47024bb96ca3189e9d09ac7348977b Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Thu, 9 May 2024 17:40:30 -0600 Subject: [PATCH 07/20] Creates 'kb entries' OAS and standardize routes and constants --- .../kbn-elastic-assistant-common/constants.ts | 6 + .../schemas/common_attributes.schema.yaml | 44 +++++ .../impl/schemas/index.ts | 5 +- .../bulk_crud_knowledge_base_route.gen.ts | 117 ++++++++++++ ...bulk_crud_knowledge_base_route.schema.yaml | 175 ++++++++++++++++++ .../knowledge_base/common_attributes.gen.ts | 139 ++++++++++++++ .../common_attributes.schema.yaml | 105 +++++++++++ .../crud_knowledge_base_route.gen.ts | 91 +++++++++ .../crud_knowledge_base_route.schema.yaml | 122 ++++++++++++ .../impl/assistant/api/index.tsx | 13 +- .../elastic_assistant/common/constants.ts | 3 - .../server/__mocks__/request.ts | 9 +- .../field_maps_configuration.ts | 40 ++++ .../knowledge_base/types.ts | 29 +++ .../knowledge_base/delete_knowledge_base.ts | 8 +- .../knowledge_base/entries/create_route.ts | 78 ++++++++ .../get_knowledge_base_status.ts | 4 +- .../knowledge_base/post_knowledge_base.ts | 4 +- 18 files changed, 973 insertions(+), 19 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts diff --git a/x-pack/packages/kbn-elastic-assistant-common/constants.ts b/x-pack/packages/kbn-elastic-assistant-common/constants.ts index 67a20011dffd9..bc8f4e3ab9db3 100755 --- a/x-pack/packages/kbn-elastic-assistant-common/constants.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/constants.ts @@ -8,6 +8,7 @@ export const ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION = '1'; export const ELASTIC_AI_ASSISTANT_URL = '/api/elastic_assistant'; +export const ELASTIC_AI_ASSISTANT_INTERNAL_URL = '/internal/elastic_assistant'; export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL = `${ELASTIC_AI_ASSISTANT_URL}/current_user/conversations`; export const ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID = `${ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL}/{id}`; @@ -23,3 +24,8 @@ export const ELASTIC_AI_ASSISTANT_PROMPTS_URL_FIND = `${ELASTIC_AI_ASSISTANT_PRO export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL = `${ELASTIC_AI_ASSISTANT_URL}/anonymization_fields`; export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL}/_bulk_action`; export const ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_FIND = `${ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL}/_find`; + +// TODO: Update existing 'status' endpoint to take resource as query param as to not conflict with 'entries' +export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/knowledge_base/{resource?}`; +export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL = `${ELASTIC_AI_ASSISTANT_URL}/knowledge_base/entries`; +export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION = `${ELASTIC_AI_ASSISTANT_URL}/knowledge_base/_bulk_action`; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml new file mode 100644 index 0000000000000..fc35e62e28dba --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml @@ -0,0 +1,44 @@ +openapi: 3.0.0 +info: + title: Common Elastic AI Assistant Attributes + version: 'not applicable' +paths: {} +components: + x-codegen-enabled: true + schemas: + NonEmptyString: + type: string + pattern: ^(?! *$).+$ + minLength: 1 + description: A string that is not empty and does not contain only whitespace + + UUID: + type: string + format: uuid + description: A universally unique identifier + + User: + type: object + description: Could be any string, not necessarily a UUID + properties: + id: + type: string + description: User id + name: + type: string + description: User name + + ErrorSchema: + type: object + required: + - statusCode + - error + - message + additionalProperties: false + properties: + statusCode: + type: number + error: + type: string + message: + type: string diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts index 24d484bdd06c6..6440ddb0931bc 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts @@ -37,5 +37,8 @@ export * from './conversations/find_conversations_route.gen'; // Actions Connector Schemas export * from './actions_connector/post_actions_connector_execute_route.gen'; -// KB Schemas +// Knowledge Base Schemas export * from './knowledge_base/crud_kb_route.gen'; +export * from './knowledge_base/bulk_crud_knowledge_base_route.gen'; +// export * from './knowledge_base/common_attributes.gen'; +export * from './knowledge_base/crud_knowledge_base_route.gen'; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.gen.ts new file mode 100644 index 0000000000000..9ff055e656fe3 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.gen.ts @@ -0,0 +1,117 @@ +/* + * 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 { z } from 'zod'; + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Bulk Knowledge Base Actions API endpoint + * version: 2023-10-31 + */ + +import { + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryUpdateProps, + KnowledgeBaseEntryResponse, +} from './common_attributes.gen'; + +export type KnowledgeBaseEntryBulkActionSkipReason = z.infer< + typeof KnowledgeBaseEntryBulkActionSkipReason +>; +export const KnowledgeBaseEntryBulkActionSkipReason = z.literal( + 'KNOWLEDGE_BASE_ENTRY_NOT_MODIFIED' +); + +export type KnowledgeBaseEntryBulkActionSkipResult = z.infer< + typeof KnowledgeBaseEntryBulkActionSkipResult +>; +export const KnowledgeBaseEntryBulkActionSkipResult = z.object({ + id: z.string(), + name: z.string().optional(), + skip_reason: KnowledgeBaseEntryBulkActionSkipReason, +}); + +export type KnowledgeBaseEntryDetailsInError = z.infer; +export const KnowledgeBaseEntryDetailsInError = z.object({ + id: z.string(), + name: z.string().optional(), +}); + +export type NormalizedKnowledgeBaseEntryError = z.infer; +export const NormalizedKnowledgeBaseEntryError = z.object({ + message: z.string(), + statusCode: z.number().int(), + err_code: z.string().optional(), + knowledgeBaseEntries: z.array(KnowledgeBaseEntryDetailsInError), +}); + +export type KnowledgeBaseEntryBulkCrudActionResults = z.infer< + typeof KnowledgeBaseEntryBulkCrudActionResults +>; +export const KnowledgeBaseEntryBulkCrudActionResults = z.object({ + updated: z.array(KnowledgeBaseEntryResponse), + created: z.array(KnowledgeBaseEntryResponse), + deleted: z.array(z.string()), + skipped: z.array(KnowledgeBaseEntryBulkActionSkipResult), +}); + +export type KnowledgeBaseEntryBulkCrudActionSummary = z.infer< + typeof KnowledgeBaseEntryBulkCrudActionSummary +>; +export const KnowledgeBaseEntryBulkCrudActionSummary = z.object({ + failed: z.number().int(), + skipped: z.number().int(), + succeeded: z.number().int(), + total: z.number().int(), +}); + +export type KnowledgeBaseEntryBulkCrudActionResponse = z.infer< + typeof KnowledgeBaseEntryBulkCrudActionResponse +>; +export const KnowledgeBaseEntryBulkCrudActionResponse = z.object({ + success: z.boolean().optional(), + statusCode: z.number().int().optional(), + message: z.string().optional(), + knowledgeBaseEntriesCount: z.number().int().optional(), + attributes: z.object({ + results: KnowledgeBaseEntryBulkCrudActionResults, + summary: KnowledgeBaseEntryBulkCrudActionSummary, + errors: z.array(NormalizedKnowledgeBaseEntryError).optional(), + }), +}); + +export type KnowledgeBaseEntryBulkActionBase = z.infer; +export const KnowledgeBaseEntryBulkActionBase = z.object({ + /** + * Query to filter Knowledge Base Entries + */ + query: z.string().optional(), + /** + * Array of Knowledge base Entry IDs + */ + ids: z.array(z.string()).min(1).optional(), +}); + +export type PerformKnowledgeBaseEntryBulkActionRequestBody = z.infer< + typeof PerformKnowledgeBaseEntryBulkActionRequestBody +>; +export const PerformKnowledgeBaseEntryBulkActionRequestBody = z.object({ + delete: KnowledgeBaseEntryBulkActionBase.optional(), + create: z.array(KnowledgeBaseEntryCreateProps).optional(), + update: z.array(KnowledgeBaseEntryUpdateProps).optional(), +}); +export type PerformKnowledgeBaseEntryBulkActionRequestBodyInput = z.input< + typeof PerformKnowledgeBaseEntryBulkActionRequestBody +>; + +export type PerformKnowledgeBaseEntryBulkActionResponse = z.infer< + typeof PerformKnowledgeBaseEntryBulkActionResponse +>; +export const PerformKnowledgeBaseEntryBulkActionResponse = KnowledgeBaseEntryBulkCrudActionResponse; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml new file mode 100644 index 0000000000000..b52be7cd8055b --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml @@ -0,0 +1,175 @@ +openapi: 3.0.0 +info: + title: Bulk Knowledge Base Actions API endpoint + version: '2023-10-31' +paths: + /api/elastic_assistant/knowledge_base/entries/_bulk_action: + post: + operationId: PerformKnowledgeBaseEntryBulkAction + x-codegen-enabled: true + summary: Applies a bulk action to multiple Knowledge Base Entries + description: The bulk action is applied to all Knowledge Base Entries that match the filter or to the list of Knowledge Base Entries by their IDs + tags: + - Knowledge Base Entries Bulk API + requestBody: + content: + application/json: + schema: + type: object + properties: + delete: + $ref: '#/components/schemas/KnowledgeBaseEntryBulkActionBase' + create: + type: array + items: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryCreateProps' + update: + type: array + items: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryUpdateProps' + responses: + 200: + description: Successful bulk operation request + content: + application/json: + schema: + $ref: '#/components/schemas/KnowledgeBaseEntryBulkCrudActionResponse' + 400: + description: Generic Error + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/ErrorSchema' + +components: + schemas: + KnowledgeBaseEntryBulkActionSkipReason: + type: string + enum: + - KNOWLEDGE_BASE_ENTRY_NOT_MODIFIED + + KnowledgeBaseEntryBulkActionSkipResult: + type: object + properties: + id: + type: string + name: + type: string + skip_reason: + $ref: '#/components/schemas/KnowledgeBaseEntryBulkActionSkipReason' + required: + - id + - skip_reason + + KnowledgeBaseEntryDetailsInError: + type: object + properties: + id: + type: string + name: + type: string + required: + - id + + NormalizedKnowledgeBaseEntryError: + type: object + properties: + message: + type: string + statusCode: + type: integer + err_code: + type: string + knowledgeBaseEntries: + type: array + items: + $ref: '#/components/schemas/KnowledgeBaseEntryDetailsInError' + required: + - message + - statusCode + - knowledgeBaseEntries + + KnowledgeBaseEntryBulkCrudActionResults: + type: object + properties: + updated: + type: array + items: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse' + created: + type: array + items: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse' + deleted: + type: array + items: + type: string + skipped: + type: array + items: + $ref: '#/components/schemas/KnowledgeBaseEntryBulkActionSkipResult' + required: + - updated + - created + - deleted + - skipped + + KnowledgeBaseEntryBulkCrudActionSummary: + type: object + properties: + failed: + type: integer + skipped: + type: integer + succeeded: + type: integer + total: + type: integer + required: + - failed + - skipped + - succeeded + - total + + KnowledgeBaseEntryBulkCrudActionResponse: + type: object + properties: + success: + type: boolean + statusCode: + type: integer + message: + type: string + knowledgeBaseEntriesCount: + type: integer + attributes: + type: object + properties: + results: + $ref: '#/components/schemas/KnowledgeBaseEntryBulkCrudActionResults' + summary: + $ref: '#/components/schemas/KnowledgeBaseEntryBulkCrudActionSummary' + errors: + type: array + items: + $ref: '#/components/schemas/NormalizedKnowledgeBaseEntryError' + required: + - results + - summary + required: + - attributes + + + KnowledgeBaseEntryBulkActionBase: + x-inline: true + type: object + properties: + query: + type: string + description: Query to filter Knowledge Base Entries + ids: + type: array + description: Array of Knowledge base Entry IDs + minItems: 1 + items: + type: string diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts new file mode 100644 index 0000000000000..6f298e4d360ac --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts @@ -0,0 +1,139 @@ +/* + * 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 { z } from 'zod'; + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Common Knowledge Base Attributes + * version: not applicable + */ + +/** + * A string that is not empty and does not contain only whitespace + */ +export type NonEmptyString = z.infer; +export const NonEmptyString = z + .string() + .min(1) + .regex(/^(?! *$).+$/); + +/** + * A universally unique identifier + */ +export type UUID = z.infer; +export const UUID = z.string().uuid(); + +/** + * Could be any string, not necessarily a UUID + */ +export type User = z.infer; +export const User = z.object({ + /** + * User id + */ + id: z.string().optional(), + /** + * User name + */ + name: z.string().optional(), +}); + +/** + * Metadata about an Knowledge Base Entry + */ +export type Metadata = z.infer; +export const Metadata = z.object({ + /** + * Knowledge Base resource name + */ + kbResource: z.string(), + /** + * Original text content + */ + source: z.string(), + /** + * Whether or not this resource should always be included + */ + required: z.boolean(), +}); + +/** + * Object containing Metadata.source embeddings and modelId used to create the embeddings + */ +export type Vector = z.infer; +export const Vector = z.object({ + /** + * ID of the model used to create the embeddings + */ + modelId: z.string(), + /** + * Tokens with their corresponding values + */ + tokens: z.object({}).catchall(z.number()), +}); + +export type ErrorSchema = z.infer; +export const ErrorSchema = z + .object({ + statusCode: z.number(), + error: z.string(), + message: z.string(), + }) + .strict(); + +export type KnowledgeBaseEntryResponse = z.infer; +export const KnowledgeBaseEntryResponse = z.object({ + timestamp: NonEmptyString.optional(), + id: z.union([UUID, NonEmptyString]), + /** + * Time the Knowledge Base Entry was created + */ + createdAt: z.string(), + /** + * User who created the Knowledge Base Entry + */ + createdBy: z.string().optional(), + /** + * Time the Knowledge Base Entry was last updated + */ + updatedAt: z.string().optional(), + /** + * User who last updated the Knowledge Base Entry + */ + updatedBy: z.string().optional(), + users: z.array(User), + /** + * Metadata about the Knowledge Base Entry + */ + metadata: Metadata.optional(), + /** + * Kibana space + */ + namespace: z.string(), + vector: Vector.optional(), +}); + +export type KnowledgeBaseEntryUpdateProps = z.infer; +export const KnowledgeBaseEntryUpdateProps = z.object({ + id: z.union([UUID, NonEmptyString]), + /** + * Metadata about the Knowledge Base Entry + */ + metadata: Metadata.optional(), +}); + +export type KnowledgeBaseEntryCreateProps = z.infer; +export const KnowledgeBaseEntryCreateProps = z.object({ + /** + * Metadata about the Knowledge Base Entry + */ + metadata: Metadata, +}); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml new file mode 100644 index 0000000000000..8df97adc1a809 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml @@ -0,0 +1,105 @@ +openapi: 3.0.0 +info: + title: Common Knowledge Base Attributes + version: 'not applicable' +paths: {} +components: + x-codegen-enabled: true + schemas: + + Metadata: + type: object + description: Metadata about an Knowledge Base Entry + required: + - 'kbResource' + - 'source' + - 'required' + properties: + kbResource: + type: string + description: Knowledge Base resource name + source: + type: string + description: Original text content + required: + type: boolean + description: Whether or not this resource should always be included + + Vector: + type: object + description: Object containing Metadata.source embeddings and modelId used to create the embeddings + required: + - 'modelId' + - 'tokens' + properties: + modelId: + type: string + description: ID of the model used to create the embeddings + tokens: + type: object + additionalProperties: + type: number + description: Tokens with their corresponding values + + KnowledgeBaseEntryResponse: + type: object + required: + - id + - createdAt + - users + - namespace + properties: + 'timestamp': + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + id: + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + createdAt: + description: Time the Knowledge Base Entry was created + type: string + createdBy: + description: User who created the Knowledge Base Entry + type: string + updatedAt: + description: Time the Knowledge Base Entry was last updated + type: string + updatedBy: + description: User who last updated the Knowledge Base Entry + type: string + users: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/User' + metadata: + $ref: '#/components/schemas/Metadata' + description: Metadata about the Knowledge Base Entry + namespace: + type: string + description: Kibana space + vector: + $ref: '#/components/schemas/Vector' + + KnowledgeBaseEntryUpdateProps: + type: object + required: + - id + properties: + id: + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + metadata: + $ref: '#/components/schemas/Metadata' + description: Metadata about the Knowledge Base Entry + + KnowledgeBaseEntryCreateProps: + type: object + required: + - metadata + properties: + metadata: + $ref: '#/components/schemas/Metadata' + description: Metadata about the Knowledge Base Entry + + diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts new file mode 100644 index 0000000000000..92803577ae3f0 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts @@ -0,0 +1,91 @@ +/* + * 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 { z } from 'zod'; + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Manage Knowledge Base Entries API endpoint + * version: 2023-10-31 + */ + +import { + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryResponse, + UUID, + KnowledgeBaseEntryUpdateProps, +} from './common_attributes.gen'; + +export type CreateKnowledgeBaseEntryRequestBody = z.infer< + typeof CreateKnowledgeBaseEntryRequestBody +>; +export const CreateKnowledgeBaseEntryRequestBody = KnowledgeBaseEntryCreateProps; +export type CreateKnowledgeBaseEntryRequestBodyInput = z.input< + typeof CreateKnowledgeBaseEntryRequestBody +>; + +export type CreateKnowledgeBaseEntryResponse = z.infer; +export const CreateKnowledgeBaseEntryResponse = KnowledgeBaseEntryResponse; + +export type DeleteKnowledgeBaseEntryRequestParams = z.infer< + typeof DeleteKnowledgeBaseEntryRequestParams +>; +export const DeleteKnowledgeBaseEntryRequestParams = z.object({ + /** + * The Knowledge Base Entry's `id` value + */ + id: UUID, +}); +export type DeleteKnowledgeBaseEntryRequestParamsInput = z.input< + typeof DeleteKnowledgeBaseEntryRequestParams +>; + +export type DeleteKnowledgeBaseEntryResponse = z.infer; +export const DeleteKnowledgeBaseEntryResponse = KnowledgeBaseEntryResponse; + +export type ReadKnowledgeBaseEntryRequestParams = z.infer< + typeof ReadKnowledgeBaseEntryRequestParams +>; +export const ReadKnowledgeBaseEntryRequestParams = z.object({ + /** + * The Knowledge Base Entry's `id` value. + */ + id: UUID, +}); +export type ReadKnowledgeBaseEntryRequestParamsInput = z.input< + typeof ReadKnowledgeBaseEntryRequestParams +>; + +export type ReadKnowledgeBaseEntryResponse = z.infer; +export const ReadKnowledgeBaseEntryResponse = KnowledgeBaseEntryResponse; + +export type UpdateKnowledgeBaseEntryRequestParams = z.infer< + typeof UpdateKnowledgeBaseEntryRequestParams +>; +export const UpdateKnowledgeBaseEntryRequestParams = z.object({ + /** + * The Knowledge Base Entry's `id` value + */ + id: UUID, +}); +export type UpdateKnowledgeBaseEntryRequestParamsInput = z.input< + typeof UpdateKnowledgeBaseEntryRequestParams +>; + +export type UpdateKnowledgeBaseEntryRequestBody = z.infer< + typeof UpdateKnowledgeBaseEntryRequestBody +>; +export const UpdateKnowledgeBaseEntryRequestBody = KnowledgeBaseEntryUpdateProps; +export type UpdateKnowledgeBaseEntryRequestBodyInput = z.input< + typeof UpdateKnowledgeBaseEntryRequestBody +>; + +export type UpdateKnowledgeBaseEntryResponse = z.infer; +export const UpdateKnowledgeBaseEntryResponse = KnowledgeBaseEntryResponse; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml new file mode 100644 index 0000000000000..681581a227ae1 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml @@ -0,0 +1,122 @@ +openapi: 3.0.0 +info: + title: Manage Knowledge Base Entries API endpoint + version: '2023-10-31' +paths: + /api/elastic_assistant/knowledge_base/entries: + post: + operationId: CreateKnowledgeBaseEntry + x-codegen-enabled: true + description: Create a Knowledge Base Entry + summary: Create a Knowledge Base Entry + tags: + - Knowledge Base Entries API + requestBody: + required: true + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryCreateProps' + responses: + 200: + description: Successful request returning Knowledge Base Entries + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse' + 400: + description: Generic Error + content: + application/json: + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + + /api/elastic_assistant/knowledge_base/entries/{id}: + get: + operationId: ReadKnowledgeBaseEntry + x-codegen-enabled: true + description: Read a Knowledge Base Entry + summary: Read a Knowledge Base Entry + tags: + - Knowledge Base Entries API + parameters: + - name: id + in: path + required: true + description: The Knowledge Base Entry's `id` value. + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + responses: + 200: + description: Successful request returning a Knowledge Base Entry + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse' + 400: + description: Generic Error + content: + application/json: + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + put: + operationId: UpdateKnowledgeBaseEntry + x-codegen-enabled: true + description: Update a Knowledge Base Entry + summary: Update a Knowledge Base Entry + tags: + - Knowledge Base Entries API + parameters: + - name: id + in: path + required: true + description: The Knowledge Base Entry's `id` value + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + requestBody: + required: true + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryUpdateProps' + responses: + 200: + description: Successful request returning the updated Knowledge Base Entry + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse' + 400: + description: Generic Error + content: + application/json: + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + delete: + operationId: DeleteKnowledgeBaseEntry + x-codegen-enabled: true + description: Deletes a single Knowledge Base Entry using the `id` field + summary: Deletes a single Knowledge Base Entry using the `id` field + tags: + - Knowledge Base Entries API + parameters: + - name: id + in: path + required: true + description: The Knowledge Base Entry's `id` value + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + responses: + 200: + description: Successful request returning the deleted Knowledge Base Entry + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryResponse' + 400: + description: Generic Error + content: + application/json: + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx index ae465cdf27ba9..524f0ca95f99c 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx @@ -7,7 +7,12 @@ import { HttpSetup } from '@kbn/core/public'; import { IHttpFetchError } from '@kbn/core-http-browser'; -import { API_VERSIONS, ApiConfig, Replacements } from '@kbn/elastic-assistant-common'; +import { + API_VERSIONS, + ApiConfig, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, + Replacements, +} from '@kbn/elastic-assistant-common'; import { API_ERROR } from '../translations'; import { getOptionalRequestParams } from '../helpers'; import { TraceOptions } from '../types'; @@ -209,7 +214,7 @@ export const getKnowledgeBaseStatus = async ({ signal, }: GetKnowledgeBaseStatusParams): Promise => { try { - const path = `/internal/elastic_assistant/knowledge_base/${resource || ''}`; + const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); const response = await http.fetch(path, { method: 'GET', signal, @@ -248,7 +253,7 @@ export const postKnowledgeBase = async ({ signal, }: PostKnowledgeBaseParams): Promise => { try { - const path = `/internal/elastic_assistant/knowledge_base/${resource || ''}`; + const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); const response = await http.fetch(path, { method: 'POST', signal, @@ -287,7 +292,7 @@ export const deleteKnowledgeBase = async ({ signal, }: DeleteKnowledgeBaseParams): Promise => { try { - const path = `/internal/elastic_assistant/knowledge_base/${resource || ''}`; + const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); const response = await http.fetch(path, { method: 'DELETE', signal, diff --git a/x-pack/plugins/elastic_assistant/common/constants.ts b/x-pack/plugins/elastic_assistant/common/constants.ts index 77fd99c7d5eb6..5c5233eaaa6c8 100755 --- a/x-pack/plugins/elastic_assistant/common/constants.ts +++ b/x-pack/plugins/elastic_assistant/common/constants.ts @@ -15,9 +15,6 @@ export const POST_ACTIONS_CONNECTOR_EXECUTE = `${BASE_PATH}/actions/connector/{c // Attack discovery export const ATTACK_DISCOVERY = `${BASE_PATH}/attack_discovery`; -// Knowledge Base -export const KNOWLEDGE_BASE = `${BASE_PATH}/knowledge_base/{resource?}`; - // Model Evaluation export const EVALUATE = `${BASE_PATH}/evaluate`; diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts index 671cb1b2f0a88..0850938633322 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts @@ -5,7 +5,7 @@ * 2.0. */ import { httpServerMock } from '@kbn/core/server/mocks'; -import { CAPABILITIES, EVALUATE, KNOWLEDGE_BASE } from '../../common/constants'; +import { CAPABILITIES, EVALUATE } from '../../common/constants'; import { ConversationCreateProps, ConversationUpdateProps, @@ -16,6 +16,7 @@ import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID, ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID_MESSAGES, ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_FIND, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION, ELASTIC_AI_ASSISTANT_PROMPTS_URL_FIND, PostEvaluateRequestBodyInput, @@ -42,21 +43,21 @@ export const requestMock = { export const getGetKnowledgeBaseStatusRequest = (resource?: string) => requestMock.create({ method: 'get', - path: KNOWLEDGE_BASE, + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, query: { resource }, }); export const getPostKnowledgeBaseRequest = (resource?: string) => requestMock.create({ method: 'post', - path: KNOWLEDGE_BASE, + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, query: { resource }, }); export const getDeleteKnowledgeBaseRequest = (resource?: string) => requestMock.create({ method: 'delete', - path: KNOWLEDGE_BASE, + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, query: { resource }, }); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index e334427f9d9f5..2a28f2467087c 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -12,6 +12,46 @@ export const knowledgeBaseFieldMap: FieldMap = { array: false, required: false, }, + id: { + type: 'keyword', + array: false, + required: true, + }, + created_at: { + type: 'date', + array: false, + required: false, + }, + created_by: { + type: 'keyword', + array: false, + required: false, + }, + updated_at: { + type: 'date', + array: false, + required: false, + }, + updated_by: { + type: 'keyword', + array: false, + required: false, + }, + users: { + type: 'nested', + array: true, + required: false, + }, + 'users.id': { + type: 'keyword', + array: false, + required: true, + }, + 'users.name': { + type: 'keyword', + array: false, + required: false, + }, metadata: { type: 'object', array: false, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts new file mode 100644 index 0000000000000..c7dd9a527da5c --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts @@ -0,0 +1,29 @@ +/* + * 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. + */ + +export interface EsKnowledgeBaseEntrySchema { + '@timestamp': string; + id: string; + created_at: string; + created_by: string; + updated_at: string; + updated_by: string; + users?: Array<{ + id?: string; + name?: string; + }>; + metadata?: { + kbResource?: string; + source?: string; + required?: boolean; + }; + namespace: string; + vector?: { + tokens?: number[]; + model_id?: string; + }; +} diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts index 17b01fdb5f9aa..e3ba4f7b6156d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts @@ -8,7 +8,10 @@ import { IRouter, KibanaRequest } from '@kbn/core/server'; import { transformError } from '@kbn/securitysolution-es-utils'; -import { ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION } from '@kbn/elastic-assistant-common'; +import { + ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, +} from '@kbn/elastic-assistant-common'; import { DeleteKnowledgeBaseRequestParams, DeleteKnowledgeBaseResponse, @@ -16,7 +19,6 @@ import { import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantRequestHandlerContext } from '../../types'; -import { KNOWLEDGE_BASE } from '../../../common/constants'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants'; @@ -30,7 +32,7 @@ export const deleteKnowledgeBaseRoute = ( router.versioned .delete({ access: 'internal', - path: KNOWLEDGE_BASE, + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, options: { tags: ['access:elasticAssistant'], }, diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts new file mode 100644 index 0000000000000..434014adaaea4 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts @@ -0,0 +1,78 @@ +/* + * 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 type { IKibanaResponse } from '@kbn/core/server'; +import { transformError } from '@kbn/securitysolution-es-utils'; +import { + API_VERSIONS, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL, +} from '@kbn/elastic-assistant-common'; +import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; +import { + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryResponse, +} from '@kbn/elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen'; +import { ElasticAssistantPluginRouter } from '../../../types'; +import { buildResponse } from '../../utils'; +import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../../helpers'; + +export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRouter): void => { + router.versioned + .post({ + access: 'public', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL, + + options: { + tags: ['access:elasticAssistant'], + }, + }) + .addVersion( + { + version: API_VERSIONS.public.v1, + validate: { + request: { + body: buildRouteValidationWithZod(KnowledgeBaseEntryCreateProps), + }, + }, + }, + async (context, request, response): Promise> => { + const assistantResponse = buildResponse(response); + try { + const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); + const license = ctx.licensing.license; + if (!hasAIAssistantLicense(license)) { + return response.forbidden({ + body: { + message: UPGRADE_LICENSE_MESSAGE, + }, + }); + } + // const dataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient( + // false + // ); + const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); + if (authenticatedUser == null) { + return assistantResponse.error({ + body: `Authenticated user not found`, + statusCode: 401, + }); + } + + return assistantResponse.error({ + body: `knowledge base entry was not created`, + statusCode: 400, + }); + } catch (err) { + const error = transformError(err as Error); + return assistantResponse.error({ + body: error.message, + statusCode: error.statusCode, + }); + } + } + ); +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index cb13249423406..ae8496bcd54a0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -9,6 +9,7 @@ import { transformError } from '@kbn/securitysolution-es-utils'; import { ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, ReadKnowledgeBaseRequestParams, ReadKnowledgeBaseResponse, } from '@kbn/elastic-assistant-common'; @@ -17,7 +18,6 @@ import { KibanaRequest } from '@kbn/core/server'; import { getKbResource } from './get_kb_resource'; import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantPluginRouter, GetElser } from '../../types'; -import { KNOWLEDGE_BASE } from '../../../common/constants'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants'; @@ -34,7 +34,7 @@ export const getKnowledgeBaseStatusRoute = ( router.versioned .get({ access: 'internal', - path: KNOWLEDGE_BASE, + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, options: { tags: ['access:elasticAssistant'], }, diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index ffcfd63059bba..752de7f842b6d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -11,12 +11,12 @@ import { ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, CreateKnowledgeBaseRequestParams, CreateKnowledgeBaseResponse, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { IKibanaResponse, KibanaRequest } from '@kbn/core/server'; import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantPluginRouter, GetElser } from '../../types'; -import { KNOWLEDGE_BASE } from '../../../common/constants'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants'; import { getKbResource } from './get_kb_resource'; @@ -35,7 +35,7 @@ export const postKnowledgeBaseRoute = ( router.versioned .post({ access: 'internal', - path: KNOWLEDGE_BASE, + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, options: { tags: ['access:elasticAssistant'], }, From acdc0e717369894ed719bb7d58e6ad91b4e176ad Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 10 May 2024 11:08:04 -0600 Subject: [PATCH 08/20] Plumb through kbDataClient create and add kbDataClient support to EsStore --- ...ost_actions_connector_execute_route.gen.ts | 3 +- ...ctions_connector_execute_route.schema.yaml | 2 +- ...ulk_crud_anonymization_fields_route.gen.ts | 2 +- ...rud_anonymization_fields_route.schema.yaml | 5 +- .../impl/schemas/common_attributes.gen.ts | 47 +++++++++ .../schemas/common_attributes.schema.yaml | 14 --- .../conversations/common_attributes.gen.ts | 30 +----- .../common_attributes.schema.yaml | 39 ++------ .../crud_conversation_route.gen.ts | 2 +- .../crud_conversation_route.schema.yaml | 14 +-- .../impl/schemas/index.ts | 5 +- ...bulk_crud_knowledge_base_route.schema.yaml | 2 +- .../knowledge_base/common_attributes.gen.ts | 58 ++++------- .../common_attributes.schema.yaml | 27 +++++- .../crud_knowledge_base_route.gen.ts | 2 +- .../crud_knowledge_base_route.schema.yaml | 8 +- .../prompts/bulk_crud_prompts_route.gen.ts | 2 +- .../bulk_crud_prompts_route.schema.yaml | 7 +- .../create_knowledge_base_entry.ts | 83 ++++++++++++++++ .../field_maps_configuration.ts | 5 + .../get_knowledge_base_entry.ts | 80 ++++++++++++++++ .../knowledge_base/index.ts | 74 +++++++++++--- .../knowledge_base/transforms.ts | 96 +++++++++++++++++++ .../knowledge_base/types.ts | 35 ++++++- .../elasticsearch_store.ts | 48 +++++++++- .../knowledge_base/post_knowledge_base.ts | 29 ++++-- 26 files changed, 546 insertions(+), 173 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.gen.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.ts diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts index be9ed538cda7f..1405a3612291e 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts @@ -16,7 +16,8 @@ import { z } from 'zod'; * version: 1 */ -import { UUID, Replacements } from '../conversations/common_attributes.gen'; +import { UUID } from '../common_attributes.gen'; +import { Replacements } from '../conversations/common_attributes.gen'; export type ExecuteConnectorRequestParams = z.infer; export const ExecuteConnectorRequestParams = z.object({ diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml index d8bb1396746ed..1372619492eb8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml @@ -31,7 +31,7 @@ paths: - subAction properties: conversationId: - $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' message: type: string model: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts index bc2cb93bb30c3..2af794d04811d 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: 2023-10-31 */ -import { UUID, NonEmptyString } from '../conversations/common_attributes.gen'; +import { UUID, NonEmptyString } from '../common_attributes.gen'; export type BulkActionSkipReason = z.infer; export const BulkActionSkipReason = z.literal('ANONYMIZATION_FIELD_NOT_MODIFIED'); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml index e07e492b5fdbe..68a94c189f326 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml @@ -103,9 +103,9 @@ components: - field properties: id: - $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' 'timestamp': - $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' field: type: string allowed: @@ -232,4 +232,3 @@ components: type: boolean anonymized: type: boolean - \ No newline at end of file diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.gen.ts new file mode 100644 index 0000000000000..d98ee02af9ce4 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.gen.ts @@ -0,0 +1,47 @@ +/* + * 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 { z } from 'zod'; + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Common Elastic AI Assistant Attributes + * version: not applicable + */ + +/** + * A string that is not empty and does not contain only whitespace + */ +export type NonEmptyString = z.infer; +export const NonEmptyString = z + .string() + .min(1) + .regex(/^(?! *$).+$/); + +/** + * A universally unique identifier + */ +export type UUID = z.infer; +export const UUID = z.string().uuid(); + +/** + * Could be any string, not necessarily a UUID + */ +export type User = z.infer; +export const User = z.object({ + /** + * User id + */ + id: z.string().optional(), + /** + * User name + */ + name: z.string().optional(), +}); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml index fc35e62e28dba..5c580c52281ad 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/common_attributes.schema.yaml @@ -28,17 +28,3 @@ components: type: string description: User name - ErrorSchema: - type: object - required: - - statusCode - - error - - message - additionalProperties: false - properties: - statusCode: - type: number - error: - type: string - message: - type: string diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts index 808cf88fcec7c..37d3fc7f38098 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts @@ -16,35 +16,7 @@ import { z } from 'zod'; * version: not applicable */ -/** - * A string that is not empty and does not contain only whitespace - */ -export type NonEmptyString = z.infer; -export const NonEmptyString = z - .string() - .min(1) - .regex(/^(?! *$).+$/); - -/** - * A universally unique identifier - */ -export type UUID = z.infer; -export const UUID = z.string().uuid(); - -/** - * Could be any string, not necessarily a UUID - */ -export type User = z.infer; -export const User = z.object({ - /** - * User id. - */ - id: z.string().optional(), - /** - * User name. - */ - name: z.string().optional(), -}); +import { NonEmptyString, UUID, User } from '../common_attributes.gen'; /** * trace Data diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml index 3f2827b348004..74817a8b7d75b 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml @@ -6,27 +6,6 @@ paths: {} components: x-codegen-enabled: true schemas: - NonEmptyString: - type: string - pattern: ^(?! *$).+$ - minLength: 1 - description: A string that is not empty and does not contain only whitespace - - UUID: - type: string - format: uuid - description: A universally unique identifier - - User: - type: object - description: Could be any string, not necessarily a UUID - properties: - id: - type: string - description: User id. - name: - type: string - description: User name. TraceData: type: object @@ -97,7 +76,7 @@ components: $ref: '#/components/schemas/MessageRole' description: Message role. timestamp: - $ref: '#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' description: The timestamp message was sent or received. isError: type: boolean @@ -135,7 +114,7 @@ components: type: string description: Summary text of the conversation over time. timestamp: - $ref: '#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' description: The timestamp summary was updated. public: type: boolean @@ -151,7 +130,7 @@ components: additionalProperties: false properties: id: - $ref: '#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' error: type: object required: @@ -176,8 +155,8 @@ components: properties: id: oneOf: - - $ref: '#/components/schemas/UUID' - - $ref: '#/components/schemas/NonEmptyString' + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' title: type: string description: The conversation title. @@ -187,7 +166,7 @@ components: summary: $ref: '#/components/schemas/ConversationSummary' 'timestamp': - $ref: '#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' updatedAt: description: The last time conversation was updated. type: string @@ -199,7 +178,7 @@ components: users: type: array items: - $ref: '#/components/schemas/User' + $ref: '../common_attributes.schema.yaml#/components/schemas/User' messages: type: array items: @@ -225,8 +204,8 @@ components: properties: id: oneOf: - - $ref: '#/components/schemas/UUID' - - $ref: '#/components/schemas/NonEmptyString' + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' title: type: string description: The conversation title. diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts index ed5f1b8f057c6..e3a2be803db7a 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts @@ -19,10 +19,10 @@ import { z } from 'zod'; import { ConversationCreateProps, ConversationResponse, - UUID, ConversationUpdateProps, ConversationMessageCreateProps, } from './common_attributes.gen'; +import { UUID } from '../common_attributes.gen'; export type AppendConversationMessageRequestParams = z.infer< typeof AppendConversationMessageRequestParams diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml index a7f08659e76e3..5ae921be7a36f 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml @@ -37,7 +37,7 @@ paths: type: string message: type: string - + /api/elastic_assistant/conversations/{id}: get: operationId: ReadConversation @@ -52,7 +52,7 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: './common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' responses: 200: description: Indicates a successful call. @@ -86,7 +86,7 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: './common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' requestBody: required: true content: @@ -126,7 +126,7 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: './common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' responses: 200: description: Indicates a successful call. @@ -147,7 +147,7 @@ paths: type: string message: type: string - + /api/elastic_assistant/conversations/{id}/messages: post: operationId: AppendConversationMessage @@ -162,7 +162,7 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: './common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' requestBody: required: true content: @@ -188,4 +188,4 @@ paths: error: type: string message: - type: string \ No newline at end of file + type: string diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts index 6440ddb0931bc..c9c2d2a8be3c0 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts @@ -18,6 +18,9 @@ export const API_VERSIONS = { export const PUBLIC_API_ACCESS = 'public'; export const INTERNAL_API_ACCESS = 'internal'; +// Common Schemas +export * from './common_attributes.gen'; + // Attack discovery Schemas export * from './attack_discovery/post_attack_discovery_route.gen'; @@ -40,5 +43,5 @@ export * from './actions_connector/post_actions_connector_execute_route.gen'; // Knowledge Base Schemas export * from './knowledge_base/crud_kb_route.gen'; export * from './knowledge_base/bulk_crud_knowledge_base_route.gen'; -// export * from './knowledge_base/common_attributes.gen'; +export * from './knowledge_base/common_attributes.gen'; export * from './knowledge_base/crud_knowledge_base_route.gen'; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml index b52be7cd8055b..f8a2ee49d399a 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml @@ -39,7 +39,7 @@ paths: content: application/json: schema: - $ref: './common_attributes.schema.yaml#/components/schemas/ErrorSchema' + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryErrorSchema' components: schemas: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts index 6f298e4d360ac..805b9a06ef074 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts @@ -16,35 +16,16 @@ import { z } from 'zod'; * version: not applicable */ -/** - * A string that is not empty and does not contain only whitespace - */ -export type NonEmptyString = z.infer; -export const NonEmptyString = z - .string() - .min(1) - .regex(/^(?! *$).+$/); - -/** - * A universally unique identifier - */ -export type UUID = z.infer; -export const UUID = z.string().uuid(); +import { NonEmptyString, UUID, User } from '../common_attributes.gen'; -/** - * Could be any string, not necessarily a UUID - */ -export type User = z.infer; -export const User = z.object({ - /** - * User id - */ - id: z.string().optional(), - /** - * User name - */ - name: z.string().optional(), -}); +export type KnowledgeBaseEntryErrorSchema = z.infer; +export const KnowledgeBaseEntryErrorSchema = z + .object({ + statusCode: z.number(), + error: z.string(), + message: z.string(), + }) + .strict(); /** * Metadata about an Knowledge Base Entry @@ -56,7 +37,7 @@ export const Metadata = z.object({ */ kbResource: z.string(), /** - * Original text content + * Original text content source */ source: z.string(), /** @@ -66,7 +47,7 @@ export const Metadata = z.object({ }); /** - * Object containing Metadata.source embeddings and modelId used to create the embeddings + * Object containing Knowledge Base Entry text embeddings and modelId used to create the embeddings */ export type Vector = z.infer; export const Vector = z.object({ @@ -80,15 +61,6 @@ export const Vector = z.object({ tokens: z.object({}).catchall(z.number()), }); -export type ErrorSchema = z.infer; -export const ErrorSchema = z - .object({ - statusCode: z.number(), - error: z.string(), - message: z.string(), - }) - .strict(); - export type KnowledgeBaseEntryResponse = z.infer; export const KnowledgeBaseEntryResponse = z.object({ timestamp: NonEmptyString.optional(), @@ -118,6 +90,10 @@ export const KnowledgeBaseEntryResponse = z.object({ * Kibana space */ namespace: z.string(), + /** + * Knowledge Base Entry content + */ + text: z.string(), vector: Vector.optional(), }); @@ -136,4 +112,8 @@ export const KnowledgeBaseEntryCreateProps = z.object({ * Metadata about the Knowledge Base Entry */ metadata: Metadata, + /** + * Knowledge Base Entry content + */ + text: z.string(), }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml index 8df97adc1a809..763a2acf34680 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml @@ -7,6 +7,21 @@ components: x-codegen-enabled: true schemas: + KnowledgeBaseEntryErrorSchema: + type: object + required: + - statusCode + - error + - message + additionalProperties: false + properties: + statusCode: + type: number + error: + type: string + message: + type: string + Metadata: type: object description: Metadata about an Knowledge Base Entry @@ -20,14 +35,14 @@ components: description: Knowledge Base resource name source: type: string - description: Original text content + description: Original text content source required: type: boolean description: Whether or not this resource should always be included Vector: type: object - description: Object containing Metadata.source embeddings and modelId used to create the embeddings + description: Object containing Knowledge Base Entry text embeddings and modelId used to create the embeddings required: - 'modelId' - 'tokens' @@ -48,6 +63,7 @@ components: - createdAt - users - namespace + - text properties: 'timestamp': $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' @@ -77,6 +93,9 @@ components: namespace: type: string description: Kibana space + text: + type: string + description: Knowledge Base Entry content vector: $ref: '#/components/schemas/Vector' @@ -97,9 +116,13 @@ components: type: object required: - metadata + - text properties: metadata: $ref: '#/components/schemas/Metadata' description: Metadata about the Knowledge Base Entry + text: + type: string + description: Knowledge Base Entry content diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts index 92803577ae3f0..d69ae81ea13e8 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts @@ -19,9 +19,9 @@ import { z } from 'zod'; import { KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, - UUID, KnowledgeBaseEntryUpdateProps, } from './common_attributes.gen'; +import { UUID } from '../common_attributes.gen'; export type CreateKnowledgeBaseEntryRequestBody = z.infer< typeof CreateKnowledgeBaseEntryRequestBody diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml index 681581a227ae1..e8fbe9eac5f64 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml @@ -29,7 +29,7 @@ paths: content: application/json: schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryErrorSchema' /api/elastic_assistant/knowledge_base/entries/{id}: get: @@ -58,7 +58,7 @@ paths: content: application/json: schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryErrorSchema' put: operationId: UpdateKnowledgeBaseEntry x-codegen-enabled: true @@ -91,7 +91,7 @@ paths: content: application/json: schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryErrorSchema' delete: operationId: DeleteKnowledgeBaseEntry x-codegen-enabled: true @@ -118,5 +118,5 @@ paths: content: application/json: schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/ErrorSchema' + $ref: './common_attributes.schema.yaml#/components/schemas/KnowledgeBaseEntryErrorSchema' diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts index ce54c6a41fecc..9c5c1dd4f31ff 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: 2023-10-31 */ -import { UUID, NonEmptyString, User } from '../conversations/common_attributes.gen'; +import { UUID, NonEmptyString, User } from '../common_attributes.gen'; export type BulkActionSkipReason = z.infer; export const BulkActionSkipReason = z.literal('PROMPT_FIELD_NOT_MODIFIED'); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml index 2f6a419d9bf2c..f8938e6a653b1 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml @@ -105,9 +105,9 @@ components: - content properties: id: - $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/UUID' + $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' 'timestamp': - $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' name: type: string promptType: @@ -131,7 +131,7 @@ components: users: type: array items: - $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/User' + $ref: '../common_attributes.schema.yaml#/components/schemas/User' namespace: type: string description: Kibana space @@ -256,4 +256,3 @@ components: type: boolean isShared: type: boolean - \ No newline at end of file diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts new file mode 100644 index 0000000000000..5430bf597ebe7 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts @@ -0,0 +1,83 @@ +/* + * 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 { v4 as uuidv4 } from 'uuid'; +import { ElasticsearchClient, Logger } from '@kbn/core/server'; + +import { + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryResponse, +} from '@kbn/elastic-assistant-common'; +import { AuthenticatedUser } from '@kbn/security-plugin-types-common'; +import { getKnowledgeBaseEntry } from './get_knowledge_base_entry'; +import { CreateKnowledgeBaseEntrySchema } from './types'; + +export interface CreateKnowledgeBaseEntryParams { + esClient: ElasticsearchClient; + logger: Logger; + knowledgeBaseIndex: string; + spaceId: string; + user: AuthenticatedUser; + knowledgeBaseEntry: KnowledgeBaseEntryCreateProps; +} + +export const createKnowledgeBaseEntry = async ({ + esClient, + knowledgeBaseIndex, + spaceId, + user, + knowledgeBaseEntry, + logger, +}: CreateKnowledgeBaseEntryParams): Promise => { + const createdAt = new Date().toISOString(); + const body = transformToCreateSchema(createdAt, spaceId, user, knowledgeBaseEntry); + try { + const response = await esClient.create({ + body, + id: uuidv4(), + index: knowledgeBaseIndex, + refresh: 'wait_for', + }); + + return await getKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex, + id: response._id, + logger, + user, + }); + } catch (err) { + logger.error( + `Error creating Knowledge Base Entry: ${err} with kbResource: ${knowledgeBaseEntry.metadata.kbResource}` + ); + throw err; + } +}; + +export const transformToCreateSchema = ( + createdAt: string, + spaceId: string, + user: AuthenticatedUser, + { metadata, text }: KnowledgeBaseEntryCreateProps +): CreateKnowledgeBaseEntrySchema => { + return { + '@timestamp': createdAt, + created_at: createdAt, + created_by: user.profile_uid ?? 'unknown', + updated_at: createdAt, + updated_by: user.profile_uid ?? 'unknown', + users: [ + { + id: user.profile_uid, + name: user.username, + }, + ], + namespace: spaceId, + metadata, + text, + }; +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index 2a28f2467087c..16b5b35a9b2bb 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -72,6 +72,11 @@ export const knowledgeBaseFieldMap: FieldMap = { array: false, required: false, }, + text: { + type: 'text', + array: false, + required: true, + }, vector: { type: 'object', array: false, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts new file mode 100644 index 0000000000000..ee39a69e50d03 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts @@ -0,0 +1,80 @@ +/* + * 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 { ElasticsearchClient, Logger } from '@kbn/core/server'; +import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; +import { AuthenticatedUser } from '@kbn/security-plugin/common'; +import { EsKnowledgeBaseEntrySchema } from './types'; +import { transformESSearchToKnowledgeBaseEntry } from './transforms'; + +export interface GetKnowledgeBaseEntryParams { + esClient: ElasticsearchClient; + logger: Logger; + knowledgeBaseIndex: string; + id: string; + user?: AuthenticatedUser | null; +} + +export const getKnowledgeBaseEntry = async ({ + esClient, + logger, + knowledgeBaseIndex, + id, + user, +}: GetKnowledgeBaseEntryParams): Promise => { + const filterByUser = user + ? [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ] + : []; + try { + const response = await esClient.search({ + query: { + bool: { + must: [ + { + bool: { + should: [ + { + term: { + _id: id, + }, + }, + ], + }, + }, + ...filterByUser, + ], + }, + }, + _source: true, + ignore_unavailable: true, + index: knowledgeBaseIndex, + seq_no_primary_term: true, + }); + const conversation = transformESSearchToKnowledgeBaseEntry(response); + return conversation[0] ?? null; + } catch (err) { + logger.error(`Error fetching conversation: ${err} with id: ${id}`); + throw err; + } +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index ce920da4cbf83..b3e3271477ab3 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -12,17 +12,24 @@ import { import { AuthenticatedUser } from '@kbn/core-security-common'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import type { KibanaRequest } from '@kbn/core-http-server'; +import type { Document } from 'langchain/document'; import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; +import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; import { GetElser } from '../../types'; +import { transformToCreateSchema } from './create_knowledge_base_entry'; +import { EsKnowledgeBaseEntrySchema } from './types'; +import { transformESSearchToKnowledgeBaseEntry } from './transforms'; +import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ml: MlPluginSetup; getElserId: GetElser; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { + // @ts-ignore private isInstallingElser: boolean = false; constructor(public readonly options: KnowledgeBaseDataClientParams) { @@ -76,44 +83,83 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }: { esStore: ElasticsearchStore; }): Promise => { - if (this.isInstallingElser) { - return; - } // TODO: Before automatically installing ELSER in the background, we should perform deployment resource checks // Note: ESS only, as Serverless can always auto-install if `productTier === complete` // See ml-team issue for providing 'dry run' flag to perform these checks: https://github.com/elastic/ml-team/issues/1208 - this.isInstallingElser = true; const elserId = await this.options.getElserId(); const isInstalled = await this.isModelInstalled(elserId); if (isInstalled) { - this.options.logger.debug(`ELSER model '${elserId}' is already installed`); - this.options.logger.debug(`Loading KB docs!`); - const loadedKnowledgeBase = await loadESQL(esStore, this.options.logger); - this.options.logger.debug(`${loadedKnowledgeBase}`); this.isInstallingElser = false; + this.options.logger.debug(`ELSER model '${elserId}' is already installed`); + + const esqlExists = (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0; + if (esqlExists) { + this.options.logger.debug(`Kb docs already loaded!`); + return; + } else { + this.options.logger.debug(`Loading KB docs!`); + const loadedKnowledgeBase = await loadESQL(esStore, this.options.logger); + this.options.logger.debug(`${loadedKnowledgeBase}`); + this.isInstallingElser = false; + } return; } try { + this.isInstallingElser = true; const elserResponse = await this.options.ml .trainedModelsProvider({} as KibanaRequest, {} as SavedObjectsClientContract) .installElasticModel(elserId); - // const esClient = await this.options.elasticsearchClientPromise; - this.options.logger.debug(`elser response:\n: ${JSON.stringify(elserResponse, null, 2)}`); } catch (e) { this.options.logger.error(`Error setting up ELSER model: ${e.message}`); } }; - public addKnowledgeBaseResource = async ({ - document, + /** + * Adds LangChain Documents to the knowledge base + * + * @param documents + * @param authenticatedUser + */ + public addKnowledgeBaseDocuments = async ({ + documents, authenticatedUser, }: { - document: Document; + documents: Document[]; authenticatedUser: AuthenticatedUser; - }): Promise => {}; + }): Promise => { + const writer = await this.getWriter(); + const changedAt = new Date().toISOString(); + // @ts-ignore + const { errors, docs_created: docsCreated } = await writer.bulk({ + documentsToCreate: documents.map((doc) => + transformToCreateSchema(changedAt, this.spaceId, authenticatedUser, { + // TODO: Update the LangChain Document Metadata type extension + metadata: { + kbResource: doc.metadata.kbResourcer ?? 'unknown', + required: doc.metadata.required ?? false, + source: doc.metadata.source ?? 'unknown', + }, + text: doc.pageContent, + }) + ), + authenticatedUser, + }); + const created = + docsCreated.length > 0 + ? await this.findDocuments({ + page: 1, + perPage: 100, + filter: docsCreated.map((c) => `_id:${c}`).join(' OR '), + }) + : undefined; + this.options.logger.debug(`created: ${created}`); + this.options.logger.debug(`errors: ${errors}`); + + return created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : []; + }; } diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.ts new file mode 100644 index 0000000000000..f185c5ba8fdc2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.ts @@ -0,0 +1,96 @@ +/* + * 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 { estypes } from '@elastic/elasticsearch'; +import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; +import { EsKnowledgeBaseEntrySchema } from './types'; + +export const transformESSearchToKnowledgeBaseEntry = ( + response: estypes.SearchResponse +): KnowledgeBaseEntryResponse[] => { + return response.hits.hits + .filter((hit) => hit._source !== undefined) + .map((hit) => { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const kbEntrySchema = hit._source!; + const kbEntry: KnowledgeBaseEntryResponse = { + timestamp: kbEntrySchema['@timestamp'], + id: hit._id, + createdAt: kbEntrySchema.created_at, + createdBy: kbEntrySchema.created_by, + updatedAt: kbEntrySchema.updated_at, + updatedBy: kbEntrySchema.updated_by, + users: + kbEntrySchema.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + ...(kbEntrySchema.metadata + ? { + metadata: { + kbResource: kbEntrySchema.metadata.kbResource, + source: kbEntrySchema.metadata.source, + required: kbEntrySchema.metadata.required, + }, + } + : {}), + namespace: kbEntrySchema.namespace, + text: kbEntrySchema.text, + ...(kbEntrySchema.vector + ? { + vector: { + modelId: kbEntrySchema.vector.model_id, + tokens: kbEntrySchema.vector.tokens, + }, + } + : {}), + }; + + return kbEntry; + }); +}; + +export const transformESToKnowledgeBase = ( + response: EsKnowledgeBaseEntrySchema[] +): KnowledgeBaseEntryResponse[] => { + return response.map((kbEntrySchema) => { + const kbEntry: KnowledgeBaseEntryResponse = { + timestamp: kbEntrySchema['@timestamp'], + id: kbEntrySchema.id, + createdAt: kbEntrySchema.created_at, + createdBy: kbEntrySchema.created_by, + updatedAt: kbEntrySchema.updated_at, + updatedBy: kbEntrySchema.updated_by, + users: + kbEntrySchema.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + ...(kbEntrySchema.metadata + ? { + metadata: { + kbResource: kbEntrySchema.metadata.kbResource, + source: kbEntrySchema.metadata.source, + required: kbEntrySchema.metadata.required, + }, + } + : {}), + namespace: kbEntrySchema.namespace, + text: kbEntrySchema.text, + ...(kbEntrySchema.vector + ? { + vector: { + modelId: kbEntrySchema.vector.model_id, + tokens: kbEntrySchema.vector.tokens, + }, + } + : {}), + }; + + return kbEntry; + }); +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts index c7dd9a527da5c..b3180d80223ce 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts @@ -17,13 +17,38 @@ export interface EsKnowledgeBaseEntrySchema { name?: string; }>; metadata?: { - kbResource?: string; - source?: string; - required?: boolean; + kbResource: string; + source: string; + required: boolean; }; namespace: string; + text: string; vector?: { - tokens?: number[]; - model_id?: string; + tokens: Record; + model_id: string; + }; +} + +export interface CreateKnowledgeBaseEntrySchema { + '@timestamp'?: string; + id?: string | undefined; + created_at: string; + created_by: string; + updated_at: string; + updated_by: string; + users: Array<{ + id?: string; + name?: string; + }>; + metadata?: { + kbResource: string; + source: string; + required: boolean; + }; + namespace: string; + text: string; + vector?: { + tokens: Record; + model_id: string; }; } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts index d6987eae089e0..0fa9f46c3c28e 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts @@ -5,7 +5,12 @@ * 2.0. */ -import { type AnalyticsServiceSetup, ElasticsearchClient, Logger } from '@kbn/core/server'; +import { + type AnalyticsServiceSetup, + AuthenticatedUser, + ElasticsearchClient, + Logger, +} from '@kbn/core/server'; import { MappingTypeMapping, MlTrainedModelDeploymentNodesStats, @@ -34,6 +39,7 @@ import { KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT, KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT, } from '../../telemetry/event_based_telemetry'; +import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; interface CreatePipelineParams { id?: string; @@ -64,8 +70,9 @@ export const TERMS_QUERY_SIZE = 10000; export class ElasticsearchStore extends VectorStore { declare FilterType: QueryDslQueryContainer; - // Note: convert to { Client } from '@elastic/elasticsearch' for langchain contribution (removing Kibana dependency) + private readonly authenticatedUser: AuthenticatedUser | undefined; private readonly esClient: ElasticsearchClient; + private readonly kbDataClient: AIAssistantKnowledgeBaseDataClient | undefined; private readonly index: string; private readonly logger: Logger; private readonly telemetry: AnalyticsServiceSetup; @@ -82,15 +89,21 @@ export class ElasticsearchStore extends VectorStore { logger: Logger, telemetry: AnalyticsServiceSetup, model?: string, - kbResource?: string | undefined + kbResource?: string | undefined, + kbDataClient?: AIAssistantKnowledgeBaseDataClient, + authenticatedUser?: AuthenticatedUser ) { super(new ElasticsearchEmbeddings(logger), { esClient, index }); this.esClient = esClient; - this.index = index ?? KNOWLEDGE_BASE_INDEX_PATTERN; + this.index = kbDataClient + ? kbDataClient.options.indexPatternsResourceName + : index ?? KNOWLEDGE_BASE_INDEX_PATTERN; this.logger = logger; this.telemetry = telemetry; this.model = model ?? '.elser_model_2'; this.kbResource = kbResource ?? ESQL_RESOURCE; + this.kbDataClient = kbDataClient; + this.authenticatedUser = authenticatedUser; } /** @@ -105,6 +118,12 @@ export class ElasticsearchStore extends VectorStore { documents: Document[], options?: Record ): Promise => { + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled + // Once removed replace addDocuments() w/ addDocumentsViaDataClient() + if (this.kbDataClient != null) { + return this.addDocumentsViaDataClient(documents, options); + } + const pipelineExists = await this.pipelineExists(); if (!pipelineExists) { await this.createPipeline(); @@ -135,6 +154,27 @@ export class ElasticsearchStore extends VectorStore { } }; + addDocumentsViaDataClient = async ( + documents: Document[], + options?: Record + ): Promise => { + if (!this.kbDataClient || !this.authenticatedUser) { + this.logger.error('No kbDataClient or authenticatedUser provided'); + return []; + } + + try { + const response = await this.kbDataClient.addKnowledgeBaseDocuments({ + documents, + authenticatedUser: this.authenticatedUser, + }); + return response.map((doc) => doc.id); + } catch (e) { + this.logger.error(`Error loading data into KB\n ${e}`); + return []; + } + }; + /** * Add vectors to the store. Returns a list of document IDs. * diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 752de7f842b6d..c3d55756488d8 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -59,6 +59,15 @@ export const postKnowledgeBaseRoute = ( const logger = assistantContext.logger; const telemetry = assistantContext.telemetry; const elserId = await getElser(); + const core = await context.core; + const esClient = core.elasticsearch.client.asInternalUser; + const authenticatedUser = assistantContext.getCurrentUser(); + if (authenticatedUser == null) { + return response.custom({ + body: `Authenticated user not found`, + statusCode: 401, + }); + } const pluginName = getPluginNameFromRequest({ request, @@ -68,33 +77,33 @@ export const postKnowledgeBaseRoute = ( const enableKnowledgeBaseByDefault = assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled if (enableKnowledgeBaseByDefault) { - // Go ahead and do that now, right here's the place :) const knowledgeBaseDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); - // Temporarily get esClient for current user until `kibana_system` user has `inference_admin` role - // See https://github.com/elastic/elasticsearch/pull/108262 - const esClient = (await context.core).elasticsearch.client.asCurrentUser; + if (!knowledgeBaseDataClient) { + return response.custom({ body: { success: false }, statusCode: 500 }); + } - // + // Continue to use esStore for loading esql docs until `semantic_text` is available and we can test the new chunking strategy const esStore = new ElasticsearchStore( esClient, - `.kibana-elastic-ai-assistant-knowledge-base-default`, + knowledgeBaseDataClient.options.indexPatternsResourceName, logger, telemetry, elserId, - getKbResource(request) + getKbResource(request), + knowledgeBaseDataClient, + authenticatedUser ); // - await knowledgeBaseDataClient?.setupKnowledgeBase({ esStore }); + await knowledgeBaseDataClient.setupKnowledgeBase({ esStore }); return response.ok({ body: { success: true } }); } try { - const core = await context.core; - const esClient = core.elasticsearch.client.asInternalUser; const kbResource = getKbResource(request); const esStore = new ElasticsearchStore( esClient, From 49986eab76c5100546fdf633c18f69afd2fbbc1a Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Mon, 13 May 2024 14:49:17 -0600 Subject: [PATCH 09/20] Adds ingest pipeline to index template and fixes initial setup and status checks when FF is enabled --- .../knowledge_base/index.ts | 2 +- .../server/ai_assistant_service/index.ts | 23 +++++++++-- .../elasticsearch_store.ts | 3 +- .../get_knowledge_base_status.test.ts | 9 +++++ .../get_knowledge_base_status.ts | 40 ++++++++++++++++++- .../post_knowledge_base.test.ts | 9 +++++ .../knowledge_base/post_knowledge_base.ts | 3 +- 7 files changed, 81 insertions(+), 8 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index b3e3271477ab3..5747afb4703eb 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -153,7 +153,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { docsCreated.length > 0 ? await this.findDocuments({ page: 1, - perPage: 100, + perPage: 10000, filter: docsCreated.map((c) => `_id:${c}`).join(' OR '), }) : undefined; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index fb6f0faf33c82..00abfa8efeea9 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -122,6 +122,19 @@ export class AIAssistantService { newDataStream.setIndexTemplate({ name: this.resourceNames.indexTemplate[resource], componentTemplateRefs: [this.resourceNames.componentTemplate[resource]], + // Apply `default_pipeline` if pipeline exists for resource + ...(resource in this.resourceNames.pipelines + ? { + template: { + settings: { + 'index.default_pipeline': + this.resourceNames.pipelines[ + resource as keyof typeof this.resourceNames.pipelines + ], + }, + }, + } + : {}), }); return newDataStream; @@ -152,7 +165,9 @@ export class AIAssistantService { id: this.resourceNames.pipelines.knowledgeBase, }); if (!pipelineCreated) { - this.options.logger.debug('Installing ingest pipeline'); + this.options.logger.debug( + `Installing ingest pipeline - ${this.resourceNames.pipelines.knowledgeBase}` + ); const response = await createPipeline({ esClient, id: this.resourceNames.pipelines.knowledgeBase, @@ -161,7 +176,9 @@ export class AIAssistantService { this.options.logger.debug(`Installed ingest pipeline: ${response}`); } else { - this.options.logger.debug('Ingest pipeline already exists'); + this.options.logger.debug( + `Ingest pipeline already exists - ${this.resourceNames.pipelines.knowledgeBase}` + ); } } @@ -303,7 +320,7 @@ export class AIAssistantService { elasticsearchClientPromise: this.options.elasticsearchClientPromise, spaceId: opts.spaceId, kibanaVersion: this.options.kibanaVersion, - indexPatternsResourceName: this.resourceNames.aliases.conversations, + indexPatternsResourceName: this.resourceNames.aliases.knowledgeBase, currentUser: opts.currentUser, ml: this.options.ml, getElserId: this.getElserId, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts index 0fa9f46c3c28e..283df8c8f5079 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts @@ -70,6 +70,7 @@ export const TERMS_QUERY_SIZE = 10000; export class ElasticsearchStore extends VectorStore { declare FilterType: QueryDslQueryContainer; + // AuthenticatedUser is required for adding/retrieving documents as it is not encapsulated in the kbDataClient private readonly authenticatedUser: AuthenticatedUser | undefined; private readonly esClient: ElasticsearchClient; private readonly kbDataClient: AIAssistantKnowledgeBaseDataClient | undefined; @@ -96,7 +97,7 @@ export class ElasticsearchStore extends VectorStore { super(new ElasticsearchEmbeddings(logger), { esClient, index }); this.esClient = esClient; this.index = kbDataClient - ? kbDataClient.options.indexPatternsResourceName + ? kbDataClient.indexTemplateAndPattern.alias : index ?? KNOWLEDGE_BASE_INDEX_PATTERN; this.logger = logger; this.telemetry = telemetry; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts index e6811ac7cf3d2..ba8992d9c19bd 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.test.ts @@ -10,6 +10,7 @@ import { serverMock } from '../../__mocks__/server'; import { requestContextMock } from '../../__mocks__/request_context'; import { getGetKnowledgeBaseStatusRequest } from '../../__mocks__/request'; import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; +import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Get Knowledge Base Status Route', () => { let server: ReturnType; @@ -19,10 +20,18 @@ describe('Get Knowledge Base Status Route', () => { clients.core.elasticsearch.client = elasticsearchServiceMock.createScopedClusterClient(); const mockGetElser = jest.fn().mockResolvedValue('.elser_model_2'); + const mockUser = { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; beforeEach(() => { server = serverMock.create(); ({ context } = requestContextMock.createTools()); + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser); getKnowledgeBaseStatusRoute(server.router, mockGetElser); }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index ae8496bcd54a0..a2ff556ea64c9 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -20,6 +20,7 @@ import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantPluginRouter, GetElser } from '../../types'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; /** * Get the status of the Knowledge Base index, pipeline, and resources (collection of documents) @@ -58,7 +59,7 @@ export const getKnowledgeBaseStatusRoute = ( const esClient = (await context.core).elasticsearch.client.asInternalUser; const elserId = await getElser(); const kbResource = getKbResource(request); - const esStore = new ElasticsearchStore( + let esStore = new ElasticsearchStore( esClient, KNOWLEDGE_BASE_INDEX_PATTERN, logger, @@ -67,6 +68,43 @@ export const getKnowledgeBaseStatusRoute = ( kbResource ); + const authenticatedUser = assistantContext.getCurrentUser(); + if (authenticatedUser == null) { + return response.custom({ + body: `Authenticated user not found`, + statusCode: 401, + }); + } + + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const enableKnowledgeBaseByDefault = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; + + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled + if (enableKnowledgeBaseByDefault) { + const knowledgeBaseDataClient = + await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); + if (!knowledgeBaseDataClient) { + return response.custom({ body: { success: false }, statusCode: 500 }); + } + + // Use old status checks by overriding esStore to use kbDataClient + esStore = new ElasticsearchStore( + esClient, + knowledgeBaseDataClient.indexTemplateAndPattern.alias, + logger, + telemetry, + elserId, + kbResource, + knowledgeBaseDataClient, + authenticatedUser + ); + } + const indexExists = await esStore.indexExists(); const pipelineExists = await esStore.pipelineExists(); const modelExists = await esStore.isModelInstalled(elserId); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts index b9ab3bc4c6866..547923b5c0d17 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.test.ts @@ -10,6 +10,7 @@ import { serverMock } from '../../__mocks__/server'; import { requestContextMock } from '../../__mocks__/request_context'; import { getPostKnowledgeBaseRequest } from '../../__mocks__/request'; import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; +import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Post Knowledge Base Route', () => { let server: ReturnType; @@ -19,10 +20,18 @@ describe('Post Knowledge Base Route', () => { clients.core.elasticsearch.client = elasticsearchServiceMock.createScopedClusterClient(); const mockGetElser = jest.fn().mockResolvedValue('.elser_model_2'); + const mockUser = { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; beforeEach(() => { server = serverMock.create(); ({ context } = requestContextMock.createTools()); + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser); postKnowledgeBaseRoute(server.router, mockGetElser); }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index c3d55756488d8..510db4a669567 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -88,7 +88,7 @@ export const postKnowledgeBaseRoute = ( // Continue to use esStore for loading esql docs until `semantic_text` is available and we can test the new chunking strategy const esStore = new ElasticsearchStore( esClient, - knowledgeBaseDataClient.options.indexPatternsResourceName, + knowledgeBaseDataClient.indexTemplateAndPattern.alias, logger, telemetry, elserId, @@ -96,7 +96,6 @@ export const postKnowledgeBaseRoute = ( knowledgeBaseDataClient, authenticatedUser ); - // await knowledgeBaseDataClient.setupKnowledgeBase({ esStore }); From 5b3d1723f909de66579b0093982637397efe7d99 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 14 May 2024 11:39:04 -0600 Subject: [PATCH 10/20] Fix kb document loading and auto elser setup --- .../knowledge_base/helpers.ts | 16 ++ .../knowledge_base/index.ts | 163 ++++++++++++++---- .../lib/data_stream/documents_data_writer.ts | 4 +- .../knowledge_base/post_knowledge_base.ts | 3 +- 4 files changed, 149 insertions(+), 37 deletions(-) create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts new file mode 100644 index 0000000000000..dc7f64e1aeee1 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -0,0 +1,16 @@ +/* + * 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 { errors } from '@elastic/elasticsearch'; + +export const isModelAlreadyExistsError = (error: Error) => { + return ( + error instanceof errors.ResponseError && + (error.body.error.type === 'resource_not_found_exception' || + error.body.error.type === 'status_exception') + ); +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 5747afb4703eb..d707a3773fe04 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -15,6 +15,7 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import type { Document } from 'langchain/document'; import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; +import pRetry from 'p-retry'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; @@ -23,34 +24,105 @@ import { transformToCreateSchema } from './create_knowledge_base_entry'; import { EsKnowledgeBaseEntrySchema } from './types'; import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; +import { isModelAlreadyExistsError } from './helpers'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ml: MlPluginSetup; getElserId: GetElser; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { - // @ts-ignore - private isInstallingElser: boolean = false; + private setupInProgress: boolean = false; constructor(public readonly options: KnowledgeBaseDataClientParams) { super(options); } /** - * Checks if the provided model is installed (deployed and allocated) in Elasticsearch + * Downloads and installs ELSER model if not already installed + * + * @param soClient SavedObjectsClientContract for installing ELSER so that ML SO's are in sync + */ + private installModel = async ({ soClient }: { soClient: SavedObjectsClientContract }) => { + try { + this.setupInProgress = true; + const elserId = await this.options.getElserId(); + const elserResponse = await this.options.ml + // TODO: See about calling `installElasticModel()` as systemUser as to not require soClient + .trainedModelsProvider({} as KibanaRequest, soClient) + .installElasticModel(elserId); + + this.options.logger.debug( + `installElasticModel response:\n: ${JSON.stringify(elserResponse, null, 2)}` + ); + } catch (e) { + this.options.logger.error(`Error setting up ELSER model: ${e.message}`); + this.setupInProgress = false; + } + }; + + /** + * Returns whether ELSER is installed/ready to deploy * - * @param modelId ID of the model to check * @returns Promise indicating whether the model is installed */ - private isModelInstalled = async (modelId: string): Promise => { + private isModelInstalled = async (): Promise => { + const elserId = await this.options.getElserId(); + try { + const esClient = await this.options.elasticsearchClientPromise; + const getResponse = await esClient.ml.getTrainedModels({ + model_id: elserId, + include: 'definition_status', + }); + this.options.logger.debug( + `Model definition status:\n${JSON.stringify(getResponse.trained_model_configs[0])}` + ); + return Boolean(getResponse.trained_model_configs[0]?.fully_defined); + } catch (error) { + if (!isModelAlreadyExistsError(error)) { + this.options.logger.error(`Error deploying ELSER model '${elserId}'\n${error}`); + } + this.options.logger.debug(`Error deploying ELSER model '${elserId}', model already deployed`); + this.setupInProgress = false; + return false; + } + }; + + /** + * Deploy the ELSER model with default configuration + */ + private deployModel = async () => { + const elserId = await this.options.getElserId(); + try { + this.setupInProgress = true; + const esClient = await this.options.elasticsearchClientPromise; + await esClient.ml.startTrainedModelDeployment({ + model_id: elserId, + wait_for: 'fully_allocated', + }); + } catch (error) { + if (!isModelAlreadyExistsError(error)) { + this.options.logger.error(`Error deploying ELSER model '${elserId}'\n${error}`); + } + this.options.logger.debug(`Error deploying ELSER model '${elserId}', model already deployed`); + this.setupInProgress = false; + } + }; + + /** + * Checks if the provided model is deployed and allocated in Elasticsearch + * + * @returns Promise indicating whether the model is deployed + */ + private isModelDeployed = async (): Promise => { + const elserId = await this.options.getElserId(); const esClient = await this.options.elasticsearchClientPromise; try { const getResponse = await esClient.ml.getTrainedModelsStats({ - model_id: modelId, + model_id: elserId, }); - this.options.logger.debug(`modelId: ${modelId}`); + this.options.logger.debug(`modelId: ${elserId}`); // For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986 const isReadyESS = (stats: MlTrainedModelStats) => @@ -74,48 +146,71 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { /** * Downloads and deploys ELSER (if not already), then loads ES|QL docs * + * NOTE: Before automatically installing ELSER in the background, we should perform deployment resource checks + * Only necessary for ESS, as Serverless can always auto-install if `productTier === complete` + * See ml-team issue for providing 'dry run' flag to perform these checks: https://github.com/elastic/ml-team/issues/1208 + * * @param options - * @param options.esStore ElasticsearchStore for loading ES|QL docs + * @param options.esStore ElasticsearchStore for loading ES|QL docs via LangChain loaders + * @param options.soClient SavedObjectsClientContract for installing ELSER so that ML SO's are in sync + * * @returns Promise */ public setupKnowledgeBase = async ({ esStore, + request, + soClient, }: { esStore: ElasticsearchStore; + request: KibanaRequest; + soClient: SavedObjectsClientContract; }): Promise => { - // TODO: Before automatically installing ELSER in the background, we should perform deployment resource checks - // Note: ESS only, as Serverless can always auto-install if `productTier === complete` - // See ml-team issue for providing 'dry run' flag to perform these checks: https://github.com/elastic/ml-team/issues/1208 + if (this.setupInProgress) { + this.options.logger.debug('Setup already in progress'); + return; + } + const elserId = await this.options.getElserId(); - const isInstalled = await this.isModelInstalled(elserId); - if (isInstalled) { - this.isInstallingElser = false; - this.options.logger.debug(`ELSER model '${elserId}' is already installed`); + try { + const isInstalled = await this.isModelInstalled(); + if (!isInstalled) { + await this.installModel({ soClient }); + await pRetry( + async () => + (await this.isModelInstalled()) + ? Promise.resolve() + : Promise.reject(new Error('Model not installed')), + { minTimeout: 10000, retries: 10 } + ); + } else { + this.options.logger.debug(`ELSER model '${elserId}' is already installed`); + } - const esqlExists = (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0; - if (esqlExists) { - this.options.logger.debug(`Kb docs already loaded!`); - return; + const isDeployed = await this.isModelDeployed(); + if (!isDeployed) { + await this.deployModel(); + await pRetry( + async () => + (await this.isModelDeployed()) + ? Promise.resolve() + : Promise.reject(new Error('Model not deployed')), + { minTimeout: 2000, retries: 10 } + ); } else { + this.options.logger.debug(`ELSER model '${elserId}' is already deployed`); + } + + const kbDocsLoaded = (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0; + if (!kbDocsLoaded) { this.options.logger.debug(`Loading KB docs!`); const loadedKnowledgeBase = await loadESQL(esStore, this.options.logger); this.options.logger.debug(`${loadedKnowledgeBase}`); - this.isInstallingElser = false; + } else { + this.options.logger.debug(`KB docs already loaded!`); } - - return; - } - - try { - this.isInstallingElser = true; - const elserResponse = await this.options.ml - .trainedModelsProvider({} as KibanaRequest, {} as SavedObjectsClientContract) - .installElasticModel(elserId); - - this.options.logger.debug(`elser response:\n: ${JSON.stringify(elserResponse, null, 2)}`); } catch (e) { - this.options.logger.error(`Error setting up ELSER model: ${e.message}`); + this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); } }; @@ -157,8 +252,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { filter: docsCreated.map((c) => `_id:${c}`).join(' OR '), }) : undefined; - this.options.logger.debug(`created: ${created}`); - this.options.logger.debug(`errors: ${errors}`); + this.options.logger.debug(`created: ${created?.data.hits.hits.length ?? '0'}`); + this.options.logger.debug(`errors: ${JSON.stringify(errors, null, 2)}`); return created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : []; }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 83b4227275066..3cae3972c8bf6 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -5,7 +5,6 @@ * 2.0. */ -import { v4 as uuidV4 } from 'uuid'; import type { BulkOperationContainer, BulkOperationType, @@ -241,7 +240,8 @@ export class DocumentsDataWriter implements DocumentsDataWriter { ): Promise => { const documentCreateBody = params.documentsToCreate ? params.documentsToCreate.flatMap((document) => [ - { create: { _index: this.options.index, _id: uuidV4() } }, + // Do not pre-gen _id for bulk create operations to avoid `version_conflict_engine_exception` + { create: { _index: this.options.index } }, document, ]) : []; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 510db4a669567..0161c9bc2a939 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -61,6 +61,7 @@ export const postKnowledgeBaseRoute = ( const elserId = await getElser(); const core = await context.core; const esClient = core.elasticsearch.client.asInternalUser; + const soClient = core.savedObjects.getClient(); const authenticatedUser = assistantContext.getCurrentUser(); if (authenticatedUser == null) { return response.custom({ @@ -97,7 +98,7 @@ export const postKnowledgeBaseRoute = ( authenticatedUser ); - await knowledgeBaseDataClient.setupKnowledgeBase({ esStore }); + await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, request, soClient }); return response.ok({ body: { success: true } }); } From aa81b12647c3a94d7eb106bc00ad5e7d87fab330 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 14 May 2024 22:53:06 -0600 Subject: [PATCH 11/20] Fixes multi-step setup issue and wires up kbDataClient retrieval in post_actions_execute route --- .../knowledge_base/index.ts | 61 +++++++++---------- .../server/ai_assistant_service/index.ts | 2 +- .../lib/data_stream/documents_data_writer.ts | 1 + .../elasticsearch_store.ts | 9 ++- .../execute_custom_llm_chain/index.test.ts | 15 +++-- .../execute_custom_llm_chain/index.ts | 16 +---- .../executors/openai_functions_executor.ts | 16 +---- .../server/lib/langchain/executors/types.ts | 6 +- .../knowledge_base/post_knowledge_base.ts | 44 ++++++------- .../post_actions_connector_execute.test.ts | 2 + .../routes/post_actions_connector_execute.ts | 29 +++++++-- 11 files changed, 101 insertions(+), 100 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index d707a3773fe04..5aabf10e4882f 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -43,20 +43,16 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { * @param soClient SavedObjectsClientContract for installing ELSER so that ML SO's are in sync */ private installModel = async ({ soClient }: { soClient: SavedObjectsClientContract }) => { + const elserId = await this.options.getElserId(); + this.options.logger.debug(`Installing ELSER model '${elserId}'...`); + try { - this.setupInProgress = true; - const elserId = await this.options.getElserId(); - const elserResponse = await this.options.ml - // TODO: See about calling `installElasticModel()` as systemUser as to not require soClient + await this.options.ml + // TODO: Potentially plumb soClient through DataClient from pluginStart .trainedModelsProvider({} as KibanaRequest, soClient) .installElasticModel(elserId); - - this.options.logger.debug( - `installElasticModel response:\n: ${JSON.stringify(elserResponse, null, 2)}` - ); - } catch (e) { - this.options.logger.error(`Error setting up ELSER model: ${e.message}`); - this.setupInProgress = false; + } catch (error) { + this.options.logger.error(`Error installing ELSER model '${elserId}':\n${error}`); } }; @@ -65,24 +61,23 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { * * @returns Promise indicating whether the model is installed */ - private isModelInstalled = async (): Promise => { + public isModelInstalled = async (): Promise => { const elserId = await this.options.getElserId(); + this.options.logger.debug(`Checking if ELSER model '${elserId}' is installed...`); + try { const esClient = await this.options.elasticsearchClientPromise; const getResponse = await esClient.ml.getTrainedModels({ model_id: elserId, include: 'definition_status', }); - this.options.logger.debug( - `Model definition status:\n${JSON.stringify(getResponse.trained_model_configs[0])}` - ); return Boolean(getResponse.trained_model_configs[0]?.fully_defined); } catch (error) { if (!isModelAlreadyExistsError(error)) { - this.options.logger.error(`Error deploying ELSER model '${elserId}'\n${error}`); + this.options.logger.error( + `Error checking if ELSER model '${elserId}' is installed:\n${error}` + ); } - this.options.logger.debug(`Error deploying ELSER model '${elserId}', model already deployed`); - this.setupInProgress = false; return false; } }; @@ -92,8 +87,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { */ private deployModel = async () => { const elserId = await this.options.getElserId(); + this.options.logger.debug(`Deploying ELSER model '${elserId}'...`); try { - this.setupInProgress = true; const esClient = await this.options.elasticsearchClientPromise; await esClient.ml.startTrainedModelDeployment({ model_id: elserId, @@ -101,10 +96,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }); } catch (error) { if (!isModelAlreadyExistsError(error)) { - this.options.logger.error(`Error deploying ELSER model '${elserId}'\n${error}`); + this.options.logger.error(`Error deploying ELSER model '${elserId}':\n${error}`); } this.options.logger.debug(`Error deploying ELSER model '${elserId}', model already deployed`); - this.setupInProgress = false; } }; @@ -113,17 +107,16 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { * * @returns Promise indicating whether the model is deployed */ - private isModelDeployed = async (): Promise => { + public isModelDeployed = async (): Promise => { const elserId = await this.options.getElserId(); - const esClient = await this.options.elasticsearchClientPromise; + this.options.logger.debug(`Checking if ELSER model '${elserId}' is deployed...`); try { + const esClient = await this.options.elasticsearchClientPromise; const getResponse = await esClient.ml.getTrainedModelsStats({ model_id: elserId, }); - this.options.logger.debug(`modelId: ${elserId}`); - // For standardized way of checking deployment status see: https://github.com/elastic/elasticsearch/issues/106986 const isReadyESS = (stats: MlTrainedModelStats) => stats.deployment_stats?.state === 'started' && @@ -144,7 +137,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }; /** - * Downloads and deploys ELSER (if not already), then loads ES|QL docs + * Downloads and deploys recommended ELSER (if not already), then loads ES|QL docs * * NOTE: Before automatically installing ELSER in the background, we should perform deployment resource checks * Only necessary for ESS, as Serverless can always auto-install if `productTier === complete` @@ -158,18 +151,18 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { */ public setupKnowledgeBase = async ({ esStore, - request, soClient, }: { esStore: ElasticsearchStore; - request: KibanaRequest; soClient: SavedObjectsClientContract; }): Promise => { if (this.setupInProgress) { - this.options.logger.debug('Setup already in progress'); + this.options.logger.debug('Knowledge Base setup already in progress'); return; } + this.options.logger.debug('Starting Knowledge Base setup...'); + this.setupInProgress = true; const elserId = await this.options.getElserId(); try { @@ -181,8 +174,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { (await this.isModelInstalled()) ? Promise.resolve() : Promise.reject(new Error('Model not installed')), - { minTimeout: 10000, retries: 10 } + { minTimeout: 10000, maxTimeout: 10000, retries: 10 } ); + this.options.logger.debug(`ELSER model '${elserId}' successfully installed!`); } else { this.options.logger.debug(`ELSER model '${elserId}' is already installed`); } @@ -197,21 +191,24 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { : Promise.reject(new Error('Model not deployed')), { minTimeout: 2000, retries: 10 } ); + this.options.logger.debug(`ELSER model '${elserId}' successfully deployed!`); } else { this.options.logger.debug(`ELSER model '${elserId}' is already deployed`); } + this.options.logger.debug(`Checking if Knowledge Base docs have been loaded...`); const kbDocsLoaded = (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0; if (!kbDocsLoaded) { - this.options.logger.debug(`Loading KB docs!`); + this.options.logger.debug(`Loading KB docs...`); const loadedKnowledgeBase = await loadESQL(esStore, this.options.logger); this.options.logger.debug(`${loadedKnowledgeBase}`); } else { - this.options.logger.debug(`KB docs already loaded!`); + this.options.logger.debug(`Knowledge Base docs already loaded!`); } } catch (e) { this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); } + this.setupInProgress = false; }; /** diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 00abfa8efeea9..6f08524ef0c12 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -316,7 +316,7 @@ export class AIAssistantService { } return new AIAssistantKnowledgeBaseDataClient({ - logger: this.options.logger, + logger: this.options.logger.get('knowledgeBase'), elasticsearchClientPromise: this.options.elasticsearchClientPromise, spaceId: opts.spaceId, kibanaVersion: this.options.kibanaVersion, diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 3cae3972c8bf6..693eda3ef6bd8 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -241,6 +241,7 @@ export class DocumentsDataWriter implements DocumentsDataWriter { const documentCreateBody = params.documentsToCreate ? params.documentsToCreate.flatMap((document) => [ // Do not pre-gen _id for bulk create operations to avoid `version_conflict_engine_exception` + // TODO: Breaks Attack Discovery anonymization-field creation :( { create: { _index: this.options.index } }, document, ]) diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts index 283df8c8f5079..5c71951856009 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts @@ -96,9 +96,7 @@ export class ElasticsearchStore extends VectorStore { ) { super(new ElasticsearchEmbeddings(logger), { esClient, index }); this.esClient = esClient; - this.index = kbDataClient - ? kbDataClient.indexTemplateAndPattern.alias - : index ?? KNOWLEDGE_BASE_INDEX_PATTERN; + this.index = index ?? KNOWLEDGE_BASE_INDEX_PATTERN; this.logger = logger; this.telemetry = telemetry; this.model = model ?? '.elser_model_2'; @@ -432,6 +430,11 @@ export class ElasticsearchStore extends VectorStore { */ async isModelInstalled(modelId?: string): Promise { try { + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled + if (this.kbDataClient != null) { + return this.kbDataClient.isModelInstalled(); + } + const getResponse = await this.esClient.ml.getTrainedModelsStats({ model_id: modelId ?? this.model, }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts index 7e7d5cf561d5e..42ffad4779d1c 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts @@ -14,13 +14,15 @@ import { initializeAgentExecutorWithOptions } from 'langchain/agents'; import { mockActionResponse } from '../../../__mocks__/action_result_data'; import { langChainMessages } from '../../../__mocks__/lang_chain_messages'; -import { ESQL_RESOURCE } from '../../../routes/knowledge_base/constants'; +import { KNOWLEDGE_BASE_INDEX_PATTERN } from '../../../routes/knowledge_base/constants'; import { callAgentExecutor } from '.'; import { PassThrough, Stream } from 'stream'; import { ActionsClientChatOpenAI, ActionsClientLlm, } from '@kbn/elastic-assistant-common/impl/language_models'; +import { AgentExecutorParams } from '../executors/types'; +import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; jest.mock('@kbn/elastic-assistant-common/impl/language_models', () => ({ ActionsClientChatOpenAI: jest.fn(), @@ -85,18 +87,23 @@ const mockActions: ActionsPluginStart = {} as ActionsPluginStart; const mockLogger = loggerMock.create(); const mockTelemetry = coreMock.createSetup().analytics; const esClientMock = elasticsearchServiceMock.createScopedClusterClient().asCurrentUser; -const defaultProps = { +const esStoreMock = new ElasticsearchStore( + esClientMock, + KNOWLEDGE_BASE_INDEX_PATTERN, + mockLogger, + mockTelemetry +); +const defaultProps: AgentExecutorParams = { actions: mockActions, isEnabledKnowledgeBase: true, connectorId: mockConnectorId, esClient: esClientMock, + esStore: esStoreMock, llmType: 'openai', langChainMessages, logger: mockLogger, onNewReplacements: jest.fn(), request: mockRequest, - kbResource: ESQL_RESOURCE, - telemetry: mockTelemetry, replacements: {}, }; const executorMock = initializeAgentExecutorWithOptions as jest.Mock; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts index f8e2f2426bf89..9c1ea6d5c36d5 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts @@ -17,8 +17,6 @@ import { ActionsClientLlm, } from '@kbn/elastic-assistant-common/impl/language_models'; import { getDefaultArguments } from '@kbn/elastic-assistant-common/impl/language_models/constants'; -import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; -import { KNOWLEDGE_BASE_INDEX_PATTERN } from '../../../routes/knowledge_base/constants'; import { AgentExecutor } from '../executors/types'; import { withAssistantSpan } from '../tracers/with_assistant_span'; import { APMTracer } from '../tracers/apm_tracer'; @@ -38,9 +36,8 @@ export const callAgentExecutor: AgentExecutor = async ({ isEnabledKnowledgeBase, assistantTools = [], connectorId, - elserId, esClient, - kbResource, + esStore, langChainMessages, llmType, logger, @@ -50,7 +47,6 @@ export const callAgentExecutor: AgentExecutor = async ({ replacements, request, size, - telemetry, traceOptions, }) => { // TODO implement llmClass for bedrock streaming @@ -87,16 +83,6 @@ export const callAgentExecutor: AgentExecutor = async ({ returnMessages: true, }); - // ELSER backed ElasticsearchStore for Knowledge Base - const esStore = new ElasticsearchStore( - esClient, - KNOWLEDGE_BASE_INDEX_PATTERN, - logger, - telemetry, - elserId, - kbResource - ); - const modelExists = await esStore.isModelInstalled(); // Create a chain that uses the ELSER backed ElasticsearchStore, override k=10 for esql query generation for now diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts index 2b3d07708e2e1..2b5ef7350e628 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts @@ -11,8 +11,6 @@ import { BufferMemory, ChatMessageHistory } from 'langchain/memory'; import { ChainTool } from 'langchain/tools/chain'; import { ActionsClientLlm } from '@kbn/elastic-assistant-common/impl/language_models'; -import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; -import { KNOWLEDGE_BASE_INDEX_PATTERN } from '../../../routes/knowledge_base/constants'; import { AgentExecutor } from './types'; import { withAssistantSpan } from '../tracers/with_assistant_span'; import { APMTracer } from '../tracers/apm_tracer'; @@ -30,13 +28,11 @@ export const callOpenAIFunctionsExecutor: AgentExecutor = async ({ actions, connectorId, esClient, + esStore, langChainMessages, llmType, logger, request, - elserId, - kbResource, - telemetry, traceOptions, }) => { const llm = new ActionsClientLlm({ @@ -59,16 +55,6 @@ export const callOpenAIFunctionsExecutor: AgentExecutor = async ({ returnMessages: true, }); - // ELSER backed ElasticsearchStore for Knowledge Base - const esStore = new ElasticsearchStore( - esClient, - KNOWLEDGE_BASE_INDEX_PATTERN, - logger, - telemetry, - elserId, - kbResource - ); - const modelExists = await esStore.isModelInstalled(); if (!modelExists) { throw new Error( diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 0e7925df20281..9686360dddb95 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -11,12 +11,12 @@ import { BaseMessage } from '@langchain/core/messages'; import { Logger } from '@kbn/logging'; import { KibanaRequest, ResponseHeaders } from '@kbn/core-http-server'; import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; -import type { AnalyticsServiceSetup } from '@kbn/core-analytics-server'; import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic-assistant-common'; import { StreamFactoryReturnType } from '@kbn/ml-response-stream/server'; import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; +import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store'; export interface AgentExecutorParams { abortSignal?: AbortSignal; @@ -27,7 +27,7 @@ export interface AgentExecutorParams { assistantTools?: AssistantTool[]; connectorId: string; esClient: ElasticsearchClient; - kbResource: string | undefined; + esStore: ElasticsearchStore; langChainMessages: BaseMessage[]; llmType?: string; logger: Logger; @@ -41,9 +41,7 @@ export interface AgentExecutorParams { ) => Promise; request: KibanaRequest; size?: number; - elserId?: string; traceOptions?: TraceOptions; - telemetry: AnalyticsServiceSetup; } export interface StaticReturnType { diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 0161c9bc2a939..6206b60ea801f 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -78,32 +78,32 @@ export const postKnowledgeBaseRoute = ( const enableKnowledgeBaseByDefault = assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; - // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled - if (enableKnowledgeBaseByDefault) { - const knowledgeBaseDataClient = - await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); - if (!knowledgeBaseDataClient) { - return response.custom({ body: { success: false }, statusCode: 500 }); - } + try { + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled + if (enableKnowledgeBaseByDefault) { + const knowledgeBaseDataClient = + await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); + if (!knowledgeBaseDataClient) { + return response.custom({ body: { success: false }, statusCode: 500 }); + } - // Continue to use esStore for loading esql docs until `semantic_text` is available and we can test the new chunking strategy - const esStore = new ElasticsearchStore( - esClient, - knowledgeBaseDataClient.indexTemplateAndPattern.alias, - logger, - telemetry, - elserId, - getKbResource(request), - knowledgeBaseDataClient, - authenticatedUser - ); + // Continue to use esStore for loading esql docs until `semantic_text` is available and we can test the new chunking strategy + const esStore = new ElasticsearchStore( + esClient, + knowledgeBaseDataClient.indexTemplateAndPattern.alias, + logger, + telemetry, + elserId, + getKbResource(request), + knowledgeBaseDataClient, + authenticatedUser + ); - await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, request, soClient }); + await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, request, soClient }); - return response.ok({ body: { success: true } }); - } + return response.ok({ body: { success: true } }); + } - try { const kbResource = getKbResource(request); const esStore = new ElasticsearchStore( esClient, diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts index 6c2be6b8c7120..ab9250047f7ef 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts @@ -23,6 +23,7 @@ import { PassThrough } from 'stream'; import { getConversationResponseMock } from '../ai_assistant_data_clients/conversations/update_conversation.test'; import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock'; import { getFindAnonymizationFieldsResultWithSingleHit } from '../__mocks__/response'; +import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; const actionsClient = actionsClientMock.create(); jest.mock('../lib/build_response', () => ({ @@ -92,6 +93,7 @@ const mockContext = { getActionsClientWithRequest: jest.fn().mockResolvedValue(actionsClient), }, getRegisteredTools: jest.fn(() => []), + getRegisteredFeatures: jest.fn(() => defaultAssistantFeatures), logger: loggingSystemMock.createLogger(), telemetry: { ...coreMock.createSetup().analytics, reportEvent }, getCurrentUser: () => ({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 61326f157bb27..25b891b2ca163 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -31,7 +31,7 @@ import { POST_ACTIONS_CONNECTOR_EXECUTE } from '../../common/constants'; import { getLangChainMessages } from '../lib/langchain/helpers'; import { buildResponse } from '../lib/build_response'; import { ElasticAssistantRequestHandlerContext, GetElser } from '../types'; -import { ESQL_RESOURCE } from './knowledge_base/constants'; +import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './knowledge_base/constants'; import { callAgentExecutor } from '../lib/langchain/execute_custom_llm_chain'; import { DEFAULT_PLUGIN_NAME, @@ -41,6 +41,7 @@ import { import { getLangSmithTracer } from './evaluate/utils'; import { EsAnonymizationFieldsSchema } from '../ai_assistant_data_clients/anonymization_fields/types'; import { transformESSearchToAnonymizationFields } from '../ai_assistant_data_clients/anonymization_fields/helpers'; +import { ElasticsearchStore } from '../lib/langchain/elasticsearch_store/elasticsearch_store'; export const postActionsConnectorExecuteRoute = ( router: IRouter, @@ -302,6 +303,7 @@ export const postActionsConnectorExecuteRoute = ( defaultPluginName: DEFAULT_PLUGIN_NAME, logger, }); + const assistantTools = (await context.elasticAssistant) .getRegisteredTools(pluginName) .filter((x) => x.id !== 'attack-discovery'); // We don't (yet) support asking the assistant for NEW attack discoveries from a conversation @@ -323,6 +325,27 @@ export const postActionsConnectorExecuteRoute = ( page: 1, }); + // Create an ElasticsearchStore for KB interactions + // Setup with kbDataClient if `enableKnowledgeBaseByDefault` FF is enabled + const enableKnowledgeBaseByDefault = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; + const kbDataClient = enableKnowledgeBaseByDefault + ? (await assistantContext.getAIAssistantKnowledgeBaseDataClient(false)) ?? undefined + : undefined; + const kbIndex = + enableKnowledgeBaseByDefault && kbDataClient != null + ? kbDataClient.indexTemplateAndPattern.alias + : KNOWLEDGE_BASE_INDEX_PATTERN; + const esStore = new ElasticsearchStore( + esClient, + kbIndex, + logger, + telemetry, + elserId, + ESQL_RESOURCE, + kbDataClient + ); + const result: StreamFactoryReturnType['responseWithHeaders'] | StaticReturnType = await callAgentExecutor({ abortSignal, @@ -334,14 +357,13 @@ export const postActionsConnectorExecuteRoute = ( isEnabledKnowledgeBase: request.body.isEnabledKnowledgeBase ?? false, assistantTools, connectorId, - elserId, esClient, + esStore, isStream: // TODO implement llmClass for bedrock streaming // tracked here: https://github.com/elastic/security-team/issues/7363 request.body.subAction !== 'invokeAI' && actionTypeId === '.gen-ai', llmType: getLlmType(actionTypeId), - kbResource: ESQL_RESOURCE, langChainMessages, logger, onNewReplacements, @@ -349,7 +371,6 @@ export const postActionsConnectorExecuteRoute = ( request, replacements: request.body.replacements, size: request.body.size, - telemetry, traceOptions: { projectName: langSmithProject, tracers: getLangSmithTracer({ From be327c114f65e0f7f9c0c3ebef13c23ed7d10326 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 15 May 2024 05:01:52 +0000 Subject: [PATCH 12/20] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/plugins/elastic_assistant/tsconfig.json | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 4671ab75e544d..345949b866262 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -29,7 +29,6 @@ "@kbn/logging", "@kbn/ml-plugin", "@kbn/apm-utils", - "@kbn/core-analytics-server", "@kbn/elastic-assistant-common", "@kbn/core-http-router-server-mocks", "@kbn/data-stream-adapter", From 2b10c9ca2973136676f3ee9622c461513ec31357 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Tue, 14 May 2024 23:55:32 -0600 Subject: [PATCH 13/20] Fix remaining UI controls from old impl, and refactor authenticatedUser out of esStore --- .../knowledge_base/index.ts | 10 +++-- .../server/ai_assistant_service/index.ts | 1 + .../elasticsearch_store.ts | 44 +++++++++++-------- .../knowledge_base/delete_knowledge_base.ts | 34 +++++++++++--- .../get_knowledge_base_status.ts | 13 +----- .../knowledge_base/post_knowledge_base.ts | 12 +---- 6 files changed, 67 insertions(+), 47 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 5aabf10e4882f..57e8eb965e2fc 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -9,7 +9,6 @@ import { MlTrainedModelDeploymentNodesStats, MlTrainedModelStats, } from '@elastic/elasticsearch/lib/api/types'; -import { AuthenticatedUser } from '@kbn/core-security-common'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { Document } from 'langchain/document'; @@ -29,6 +28,7 @@ import { isModelAlreadyExistsError } from './helpers'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ml: MlPluginSetup; getElserId: GetElser; + ingestPipelineResourceName: string; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { private setupInProgress: boolean = false; @@ -219,13 +219,17 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { */ public addKnowledgeBaseDocuments = async ({ documents, - authenticatedUser, }: { documents: Document[]; - authenticatedUser: AuthenticatedUser; }): Promise => { const writer = await this.getWriter(); const changedAt = new Date().toISOString(); + const authenticatedUser = this.options.currentUser; + if (authenticatedUser == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } // @ts-ignore const { errors, docs_created: docsCreated } = await writer.bulk({ documentsToCreate: documents.map((doc) => diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 6f08524ef0c12..7847c40331426 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -321,6 +321,7 @@ export class AIAssistantService { spaceId: opts.spaceId, kibanaVersion: this.options.kibanaVersion, indexPatternsResourceName: this.resourceNames.aliases.knowledgeBase, + ingestPipelineResourceName: this.resourceNames.pipelines.knowledgeBase, currentUser: opts.currentUser, ml: this.options.ml, getElserId: this.getElserId, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts index 5c71951856009..86791cec5f5ce 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts @@ -5,12 +5,7 @@ * 2.0. */ -import { - type AnalyticsServiceSetup, - AuthenticatedUser, - ElasticsearchClient, - Logger, -} from '@kbn/core/server'; +import { type AnalyticsServiceSetup, ElasticsearchClient, Logger } from '@kbn/core/server'; import { MappingTypeMapping, MlTrainedModelDeploymentNodesStats, @@ -70,8 +65,6 @@ export const TERMS_QUERY_SIZE = 10000; export class ElasticsearchStore extends VectorStore { declare FilterType: QueryDslQueryContainer; - // AuthenticatedUser is required for adding/retrieving documents as it is not encapsulated in the kbDataClient - private readonly authenticatedUser: AuthenticatedUser | undefined; private readonly esClient: ElasticsearchClient; private readonly kbDataClient: AIAssistantKnowledgeBaseDataClient | undefined; private readonly index: string; @@ -91,8 +84,7 @@ export class ElasticsearchStore extends VectorStore { telemetry: AnalyticsServiceSetup, model?: string, kbResource?: string | undefined, - kbDataClient?: AIAssistantKnowledgeBaseDataClient, - authenticatedUser?: AuthenticatedUser + kbDataClient?: AIAssistantKnowledgeBaseDataClient ) { super(new ElasticsearchEmbeddings(logger), { esClient, index }); this.esClient = esClient; @@ -102,7 +94,6 @@ export class ElasticsearchStore extends VectorStore { this.model = model ?? '.elser_model_2'; this.kbResource = kbResource ?? ESQL_RESOURCE; this.kbDataClient = kbDataClient; - this.authenticatedUser = authenticatedUser; } /** @@ -157,15 +148,14 @@ export class ElasticsearchStore extends VectorStore { documents: Document[], options?: Record ): Promise => { - if (!this.kbDataClient || !this.authenticatedUser) { - this.logger.error('No kbDataClient or authenticatedUser provided'); + if (!this.kbDataClient) { + this.logger.error('No kbDataClient provided'); return []; } try { const response = await this.kbDataClient.addKnowledgeBaseDocuments({ documents, - authenticatedUser: this.authenticatedUser, }); return response.map((doc) => doc.id); } catch (e) { @@ -352,6 +342,13 @@ export class ElasticsearchStore extends VectorStore { * @returns Promise indicating whether the index was created */ deleteIndex = async (index?: string): Promise => { + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled + // We won't be supporting delete operations for the KB data stream going forward, so this can be removed along with the FF + if (this.kbDataClient != null) { + const response = await this.esClient.indices.deleteDataStream({ name: index ?? this.index }); + return response.acknowledged; + } + const response = await this.esClient.indices.delete({ index: index ?? this.index, }); @@ -367,8 +364,12 @@ export class ElasticsearchStore extends VectorStore { */ pipelineExists = async (pipelineId?: string): Promise => { try { + const id = + pipelineId ?? + this.kbDataClient?.options.ingestPipelineResourceName ?? + KNOWLEDGE_BASE_INGEST_PIPELINE; const response = await this.esClient.ingest.getPipeline({ - id: KNOWLEDGE_BASE_INGEST_PIPELINE, + id, }); return Object.keys(response).length > 0; } catch (e) { @@ -384,7 +385,10 @@ export class ElasticsearchStore extends VectorStore { */ createPipeline = async ({ id, description }: CreatePipelineParams = {}): Promise => { const response = await this.esClient.ingest.putPipeline({ - id: id ?? KNOWLEDGE_BASE_INGEST_PIPELINE, + id: + id ?? + this.kbDataClient?.options.ingestPipelineResourceName ?? + KNOWLEDGE_BASE_INGEST_PIPELINE, description: description ?? 'Embedding pipeline for Elastic AI Assistant ELSER Knowledge Base', processors: [ @@ -416,7 +420,10 @@ export class ElasticsearchStore extends VectorStore { */ deletePipeline = async (pipelineId?: string): Promise => { const response = await this.esClient.ingest.deletePipeline({ - id: pipelineId ?? KNOWLEDGE_BASE_INGEST_PIPELINE, + id: + pipelineId ?? + this.kbDataClient?.options.ingestPipelineResourceName ?? + KNOWLEDGE_BASE_INGEST_PIPELINE, }); return response.acknowledged; @@ -432,7 +439,8 @@ export class ElasticsearchStore extends VectorStore { try { // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled if (this.kbDataClient != null) { - return this.kbDataClient.isModelInstalled(); + // esStore.isModelInstalled() is actually checking if the model is deployed, not installed, so do that instead + return this.kbDataClient.isModelDeployed(); } const getResponse = await this.esClient.ml.getTrainedModelsStats({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts index e3ba4f7b6156d..6e81738a7376f 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/delete_knowledge_base.ts @@ -21,6 +21,8 @@ import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantRequestHandlerContext } from '../../types'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './constants'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; +import { getKbResource } from './get_kb_resource'; /** * Delete Knowledge Base index, pipeline, and resources (collection of documents) @@ -51,21 +53,43 @@ export const deleteKnowledgeBaseRoute = ( const assistantContext = await context.elasticAssistant; const logger = assistantContext.logger; const telemetry = assistantContext.telemetry; + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const enableKnowledgeBaseByDefault = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; try { - const kbResource = - request.params.resource != null - ? decodeURIComponent(request.params.resource) - : undefined; + const kbResource = getKbResource(request); const esClient = (await context.core).elasticsearch.client.asInternalUser; - const esStore = new ElasticsearchStore( + let esStore = new ElasticsearchStore( esClient, KNOWLEDGE_BASE_INDEX_PATTERN, logger, telemetry ); + // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled, only need an esStore w/ kbDataClient + if (enableKnowledgeBaseByDefault) { + const knowledgeBaseDataClient = + await assistantContext.getAIAssistantKnowledgeBaseDataClient(false); + if (!knowledgeBaseDataClient) { + return response.custom({ body: { success: false }, statusCode: 500 }); + } + esStore = new ElasticsearchStore( + esClient, + knowledgeBaseDataClient.indexTemplateAndPattern.alias, + logger, + telemetry, + 'elserId', // Not needed for delete ops + kbResource, + knowledgeBaseDataClient + ); + } + if (kbResource === ESQL_RESOURCE) { // For now, tearing down the Knowledge Base is fine, but will want to support removing specific assets based // on resource name or document query diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index a2ff556ea64c9..f1faeae17f181 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -68,14 +68,6 @@ export const getKnowledgeBaseStatusRoute = ( kbResource ); - const authenticatedUser = assistantContext.getCurrentUser(); - if (authenticatedUser == null) { - return response.custom({ - body: `Authenticated user not found`, - statusCode: 401, - }); - } - const pluginName = getPluginNameFromRequest({ request, defaultPluginName: DEFAULT_PLUGIN_NAME, @@ -87,7 +79,7 @@ export const getKnowledgeBaseStatusRoute = ( // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled if (enableKnowledgeBaseByDefault) { const knowledgeBaseDataClient = - await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); + await assistantContext.getAIAssistantKnowledgeBaseDataClient(false); if (!knowledgeBaseDataClient) { return response.custom({ body: { success: false }, statusCode: 500 }); } @@ -100,8 +92,7 @@ export const getKnowledgeBaseStatusRoute = ( telemetry, elserId, kbResource, - knowledgeBaseDataClient, - authenticatedUser + knowledgeBaseDataClient ); } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 6206b60ea801f..063ffc36eda79 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -62,13 +62,6 @@ export const postKnowledgeBaseRoute = ( const core = await context.core; const esClient = core.elasticsearch.client.asInternalUser; const soClient = core.savedObjects.getClient(); - const authenticatedUser = assistantContext.getCurrentUser(); - if (authenticatedUser == null) { - return response.custom({ - body: `Authenticated user not found`, - statusCode: 401, - }); - } const pluginName = getPluginNameFromRequest({ request, @@ -95,11 +88,10 @@ export const postKnowledgeBaseRoute = ( telemetry, elserId, getKbResource(request), - knowledgeBaseDataClient, - authenticatedUser + knowledgeBaseDataClient ); - await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, request, soClient }); + await knowledgeBaseDataClient.setupKnowledgeBase({ esStore, soClient }); return response.ok({ body: { success: true } }); } From 684488a952bdd4dce77331609f5d1b62a2ba4581 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 15 May 2024 00:46:27 -0600 Subject: [PATCH 14/20] Plumb through isSetupInProgress to status endpoints and migrate them to use OAS generated types --- .../knowledge_base/crud_kb_route.gen.ts | 2 + .../knowledge_base/crud_kb_route.schema.yaml | 4 ++ .../impl/assistant/api/index.tsx | 61 +++++++------------ .../install_knowledge_base_button.tsx | 7 ++- .../use_knowledge_base_status.tsx | 13 +--- .../packages/kbn-elastic-assistant/index.ts | 7 --- .../knowledge_base/index.ts | 12 ++-- .../server/routes/evaluate/post_evaluate.ts | 28 +++++++-- .../get_knowledge_base_status.ts | 12 ++-- 9 files changed, 73 insertions(+), 73 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts index 634cd8cb6e78b..8c85cbce32108 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts @@ -67,6 +67,8 @@ export type ReadKnowledgeBaseRequestParamsInput = z.input; export const ReadKnowledgeBaseResponse = z.object({ elser_exists: z.boolean().optional(), + esql_exists: z.boolean().optional(), index_exists: z.boolean().optional(), + is_setup_in_progress: z.boolean().optional(), pipeline_exists: z.boolean().optional(), }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml index 650a7e141ce39..16d13c24f23ea 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml @@ -60,8 +60,12 @@ paths: properties: elser_exists: type: boolean + esql_exists: + type: boolean index_exists: type: boolean + is_setup_in_progress: + type: boolean pipeline_exists: type: boolean 400: diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx index 524f0ca95f99c..94f62aa4040f2 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx @@ -10,7 +10,13 @@ import { IHttpFetchError } from '@kbn/core-http-browser'; import { API_VERSIONS, ApiConfig, + CreateKnowledgeBaseRequestParams, + CreateKnowledgeBaseResponse, + DeleteKnowledgeBaseRequestParams, + DeleteKnowledgeBaseResponse, ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, + ReadKnowledgeBaseRequestParams, + ReadKnowledgeBaseResponse, Replacements, } from '@kbn/elastic-assistant-common'; import { API_ERROR } from '../translations'; @@ -184,19 +190,6 @@ export const fetchConnectorExecuteAction = async ({ } }; -export interface GetKnowledgeBaseStatusParams { - http: HttpSetup; - resource?: string; - signal?: AbortSignal | undefined; -} - -export interface GetKnowledgeBaseStatusResponse { - elser_exists: boolean; - esql_exists?: boolean; - index_exists: boolean; - pipeline_exists: boolean; -} - /** * API call for getting the status of the Knowledge Base. Provide * a resource to include the status of that specific resource. @@ -206,13 +199,15 @@ export interface GetKnowledgeBaseStatusResponse { * @param {string} [options.resource] - Resource to get the status of, otherwise status of overall KB * @param {AbortSignal} [options.signal] - AbortSignal * - * @returns {Promise} + * @returns {Promise} */ export const getKnowledgeBaseStatus = async ({ http, resource, signal, -}: GetKnowledgeBaseStatusParams): Promise => { +}: ReadKnowledgeBaseRequestParams & { http: HttpSetup; signal?: AbortSignal | undefined }): Promise< + ReadKnowledgeBaseResponse | IHttpFetchError +> => { try { const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); const response = await http.fetch(path, { @@ -221,22 +216,12 @@ export const getKnowledgeBaseStatus = async ({ version: API_VERSIONS.internal.v1, }); - return response as GetKnowledgeBaseStatusResponse; + return response as ReadKnowledgeBaseResponse; } catch (error) { return error as IHttpFetchError; } }; -export interface PostKnowledgeBaseParams { - http: HttpSetup; - resource?: string; - signal?: AbortSignal | undefined; -} - -export interface PostKnowledgeBaseResponse { - success: boolean; -} - /** * API call for setting up the Knowledge Base. Provide a resource to set up a specific resource. * @@ -245,13 +230,16 @@ export interface PostKnowledgeBaseResponse { * @param {string} [options.resource] - Resource to be added to the KB, otherwise sets up the base KB * @param {AbortSignal} [options.signal] - AbortSignal * - * @returns {Promise} + * @returns {Promise} */ export const postKnowledgeBase = async ({ http, resource, signal, -}: PostKnowledgeBaseParams): Promise => { +}: CreateKnowledgeBaseRequestParams & { + http: HttpSetup; + signal?: AbortSignal | undefined; +}): Promise => { try { const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); const response = await http.fetch(path, { @@ -260,22 +248,12 @@ export const postKnowledgeBase = async ({ version: API_VERSIONS.internal.v1, }); - return response as PostKnowledgeBaseResponse; + return response as CreateKnowledgeBaseResponse; } catch (error) { return error as IHttpFetchError; } }; -export interface DeleteKnowledgeBaseParams { - http: HttpSetup; - resource?: string; - signal?: AbortSignal | undefined; -} - -export interface DeleteKnowledgeBaseResponse { - success: boolean; -} - /** * API call for deleting the Knowledge Base. Provide a resource to delete that specific resource. * @@ -290,7 +268,10 @@ export const deleteKnowledgeBase = async ({ http, resource, signal, -}: DeleteKnowledgeBaseParams): Promise => { +}: DeleteKnowledgeBaseRequestParams & { + http: HttpSetup; + signal?: AbortSignal | undefined; +}): Promise => { try { const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || ''); const response = await http.fetch(path, { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx index c58061a733210..9f9e08915886d 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx @@ -11,6 +11,7 @@ import { i18n } from '@kbn/i18n'; import { useAssistantContext } from '../..'; import { useSetupKnowledgeBase } from './use_setup_knowledge_base'; +import { useKnowledgeBaseStatus } from './use_knowledge_base_status'; const ESQL_RESOURCE = 'esql'; @@ -25,9 +26,11 @@ export const InstallKnowledgeBaseButton: React.FC = React.memo(() => { http, } = useAssistantContext(); - // const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); + const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http }); + const isLoading = kbStatus?.is_setup_in_progress || isSettingUpKB; + const onInstallKnowledgeBase = useCallback(() => { setupKB(ESQL_RESOURCE); }, [setupKB]); @@ -41,7 +44,7 @@ export const InstallKnowledgeBaseButton: React.FC = React.memo(() => { color="primary" data-test-subj="install-knowledge-base-button" fill - isLoading={isSettingUpKB} + isLoading={isLoading} iconType="importAction" onClick={onInstallKnowledgeBase} > diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/use_knowledge_base_status.tsx index 7ec7d227f9ef8..c03eb31581e42 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/use_knowledge_base_status.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/use_knowledge_base_status.tsx @@ -11,6 +11,7 @@ import type { HttpSetup, IHttpFetchError, ResponseErrorBody } from '@kbn/core-ht import type { IToasts } from '@kbn/core-notifications-browser'; import { i18n } from '@kbn/i18n'; import { useCallback } from 'react'; +import { ReadKnowledgeBaseResponse } from '@kbn/elastic-assistant-common'; import { getKnowledgeBaseStatus } from '../assistant/api'; const KNOWLEDGE_BASE_STATUS_QUERY_KEY = ['elastic-assistant', 'knowledge-base-status']; @@ -21,13 +22,6 @@ export interface UseKnowledgeBaseStatusParams { toasts?: IToasts; } -export interface GetKnowledgeBaseStatusResponse { - elser_exists: boolean; - esql_exists?: boolean; - index_exists: boolean; - pipeline_exists: boolean; -} - /** * Hook for getting the status of the Knowledge Base. Provide a resource name to include * the status for that specific resource within the KB. @@ -42,10 +36,7 @@ export const useKnowledgeBaseStatus = ({ http, resource, toasts, -}: UseKnowledgeBaseStatusParams): UseQueryResult< - GetKnowledgeBaseStatusResponse, - IHttpFetchError -> => { +}: UseKnowledgeBaseStatusParams): UseQueryResult => { return useQuery( KNOWLEDGE_BASE_STATUS_QUERY_KEY, async ({ signal }) => { diff --git a/x-pack/packages/kbn-elastic-assistant/index.ts b/x-pack/packages/kbn-elastic-assistant/index.ts index 56fb74af3aba4..df0ba1e8db0f9 100644 --- a/x-pack/packages/kbn-elastic-assistant/index.ts +++ b/x-pack/packages/kbn-elastic-assistant/index.ts @@ -145,13 +145,6 @@ export type { PromptContextTemplate } from './impl/assistant/prompt_context/type */ export type { QuickPrompt } from './impl/assistant/quick_prompts/types'; -/** - * Knowledge Base API Responses - */ -export type { DeleteKnowledgeBaseResponse } from './impl/assistant/api'; -export type { GetKnowledgeBaseStatusResponse } from './impl/assistant/api'; -export type { PostKnowledgeBaseResponse } from './impl/assistant/api'; - export { useFetchCurrentUserConversations } from './impl/assistant/api/conversations/use_fetch_current_user_conversations'; export * from './impl/assistant/api/conversations/bulk_update_actions_conversations'; export { getConversationById } from './impl/assistant/api/conversations/conversations'; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 57e8eb965e2fc..dd575a1834c36 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -31,12 +31,16 @@ interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ingestPipelineResourceName: string; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { - private setupInProgress: boolean = false; + private _isSetupInProgress: boolean = false; constructor(public readonly options: KnowledgeBaseDataClientParams) { super(options); } + public get isSetupInProgress() { + return this._isSetupInProgress; + } + /** * Downloads and installs ELSER model if not already installed * @@ -156,13 +160,13 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { esStore: ElasticsearchStore; soClient: SavedObjectsClientContract; }): Promise => { - if (this.setupInProgress) { + if (this._isSetupInProgress) { this.options.logger.debug('Knowledge Base setup already in progress'); return; } this.options.logger.debug('Starting Knowledge Base setup...'); - this.setupInProgress = true; + this._isSetupInProgress = true; const elserId = await this.options.getElserId(); try { @@ -208,7 +212,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } catch (e) { this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); } - this.setupInProgress = false; + this._isSetupInProgress = false; }; /** diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index fc919d10e345c..46a3953370826 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -19,7 +19,7 @@ import { } from '@kbn/elastic-assistant-common'; import { ActionsClientLlm } from '@kbn/elastic-assistant-common/impl/language_models'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; -import { ESQL_RESOURCE } from '../knowledge_base/constants'; +import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from '../knowledge_base/constants'; import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantRequestHandlerContext, GetElser } from '../../types'; import { EVALUATE } from '../../../common/constants'; @@ -37,6 +37,7 @@ import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; * and reference your specific AgentExecutor function */ import { AGENT_EXECUTOR_MAP } from '../../lib/langchain/executors'; +import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; const DEFAULT_SIZE = 20; @@ -157,6 +158,27 @@ export const postEvaluateRoute = ( }, }; + // Create an ElasticsearchStore for KB interactions + // Setup with kbDataClient if `enableKnowledgeBaseByDefault` FF is enabled + const enableKnowledgeBaseByDefault = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; + const kbDataClient = enableKnowledgeBaseByDefault + ? (await assistantContext.getAIAssistantKnowledgeBaseDataClient(false)) ?? undefined + : undefined; + const kbIndex = + enableKnowledgeBaseByDefault && kbDataClient != null + ? kbDataClient.indexTemplateAndPattern.alias + : KNOWLEDGE_BASE_INDEX_PATTERN; + const esStore = new ElasticsearchStore( + esClient, + kbIndex, + logger, + telemetry, + elserId, + ESQL_RESOURCE, + kbDataClient + ); + // Create an array of executor functions to call in batches // One for each connector/model + agent combination // Hoist `langChainMessages` so they can be batched by dataset.input in the evaluator @@ -175,14 +197,12 @@ export const postEvaluateRoute = ( assistantTools, connectorId, esClient, - elserId, + esStore, isStream: false, langChainMessages, llmType: 'openai', logger, request: skeletonRequest, - kbResource: ESQL_RESOURCE, - telemetry, traceOptions: { exampleId, projectName, diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index f1faeae17f181..7382be25c9538 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -77,23 +77,24 @@ export const getKnowledgeBaseStatusRoute = ( assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled + let isSetupInProgress = false; if (enableKnowledgeBaseByDefault) { - const knowledgeBaseDataClient = - await assistantContext.getAIAssistantKnowledgeBaseDataClient(false); - if (!knowledgeBaseDataClient) { + const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); + if (!kbDataClient) { return response.custom({ body: { success: false }, statusCode: 500 }); } // Use old status checks by overriding esStore to use kbDataClient esStore = new ElasticsearchStore( esClient, - knowledgeBaseDataClient.indexTemplateAndPattern.alias, + kbDataClient.indexTemplateAndPattern.alias, logger, telemetry, elserId, kbResource, - knowledgeBaseDataClient + kbDataClient ); + isSetupInProgress = kbDataClient.isSetupInProgress; } const indexExists = await esStore.indexExists(); @@ -103,6 +104,7 @@ export const getKnowledgeBaseStatusRoute = ( const body: ReadKnowledgeBaseResponse = { elser_exists: modelExists, index_exists: indexExists, + is_setup_in_progress: isSetupInProgress, pipeline_exists: pipelineExists, }; From 252b8e809510ddfd7c8afd61d4915e42d338c252 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 15 May 2024 10:49:40 -0600 Subject: [PATCH 15/20] Fix isSetupInProgress update not propagating --- .../install_knowledge_base_button.tsx | 11 ++++++++--- .../knowledge_base/index.ts | 12 ++++++------ .../server/ai_assistant_service/index.ts | 17 +++++++++++++---- .../knowledge_base/get_knowledge_base_status.ts | 4 +++- 4 files changed, 30 insertions(+), 14 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx index 9f9e08915886d..631bbae89b891 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx @@ -29,13 +29,18 @@ export const InstallKnowledgeBaseButton: React.FC = React.memo(() => { const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http }); - const isLoading = kbStatus?.is_setup_in_progress || isSettingUpKB; + const isSetupInProgress = kbStatus?.is_setup_in_progress || isSettingUpKB; + const isSetupComplete = + kbStatus?.elser_exists && + kbStatus?.index_exists && + kbStatus?.pipeline_exists && + kbStatus?.esql_exists; const onInstallKnowledgeBase = useCallback(() => { setupKB(ESQL_RESOURCE); }, [setupKB]); - if (!enableKnowledgeBaseByDefault) { + if (!enableKnowledgeBaseByDefault || isSetupComplete) { return null; } @@ -44,7 +49,7 @@ export const InstallKnowledgeBaseButton: React.FC = React.memo(() => { color="primary" data-test-subj="install-knowledge-base-button" fill - isLoading={isLoading} + isLoading={isSetupInProgress} iconType="importAction" onClick={onInstallKnowledgeBase} > diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index dd575a1834c36..3957bab35e914 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -24,21 +24,21 @@ import { EsKnowledgeBaseEntrySchema } from './types'; import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; import { isModelAlreadyExistsError } from './helpers'; +import { AIAssistantService } from '../../ai_assistant_service'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { + assistantService: AIAssistantService; ml: MlPluginSetup; getElserId: GetElser; ingestPipelineResourceName: string; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { - private _isSetupInProgress: boolean = false; - constructor(public readonly options: KnowledgeBaseDataClientParams) { super(options); } public get isSetupInProgress() { - return this._isSetupInProgress; + return this.options.assistantService.isKBSetupInProgress; } /** @@ -160,13 +160,13 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { esStore: ElasticsearchStore; soClient: SavedObjectsClientContract; }): Promise => { - if (this._isSetupInProgress) { + if (this.options.assistantService.isKBSetupInProgress) { this.options.logger.debug('Knowledge Base setup already in progress'); return; } this.options.logger.debug('Starting Knowledge Base setup...'); - this._isSetupInProgress = true; + this.options.assistantService.isKBSetupInProgress = true; const elserId = await this.options.getElserId(); try { @@ -212,7 +212,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } catch (e) { this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); } - this._isSetupInProgress = false; + this.options.assistantService.isKBSetupInProgress = false; }; /** diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 7847c40331426..f753df5e73a68 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -70,6 +70,7 @@ export class AIAssistantService { private anonymizationFieldsDataStream: DataStreamSpacesAdapter; private resourceInitializationHelper: ResourceInstallationHelper; private initPromise: Promise; + private _isKBSetupInProgress: boolean = false; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -108,6 +109,13 @@ export class AIAssistantService { return this.initialized; } + public get isKBSetupInProgress() { + return this._isKBSetupInProgress; + } + public set isKBSetupInProgress(inProgress: boolean) { + this._isKBSetupInProgress = inProgress; + } + private createDataStream: CreateDataStream = ({ resource, kibanaVersion, fieldMap }) => { const newDataStream = new DataStreamSpacesAdapter(this.resourceNames.aliases[resource], { kibanaVersion, @@ -316,15 +324,16 @@ export class AIAssistantService { } return new AIAssistantKnowledgeBaseDataClient({ + assistantService: this, logger: this.options.logger.get('knowledgeBase'), + currentUser: opts.currentUser, elasticsearchClientPromise: this.options.elasticsearchClientPromise, - spaceId: opts.spaceId, - kibanaVersion: this.options.kibanaVersion, indexPatternsResourceName: this.resourceNames.aliases.knowledgeBase, ingestPipelineResourceName: this.resourceNames.pipelines.knowledgeBase, - currentUser: opts.currentUser, - ml: this.options.ml, getElserId: this.getElserId, + kibanaVersion: this.options.kibanaVersion, + ml: this.options.ml, + spaceId: opts.spaceId, }); } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts index 7382be25c9538..a8d2e96227d60 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts @@ -79,7 +79,9 @@ export const getKnowledgeBaseStatusRoute = ( // Code path for when `assistantKnowledgeBaseByDefault` FF is enabled let isSetupInProgress = false; if (enableKnowledgeBaseByDefault) { - const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient(true); + const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient( + false + ); if (!kbDataClient) { return response.custom({ body: { success: false }, statusCode: 500 }); } From ddf93a8dd1629121964a474f37b7251663519f72 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 15 May 2024 11:02:26 -0600 Subject: [PATCH 16/20] Relax id validation from UUID to UUID or NonEmptyString to support ES generated id's --- .../post_actions_connector_execute_route.gen.ts | 4 ++-- ...t_actions_connector_execute_route.schema.yaml | 4 +++- .../bulk_crud_anonymization_fields_route.gen.ts | 2 +- ...k_crud_anonymization_fields_route.schema.yaml | 4 +++- .../conversations/common_attributes.gen.ts | 2 +- .../conversations/common_attributes.schema.yaml | 4 +++- .../conversations/crud_conversation_route.gen.ts | 10 +++++----- .../crud_conversation_route.schema.yaml | 16 ++++++++++++---- .../crud_knowledge_base_route.gen.ts | 8 ++++---- .../crud_knowledge_base_route.schema.yaml | 12 +++++++++--- .../prompts/bulk_crud_prompts_route.gen.ts | 2 +- .../prompts/bulk_crud_prompts_route.schema.yaml | 4 +++- 12 files changed, 47 insertions(+), 25 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts index 1405a3612291e..3e33a9e43878c 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: 1 */ -import { UUID } from '../common_attributes.gen'; +import { UUID, NonEmptyString } from '../common_attributes.gen'; import { Replacements } from '../conversations/common_attributes.gen'; export type ExecuteConnectorRequestParams = z.infer; @@ -30,7 +30,7 @@ export type ExecuteConnectorRequestParamsInput = z.input; export const ExecuteConnectorRequestBody = z.object({ - conversationId: UUID.optional(), + conversationId: z.union([UUID, NonEmptyString]).optional(), message: z.string().optional(), model: z.string().optional(), subAction: z.enum(['invokeAI', 'invokeStream']), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml index 1372619492eb8..3b3f3458360c2 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml @@ -31,7 +31,9 @@ paths: - subAction properties: conversationId: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' message: type: string model: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts index 2af794d04811d..786e300f4e501 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts @@ -44,7 +44,7 @@ export const NormalizedAnonymizationFieldError = z.object({ export type AnonymizationFieldResponse = z.infer; export const AnonymizationFieldResponse = z.object({ - id: UUID, + id: z.union([UUID, NonEmptyString]), timestamp: NonEmptyString.optional(), field: z.string(), allowed: z.boolean().optional(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml index 68a94c189f326..0502d9f51eee7 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml @@ -103,7 +103,9 @@ components: - field properties: id: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' 'timestamp': $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' field: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts index 37d3fc7f38098..e2a523822b18b 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts @@ -152,7 +152,7 @@ export const ConversationSummary = z.object({ export type ErrorSchema = z.infer; export const ErrorSchema = z .object({ - id: UUID.optional(), + id: z.union([UUID, NonEmptyString]).optional(), error: z.object({ status_code: z.number().int().min(400), message: z.string(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml index 74817a8b7d75b..a226ae66a12dc 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml @@ -130,7 +130,9 @@ components: additionalProperties: false properties: id: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' error: type: object required: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts index e3a2be803db7a..bbd743166cd75 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts @@ -22,7 +22,7 @@ import { ConversationUpdateProps, ConversationMessageCreateProps, } from './common_attributes.gen'; -import { UUID } from '../common_attributes.gen'; +import { UUID, NonEmptyString } from '../common_attributes.gen'; export type AppendConversationMessageRequestParams = z.infer< typeof AppendConversationMessageRequestParams @@ -31,7 +31,7 @@ export const AppendConversationMessageRequestParams = z.object({ /** * The conversation's `id` value. */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type AppendConversationMessageRequestParamsInput = z.input< typeof AppendConversationMessageRequestParams @@ -60,7 +60,7 @@ export const DeleteConversationRequestParams = z.object({ /** * The conversation's `id` value. */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type DeleteConversationRequestParamsInput = z.input; @@ -72,7 +72,7 @@ export const ReadConversationRequestParams = z.object({ /** * The conversation's `id` value. */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type ReadConversationRequestParamsInput = z.input; @@ -84,7 +84,7 @@ export const UpdateConversationRequestParams = z.object({ /** * The conversation's `id` value. */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type UpdateConversationRequestParamsInput = z.input; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml index 5ae921be7a36f..1daa3e7edd36c 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml @@ -52,7 +52,9 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Indicates a successful call. @@ -86,7 +88,9 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: @@ -126,7 +130,9 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Indicates a successful call. @@ -162,7 +168,9 @@ paths: required: true description: The conversation's `id` value. schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts index d69ae81ea13e8..f5961344d9293 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts @@ -21,7 +21,7 @@ import { KnowledgeBaseEntryResponse, KnowledgeBaseEntryUpdateProps, } from './common_attributes.gen'; -import { UUID } from '../common_attributes.gen'; +import { UUID, NonEmptyString } from '../common_attributes.gen'; export type CreateKnowledgeBaseEntryRequestBody = z.infer< typeof CreateKnowledgeBaseEntryRequestBody @@ -41,7 +41,7 @@ export const DeleteKnowledgeBaseEntryRequestParams = z.object({ /** * The Knowledge Base Entry's `id` value */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type DeleteKnowledgeBaseEntryRequestParamsInput = z.input< typeof DeleteKnowledgeBaseEntryRequestParams @@ -57,7 +57,7 @@ export const ReadKnowledgeBaseEntryRequestParams = z.object({ /** * The Knowledge Base Entry's `id` value. */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type ReadKnowledgeBaseEntryRequestParamsInput = z.input< typeof ReadKnowledgeBaseEntryRequestParams @@ -73,7 +73,7 @@ export const UpdateKnowledgeBaseEntryRequestParams = z.object({ /** * The Knowledge Base Entry's `id` value */ - id: UUID, + id: z.union([UUID, NonEmptyString]), }); export type UpdateKnowledgeBaseEntryRequestParamsInput = z.input< typeof UpdateKnowledgeBaseEntryRequestParams diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml index e8fbe9eac5f64..e55333bacc9ac 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml @@ -45,7 +45,9 @@ paths: required: true description: The Knowledge Base Entry's `id` value. schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Successful request returning a Knowledge Base Entry @@ -72,7 +74,9 @@ paths: required: true description: The Knowledge Base Entry's `id` value schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: @@ -105,7 +109,9 @@ paths: required: true description: The Knowledge Base Entry's `id` value schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Successful request returning the deleted Knowledge Base Entry diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts index 9c5c1dd4f31ff..b4d0b42cdac7c 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts @@ -44,7 +44,7 @@ export const NormalizedPromptError = z.object({ export type PromptResponse = z.infer; export const PromptResponse = z.object({ - id: UUID, + id: z.union([UUID, NonEmptyString]), timestamp: NonEmptyString.optional(), name: z.string(), promptType: z.string(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml index f8938e6a653b1..d8df6913ed8ec 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml @@ -105,7 +105,9 @@ components: - content properties: id: - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + oneOf: + - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' + - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' 'timestamp': $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' name: From 7209e09230a887af9c8bc6670bfaee35258cdd33 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 17 May 2024 13:05:40 -0600 Subject: [PATCH 17/20] Update kb entry naming on naming data client get --- .../knowledge_base/get_knowledge_base_entry.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts index ee39a69e50d03..15bd7bfb910ba 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts @@ -71,10 +71,10 @@ export const getKnowledgeBaseEntry = async ({ index: knowledgeBaseIndex, seq_no_primary_term: true, }); - const conversation = transformESSearchToKnowledgeBaseEntry(response); - return conversation[0] ?? null; + const knowledgeBaseEntry = transformESSearchToKnowledgeBaseEntry(response); + return knowledgeBaseEntry[0] ?? null; } catch (err) { - logger.error(`Error fetching conversation: ${err} with id: ${id}`); + logger.error(`Error fetching knowledge base entry: ${err} with id: ${id}`); throw err; } }; From 311b6b2e4edbfb093301dcc5f13956952af11797 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 17 May 2024 17:07:04 -0600 Subject: [PATCH 18/20] Comments from PR review --- .../get_knowledge_base_entry.ts | 34 +++++++++---------- .../knowledge_base/index.ts | 3 +- .../server/ai_assistant_service/helpers.ts | 18 ++++++---- .../lib/data_stream/documents_data_writer.ts | 1 - .../knowledge_base/entries/create_route.ts | 4 +-- 5 files changed, 29 insertions(+), 31 deletions(-) diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts index 15bd7bfb910ba..c021e022765d0 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.ts @@ -16,7 +16,7 @@ export interface GetKnowledgeBaseEntryParams { logger: Logger; knowledgeBaseIndex: string; id: string; - user?: AuthenticatedUser | null; + user: AuthenticatedUser; } export const getKnowledgeBaseEntry = async ({ @@ -26,26 +26,24 @@ export const getKnowledgeBaseEntry = async ({ id, user, }: GetKnowledgeBaseEntryParams): Promise => { - const filterByUser = user - ? [ - { - nested: { - path: 'users', - query: { - bool: { - must: [ - { - match: user.profile_uid - ? { 'users.id': user.profile_uid } - : { 'users.name': user.username }, - }, - ], + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, }, - }, + ], }, }, - ] - : []; + }, + }, + ]; try { const response = await esClient.search({ query: { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 3957bab35e914..469b039bca84a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -204,8 +204,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { const kbDocsLoaded = (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0; if (!kbDocsLoaded) { this.options.logger.debug(`Loading KB docs...`); - const loadedKnowledgeBase = await loadESQL(esStore, this.options.logger); - this.options.logger.debug(`${loadedKnowledgeBase}`); + await loadESQL(esStore, this.options.logger); } else { this.options.logger.debug(`Knowledge Base docs already loaded!`); } diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts index 732339b18ded3..8503494cf1535 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts @@ -78,14 +78,18 @@ export const createPipeline = async ({ id, modelId, }: CreatePipelineParams): Promise => { - const response = await esClient.ingest.putPipeline( - knowledgeBaseIngestPipeline({ - id, - modelId, - }) - ); + try { + const response = await esClient.ingest.putPipeline( + knowledgeBaseIngestPipeline({ + id, + modelId, + }) + ); - return response.acknowledged; + return response.acknowledged; + } catch (e) { + return false; + } }; interface DeletePipelineParams { diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 693eda3ef6bd8..3cae3972c8bf6 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -241,7 +241,6 @@ export class DocumentsDataWriter implements DocumentsDataWriter { const documentCreateBody = params.documentsToCreate ? params.documentsToCreate.flatMap((document) => [ // Do not pre-gen _id for bulk create operations to avoid `version_conflict_engine_exception` - // TODO: Breaks Attack Discovery anonymization-field creation :( { create: { _index: this.options.index } }, document, ]) diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts index 434014adaaea4..0c31974a20785 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts @@ -51,9 +51,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout }, }); } - // const dataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient( - // false - // ); + const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); if (authenticatedUser == null) { return assistantResponse.error({ From 00468fa5cc0537d166f69b1135185654838e3cd0 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Fri, 17 May 2024 17:41:30 -0600 Subject: [PATCH 19/20] Remaining comments from PR review --- .../install_knowledge_base_button.tsx | 3 ++- .../knowledge_base/knowledge_base_settings.tsx | 15 +++++++++++---- .../knowledge_base/index.ts | 12 ++++++------ .../server/ai_assistant_service/index.ts | 13 +++++++------ .../routes/knowledge_base/post_knowledge_base.ts | 7 +++++++ 5 files changed, 33 insertions(+), 17 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx index 631bbae89b891..f5a82fd02c55d 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/install_knowledge_base_button.tsx @@ -24,10 +24,11 @@ export const InstallKnowledgeBaseButton: React.FC = React.memo(() => { const { assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault }, http, + toasts, } = useAssistantContext(); const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); - const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http }); + const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts }); const isSetupInProgress = kbStatus?.is_setup_in_progress || isSettingUpKB; const isSetupComplete = diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx index 6c0e2a8f5e244..8c83d1f3403e8 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx @@ -35,7 +35,8 @@ import { useKnowledgeBaseStatus } from './use_knowledge_base_status'; import { useSetupKnowledgeBase } from './use_setup_knowledge_base'; const ESQL_RESOURCE = 'esql'; -const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-kb'; +const KNOWLEDGE_BASE_INDEX_PATTERN_OLD = '.kibana-elastic-ai-assistant-kb'; +const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-knowledge-base-(SPACE)'; interface Props { knowledgeBase: KnowledgeBaseConfig; @@ -63,9 +64,11 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const isElserEnabled = kbStatus?.elser_exists ?? false; const isKnowledgeBaseEnabled = (kbStatus?.index_exists && kbStatus?.pipeline_exists) ?? false; const isESQLEnabled = kbStatus?.esql_exists ?? false; + const isSetupInProgress = kbStatus?.is_setup_in_progress ?? false; // Resource availability state - const isLoadingKb = isLoading || isFetching || isSettingUpKB || isDeletingUpKB; + const isLoadingKb = + isLoading || isFetching || isSettingUpKB || isDeletingUpKB || isSetupInProgress; const isKnowledgeBaseAvailable = knowledgeBase.isEnabledKnowledgeBase && kbStatus?.elser_exists; const isESQLAvailable = knowledgeBase.isEnabledKnowledgeBase && isKnowledgeBaseAvailable && isKnowledgeBaseEnabled; @@ -160,7 +163,11 @@ export const KnowledgeBaseSettings: React.FC = React.memo( const knowledgeBaseDescription = useMemo(() => { return isKnowledgeBaseEnabled ? ( - {i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED(KNOWLEDGE_BASE_INDEX_PATTERN)}{' '} + {i18n.KNOWLEDGE_BASE_DESCRIPTION_INSTALLED( + enableKnowledgeBaseByDefault + ? KNOWLEDGE_BASE_INDEX_PATTERN + : KNOWLEDGE_BASE_INDEX_PATTERN_OLD + )}{' '} {knowledgeBaseActionButton} ) : ( @@ -168,7 +175,7 @@ export const KnowledgeBaseSettings: React.FC = React.memo( {i18n.KNOWLEDGE_BASE_DESCRIPTION} {knowledgeBaseActionButton} ); - }, [isKnowledgeBaseEnabled, knowledgeBaseActionButton]); + }, [enableKnowledgeBaseByDefault, isKnowledgeBaseEnabled, knowledgeBaseActionButton]); ////////////////////////////////////////////////////////////////////////////////////////// // ESQL Resource diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 469b039bca84a..771ec35c07c51 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -24,13 +24,13 @@ import { EsKnowledgeBaseEntrySchema } from './types'; import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; import { isModelAlreadyExistsError } from './helpers'; -import { AIAssistantService } from '../../ai_assistant_service'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { - assistantService: AIAssistantService; ml: MlPluginSetup; getElserId: GetElser; + getIsKBSetupInProgress: () => boolean; ingestPipelineResourceName: string; + setIsKBSetupInProgress: (isInProgress: boolean) => void; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { constructor(public readonly options: KnowledgeBaseDataClientParams) { @@ -38,7 +38,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } public get isSetupInProgress() { - return this.options.assistantService.isKBSetupInProgress; + return this.options.getIsKBSetupInProgress(); } /** @@ -160,13 +160,13 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { esStore: ElasticsearchStore; soClient: SavedObjectsClientContract; }): Promise => { - if (this.options.assistantService.isKBSetupInProgress) { + if (this.options.getIsKBSetupInProgress()) { this.options.logger.debug('Knowledge Base setup already in progress'); return; } this.options.logger.debug('Starting Knowledge Base setup...'); - this.options.assistantService.isKBSetupInProgress = true; + this.options.setIsKBSetupInProgress(true); const elserId = await this.options.getElserId(); try { @@ -211,7 +211,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } catch (e) { this.options.logger.error(`Error setting up Knowledge Base: ${e.message}`); } - this.options.assistantService.isKBSetupInProgress = false; + this.options.setIsKBSetupInProgress(false); }; /** diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index f753df5e73a68..619d9e9bca256 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -70,7 +70,7 @@ export class AIAssistantService { private anonymizationFieldsDataStream: DataStreamSpacesAdapter; private resourceInitializationHelper: ResourceInstallationHelper; private initPromise: Promise; - private _isKBSetupInProgress: boolean = false; + private isKBSetupInProgress: boolean = false; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -109,11 +109,11 @@ export class AIAssistantService { return this.initialized; } - public get isKBSetupInProgress() { - return this._isKBSetupInProgress; + public getIsKBSetupInProgress() { + return this.isKBSetupInProgress; } - public set isKBSetupInProgress(inProgress: boolean) { - this._isKBSetupInProgress = inProgress; + public setIsKBSetupInProgress(isInProgress: boolean) { + this.isKBSetupInProgress = isInProgress; } private createDataStream: CreateDataStream = ({ resource, kibanaVersion, fieldMap }) => { @@ -324,15 +324,16 @@ export class AIAssistantService { } return new AIAssistantKnowledgeBaseDataClient({ - assistantService: this, logger: this.options.logger.get('knowledgeBase'), currentUser: opts.currentUser, elasticsearchClientPromise: this.options.elasticsearchClientPromise, indexPatternsResourceName: this.resourceNames.aliases.knowledgeBase, ingestPipelineResourceName: this.resourceNames.pipelines.knowledgeBase, getElserId: this.getElserId, + getIsKBSetupInProgress: this.getIsKBSetupInProgress.bind(this), kibanaVersion: this.options.kibanaVersion, ml: this.options.ml, + setIsKBSetupInProgress: this.setIsKBSetupInProgress.bind(this), spaceId: opts.spaceId, }); } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts index 063ffc36eda79..bc511d99eb63d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/post_knowledge_base.ts @@ -23,6 +23,10 @@ import { getKbResource } from './get_kb_resource'; import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader'; import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; +// Since we're awaiting on ELSER setup, this could take a bit (especially if ML needs to autoscale) +// Consider just returning if attempt was successful, and switch to client polling +const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes + /** * Load Knowledge Base index, pipeline, and resources (collection of documents) * @param router @@ -38,6 +42,9 @@ export const postKnowledgeBaseRoute = ( path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, options: { tags: ['access:elasticAssistant'], + timeout: { + idleSocket: ROUTE_HANDLER_TIMEOUT, + }, }, }) .addVersion( From a1f2820cf45b117f11180011bb7248ffb8bb7e7c Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Mon, 20 May 2024 10:07:11 -0600 Subject: [PATCH 20/20] Align id schema to NonEmptyString --- .../post_actions_connector_execute_route.gen.ts | 4 ++-- ...t_actions_connector_execute_route.schema.yaml | 4 +--- .../bulk_crud_anonymization_fields_route.gen.ts | 4 ++-- ...k_crud_anonymization_fields_route.schema.yaml | 4 +--- .../conversations/common_attributes.gen.ts | 8 ++++---- .../conversations/common_attributes.schema.yaml | 12 +++--------- .../conversations/crud_conversation_route.gen.ts | 10 +++++----- .../crud_conversation_route.schema.yaml | 16 ++++------------ .../knowledge_base/common_attributes.gen.ts | 6 +++--- .../knowledge_base/common_attributes.schema.yaml | 8 ++------ .../crud_knowledge_base_route.gen.ts | 8 ++++---- .../crud_knowledge_base_route.schema.yaml | 12 +++--------- .../prompts/bulk_crud_prompts_route.gen.ts | 4 ++-- .../prompts/bulk_crud_prompts_route.schema.yaml | 4 +--- 14 files changed, 37 insertions(+), 67 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts index 3e33a9e43878c..1b855650ddcfc 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: 1 */ -import { UUID, NonEmptyString } from '../common_attributes.gen'; +import { NonEmptyString } from '../common_attributes.gen'; import { Replacements } from '../conversations/common_attributes.gen'; export type ExecuteConnectorRequestParams = z.infer; @@ -30,7 +30,7 @@ export type ExecuteConnectorRequestParamsInput = z.input; export const ExecuteConnectorRequestBody = z.object({ - conversationId: z.union([UUID, NonEmptyString]).optional(), + conversationId: NonEmptyString.optional(), message: z.string().optional(), model: z.string().optional(), subAction: z.enum(['invokeAI', 'invokeStream']), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml index 3b3f3458360c2..9ea3c4107c7d1 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml @@ -31,9 +31,7 @@ paths: - subAction properties: conversationId: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' message: type: string model: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts index 786e300f4e501..d25c0198aaa1a 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: 2023-10-31 */ -import { UUID, NonEmptyString } from '../common_attributes.gen'; +import { NonEmptyString } from '../common_attributes.gen'; export type BulkActionSkipReason = z.infer; export const BulkActionSkipReason = z.literal('ANONYMIZATION_FIELD_NOT_MODIFIED'); @@ -44,7 +44,7 @@ export const NormalizedAnonymizationFieldError = z.object({ export type AnonymizationFieldResponse = z.infer; export const AnonymizationFieldResponse = z.object({ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, timestamp: NonEmptyString.optional(), field: z.string(), allowed: z.boolean().optional(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml index 0502d9f51eee7..663dafe763303 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml @@ -103,9 +103,7 @@ components: - field properties: id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' 'timestamp': $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' field: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts index e2a523822b18b..a8637be38c146 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: not applicable */ -import { NonEmptyString, UUID, User } from '../common_attributes.gen'; +import { NonEmptyString, User } from '../common_attributes.gen'; /** * trace Data @@ -152,7 +152,7 @@ export const ConversationSummary = z.object({ export type ErrorSchema = z.infer; export const ErrorSchema = z .object({ - id: z.union([UUID, NonEmptyString]).optional(), + id: NonEmptyString.optional(), error: z.object({ status_code: z.number().int().min(400), message: z.string(), @@ -162,7 +162,7 @@ export const ErrorSchema = z export type ConversationResponse = z.infer; export const ConversationResponse = z.object({ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, /** * The conversation title. */ @@ -207,7 +207,7 @@ export const ConversationResponse = z.object({ export type ConversationUpdateProps = z.infer; export const ConversationUpdateProps = z.object({ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, /** * The conversation title. */ diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml index a226ae66a12dc..49aaaa5663a1c 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/common_attributes.schema.yaml @@ -130,9 +130,7 @@ components: additionalProperties: false properties: id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' error: type: object required: @@ -156,9 +154,7 @@ components: - category properties: id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' title: type: string description: The conversation title. @@ -205,9 +201,7 @@ components: - id properties: id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' title: type: string description: The conversation title. diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts index bbd743166cd75..58dc70f38789b 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.gen.ts @@ -22,7 +22,7 @@ import { ConversationUpdateProps, ConversationMessageCreateProps, } from './common_attributes.gen'; -import { UUID, NonEmptyString } from '../common_attributes.gen'; +import { NonEmptyString } from '../common_attributes.gen'; export type AppendConversationMessageRequestParams = z.infer< typeof AppendConversationMessageRequestParams @@ -31,7 +31,7 @@ export const AppendConversationMessageRequestParams = z.object({ /** * The conversation's `id` value. */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type AppendConversationMessageRequestParamsInput = z.input< typeof AppendConversationMessageRequestParams @@ -60,7 +60,7 @@ export const DeleteConversationRequestParams = z.object({ /** * The conversation's `id` value. */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type DeleteConversationRequestParamsInput = z.input; @@ -72,7 +72,7 @@ export const ReadConversationRequestParams = z.object({ /** * The conversation's `id` value. */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type ReadConversationRequestParamsInput = z.input; @@ -84,7 +84,7 @@ export const UpdateConversationRequestParams = z.object({ /** * The conversation's `id` value. */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type UpdateConversationRequestParamsInput = z.input; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml index 1daa3e7edd36c..1a4f490a202c7 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/conversations/crud_conversation_route.schema.yaml @@ -52,9 +52,7 @@ paths: required: true description: The conversation's `id` value. schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Indicates a successful call. @@ -88,9 +86,7 @@ paths: required: true description: The conversation's `id` value. schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: @@ -130,9 +126,7 @@ paths: required: true description: The conversation's `id` value. schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Indicates a successful call. @@ -168,9 +162,7 @@ paths: required: true description: The conversation's `id` value. schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts index 805b9a06ef074..0d44cbe51e320 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: not applicable */ -import { NonEmptyString, UUID, User } from '../common_attributes.gen'; +import { NonEmptyString, User } from '../common_attributes.gen'; export type KnowledgeBaseEntryErrorSchema = z.infer; export const KnowledgeBaseEntryErrorSchema = z @@ -64,7 +64,7 @@ export const Vector = z.object({ export type KnowledgeBaseEntryResponse = z.infer; export const KnowledgeBaseEntryResponse = z.object({ timestamp: NonEmptyString.optional(), - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, /** * Time the Knowledge Base Entry was created */ @@ -99,7 +99,7 @@ export const KnowledgeBaseEntryResponse = z.object({ export type KnowledgeBaseEntryUpdateProps = z.infer; export const KnowledgeBaseEntryUpdateProps = z.object({ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, /** * Metadata about the Knowledge Base Entry */ diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml index 763a2acf34680..a9a9794852953 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml @@ -68,9 +68,7 @@ components: 'timestamp': $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' createdAt: description: Time the Knowledge Base Entry was created type: string @@ -105,9 +103,7 @@ components: - id properties: id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' metadata: $ref: '#/components/schemas/Metadata' description: Metadata about the Knowledge Base Entry diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts index f5961344d9293..92523a43b8e76 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts @@ -21,7 +21,7 @@ import { KnowledgeBaseEntryResponse, KnowledgeBaseEntryUpdateProps, } from './common_attributes.gen'; -import { UUID, NonEmptyString } from '../common_attributes.gen'; +import { NonEmptyString } from '../common_attributes.gen'; export type CreateKnowledgeBaseEntryRequestBody = z.infer< typeof CreateKnowledgeBaseEntryRequestBody @@ -41,7 +41,7 @@ export const DeleteKnowledgeBaseEntryRequestParams = z.object({ /** * The Knowledge Base Entry's `id` value */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type DeleteKnowledgeBaseEntryRequestParamsInput = z.input< typeof DeleteKnowledgeBaseEntryRequestParams @@ -57,7 +57,7 @@ export const ReadKnowledgeBaseEntryRequestParams = z.object({ /** * The Knowledge Base Entry's `id` value. */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type ReadKnowledgeBaseEntryRequestParamsInput = z.input< typeof ReadKnowledgeBaseEntryRequestParams @@ -73,7 +73,7 @@ export const UpdateKnowledgeBaseEntryRequestParams = z.object({ /** * The Knowledge Base Entry's `id` value */ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, }); export type UpdateKnowledgeBaseEntryRequestParamsInput = z.input< typeof UpdateKnowledgeBaseEntryRequestParams diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml index e55333bacc9ac..6db7da89f55e5 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml @@ -45,9 +45,7 @@ paths: required: true description: The Knowledge Base Entry's `id` value. schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Successful request returning a Knowledge Base Entry @@ -74,9 +72,7 @@ paths: required: true description: The Knowledge Base Entry's `id` value schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: @@ -109,9 +105,7 @@ paths: required: true description: The Knowledge Base Entry's `id` value schema: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' responses: 200: description: Successful request returning the deleted Knowledge Base Entry diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts index b4d0b42cdac7c..234b1157b5f71 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen.ts @@ -16,7 +16,7 @@ import { z } from 'zod'; * version: 2023-10-31 */ -import { UUID, NonEmptyString, User } from '../common_attributes.gen'; +import { NonEmptyString, User } from '../common_attributes.gen'; export type BulkActionSkipReason = z.infer; export const BulkActionSkipReason = z.literal('PROMPT_FIELD_NOT_MODIFIED'); @@ -44,7 +44,7 @@ export const NormalizedPromptError = z.object({ export type PromptResponse = z.infer; export const PromptResponse = z.object({ - id: z.union([UUID, NonEmptyString]), + id: NonEmptyString, timestamp: NonEmptyString.optional(), name: z.string(), promptType: z.string(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml index d8df6913ed8ec..355583ae86667 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.schema.yaml @@ -105,9 +105,7 @@ components: - content properties: id: - oneOf: - - $ref: '../common_attributes.schema.yaml#/components/schemas/UUID' - - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' 'timestamp': $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' name: