diff --git a/x-pack/plugins/elastic_assistant/kibana.jsonc b/x-pack/plugins/elastic_assistant/kibana.jsonc index 435ec0b916d01..2b06f1e0db65e 100644 --- a/x-pack/plugins/elastic_assistant/kibana.jsonc +++ b/x-pack/plugins/elastic_assistant/kibana.jsonc @@ -17,9 +17,11 @@ "ml", "taskManager", "licensing", + "llmTasks", "inference", + "productDocBase", "spaces", "security" ] } -} \ No newline at end of file +} 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 77bd6b00105b6..3837c158ba199 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -53,6 +53,7 @@ export const createMockClients = () => { getSpaceId: jest.fn(), getCurrentUser: jest.fn(), inference: jest.fn(), + llmTasks: jest.fn(), }, savedObjectsClient: core.savedObjects.client, @@ -145,6 +146,7 @@ const createElasticAssistantRequestContextMock = ( getServerBasePath: jest.fn(), getSpaceId: jest.fn().mockReturnValue('default'), inference: { getClient: jest.fn() }, + llmTasks: { retrieveDocumentationAvailable: jest.fn(), retrieveDocumentation: jest.fn() }, core: clients.core, telemetry: clients.elasticAssistant.telemetry, }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.test.ts new file mode 100644 index 0000000000000..8cd020149c433 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/helpers.test.ts @@ -0,0 +1,79 @@ +/* + * 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 { ensureProductDocumentationInstalled } from './helpers'; +import { loggerMock } from '@kbn/logging-mocks'; + +const mockLogger = loggerMock.create(); +const mockProductDocManager = { + getStatus: jest.fn(), + install: jest.fn(), + uninstall: jest.fn(), + update: jest.fn(), +}; + +describe('helpers', () => { + describe('ensureProductDocumentationInstalled', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('should install product documentation if not installed', async () => { + mockProductDocManager.getStatus.mockResolvedValue({ status: 'uninstalled' }); + mockProductDocManager.install.mockResolvedValue(null); + + await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger); + + expect(mockProductDocManager.getStatus).toHaveBeenCalled(); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Installing product documentation for AIAssistantService' + ); + expect(mockProductDocManager.install).toHaveBeenCalled(); + expect(mockLogger.debug).toHaveBeenNthCalledWith( + 2, + 'Successfully installed product documentation for AIAssistantService' + ); + }); + + it('should not install product documentation if already installed', async () => { + mockProductDocManager.getStatus.mockResolvedValue({ status: 'installed' }); + + await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger); + + expect(mockProductDocManager.getStatus).toHaveBeenCalled(); + expect(mockProductDocManager.install).not.toHaveBeenCalled(); + expect(mockLogger.debug).not.toHaveBeenCalledWith( + 'Installing product documentation for AIAssistantService' + ); + }); + it('should log a warning if install fails', async () => { + mockProductDocManager.getStatus.mockResolvedValue({ status: 'not_installed' }); + mockProductDocManager.install.mockRejectedValue(new Error('Install failed')); + + await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger); + + expect(mockProductDocManager.getStatus).toHaveBeenCalled(); + expect(mockProductDocManager.install).toHaveBeenCalled(); + + expect(mockLogger.warn).toHaveBeenCalledWith( + 'Failed to install product documentation for AIAssistantService: Install failed' + ); + }); + + it('should log a warning if getStatus fails', async () => { + mockProductDocManager.getStatus.mockRejectedValue(new Error('Status check failed')); + + await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger); + + expect(mockProductDocManager.getStatus).toHaveBeenCalled(); + expect(mockLogger.warn).toHaveBeenCalledWith( + 'Failed to get status of product documentation installation for AIAssistantService: Status check failed' + ); + expect(mockProductDocManager.install).not.toHaveBeenCalled(); + }); + }); +}); 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 2a4ad628eb757..9067e42ca88bb 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 @@ -11,6 +11,8 @@ import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-ser import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import { DeleteByQueryRequest } from '@elastic/elasticsearch/lib/api/types'; import { i18n } from '@kbn/i18n'; +import { ProductDocBaseStartContract } from '@kbn/product-doc-base-plugin/server'; +import type { Logger } from '@kbn/logging'; import { getResourceName } from '.'; import { knowledgeBaseIngestPipeline } from '../ai_assistant_data_clients/knowledge_base/ingest_pipeline'; import { GetElser } from '../types'; @@ -141,3 +143,25 @@ const ESQL_QUERY_GENERATION_TITLE = i18n.translate( defaultMessage: 'ES|QL Query Generation', } ); + +export const ensureProductDocumentationInstalled = async ( + productDocManager: ProductDocBaseStartContract['management'], + logger: Logger +) => { + try { + const { status } = await productDocManager.getStatus(); + if (status !== 'installed') { + logger.debug(`Installing product documentation for AIAssistantService`); + try { + await productDocManager.install(); + logger.debug(`Successfully installed product documentation for AIAssistantService`); + } catch (e) { + logger.warn(`Failed to install product documentation for AIAssistantService: ${e.message}`); + } + } + } catch (e) { + logger.warn( + `Failed to get status of product documentation installation for AIAssistantService: ${e.message}` + ); + } +}; 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 4bfd4da6cfcbf..c60fe9a220482 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 @@ -122,6 +122,12 @@ describe('AI Assistant Service', () => { kibanaVersion: '8.8.0', ml, taskManager: taskManagerMock.createSetup(), + productDocManager: Promise.resolve({ + getStatus: jest.fn(), + install: jest.fn(), + update: jest.fn(), + uninstall: jest.fn(), + }), }; }); 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 233b5781ddf68..ff0f95340d466 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 @@ -12,6 +12,7 @@ import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import { Subject } from 'rxjs'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; +import { ProductDocBaseStartContract } from '@kbn/product-doc-base-plugin/server'; import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration'; import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; @@ -35,7 +36,12 @@ import { } from '../ai_assistant_data_clients/knowledge_base'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights'; -import { createGetElserId, createPipeline, pipelineExists } from './helpers'; +import { + createGetElserId, + createPipeline, + ensureProductDocumentationInstalled, + pipelineExists, +} from './helpers'; import { hasAIAssistantLicense } from '../routes/helpers'; const TOTAL_FIELDS_LIMIT = 2500; @@ -51,6 +57,7 @@ export interface AIAssistantServiceOpts { ml: MlPluginSetup; taskManager: TaskManagerSetupContract; pluginStop$: Subject; + productDocManager: Promise; } export interface CreateAIAssistantClientParams { @@ -87,6 +94,7 @@ export class AIAssistantService { private initPromise: Promise; private isKBSetupInProgress: boolean = false; private hasInitializedV2KnowledgeBase: boolean = false; + private productDocManager?: ProductDocBaseStartContract['management']; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -129,6 +137,13 @@ export class AIAssistantService { this.initPromise, this.installAndUpdateSpaceLevelResources.bind(this) ); + options.productDocManager + .then((productDocManager) => { + this.productDocManager = productDocManager; + }) + .catch((error) => { + this.options.logger.warn(`Failed to initialize productDocManager: ${error.message}`); + }); } public isInitialized() { @@ -183,6 +198,11 @@ export class AIAssistantService { this.options.logger.debug(`Initializing resources for AIAssistantService`); const esClient = await this.options.elasticsearchClientPromise; + if (this.productDocManager) { + // install product documentation without blocking other resources + void ensureProductDocumentationInstalled(this.productDocManager, this.options.logger); + } + await this.conversationsDataStream.install({ esClient, logger: this.options.logger, 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 7dea19755a686..abef39d8b2e25 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 @@ -17,6 +17,7 @@ import { PublicMethodsOf } from '@kbn/utility-types'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; import { AnalyticsServiceSetup } from '@kbn/core-analytics-server'; import { TelemetryParams } from '@kbn/langchain/server/tracers/telemetry/telemetry_tracer'; +import type { LlmTasksPluginStart } from '@kbn/llm-tasks-plugin/server'; import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; @@ -45,10 +46,11 @@ export interface AgentExecutorParams { dataClients?: AssistantDataClients; esClient: ElasticsearchClient; langChainMessages: BaseMessage[]; + llmTasks?: LlmTasksPluginStart; llmType?: string; isOssModel?: boolean; - logger: Logger; inference: InferenceServerStart; + logger: Logger; onNewReplacements?: (newReplacements: Replacements) => void; replacements: Replacements; isStream?: T; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index e9d2c1dd2618b..4ddd3eae11624 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -35,6 +35,7 @@ export const callAssistantGraph: AgentExecutor = async ({ esClient, inference, langChainMessages, + llmTasks, llmType, isOssModel, logger: parentLogger, @@ -106,6 +107,7 @@ export const callAssistantGraph: AgentExecutor = async ({ inference, isEnabledKnowledgeBase, kbDataClient: dataClients?.kbDataClient, + llmTasks, logger, onNewReplacements, replacements, diff --git a/x-pack/plugins/elastic_assistant/server/plugin.ts b/x-pack/plugins/elastic_assistant/server/plugin.ts index 110dbbc05f2a6..e93e3786b123c 100755 --- a/x-pack/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/plugins/elastic_assistant/server/plugin.ts @@ -63,6 +63,9 @@ export class ElasticAssistantPlugin elasticsearchClientPromise: core .getStartServices() .then(([{ elasticsearch }]) => elasticsearch.client.asInternalUser), + productDocManager: core + .getStartServices() + .then(([_, { productDocBase }]) => productDocBase.management), pluginStop$: this.pluginStop$, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts index 5d277abb00667..8cd2f0fd801d0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts @@ -61,6 +61,7 @@ const mockContext = { getRegisteredFeatures: jest.fn(() => defaultAssistantFeatures), logger: loggingSystemMock.createLogger(), telemetry: { ...coreMock.createSetup().analytics, reportEvent }, + llmTasks: { retrieveDocumentationAvailable: jest.fn(), retrieveDocumentation: jest.fn() }, getCurrentUser: () => ({ username: 'user', email: 'email', diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts index 35b4999a30249..cf1cff3d6201d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts @@ -69,6 +69,8 @@ export const chatCompleteRoute = ( try { telemetry = ctx.elasticAssistant.telemetry; const inference = ctx.elasticAssistant.inference; + const productDocsAvailable = + (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false; // Perform license and authenticated user checks const checkResponse = performChecks({ @@ -217,6 +219,7 @@ export const chatCompleteRoute = ( response, telemetry, responseLanguage: request.body.responseLanguage, + ...(productDocsAvailable ? { llmTasks: ctx.elasticAssistant.llmTasks } : {}), }); } catch (err) { const error = transformError(err as Error); 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 4e4b7e5fcd251..3c61fe75f6d00 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 @@ -150,6 +150,8 @@ export const postEvaluateRoute = ( const esClient = ctx.core.elasticsearch.client.asCurrentUser; const inference = ctx.elasticAssistant.inference; + const productDocsAvailable = + (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false; // Data clients const anonymizationFieldsDataClient = @@ -280,6 +282,7 @@ export const postEvaluateRoute = ( connectorId: connector.id, size, telemetry: ctx.elasticAssistant.telemetry, + ...(productDocsAvailable ? { llmTasks: ctx.elasticAssistant.llmTasks } : {}), }; const tools: StructuredTool[] = assistantTools.flatMap( diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts index 23ec7011be5b7..25d4ce1a2ec45 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -29,6 +29,7 @@ import { ActionsClient } from '@kbn/actions-plugin/server'; import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; +import type { LlmTasksPluginStart } from '@kbn/llm-tasks-plugin/server'; import { INVOKE_ASSISTANT_SUCCESS_EVENT } from '../lib/telemetry/event_based_telemetry'; import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; import { FindResponse } from '../ai_assistant_data_clients/find'; @@ -215,6 +216,7 @@ export interface LangChainExecuteParams { telemetry: AnalyticsServiceSetup; actionTypeId: string; connectorId: string; + llmTasks?: LlmTasksPluginStart; inference: InferenceServerStart; isOssModel?: boolean; conversationId?: string; @@ -246,6 +248,7 @@ export const langChainExecute = async ({ isOssModel, context, actionsClient, + llmTasks, inference, request, logger, @@ -301,6 +304,7 @@ export const langChainExecute = async ({ conversationId, connectorId, esClient, + llmTasks, inference, isStream, llmType: getLlmType(actionTypeId), 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 a7abac27dac6f..9f4d0beb3caff 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 @@ -67,6 +67,7 @@ const mockContext = { actions: { getActionsClientWithRequest: jest.fn().mockResolvedValue(actionsClient), }, + llmTasks: { retrieveDocumentationAvailable: jest.fn(), retrieveDocumentation: jest.fn() }, getRegisteredTools: jest.fn(() => []), getRegisteredFeatures: jest.fn(() => defaultAssistantFeatures), logger: loggingSystemMock.createLogger(), 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 43264a6c1f54b..55c23629c5de1 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 @@ -97,6 +97,8 @@ export const postActionsConnectorExecuteRoute = ( // get the actions plugin start contract from the request context: const actions = ctx.elasticAssistant.actions; const inference = ctx.elasticAssistant.inference; + const productDocsAvailable = + (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false; const actionsClient = await actions.getActionsClientWithRequest(request); const connectors = await actionsClient.getBulk({ ids: [connectorId] }); const connector = connectors.length > 0 ? connectors[0] : undefined; @@ -150,6 +152,7 @@ export const postActionsConnectorExecuteRoute = ( response, telemetry, systemPrompt, + ...(productDocsAvailable ? { llmTasks: ctx.elasticAssistant.llmTasks } : {}), }); } catch (err) { logger.error(err); 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 ef921d7c91a28..30045b3da8ad9 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 @@ -78,7 +78,7 @@ export class RequestContextFactory implements IRequestContextFactory { getRegisteredFeatures: (pluginName: string) => { return appContextService.getRegisteredFeatures(pluginName); }, - + llmTasks: startPlugins.llmTasks, inference: startPlugins.inference, telemetry: core.analytics, diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 93f35d11eb877..6158acde679ff 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -20,6 +20,7 @@ import type { Logger, SecurityServiceStart, } from '@kbn/core/server'; +import type { LlmTasksPluginStart } from '@kbn/llm-tasks-plugin/server'; import { type MlPluginSetup } from '@kbn/ml-plugin/server'; import { DynamicStructuredTool, Tool } from '@langchain/core/tools'; import { SpacesPluginSetup, SpacesPluginStart } from '@kbn/spaces-plugin/server'; @@ -46,6 +47,7 @@ import { } from '@kbn/langchain/server'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; +import { ProductDocBaseStartContract } from '@kbn/product-doc-base-plugin/server'; import type { GetAIAssistantKnowledgeBaseDataClientParams } from './ai_assistant_data_clients/knowledge_base'; import { AttackDiscoveryDataClient } from './lib/attack_discovery/persistence'; import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations'; @@ -111,10 +113,12 @@ export interface ElasticAssistantPluginSetupDependencies { } export interface ElasticAssistantPluginStartDependencies { actions: ActionsPluginStart; + llmTasks: LlmTasksPluginStart; inference: InferenceServerStart; spaces?: SpacesPluginStart; security: SecurityServiceStart; licensing: LicensingPluginStart; + productDocBase: ProductDocBaseStartContract; } export interface ElasticAssistantApiRequestHandlerContext { @@ -134,6 +138,7 @@ export interface ElasticAssistantApiRequestHandlerContext { getDefendInsightsDataClient: () => Promise; getAIAssistantPromptsDataClient: () => Promise; getAIAssistantAnonymizationFieldsDataClient: () => Promise; + llmTasks: LlmTasksPluginStart; inference: InferenceServerStart; telemetry: AnalyticsServiceSetup; } @@ -230,6 +235,7 @@ export interface AssistantToolParams { kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; llm?: ActionsClientLlm | AssistantToolLlm; + llmTasks?: LlmTasksPluginStart; isOssModel?: boolean; logger: Logger; onNewReplacements?: (newReplacements: Replacements) => void; diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index 52ed30dde67f8..5b9a7cb3466db 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -50,7 +50,9 @@ "@kbn/zod", "@kbn/inference-plugin", "@kbn/data-views-plugin", - "@kbn/core-analytics-server" + "@kbn/core-analytics-server", + "@kbn/llm-tasks-plugin", + "@kbn/product-doc-base-plugin" ], "exclude": [ "target/**/*", diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index f7824e688afe2..dc32e01335b30 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -5,8 +5,7 @@ * 2.0. */ -import type { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; - +import { PRODUCT_DOCUMENTATION_TOOL } from './product_docs/product_documentation_tool'; import { NL_TO_ESQL_TOOL } from './esql/nl_to_esql_tool'; import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool'; @@ -15,12 +14,13 @@ import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_r import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; import { SECURITY_LABS_KNOWLEDGE_BASE_TOOL } from './security_labs/security_labs_tool'; -export const assistantTools: AssistantTool[] = [ +export const assistantTools = [ ALERT_COUNTS_TOOL, DEFEND_INSIGHTS_TOOL, - NL_TO_ESQL_TOOL, KNOWLEDGE_BASE_RETRIEVAL_TOOL, KNOWLEDGE_BASE_WRITE_TOOL, + NL_TO_ESQL_TOOL, OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL, + PRODUCT_DOCUMENTATION_TOOL, SECURITY_LABS_KNOWLEDGE_BASE_TOOL, ]; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts new file mode 100644 index 0000000000000..d8d7e5995c92e --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts @@ -0,0 +1,93 @@ +/* + * 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 { RetrievalQAChain } from 'langchain/chains'; +import type { DynamicStructuredTool, DynamicTool } from '@langchain/core/tools'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; +import { loggerMock } from '@kbn/logging-mocks'; +import { PRODUCT_DOCUMENTATION_TOOL } from './product_documentation_tool'; +import type { LlmTasksPluginStart } from '@kbn/llm-tasks-plugin/server'; + +describe('ProductDocumentationTool', () => { + const chain = {} as RetrievalQAChain; + const esClient = { + search: jest.fn().mockResolvedValue({}), + } as unknown as ElasticsearchClient; + const request = {} as unknown as KibanaRequest; + const logger = loggerMock.create(); + const retrieveDocumentation = jest.fn(); + const llmTasks = { + retrieveDocumentation, + retrieveDocumentationAvailable: jest.fn(), + } as LlmTasksPluginStart; + const connectorId = 'fake-connector'; + const defaultArgs = { + chain, + esClient, + logger, + request, + llmTasks, + connectorId, + isEnabledKnowledgeBase: true, + }; + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('isSupported', () => { + it('returns true if connectorId and llmTasks have values', () => { + expect(PRODUCT_DOCUMENTATION_TOOL.isSupported(defaultArgs)).toBe(true); + }); + }); + + describe('getTool', () => { + it('should return a tool as expected when all required values are present', () => { + const tool = PRODUCT_DOCUMENTATION_TOOL.getTool(defaultArgs) as DynamicTool; + expect(tool.name).toEqual('ProductDocumentationTool'); + expect(tool.tags).toEqual(['product-documentation']); + }); + + it('returns null if llmTasks plugin is not provided', () => { + const tool = PRODUCT_DOCUMENTATION_TOOL.getTool({ + ...defaultArgs, + llmTasks: undefined, + }); + + expect(tool).toBeNull(); + }); + + it('returns null if connectorId is not provided', () => { + const tool = PRODUCT_DOCUMENTATION_TOOL.getTool({ + ...defaultArgs, + connectorId: undefined, + }); + + expect(tool).toBeNull(); + }); + }); + describe('DynamicStructuredTool', () => { + beforeEach(() => { + retrieveDocumentation.mockResolvedValue({ documents: [] }); + }); + it('the tool invokes retrieveDocumentation', async () => { + const tool = PRODUCT_DOCUMENTATION_TOOL.getTool(defaultArgs) as DynamicStructuredTool; + + await tool.func({ query: 'What is Kibana Security?', product: 'kibana' }); + + expect(retrieveDocumentation).toHaveBeenCalledWith({ + searchTerm: 'What is Kibana Security?', + products: ['kibana'], + max: 3, + connectorId: 'fake-connector', + request, + functionCalling: 'native', + }); + }); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts new file mode 100644 index 0000000000000..071a435e1311f --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts @@ -0,0 +1,79 @@ +/* + * 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 { DynamicStructuredTool } from '@langchain/core/tools'; + +import { z } from '@kbn/zod'; +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import { APP_UI_ID } from '../../../../common'; + +const toolDetails = { + description: + 'Use this tool to retrieve documentation about Elastic products. You can retrieve documentation about the Elastic stack, such as Kibana and Elasticsearch, or for Elastic solutions, such as Elastic Security, Elastic Observability or Elastic Enterprise Search.', + id: 'product-documentation-tool', + name: 'ProductDocumentationTool', +}; +export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { + ...toolDetails, + sourceRegister: APP_UI_ID, + isSupported: (params: AssistantToolParams): params is AssistantToolParams => { + return params.llmTasks != null && params.connectorId != null; + }, + getTool(params: AssistantToolParams) { + if (!this.isSupported(params)) return null; + + const { connectorId, llmTasks, request } = params as AssistantToolParams; + + // This check is here in order to satisfy TypeScript + if (llmTasks == null || connectorId == null) return null; + + return new DynamicStructuredTool({ + name: toolDetails.name, + description: toolDetails.description, + schema: z.object({ + query: z.string().describe( + `The query to use to retrieve documentation + Examples: + - "How to enable TLS for Elasticsearch?" + - "What is Kibana Security?"` + ), + product: z + .enum(['kibana', 'elasticsearch', 'observability', 'security']) + .describe( + `If specified, will filter the products to retrieve documentation for + Possible options are: + - "kibana": Kibana product + - "elasticsearch": Elasticsearch product + - "observability": Elastic Observability solution + - "security": Elastic Security solution + If not specified, will search against all products + ` + ) + .optional(), + }), + func: async ({ query, product }) => { + const response = await llmTasks.retrieveDocumentation({ + searchTerm: query, + products: product ? [product] : undefined, + max: 3, + connectorId, + request, + // o11y specific parameter, hardcode to native as we do not utilize the other value (simulated) + functionCalling: 'native', + }); + + return { + content: { + documents: response.documents, + }, + }; + }, + tags: ['product-documentation'], + // TODO: Remove after ZodAny is fixed https://github.com/langchain-ai/langchainjs/blob/main/langchain-core/src/tools.ts + }) as unknown as DynamicStructuredTool; + }, +}; diff --git a/x-pack/plugins/security_solution/tsconfig.json b/x-pack/plugins/security_solution/tsconfig.json index 4ed7e1cbdd35f..4d11804796e2b 100644 --- a/x-pack/plugins/security_solution/tsconfig.json +++ b/x-pack/plugins/security_solution/tsconfig.json @@ -235,6 +235,7 @@ "@kbn/discover-shared-plugin", "@kbn/react-hooks", "@kbn/index-adapter", - "@kbn/core-http-server-utils" + "@kbn/core-http-server-utils", + "@kbn/llm-tasks-plugin" ] }