diff --git a/.buildkite/ftr_security_serverless_configs.yml b/.buildkite/ftr_security_serverless_configs.yml index a6aca2fe7750e0..89ebb4aa12cd4c 100644 --- a/.buildkite/ftr_security_serverless_configs.yml +++ b/.buildkite/ftr_security_serverless_configs.yml @@ -77,6 +77,8 @@ enabled: - x-pack/test/security_solution_api_integration/test_suites/detections_response/user_roles/trial_license_complete_tier/configs/serverless.config.ts - x-pack/test/security_solution_api_integration/test_suites/genai/nlp_cleanup_task/trial_license_complete_tier/configs/serverless.config.ts - x-pack/test/security_solution_api_integration/test_suites/genai/nlp_cleanup_task/basic_license_essentials_tier/configs/serverless.config.ts + - x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts + - x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/serverless.config.ts - x-pack/test/security_solution_api_integration/test_suites/entity_analytics/risk_engine/trial_license_complete_tier/configs/serverless.config.ts - x-pack/test/security_solution_api_integration/test_suites/entity_analytics/risk_engine/basic_license_essentials_tier/configs/serverless.config.ts - x-pack/test/security_solution_api_integration/test_suites/lists_and_exception_lists/exception_lists_items/trial_license_complete_tier/configs/serverless.config.ts diff --git a/packages/kbn-cli-dev-mode/src/watcher.ts b/packages/kbn-cli-dev-mode/src/watcher.ts index 6f7f70466ae251..3c9763e0543aab 100644 --- a/packages/kbn-cli-dev-mode/src/watcher.ts +++ b/packages/kbn-cli-dev-mode/src/watcher.ts @@ -95,7 +95,7 @@ export class Watcher { ignore: [ '**/{node_modules,target,public,coverage,__*__}/**', '**/*.{test,spec,story,stories}.*', - '**/*.{md,sh,txt}', + '**/*.{http,md,sh,txt}', '**/debug.log', ], } diff --git a/src/dev/build/tasks/copy_legacy_source_task.ts b/src/dev/build/tasks/copy_legacy_source_task.ts index c4005bded29b74..2a377c42264de4 100644 --- a/src/dev/build/tasks/copy_legacy_source_task.ts +++ b/src/dev/build/tasks/copy_legacy_source_task.ts @@ -38,9 +38,10 @@ export const CopyLegacySource: Task = { '!**/jest*', '!**/*.{story,stories}.{js,ts}', '!**/{test_mocks,stubs}.ts', - '!**/*.{scss,console,d.ts,sh,md,mdx,asciidoc,docnav.json}', + '!**/*.{scss,console,d.ts,sh,md,mdx,asciidoc,docnav.json,http}', '!**/*.{test,test.mocks,mock,mocks,spec}.*', '!**/{packages,dev_docs,docs,public,__stories__,storybook,.storybook,ftr_e2e,e2e,scripts,test,tests,test_resources,test_data,__tests__,manual_tests,__jest__,__snapshots__,__mocks__,mock_responses,mocks,fixtures,__fixtures__,cypress,integration_tests}/**', + '!**/http-client.env.json', // explicitly exclude every package directory outside of the root packages dir ...getPackages(config.resolveFromRepo('.')).flatMap((p) => diff --git a/src/dev/precommit_hook/casing_check_config.js b/src/dev/precommit_hook/casing_check_config.js index 1b85b8fb40b5f9..9feec923b869a4 100644 --- a/src/dev/precommit_hook/casing_check_config.js +++ b/src/dev/precommit_hook/casing_check_config.js @@ -83,6 +83,9 @@ export const IGNORE_FILE_GLOBS = [ // ecs templates '**/ecs/fields/**/*', + + // Support for including http-client.env.json configurations + '**/http-client.env.json', ]; /** diff --git a/x-pack/packages/kbn-elastic-assistant-common/.gitignore b/x-pack/packages/kbn-elastic-assistant-common/.gitignore new file mode 100644 index 00000000000000..05014e534ec88a --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/.gitignore @@ -0,0 +1,2 @@ +# Private configs for .http files +http-client.private.env.json \ No newline at end of file diff --git a/x-pack/packages/kbn-elastic-assistant-common/env/README.md b/x-pack/packages/kbn-elastic-assistant-common/env/README.md new file mode 100644 index 00000000000000..b90d83aaca8123 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/env/README.md @@ -0,0 +1,33 @@ +With https://github.com/elastic/kibana/pull/186566, we've introduced a few sample `*.http` files for easier development/testing. These files are supported out of the box in JetBrains IDE's or in VSCode with the [httpyac](https://httpyac.github.io/) (and many other) extensions. Since the configuration for these files includes a `-` in the name, a few @elastic/kibana-operations files have been updated to exclude them from checks and being included in the distribution. + +You can read more about `http` files [here](https://www.jetbrains.com/help/webstorm/http-client-in-product-code-editor.html) and for the spec see this repo [here](https://github.com/JetBrains/http-request-in-editor-spec/blob/master/spec.md). If we find these useful, we could add support to our [OpenAPI Generator](https://openapi-generator.tech/docs/generators/jetbrains-http-client) to create these automatically. They currently live co-located next to the OAS and generated schema files here: + +``` +x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.http +x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/crud_knowledge_base_entries_route.http +``` + +and the main config here in this directory: + +``` +x-pack/packages/kbn-elastic-assistant-common/env/http-client.env.json +``` + +The `x-pack/packages/kbn-elastic-assistant-common/.gitignore` has been updated to ignore `http-client.private.env.json` files locally, which is how you can override the config as you'd like. This is helpful to add variables like `basePath` as below: + +``` +{ + "dev": { + "basePath": "/kbn" + } +} +``` + +To use them, just open the corresponding `*.http` for the API you want to test, and click `Send`, and the response will open in another tab. Here is what that looks like for creating one of the new `IndexEntry` KB documents that have been introduced in the initial PR: + +

+ +

+ + +For continuing this effort, https://github.com/elastic/kibana/issues/192386 has been created. diff --git a/x-pack/packages/kbn-elastic-assistant-common/env/http-client.env.json b/x-pack/packages/kbn-elastic-assistant-common/env/http-client.env.json new file mode 100644 index 00000000000000..3f48e01e099fac --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/env/http-client.env.json @@ -0,0 +1,15 @@ +{ + "dev": { + "host": "localhost", + "port": "5601", + "basePath": "", + "elasticApiVersion": "1", + "auth": { + "username": "elastic", + "password": "changeme" + }, + "appContext": { + "security": "%7B%22type%22%3A%22application%22%2C%22name%22%3A%22securitySolutionUI%22%2C%22url%22%3A%22%2Fkbn%2Fapp%2Fsecurity%22%7D" + } + } +} 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 c6c6c758a8cf6e..6304bfa4786cf8 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 @@ -49,10 +49,10 @@ 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/crud_knowledge_base_route.gen'; -export * from './knowledge_base/find_knowledge_base_entries_route.gen'; +export * from './knowledge_base/entries/bulk_crud_knowledge_base_entries_route.gen'; +export * from './knowledge_base/entries/common_attributes.gen'; +export * from './knowledge_base/entries/crud_knowledge_base_entries_route.gen'; +export * from './knowledge_base/entries/find_knowledge_base_entries_route.gen'; export * from './prompts/find_prompts_route.gen'; export { PromptResponse, PromptTypeEnum } from './prompts/bulk_crud_prompts_route.gen'; 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 deleted file mode 100644 index 2736baf8a135d5..00000000000000 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts +++ /dev/null @@ -1,119 +0,0 @@ -/* - * 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. - */ - -/* - * 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 - */ - -import { z } from '@kbn/zod'; - -import { NonEmptyString, User } from '../common_attributes.gen'; - -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 - */ -export type Metadata = z.infer; -export const Metadata = z.object({ - /** - * Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc - */ - kbResource: z.string(), - /** - * Source document name or filepath - */ - source: z.string(), - /** - * Whether or not this resource should always be included - */ - required: z.boolean(), -}); - -/** - * Object containing Knowledge Base Entry text 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 KnowledgeBaseEntryResponse = z.infer; -export const KnowledgeBaseEntryResponse = z.object({ - timestamp: NonEmptyString.optional(), - id: 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(), - /** - * Knowledge Base Entry content - */ - text: z.string(), - vector: Vector.optional(), -}); - -export type KnowledgeBaseEntryUpdateProps = z.infer; -export const KnowledgeBaseEntryUpdateProps = z.object({ - id: 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, - /** - * 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 deleted file mode 100644 index 53c69426b69bd9..00000000000000 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml +++ /dev/null @@ -1,123 +0,0 @@ -openapi: 3.0.0 -info: - title: Common Knowledge Base Attributes - version: 'not applicable' -paths: {} -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 - required: - - 'kbResource' - - 'source' - - 'required' - properties: - kbResource: - type: string - description: Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc - source: - type: string - description: Source document name or filepath - required: - type: boolean - description: Whether or not this resource should always be included - - Vector: - type: object - description: Object containing Knowledge Base Entry text 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 - - text - properties: - 'timestamp': - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' - id: - $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 - text: - type: string - description: Knowledge Base Entry content - vector: - $ref: '#/components/schemas/Vector' - - KnowledgeBaseEntryUpdateProps: - type: object - required: - - id - properties: - id: - $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 - - 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/bulk_crud_knowledge_base_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.gen.ts similarity index 100% rename from x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.gen.ts rename to x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.gen.ts diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.http b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.http new file mode 100644 index 00000000000000..d8ad084a792a10 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.http @@ -0,0 +1,13 @@ +### Empty Bulk Action +POST http://{{host}}:{{port}}{{basePath}}/internal/elastic_assistant/knowledge_base/entries/_bulk_action +kbn-xsrf: "true" +Content-Type: application/json +Elastic-Api-Version: {{elasticApiVersion}} +Authorization: Basic {{auth.username}} {{auth.password}} +X-Kbn-Context: {{appContext.security}} + +{ + "delete": {}, + "create": [], + "update": [] +} 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/entries/bulk_crud_knowledge_base_entries_route.schema.yaml similarity index 100% rename from x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/bulk_crud_knowledge_base_route.schema.yaml rename to x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/bulk_crud_knowledge_base_entries_route.schema.yaml diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts new file mode 100644 index 00000000000000..1af5c46b1c1301 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts @@ -0,0 +1,248 @@ +/* + * 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. + */ + +/* + * 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 + */ + +import { z } from '@kbn/zod'; + +import { User, NonEmptyString } from '../../common_attributes.gen'; + +/** + * Array of objects defining the input schema, allowing the LLM to extract structured data to be used in retrieval + */ +export type InputSchema = z.infer; +export const InputSchema = z.array( + z.object({ + /** + * Name of the field + */ + fieldName: z.string(), + /** + * Type of the field + */ + fieldType: z.string(), + /** + * Description of the field + */ + description: z.string(), + }) +); + +export type KnowledgeBaseEntryErrorSchema = z.infer; +export const KnowledgeBaseEntryErrorSchema = z + .object({ + statusCode: z.number(), + error: z.string(), + message: z.string(), + }) + .strict(); + +/** + * Metadata about a Knowledge Base Entry + */ +export type Metadata = z.infer; +export const Metadata = z.object({ + /** + * Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc + */ + kbResource: z.string(), + /** + * Source document name or filepath + */ + source: z.string(), + /** + * Whether this resource should always be included + */ + required: z.boolean(), +}); + +/** + * Object containing Knowledge Base Entry text 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 BaseRequiredFields = z.infer; +export const BaseRequiredFields = z.object({ + /** + * Name of the Knowledge Base Entry + */ + name: z.string(), +}); + +export type BaseDefaultableFields = z.infer; +export const BaseDefaultableFields = z.object({ + /** + * Kibana Space, defaults to 'default' space + */ + namespace: z.string().optional(), + /** + * Users who have access to the Knowledge Base Entry, defaults to current user. Empty array provides access to all users. + */ + users: z.array(User).optional(), +}); + +export type BaseCreateProps = z.infer; +export const BaseCreateProps = BaseRequiredFields.merge(BaseDefaultableFields); + +export type BaseUpdateProps = z.infer; +export const BaseUpdateProps = BaseCreateProps.partial(); + +export type BaseResponseProps = z.infer; +export const BaseResponseProps = BaseRequiredFields.merge(BaseDefaultableFields.required()); + +export type ResponseFields = z.infer; +export const ResponseFields = z.object({ + id: NonEmptyString, + /** + * Time the Knowledge Base Entry was created + */ + createdAt: z.string(), + /** + * User who created the Knowledge Base Entry + */ + createdBy: z.string(), + /** + * Time the Knowledge Base Entry was last updated + */ + updatedAt: z.string(), + /** + * User who last updated the Knowledge Base Entry + */ + updatedBy: z.string(), +}); + +export type SharedResponseProps = z.infer; +export const SharedResponseProps = BaseResponseProps.merge(ResponseFields); + +export type DocumentEntryType = z.infer; +export const DocumentEntryType = z.literal('document'); + +export type DocumentEntryRequiredFields = z.infer; +export const DocumentEntryRequiredFields = z.object({ + /** + * Entry type + */ + type: z.literal('document'), + /** + * Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc + */ + kbResource: z.string(), + /** + * Source document name or filepath + */ + source: z.string(), + /** + * Knowledge Base Entry content + */ + text: z.string(), +}); + +export type DocumentEntryOptionalFields = z.infer; +export const DocumentEntryOptionalFields = z.object({ + /** + * Whether this resource should always be included, defaults to false + */ + required: z.boolean().optional(), + vector: Vector.optional(), +}); + +export type DocumentEntryCreateFields = z.infer; +export const DocumentEntryCreateFields = BaseCreateProps.merge(DocumentEntryRequiredFields).merge( + DocumentEntryOptionalFields +); + +export type DocumentEntryUpdateFields = z.infer; +export const DocumentEntryUpdateFields = BaseUpdateProps.merge(DocumentEntryCreateFields); + +export type DocumentEntryResponseFields = z.infer; +export const DocumentEntryResponseFields = DocumentEntryRequiredFields.merge( + DocumentEntryOptionalFields +); + +export type DocumentEntry = z.infer; +export const DocumentEntry = SharedResponseProps.merge(DocumentEntryResponseFields); + +export type IndexEntryType = z.infer; +export const IndexEntryType = z.literal('index'); + +export type IndexEntryRequiredFields = z.infer; +export const IndexEntryRequiredFields = z.object({ + /** + * Entry type + */ + type: z.literal('index'), + /** + * Index or Data Stream to query for Knowledge Base content + */ + index: z.string(), + /** + * Field to query for Knowledge Base content + */ + field: z.string(), + /** + * Description for when this index or data stream should be queried for Knowledge Base content. Passed to the LLM as a tool description + */ + description: z.string(), + /** + * Description of query field used to fetch Knowledge Base content. Passed to the LLM as part of the tool input schema + */ + queryDescription: z.string(), +}); + +export type IndexEntryOptionalFields = z.infer; +export const IndexEntryOptionalFields = z.object({ + inputSchema: InputSchema.optional(), + /** + * Fields to extract from the query result, defaults to all fields if not provided or empty + */ + outputFields: z.array(z.string()).optional(), +}); + +export type IndexEntryCreateFields = z.infer; +export const IndexEntryCreateFields = + BaseCreateProps.merge(IndexEntryRequiredFields).merge(IndexEntryOptionalFields); + +export type IndexEntryUpdateFields = z.infer; +export const IndexEntryUpdateFields = BaseUpdateProps.merge(IndexEntryCreateFields); + +export type IndexEntryResponseFields = z.infer; +export const IndexEntryResponseFields = IndexEntryRequiredFields.merge(IndexEntryOptionalFields); + +export type IndexEntry = z.infer; +export const IndexEntry = SharedResponseProps.merge(IndexEntryResponseFields); + +export type KnowledgeBaseEntryCreateProps = z.infer; +export const KnowledgeBaseEntryCreateProps = z.discriminatedUnion('type', [ + DocumentEntryCreateFields, + IndexEntryCreateFields, +]); + +export type KnowledgeBaseEntryUpdateProps = z.infer; +export const KnowledgeBaseEntryUpdateProps = z.discriminatedUnion('type', [ + DocumentEntryUpdateFields, + IndexEntryUpdateFields, +]); + +export type KnowledgeBaseEntryResponse = z.infer; +export const KnowledgeBaseEntryResponse = z.discriminatedUnion('type', [DocumentEntry, IndexEntry]); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml new file mode 100644 index 00000000000000..c1c551059f04b1 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml @@ -0,0 +1,302 @@ +openapi: 3.0.0 +info: + title: Common Knowledge Base Attributes + version: "not applicable" +paths: {} +components: + x-codegen-enabled: true + schemas: + InputSchema: + type: array + description: Array of objects defining the input schema, allowing the LLM to extract structured data to be used in retrieval + items: + type: object + required: + - fieldName + - fieldType + - description + properties: + fieldName: + type: string + description: Name of the field + fieldType: + type: string + description: Type of the field + description: + type: string + description: Description of the field + + 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 a Knowledge Base Entry + required: + - "kbResource" + - "source" + - "required" + properties: + kbResource: + type: string + description: Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc + source: + type: string + description: Source document name or filepath + required: + type: boolean + description: Whether this resource should always be included + + Vector: + type: object + description: Object containing Knowledge Base Entry text 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 + + ########### + # Base Entry + ########### + BaseRequiredFields: + x-inline: true + type: object + properties: + name: + type: string + description: Name of the Knowledge Base Entry + required: + - name + + BaseDefaultableFields: + x-inline: true + type: object + properties: + namespace: + type: string + description: Kibana Space, defaults to 'default' space + users: + type: array + description: Users who have access to the Knowledge Base Entry, defaults to current user. Empty array provides access to all users. + items: + $ref: "../../common_attributes.schema.yaml#/components/schemas/User" + + BaseCreateProps: + x-inline: true + allOf: + - $ref: "#/components/schemas/BaseRequiredFields" + - $ref: "#/components/schemas/BaseDefaultableFields" + + BaseUpdateProps: + x-inline: true + allOf: + - $ref: "#/components/schemas/BaseCreateProps" + x-modify: partial + + BaseResponseProps: + x-inline: true + allOf: + - $ref: "#/components/schemas/BaseRequiredFields" + - $ref: "#/components/schemas/BaseDefaultableFields" + x-modify: required + + ResponseFields: + type: object + properties: + id: + $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 + required: + - id + - createdAt + - createdBy + - updatedAt + - updatedBy + + SharedResponseProps: + x-inline: true + allOf: + - $ref: "#/components/schemas/BaseResponseProps" + - $ref: "#/components/schemas/ResponseFields" + + ########### + # Document Knowledge Base Entry + ########### + DocumentEntryType: + type: string + enum: + - document + + DocumentEntryRequiredFields: + type: object + properties: + type: + type: string + enum: [document] + description: Entry type + kbResource: + type: string + description: Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc + source: + type: string + description: Source document name or filepath + text: + type: string + description: Knowledge Base Entry content + + required: + - type + - kbResource + - source + - text + + DocumentEntryOptionalFields: + type: object + properties: + required: + type: boolean + description: Whether this resource should always be included, defaults to false + vector: + $ref: "#/components/schemas/Vector" + + DocumentEntryCreateFields: + allOf: + - $ref: "#/components/schemas/BaseCreateProps" + - $ref: "#/components/schemas/DocumentEntryRequiredFields" + - $ref: "#/components/schemas/DocumentEntryOptionalFields" + + DocumentEntryUpdateFields: + allOf: + - $ref: "#/components/schemas/BaseUpdateProps" + - $ref: "#/components/schemas/DocumentEntryCreateFields" + + DocumentEntryResponseFields: + allOf: + - $ref: "#/components/schemas/DocumentEntryRequiredFields" + - $ref: "#/components/schemas/DocumentEntryOptionalFields" + + DocumentEntry: + allOf: + - $ref: "#/components/schemas/SharedResponseProps" + - $ref: "#/components/schemas/DocumentEntryResponseFields" + + ########### + # Index Knowledge Base Entry + ########### + IndexEntryType: + type: string + enum: + - index + + IndexEntryRequiredFields: + type: object + properties: + type: + type: string + enum: [index] + description: Entry type + index: + type: string + description: Index or Data Stream to query for Knowledge Base content + field: + type: string + description: Field to query for Knowledge Base content + description: + type: string + description: Description for when this index or data stream should be queried for Knowledge Base content. Passed to the LLM as a tool description + queryDescription: + type: string + description: Description of query field used to fetch Knowledge Base content. Passed to the LLM as part of the tool input schema + required: + - type + - index + - field + - description + - queryDescription + + IndexEntryOptionalFields: + type: object + properties: + inputSchema: + $ref: "#/components/schemas/InputSchema" + outputFields: + type: array + description: Fields to extract from the query result, defaults to all fields if not provided or empty + items: + type: string + + IndexEntryCreateFields: + allOf: + - $ref: "#/components/schemas/BaseCreateProps" + - $ref: "#/components/schemas/IndexEntryRequiredFields" + - $ref: "#/components/schemas/IndexEntryOptionalFields" + + IndexEntryUpdateFields: + allOf: + - $ref: "#/components/schemas/BaseUpdateProps" + - $ref: "#/components/schemas/IndexEntryCreateFields" + + IndexEntryResponseFields: + allOf: + - $ref: "#/components/schemas/IndexEntryRequiredFields" + - $ref: "#/components/schemas/IndexEntryOptionalFields" + + IndexEntry: + allOf: + - $ref: "#/components/schemas/SharedResponseProps" + - $ref: "#/components/schemas/IndexEntryResponseFields" + + ################### + # Combined Props + ################### + KnowledgeBaseEntryCreateProps: + discriminator: + propertyName: type + anyOf: + - $ref: "#/components/schemas/DocumentEntryCreateFields" + - $ref: "#/components/schemas/IndexEntryCreateFields" + + KnowledgeBaseEntryUpdateProps: + discriminator: + propertyName: type + anyOf: + - $ref: "#/components/schemas/DocumentEntryUpdateFields" + - $ref: "#/components/schemas/IndexEntryUpdateFields" + + KnowledgeBaseEntryResponse: + discriminator: + propertyName: type + anyOf: + - $ref: "#/components/schemas/DocumentEntry" + - $ref: "#/components/schemas/IndexEntry" 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/entries/crud_knowledge_base_entries_route.gen.ts similarity index 98% rename from x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.gen.ts rename to x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/crud_knowledge_base_entries_route.gen.ts index c73fb3ba7b7cf8..1650f4ddc3cc62 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/entries/crud_knowledge_base_entries_route.gen.ts @@ -21,7 +21,7 @@ import { KnowledgeBaseEntryResponse, KnowledgeBaseEntryUpdateProps, } from './common_attributes.gen'; -import { NonEmptyString } from '../common_attributes.gen'; +import { NonEmptyString } 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/entries/crud_knowledge_base_entries_route.http b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/crud_knowledge_base_entries_route.http new file mode 100644 index 00000000000000..88ad6555da50ff --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/crud_knowledge_base_entries_route.http @@ -0,0 +1,49 @@ +### Create Document Entry +POST http://{{host}}:{{port}}{{basePath}}/internal/elastic_assistant/knowledge_base/entries +kbn-xsrf: "true" +Content-Type: application/json +Elastic-Api-Version: {{elasticApiVersion}} +Authorization: Basic {{auth.username}} {{auth.password}} +X-Kbn-Context: {{appContext.security}} + +{ + "type": "document", + "name": "Favorites", + "kbResource": "user", + "source": "api", + "required": true, + "text": "My favorite food is Dan Bing" +} + +### Create Index Entry +POST http://{{host}}:{{port}}{{basePath}}/internal/elastic_assistant/knowledge_base/entries +kbn-xsrf: "true" +Content-Type: application/json +Elastic-Api-Version: {{elasticApiVersion}} +Authorization: Basic {{auth.username}} {{auth.password}} +X-Kbn-Context: {{appContext.security}} + +{ + "type": "index", + "name": "SpongBotSlackConnector", + "namespace": "default", + "index": "spongbot-slack", + "field": "semantic_text", + "description": "Use this index to search for the user's Slack messages.", + "queryDescription": + "The free text search that the user wants to perform over this dataset. So if asking \"what are my slack messages from last week about failed tests\", the query would be \"A test has failed! failing test failed test\"", + "inputSchema": [ + { + "fieldName": "author", + "fieldType": "string", + "description": "The author of the message. So if asking for recent messages from Stan, you would provide 'Stan' as the author." + } + ], + "outputFields": ["author", "text", "timestamp"] +} + + + + + + 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/entries/crud_knowledge_base_entries_route.schema.yaml similarity index 93% rename from x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_knowledge_base_route.schema.yaml rename to x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/crud_knowledge_base_entries_route.schema.yaml index 32e66efffc13ca..7479b5cca8225e 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/entries/crud_knowledge_base_entries_route.schema.yaml @@ -3,7 +3,7 @@ info: title: Manage Knowledge Base Entries API endpoint version: '1' paths: - /internnal/elastic_assistant/knowledge_base/entries: + /internal/elastic_assistant/knowledge_base/entries: post: x-codegen-enabled: true # This API is still behind the `assistantKnowledgeBaseByDefault` feature flag @@ -51,7 +51,7 @@ paths: required: true description: The Knowledge Base Entry's `id` value. schema: - $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 @@ -81,7 +81,7 @@ paths: required: true description: The Knowledge Base Entry's `id` value schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + $ref: '../../common_attributes.schema.yaml#/components/schemas/NonEmptyString' requestBody: required: true content: @@ -117,7 +117,7 @@ paths: required: true description: The Knowledge Base Entry's `id` value schema: - $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/knowledge_base/find_knowledge_base_entries_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.gen.ts similarity index 97% rename from x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/find_knowledge_base_entries_route.gen.ts rename to x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.gen.ts index cb5ae50ec81f00..81930d69bb2653 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/find_knowledge_base_entries_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.gen.ts @@ -17,7 +17,7 @@ import { z } from '@kbn/zod'; import { ArrayFromString } from '@kbn/zod-helpers'; -import { SortOrder } from '../common_attributes.gen'; +import { SortOrder } from '../../common_attributes.gen'; import { KnowledgeBaseEntryResponse } from './common_attributes.gen'; export type FindKnowledgeBaseEntriesSortField = z.infer; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.http b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.http new file mode 100644 index 00000000000000..3d90053767dfef --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.http @@ -0,0 +1,6 @@ +### Find all knowledge base entries +GET http://{{host}}:{{port}}{{basePath}}/internal/elastic_assistant/knowledge_base/entries/_find +Elastic-Api-Version: {{elasticApiVersion}} +Authorization: Basic {{auth.username}} {{auth.password}} +X-Kbn-Context: {{appContext.security}} + diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/find_knowledge_base_entries_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.schema.yaml similarity index 97% rename from x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/find_knowledge_base_entries_route.schema.yaml rename to x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.schema.yaml index cf88de73df1dbe..8794a94b0efc9d 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/find_knowledge_base_entries_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/find_knowledge_base_entries_route.schema.yaml @@ -39,7 +39,7 @@ paths: description: Sort order required: false schema: - $ref: '../common_attributes.schema.yaml#/components/schemas/SortOrder' + $ref: '../../common_attributes.schema.yaml#/components/schemas/SortOrder' - name: 'page' in: query description: Page number diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/mocks.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/mocks.ts new file mode 100644 index 00000000000000..a4828c2f24d913 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/mocks.ts @@ -0,0 +1,27 @@ +/* + * 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 { IndexEntryCreateFields } from './common_attributes.gen'; + +export const indexEntryMock: IndexEntryCreateFields = { + type: 'index', + name: 'SpongBotSlackConnector', + namespace: 'default', + index: 'spongbot', + field: 'semantic_text', + description: "Use this index to search for the user's Slack messages.", + queryDescription: + 'The free text search that the user wants to perform over this dataset. So if asking "what are my slack messages from last week about failed tests", the query would be "A test has failed! failing test failed test".', + inputSchema: [ + { + fieldName: 'author', + fieldType: 'string', + description: + "The author of the message. So if asking for recent messages from Stan, you would provide 'Stan' as the author.", + }, + ], +}; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts index aa1d72d755ffa4..43d70af7dd2558 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.ts @@ -10,6 +10,7 @@ import { useQuery, useQueryClient } from '@tanstack/react-query'; import { API_VERSIONS, ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, + FindKnowledgeBaseEntriesRequestQuery, FindKnowledgeBaseEntriesResponse, } from '@kbn/elastic-assistant-common'; @@ -17,9 +18,22 @@ import { useCallback } from 'react'; export interface UseKnowledgeBaseEntriesParams { http: HttpSetup; + query: FindKnowledgeBaseEntriesRequestQuery; signal?: AbortSignal | undefined; } +const defaultQuery: FindKnowledgeBaseEntriesRequestQuery = { + page: 1, + per_page: 100, +}; + +export const KNOWLEDGE_BASE_ENTRY_QUERY_KEY = [ + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, + defaultQuery.page, + defaultQuery.per_page, + API_VERSIONS.internal.v1, +]; + /** * Hook for fetching Knowledge Base Entries. * @@ -27,24 +41,16 @@ export interface UseKnowledgeBaseEntriesParams { * * @param {Object} options - The options object. * @param {HttpSetup} options.http - HttpSetup - * @param {Function} [options.onFetch] - transformation function for kb entries fetch result + * @param {Function} [options.query] - Query params to include, like filters, pagination, etc. * @param {AbortSignal} [options.signal] - AbortSignal * - * @returns {useQuery} hook for fetching Knowledge Base Entries + * @returns hook for fetching Knowledge Base Entries */ -const query = { - page: 1, - perPage: 100, -}; - -export const KNOWLEDGE_BASE_ENTRY_QUERY_KEY = [ - ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, - query.page, - query.perPage, - API_VERSIONS.internal.v1, -]; - -export const useKnowledgeBaseEntries = ({ http, signal }: UseKnowledgeBaseEntriesParams) => +export const useKnowledgeBaseEntries = ({ + http, + query = defaultQuery, + signal, +}: UseKnowledgeBaseEntriesParams) => useQuery( KNOWLEDGE_BASE_ENTRY_QUERY_KEY, async () => 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 be3965fd1c26b7..adface75dbced5 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -123,10 +123,10 @@ const createElasticAssistantRequestContextMock = ( () => clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient ) as unknown as jest.MockInstance< Promise, - [], + [boolean | undefined], unknown > & - (() => Promise), + ((v2KnowledgeBaseEnabled?: boolean) => Promise), getCurrentUser: jest.fn(), getServerBasePath: jest.fn(), getSpaceId: jest.fn(), 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 index 97060a0f48aa61..304a5a54a737fb 100644 --- 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 @@ -51,7 +51,7 @@ export const createKnowledgeBaseEntry = async ({ }); } catch (err) { logger.error( - `Error creating Knowledge Base Entry: ${err} with kbResource: ${knowledgeBaseEntry.metadata.kbResource}` + `Error creating Knowledge Base Entry: ${err} with kbResource: ${knowledgeBaseEntry.name}` ); throw err; } @@ -61,22 +61,37 @@ export const transformToCreateSchema = ( createdAt: string, spaceId: string, user: AuthenticatedUser, - { metadata, text }: KnowledgeBaseEntryCreateProps + entry: KnowledgeBaseEntryCreateProps ): CreateKnowledgeBaseEntrySchema => { - return { + const base = { '@timestamp': createdAt, created_at: createdAt, created_by: user.profile_uid ?? 'unknown', updated_at: createdAt, updated_by: user.profile_uid ?? 'unknown', + namespace: spaceId, users: [ { id: user.profile_uid, name: user.username, }, ], - namespace: spaceId, - metadata, - text, }; + + if (entry.type === 'index') { + const { inputSchema, outputFields, queryDescription, ...restEntry } = entry; + return { + ...base, + ...restEntry, + query_description: queryDescription, + input_schema: + entry.inputSchema?.map((schema) => ({ + field_name: schema.fieldName, + field_type: schema.fieldType, + description: schema.description, + })) ?? undefined, + output_fields: outputFields ?? undefined, + }; + } + return { ...base, ...entry, vector: undefined }; }; 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 16b5b35a9b2bb9..0712664bbfeedb 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 @@ -88,3 +88,141 @@ export const knowledgeBaseFieldMap: FieldMap = { required: false, }, } as const; + +export const knowledgeBaseFieldMapV2: FieldMap = { + // Base fields + '@timestamp': { + type: 'date', + 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, + }, + name: { + type: 'keyword', + array: false, + required: false, + }, + // Discriminator: 'document' | 'index' + type: { + type: 'keyword', + array: false, + required: true, + }, + // Document Entry fields + kb_resource: { + type: 'keyword', + array: false, + required: false, + }, + required: { + type: 'boolean', + array: false, + required: false, + }, + source: { + type: 'keyword', + array: false, + required: false, + }, + text: { + type: 'text', + array: false, + required: false, + }, + // Embeddings field + vector: { + type: 'object', + array: false, + required: false, + }, + 'vector.tokens': { + type: 'rank_features', + array: false, + required: false, + }, + // Index Entry fields + index: { + type: 'keyword', + array: false, + required: false, + }, + field: { + type: 'keyword', + array: false, + required: false, + }, + description: { + type: 'text', + array: false, + required: false, + }, + query_description: { + type: 'text', + array: false, + required: false, + }, + input_schema: { + type: 'nested', + array: true, + required: false, + }, + 'input_schema.field_name': { + type: 'keyword', + array: false, + required: true, + }, + 'input_schema.field_type': { + type: 'keyword', + array: false, + required: true, + }, + 'input_schema.description': { + type: 'text', + array: false, + required: true, + }, + output_fields: { + type: 'keyword', + array: true, + required: false, + }, +} as const; 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 index 839ac3e559ba12..f0d87b8b14cffc 100644 --- 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 @@ -5,9 +5,13 @@ * 2.0. */ +import { z } from '@kbn/zod'; +import { DynamicStructuredTool } from '@langchain/core/tools'; import { errors } from '@elastic/elasticsearch'; -import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types'; +import { QueryDslQueryContainer, SearchRequest } from '@elastic/elasticsearch/lib/api/types'; import { AuthenticatedUser } from '@kbn/core-security-common'; +import { IndexEntry } from '@kbn/elastic-assistant-common'; +import { ElasticsearchClient, Logger } from '@kbn/core/server'; export const isModelAlreadyExistsError = (error: Error) => { return ( @@ -18,7 +22,7 @@ export const isModelAlreadyExistsError = (error: Error) => { }; /** - * Returns an Elasticsearch query DSL that performs a vector search against the Knowledge Base for the given query/user/filter. + * Returns an Elasticsearch query DSL that performs a vector search against the Knowledge Base for the given query/user/filter. Searches only for DocumentEntries, not IndexEntries as they have no content. * * @param filter - Optional filter to apply to the search * @param kbResource - Specific resource tag to filter for, e.g. 'esql' or 'user' @@ -26,6 +30,7 @@ export const isModelAlreadyExistsError = (error: Error) => { * @param query - The search query provided by the user * @param required - Whether to only include required entries * @param user - The authenticated user + * @param v2KnowledgeBaseEnabled whether the new v2 KB is enabled * @returns */ export const getKBVectorSearchQuery = ({ @@ -35,6 +40,7 @@ export const getKBVectorSearchQuery = ({ query, required, user, + v2KnowledgeBaseEnabled = false, }: { filter?: QueryDslQueryContainer | undefined; kbResource?: string | undefined; @@ -42,12 +48,15 @@ export const getKBVectorSearchQuery = ({ query: string; required?: boolean | undefined; user: AuthenticatedUser; + v2KnowledgeBaseEnabled: boolean; }): QueryDslQueryContainer => { + const kbResourceKey = v2KnowledgeBaseEnabled ? 'kb_resource' : 'metadata.kbResource'; + const requiredKey = v2KnowledgeBaseEnabled ? 'required' : 'metadata.required'; const resourceFilter = kbResource ? [ { term: { - 'metadata.kbResource': kbResource, + [kbResourceKey]: kbResource, }, }, ] @@ -56,7 +65,7 @@ export const getKBVectorSearchQuery = ({ ? [ { term: { - 'metadata.required': required, + [requiredKey]: required, }, }, ] @@ -100,3 +109,110 @@ export const getKBVectorSearchQuery = ({ }, }; }; + +/** + * Returns a StructuredTool for a given IndexEntry + */ +export const getStructuredToolForIndexEntry = ({ + indexEntry, + esClient, + logger, + elserId, +}: { + indexEntry: IndexEntry; + esClient: ElasticsearchClient; + logger: Logger; + elserId: string; +}): DynamicStructuredTool => { + const inputSchema = indexEntry.inputSchema?.reduce((prev, input) => { + const fieldType = + input.fieldType === 'string' + ? z.string() + : input.fieldType === 'number' + ? z.number() + : input.fieldType === 'boolean' + ? z.boolean() + : z.any(); + return { ...prev, [input.fieldName]: fieldType.describe(input.description) }; + }, {}); + + return new DynamicStructuredTool({ + name: indexEntry.name.replaceAll(' ', ''), // Tool names cannot contain spaces, further sanitization possibly needed + description: indexEntry.description, + schema: z.object({ + query: z.string().describe(indexEntry.queryDescription), + ...inputSchema, + }), + func: async (input, _, cbManager) => { + logger.debug( + () => `Generated ${indexEntry.name} Tool:input\n ${JSON.stringify(input, null, 2)}` + ); + + // Generate filters for inputSchema fields + const filter = + indexEntry.inputSchema?.reduce((prev, i) => { + return [ + ...prev, + // @ts-expect-error Possible to override types with dynamic input schema? + { term: { [`${i.fieldName}`]: input?.[i.fieldName] } }, + ]; + }, [] as Array<{ term: { [key: string]: string } }>) ?? []; + + const params: SearchRequest = { + index: indexEntry.index, + size: 10, + retriever: { + standard: { + query: { + nested: { + path: 'semantic_text.inference.chunks', + query: { + sparse_vector: { + inference_id: elserId, + field: `${indexEntry.field}.inference.chunks.embeddings`, + query: input.query, + }, + }, + inner_hits: { + size: 2, + name: `${indexEntry.name}.${indexEntry.field}`, + _source: [`${indexEntry.field}.inference.chunks.text`], + }, + }, + }, + filter, + }, + }, + }; + + try { + const result = await esClient.search(params); + + const kbDocs = result.hits.hits.map((hit) => { + if (indexEntry.outputFields && indexEntry.outputFields.length > 0) { + return indexEntry.outputFields.reduce((prev, field) => { + // @ts-expect-error + return { ...prev, [field]: hit._source[field] }; + }, {}); + } + return { + text: (hit._source as { text: string }).text, + }; + }); + + logger.debug(() => `Similarity Search Params:\n ${JSON.stringify(params)}`); + logger.debug(() => `Similarity Search Results:\n ${JSON.stringify(result)}`); + logger.debug(() => `Similarity Text Extract Results:\n ${JSON.stringify(kbDocs)}`); + + return `###\nBelow are all relevant documents in JSON format:\n${JSON.stringify( + kbDocs + )}\n###`; + } catch (e) { + logger.error(`Error performing IndexEntry KB Similarity Search: ${e.message}`); + return `I'm sorry, but I was unable to find any information in the knowledge base. Perhaps this error would be useful to deliver to the user. Be sure to print it below your response and in a codeblock so it is rendered nicely: ${e.message}`; + } + }, + tags: ['knowledge-base'], + // 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/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 5fff839b497bba..38852dc8a91fc8 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 @@ -14,21 +14,30 @@ import type { KibanaRequest } from '@kbn/core-http-server'; import { Document } from 'langchain/document'; import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import { + DocumentEntryType, + IndexEntry, KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, Metadata, } from '@kbn/elastic-assistant-common'; import pRetry from 'p-retry'; import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey'; +import { StructuredTool } from '@langchain/core/tools'; +import { ElasticsearchClient } from '@kbn/core/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'; +import { AssistantToolParams, GetElser } from '../../types'; import { createKnowledgeBaseEntry, transformToCreateSchema } from './create_knowledge_base_entry'; -import { EsKnowledgeBaseEntrySchema } from './types'; +import { EsDocumentEntry, EsIndexEntry, EsKnowledgeBaseEntrySchema } from './types'; import { transformESSearchToKnowledgeBaseEntry } from './transforms'; import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants'; -import { getKBVectorSearchQuery, isModelAlreadyExistsError } from './helpers'; +import { + getKBVectorSearchQuery, + getStructuredToolForIndexEntry, + isModelAlreadyExistsError, +} from './helpers'; +import { getKBUserFilter } from '../../routes/knowledge_base/entries/utils'; interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ml: MlPluginSetup; @@ -36,6 +45,7 @@ interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { getIsKBSetupInProgress: () => boolean; ingestPipelineResourceName: string; setIsKBSetupInProgress: (isInProgress: boolean) => void; + v2KnowledgeBaseEnabled: boolean; } export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { constructor(public readonly options: KnowledgeBaseDataClientParams) { @@ -46,6 +56,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { return this.options.getIsKBSetupInProgress(); } + public get isV2KnowledgeBaseEnabled() { + return this.options.v2KnowledgeBaseEnabled; + } + /** * Returns whether setup of the Knowledge Base can be performed (essentially an ML features check) * @@ -254,18 +268,30 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { '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) => - transformToCreateSchema(changedAt, this.spaceId, authenticatedUser, { - metadata: { - kbResource: doc.metadata.kbResource ?? 'unknown', - required: doc.metadata.required ?? false, - source: doc.metadata.source ?? 'unknown', - }, + documentsToCreate: documents.map((doc) => { + // v1 schema has metadata nested in a `metadata` object, and kbResource vs kb_resource + const body = this.options.v2KnowledgeBaseEnabled + ? { + kb_resource: doc.metadata.kbResource ?? 'unknown', + required: doc.metadata.required ?? false, + source: doc.metadata.source ?? 'unknown', + } + : { + metadata: { + kbResource: doc.metadata.kbResource ?? 'unknown', + required: doc.metadata.required ?? false, + source: doc.metadata.source ?? 'unknown', + }, + }; + // @ts-ignore Transform only explicitly supports v2 schema, but technically still supports v1 + return transformToCreateSchema(changedAt, this.spaceId, authenticatedUser, { + type: DocumentEntryType.value, + name: 'unknown', text: doc.pageContent, - }) - ), + ...body, + }); + }), authenticatedUser, }); const created = @@ -285,7 +311,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { /** * Performs similarity search to retrieve LangChain Documents from the knowledge base */ - public getKnowledgeBaseDocuments = async ({ + public getKnowledgeBaseDocumentEntries = async ({ filter, kbResource, query, @@ -313,22 +339,30 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { query, required, user, + v2KnowledgeBaseEnabled: this.options.v2KnowledgeBaseEnabled, }); try { - const result = await esClient.search({ + const result = await esClient.search({ index: this.indexTemplateAndPattern.alias, size: 10, query: vectorSearchQuery, }); - const results = result.hits.hits.map( - (hit) => - new Document({ - pageContent: hit?._source?.text ?? '', - metadata: hit?._source?.metadata ?? {}, - }) - ); + const results = result.hits.hits.map((hit) => { + const metadata = this.options.v2KnowledgeBaseEnabled + ? { + source: hit?._source?.source, + required: hit?._source?.required, + kbResource: hit?._source?.kb_resource, + } + : // @ts-ignore v1 schema has metadata nested in a `metadata` object and kbResource vs kb_resource + hit?._source?.metadata ?? {}; + return new Document({ + pageContent: hit?._source?.text ?? '', + metadata, + }); + }); this.options.logger.debug( () => @@ -364,7 +398,6 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' ); } - this.options.logger.debug( () => `Creating Knowledge Base Entry:\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}` ); @@ -379,4 +412,58 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { knowledgeBaseEntry, }); }; + + /** + * Returns AssistantTools for any 'relevant' KB IndexEntries that exist in the knowledge base + */ + public getAssistantTools = async ({ + assistantToolParams, + esClient, + }: { + assistantToolParams: AssistantToolParams; + esClient: ElasticsearchClient; + }): Promise => { + const user = this.options.currentUser; + if (user == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } + + try { + const elserId = await this.options.getElserId(); + const userFilter = getKBUserFilter(user); + const results = await this.findDocuments({ + // Note: This is a magic number to set some upward bound as to not blow the context with too + // many registered tools. As discussed in review, this will initially be mitigated by caps on + // the IndexEntries field lengths, context trimming at the graph layer (before compilation), + // and eventually some sort of tool discovery sub-graph or generic retriever to scale tool usage. + perPage: 23, + page: 1, + sortField: 'created_at', + sortOrder: 'asc', + filter: `${userFilter}${` AND type:index`}`, // TODO: Support global tools (no user filter), and filter by space as well + }); + this.options.logger.debug( + `kbDataClient.getAssistantTools() - results:\n${JSON.stringify(results, null, 2)}` + ); + + if (results) { + const entries = transformESSearchToKnowledgeBaseEntry(results.data) as IndexEntry[]; + return entries.map((indexEntry) => { + return getStructuredToolForIndexEntry({ + indexEntry, + esClient, + logger: this.options.logger, + elserId, + }); + }); + } + } catch (e) { + this.options.logger.error(`kbDataClient.getAssistantTools() - Failed to fetch IndexEntries`); + return []; + } + + return []; + }; } 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 index 170fa0342f9d9e..e11840b94e6602 100644 --- 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 @@ -5,31 +5,21 @@ * 2.0. */ -import { IngestPutPipelineRequest } from '@elastic/elasticsearch/lib/api/types'; - -export const knowledgeBaseIngestPipeline = ({ - id, - modelId, -}: { - id: string; - modelId: string; -}): IngestPutPipelineRequest => ({ +// TODO: Ensure old pipeline is updated/replaced +export const knowledgeBaseIngestPipeline = ({ id, modelId }: { id: string; modelId: string }) => ({ id, description: 'Embedding pipeline for Elastic AI Assistant ELSER Knowledge Base', processors: [ { inference: { + if: 'ctx?.text != null', model_id: modelId, - target_field: 'vector', - field_map: { - text: 'text_field', - }, - inference_config: { - // @ts-expect-error - text_expansion: { - results_field: 'tokens', + input_output: [ + { + input_field: 'text', + output_field: 'vector.tokens', }, - }, + ], }, }, ], 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 index 475f9f880ee138..16ef4ffb0595ed 100644 --- 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 @@ -6,8 +6,14 @@ */ import { estypes } from '@elastic/elasticsearch'; -import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; -import { EsKnowledgeBaseEntrySchema } from './types'; +import { + DocumentEntry, + DocumentEntryType, + IndexEntry, + IndexEntryType, + KnowledgeBaseEntryResponse, +} from '@kbn/elastic-assistant-common'; +import { EsKnowledgeBaseEntrySchema, LegacyEsKnowledgeBaseEntrySchema } from './types'; export const transformESSearchToKnowledgeBaseEntry = ( response: estypes.SearchResponse @@ -17,41 +23,11 @@ export const transformESSearchToKnowledgeBaseEntry = ( .map((hit) => { // eslint-disable-next-line @typescript-eslint/no-non-null-assertion const kbEntrySchema = hit._source!; - const kbEntry: KnowledgeBaseEntryResponse = { - timestamp: kbEntrySchema['@timestamp'], + return { + ...transformEsSchemaToEntry(kbEntrySchema), // eslint-disable-next-line @typescript-eslint/no-non-null-assertion 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; }); }; @@ -59,39 +35,108 @@ 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, + return transformEsSchemaToEntry(kbEntrySchema); + }); +}; + +const transformEsSchemaToEntry = ( + esKbEntry: EsKnowledgeBaseEntrySchema +): DocumentEntry | IndexEntry => { + if (esKbEntry.type === DocumentEntryType.value) { + const documentEntry: DocumentEntry = { + id: esKbEntry.id, + createdAt: esKbEntry.created_at, + createdBy: esKbEntry.created_by, + updatedAt: esKbEntry.updated_at, + updatedBy: esKbEntry.updated_by, users: - kbEntrySchema.users?.map((user) => ({ + esKbEntry.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 + + name: esKbEntry.name, + namespace: esKbEntry.namespace, + type: esKbEntry.type, + kbResource: esKbEntry.kb_resource, + source: esKbEntry.source, + required: esKbEntry.required, + text: esKbEntry.text, + ...(esKbEntry.vector ? { vector: { - modelId: kbEntrySchema.vector.model_id, - tokens: kbEntrySchema.vector.tokens, + modelId: esKbEntry.vector.model_id, + tokens: esKbEntry.vector.tokens, }, } : {}), }; + return documentEntry; + } else if (esKbEntry.type === IndexEntryType.value) { + const indexEntry: IndexEntry = { + id: esKbEntry.id, + createdAt: esKbEntry.created_at, + createdBy: esKbEntry.created_by, + updatedAt: esKbEntry.updated_at, + updatedBy: esKbEntry.updated_by, + users: + esKbEntry.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + name: esKbEntry.name, + namespace: esKbEntry.namespace, + // Document Entry Fields + type: esKbEntry.type, + index: esKbEntry.index, + field: esKbEntry.field, + description: esKbEntry.description, + queryDescription: esKbEntry.query_description, + inputSchema: + esKbEntry.input_schema?.map((schema) => ({ + fieldName: schema.field_name, + fieldType: schema.field_type, + description: schema.description, + })) ?? [], + outputFields: esKbEntry.output_fields ?? [], + }; + return indexEntry; + } - return kbEntry; - }); + // Parse Legacy KB Entry as a DocumentEntry + return getDocumentEntryFromLegacyKbEntry(esKbEntry); +}; + +const getDocumentEntryFromLegacyKbEntry = ( + legacyEsKbDoc: LegacyEsKnowledgeBaseEntrySchema +): DocumentEntry => { + const documentEntry: DocumentEntry = { + id: legacyEsKbDoc.id, + createdAt: legacyEsKbDoc.created_at, + createdBy: legacyEsKbDoc.created_by, + updatedAt: legacyEsKbDoc.updated_at, + updatedBy: legacyEsKbDoc.updated_by, + users: + legacyEsKbDoc.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + + name: legacyEsKbDoc.text, + namespace: legacyEsKbDoc.namespace, + type: DocumentEntryType.value, + kbResource: legacyEsKbDoc.metadata?.kbResource ?? 'unknown', + source: legacyEsKbDoc.metadata?.source ?? 'unknown', + required: legacyEsKbDoc.metadata?.required ?? false, + text: legacyEsKbDoc.text, + ...(legacyEsKbDoc.vector + ? { + vector: { + modelId: legacyEsKbDoc.vector.model_id, + tokens: legacyEsKbDoc.vector.tokens, + }, + } + : {}), + }; + return documentEntry; }; 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 b3180d80223cea..ecf9260e999d27 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 @@ -5,7 +5,61 @@ * 2.0. */ -export interface EsKnowledgeBaseEntrySchema { +import type { DocumentEntryType, IndexEntryType } from '@kbn/elastic-assistant-common'; + +export type EsKnowledgeBaseEntrySchema = EsDocumentEntry | EsIndexEntry; + +export interface EsDocumentEntry { + '@timestamp': string; + id: string; + created_at: string; + created_by: string; + updated_at: string; + updated_by: string; + users?: Array<{ + id?: string; + name?: string; + }>; + name: string; + namespace: string; + type: DocumentEntryType; + kb_resource: string; + required: boolean; + source: string; + text: string; + vector?: { + tokens: Record; + model_id: string; + }; +} + +export interface EsIndexEntry { + '@timestamp': string; + id: string; + created_at: string; + created_by: string; + updated_at: string; + updated_by: string; + users?: Array<{ + id?: string; + name?: string; + }>; + name: string; + namespace: string; + type: IndexEntryType; + index: string; + field: string; + description: string; + query_description: string; + input_schema?: Array<{ + field_name: string; + field_type: string; + description: string; + }>; + output_fields?: string[]; +} + +export interface LegacyEsKnowledgeBaseEntrySchema { '@timestamp': string; id: string; created_at: string; @@ -40,15 +94,27 @@ export interface CreateKnowledgeBaseEntrySchema { id?: string; name?: string; }>; - metadata?: { - kbResource: string; - source: string; - required: boolean; - }; + name: string; namespace: string; - text: string; + type: string; + // Document Entry Fields + kb_resource?: string; + required?: boolean; + source?: string; + text?: string; vector?: { tokens: Record; model_id: string; }; + // Index Entry Fields + index?: string; + field?: string; + description?: string; + query_description?: string; + input_schema?: Array<{ + field_name: string; + field_type: string; + description: string; + }>; + output_fields?: string[]; } 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 3191e6844ea137..07da9303207123 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 @@ -81,6 +81,7 @@ export const createPipeline = async ({ return response.acknowledged; } catch (e) { + // TODO: log error or just use semantic_text already return 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 6fff5f767154b8..ebdc5accec4630 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 @@ -76,6 +76,8 @@ export class AIAssistantService { private resourceInitializationHelper: ResourceInstallationHelper; private initPromise: Promise; private isKBSetupInProgress: boolean = false; + // Temporary 'feature flag' to determine if we should initialize the new kb mappings, toggled when accessing kbDataClient + private v2KnowledgeBaseEnabled: boolean = false; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -88,7 +90,7 @@ export class AIAssistantService { this.knowledgeBaseDataStream = this.createDataStream({ resource: 'knowledgeBase', kibanaVersion: options.kibanaVersion, - fieldMap: knowledgeBaseFieldMap, + fieldMap: knowledgeBaseFieldMap, // TODO: use v2 if FF is enabled }); this.promptsDataStream = this.createDataStream({ resource: 'prompts', @@ -182,7 +184,7 @@ export class AIAssistantService { esClient, id: this.resourceNames.pipelines.knowledgeBase, }); - if (!pipelineCreated) { + if (!pipelineCreated || this.v2KnowledgeBaseEnabled) { this.options.logger.debug( `Installing ingest pipeline - ${this.resourceNames.pipelines.knowledgeBase}` ); @@ -327,8 +329,15 @@ export class AIAssistantService { } public async createAIAssistantKnowledgeBaseDataClient( - opts: CreateAIAssistantClientParams + opts: CreateAIAssistantClientParams & { v2KnowledgeBaseEnabled: 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.v2KnowledgeBaseEnabled) { + this.v2KnowledgeBaseEnabled = true; + await this.initializeResources(); + } + const res = await this.checkResourcesInstallation(opts); if (res === null) { @@ -347,6 +356,7 @@ export class AIAssistantService { ml: this.options.ml, setIsKBSetupInProgress: this.setIsKBSetupInProgress.bind(this), spaceId: opts.spaceId, + v2KnowledgeBaseEnabled: opts.v2KnowledgeBaseEnabled, }); } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.test.ts index b77e20a1030ff4..17938cbd4eb483 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.test.ts @@ -38,7 +38,7 @@ describe('addRequiredKbResourceMetadata', () => { }); }); - test('it adds the expected `required` metadata to each document', () => { + test('it adds the expected `required` metadata to each document, defaulting to true', () => { const transformedDocs = addRequiredKbResourceMetadata({ docs: mockExampleQueryDocsFromDirectoryLoader, kbResource, @@ -48,4 +48,18 @@ describe('addRequiredKbResourceMetadata', () => { expect(doc.metadata).toHaveProperty('required', true); }); }); + + test('it sets the `required` metadata to the provided value', () => { + const required = false; + + const transformedDocs = addRequiredKbResourceMetadata({ + docs: mockExampleQueryDocsFromDirectoryLoader, + kbResource, + required, + }); + + transformedDocs.forEach((doc) => { + expect(doc.metadata).toHaveProperty('required', required); + }); + }); }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.ts index 43804cc344f5b0..b85904b517e278 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/add_required_kb_resource_metadata.ts @@ -22,15 +22,17 @@ import { Document } from 'langchain/document'; export const addRequiredKbResourceMetadata = ({ docs, kbResource, + required = true, }: { docs: Array>>; kbResource: string; + required?: boolean; }): Array>> => docs.map((doc) => ({ ...doc, metadata: { ...doc.metadata, kbResource, - required: true, // indicates that the document is required for searches on the kbResource topic + required, // indicates that the document is required for searches on the kbResource topic }, })); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.test.ts index 319e1ea3eedf80..56ae404e4a2dbc 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.test.ts @@ -56,8 +56,16 @@ describe('loadESQL', () => { it('loads ES|QL docs, language files, and example queries into the Knowledge Base', async () => { expect(esStore.addDocuments).toHaveBeenCalledWith([ - ...mockEsqlDocsFromDirectoryLoader, - ...mockEsqlLanguageDocsFromDirectoryLoader, + ...addRequiredKbResourceMetadata({ + docs: mockEsqlDocsFromDirectoryLoader, + kbResource: ESQL_RESOURCE, + required: false, + }), + ...addRequiredKbResourceMetadata({ + docs: mockEsqlLanguageDocsFromDirectoryLoader, + kbResource: ESQL_RESOURCE, + required: false, + }), ...addRequiredKbResourceMetadata({ docs: mockExampleQueryDocsFromDirectoryLoader, kbResource: ESQL_RESOURCE, diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.ts index 3d4666cdf7540c..72882f771c5a1e 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/content_loaders/esql_loader.ts @@ -59,13 +59,26 @@ export const loadESQL = async (esStore: ElasticsearchStore, logger: Logger): Pro kbResource: ESQL_RESOURCE, }) as Array>; + // And make sure remaining docs have `kbResource:esql` + const docsWithMetadata = addRequiredKbResourceMetadata({ + docs, + kbResource: ESQL_RESOURCE, + required: false, + }) as Array>; + + const languageDocsWithMetadata = addRequiredKbResourceMetadata({ + docs: languageDocs, + kbResource: ESQL_RESOURCE, + required: false, + }) as Array>; + logger.info( - `Loading ${docs.length} ES|QL docs, ${languageDocs.length} language docs, and ${requiredExampleQueries.length} example queries into the Knowledge Base` + `Loading ${docsWithMetadata.length} ES|QL docs, ${languageDocsWithMetadata.length} language docs, and ${requiredExampleQueries.length} example queries into the Knowledge Base` ); const response = await esStore.addDocuments([ - ...docs, - ...languageDocs, + ...docsWithMetadata, + ...languageDocsWithMetadata, ...requiredExampleQueries, ]); 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 f17accfc85d08e..d0fe4f5097dfe9 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 @@ -126,6 +126,17 @@ export const callAssistantGraph: AgentExecutor = async ({ (tool) => tool.getTool({ ...assistantToolParams, llm: createLlmInstance() }) ?? [] ); + // If KB enabled, fetch for any KB IndexEntries and generate a tool for each + if (isEnabledKnowledgeBase && dataClients?.kbDataClient?.isV2KnowledgeBaseEnabled) { + const kbTools = await dataClients?.kbDataClient?.getAssistantTools({ + assistantToolParams, + esClient, + }); + if (kbTools) { + tools.push(...kbTools); + } + } + const agentRunnable = isOpenAI ? await createOpenAIFunctionsAgent({ llm: createLlmInstance(), 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 574ba5290ce33b..6cc8853d119dd6 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 @@ -33,7 +33,7 @@ import { RetrievalQAChain } from 'langchain/chains'; import { buildResponse } from '../../lib/build_response'; import { AssistantDataClients } from '../../lib/langchain/executors/types'; import { AssistantToolParams, ElasticAssistantRequestHandlerContext, GetElser } from '../../types'; -import { DEFAULT_PLUGIN_NAME, performChecks } from '../helpers'; +import { DEFAULT_PLUGIN_NAME, isV2KnowledgeBaseEnabled, performChecks } from '../helpers'; import { ESQL_RESOURCE } from '../knowledge_base/constants'; import { fetchLangSmithDataset } from './utils'; import { transformESSearchToAnonymizationFields } from '../../ai_assistant_data_clients/anonymization_fields/helpers'; @@ -90,6 +90,7 @@ export const postEvaluateRoute = ( const logger = assistantContext.logger.get('evaluate'); const telemetry = assistantContext.telemetry; const abortSignal = getRequestAbortedSignal(request.events.aborted$); + const v2KnowledgeBaseEnabled = isV2KnowledgeBaseEnabled({ context: ctx, request }); // Perform license, authenticated user and evaluation FF checks const checkResponse = performChecks({ @@ -155,7 +156,9 @@ export const postEvaluateRoute = ( const conversationsDataClient = (await assistantContext.getAIAssistantConversationsDataClient()) ?? undefined; const kbDataClient = - (await assistantContext.getAIAssistantKnowledgeBaseDataClient()) ?? undefined; + (await assistantContext.getAIAssistantKnowledgeBaseDataClient( + v2KnowledgeBaseEnabled + )) ?? undefined; const dataClients: AssistantDataClients = { anonymizationFieldsDataClient, conversationsDataClient, diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts index 422b33b560f8c8..fba788df5cc2ef 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -330,6 +330,8 @@ export const langChainExecute = async ({ const assistantTools = assistantContext .getRegisteredTools(pluginName) .filter((x) => x.id !== 'attack-discovery'); // We don't (yet) support asking the assistant for NEW attack discoveries from a conversation + const v2KnowledgeBaseEnabled = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; // get a scoped esClient for assistant memory const esClient = context.core.elasticsearch.client.asCurrentUser; @@ -345,7 +347,8 @@ export const langChainExecute = async ({ // Create an ElasticsearchStore for KB interactions const kbDataClient = - (await assistantContext.getAIAssistantKnowledgeBaseDataClient()) ?? undefined; + (await assistantContext.getAIAssistantKnowledgeBaseDataClient(v2KnowledgeBaseEnabled)) ?? + undefined; const bedrockChatEnabled = assistantContext.getRegisteredFeatures(pluginName).assistantBedrockChat; const esStore = new ElasticsearchStore( @@ -587,3 +590,26 @@ export const performChecks = ({ return undefined; }; + +/** + * Returns whether the v2 KB is enabled + * + * @param context - Route context + * @param request - Route KibanaRequest + + */ +export const isV2KnowledgeBaseEnabled = ({ + context, + request, +}: { + context: AwaitedProperties< + Pick + >; + request: KibanaRequest; +}): boolean => { + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); + return context.elasticAssistant.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; +}; 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 4352c866a19c1d..427bca2da3dcc1 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 @@ -22,6 +22,7 @@ import { ElasticAssistantRequestHandlerContext } from '../../types'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { ESQL_RESOURCE } from './constants'; import { getKbResource } from './get_kb_resource'; +import { isV2KnowledgeBaseEnabled } from '../helpers'; /** * Delete Knowledge Base index, pipeline, and resources (collection of documents) @@ -49,16 +50,20 @@ export const deleteKnowledgeBaseRoute = ( }, async (context, request: KibanaRequest, response) => { const resp = buildResponse(response); - const assistantContext = await context.elasticAssistant; - const logger = assistantContext.logger; + const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); + const assistantContext = ctx.elasticAssistant; + const logger = ctx.elasticAssistant.logger; const telemetry = assistantContext.telemetry; + // FF Check for V2 KB + const v2KnowledgeBaseEnabled = isV2KnowledgeBaseEnabled({ context: ctx, request }); + try { const kbResource = getKbResource(request); const esClient = (await context.core).elasticsearch.client.asInternalUser; const knowledgeBaseDataClient = - await assistantContext.getAIAssistantKnowledgeBaseDataClient(); + await assistantContext.getAIAssistantKnowledgeBaseDataClient(v2KnowledgeBaseEnabled); if (!knowledgeBaseDataClient) { return response.custom({ body: { success: false }, statusCode: 500 }); } diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts index 22543a21973ce0..af6bf559376adb 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts @@ -166,7 +166,9 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug // subscribing to completed$, because it handles both cases when request was completed and aborted. // when route is finished by timeout, aborted$ is not getting fired request.events.completed$.subscribe(() => abortController.abort()); - const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(); + const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient( + true + ); const spaceId = ctx.elasticAssistant.getSpaceId(); // Authenticated user null check completed in `performChecks()` above const authenticatedUser = ctx.elasticAssistant.getCurrentUser() as AuthenticatedUser; 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 009ad189708aa2..cf7c31d980ac8d 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 @@ -15,9 +15,7 @@ import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/ import { KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, - Metadata, -} from '@kbn/elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen'; -import type { Document } from 'langchain/document'; +} from '@kbn/elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen'; import { ElasticAssistantPluginRouter } from '../../../types'; import { buildResponse } from '../../utils'; import { performChecks } from '../../helpers'; @@ -60,15 +58,15 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout return checkResponse; } + // Check mappings and upgrade if necessary -- this route only supports v2 KB, so always `true` + const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient( + true + ); + logger.debug(() => `Creating KB Entry:\n${JSON.stringify(request.body)}`); - const documents: Array> = [ - { - metadata: request.body.metadata, - pageContent: request.body.text, - }, - ]; - const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(); - const createResponse = await kbDataClient?.addKnowledgeBaseDocuments({ documents }); + const createResponse = await kbDataClient?.createKnowledgeBaseEntry({ + knowledgeBaseEntry: request.body, + }); if (createResponse == null) { return assistantResponse.error({ @@ -76,7 +74,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout statusCode: 400, }); } - return response.ok({ body: KnowledgeBaseEntryResponse.parse(createResponse[0]) }); + return response.ok({ body: createResponse }); } catch (err) { const error = transformError(err as Error); return assistantResponse.error({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts index e939eae1eb3ad5..25d4bbb195d68c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts @@ -21,6 +21,7 @@ import { buildResponse } from '../../utils'; import { performChecks } from '../../helpers'; import { transformESSearchToKnowledgeBaseEntry } from '../../../ai_assistant_data_clients/knowledge_base/transforms'; import { EsKnowledgeBaseEntrySchema } from '../../../ai_assistant_data_clients/knowledge_base/types'; +import { getKBUserFilter } from './utils'; export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -63,16 +64,19 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout return checkResponse; } - const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(); + const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient( + true + ); const currentUser = ctx.elasticAssistant.getCurrentUser(); - + const userFilter = getKBUserFilter(currentUser); const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; + const result = await kbDataClient?.findDocuments({ perPage: query.per_page, page: query.page, sortField: query.sort_field, sortOrder: query.sort_order, - filter: `users:{ id: "${currentUser?.profile_uid}" }${additionalFilter}`, // TODO: Update filter to include non-user system entries + filter: `${userFilter}${additionalFilter}`, fields: query.fields, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/utils.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/utils.ts new file mode 100644 index 00000000000000..6ca137d64e7a3a --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/utils.ts @@ -0,0 +1,12 @@ +/* + * 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 { AuthenticatedUser } from '@kbn/core-security-common'; + +export const getKBUserFilter = (user: AuthenticatedUser | null) => { + return user?.profile_uid ? `users.id: "${user?.profile_uid}" or NOT users: *` : 'NOT users: *'; +}; 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 4907263d3713be..404e489d796066 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 } from './constants'; +import { isV2KnowledgeBaseEnabled } from '../helpers'; /** * Get the status of the Knowledge Base index, pipeline, and resources (collection of documents) @@ -50,8 +51,9 @@ export const getKnowledgeBaseStatusRoute = ( }, async (context, request: KibanaRequest, response) => { const resp = buildResponse(response); - const assistantContext = await context.elasticAssistant; - const logger = assistantContext.logger; + const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); + const assistantContext = ctx.elasticAssistant; + const logger = ctx.elasticAssistant.logger; const telemetry = assistantContext.telemetry; try { @@ -60,7 +62,12 @@ export const getKnowledgeBaseStatusRoute = ( const elserId = await getElser(); const kbResource = getKbResource(request); - const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient(); + // FF Check for V2 KB + const v2KnowledgeBaseEnabled = isV2KnowledgeBaseEnabled({ context: ctx, request }); + + const kbDataClient = await assistantContext.getAIAssistantKnowledgeBaseDataClient( + v2KnowledgeBaseEnabled + ); if (!kbDataClient) { return response.custom({ body: { success: false }, statusCode: 500 }); } @@ -89,9 +96,14 @@ export const getKnowledgeBaseStatusRoute = ( pipeline_exists: pipelineExists, }; - if (kbResource === ESQL_RESOURCE) { + if (indexExists && kbResource === ESQL_RESOURCE) { const esqlExists = - indexExists && (await esStore.similaritySearch(ESQL_DOCS_LOADED_QUERY)).length > 0; + ( + await kbDataClient.getKnowledgeBaseDocumentEntries({ + query: ESQL_DOCS_LOADED_QUERY, + required: true, + }) + ).length > 0; return response.ok({ body: { ...body, esql_exists: esqlExists } }); } 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 c6bc89da345b91..36ea1e867eb7d2 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 @@ -17,6 +17,7 @@ import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantPluginRouter, GetElser } from '../../types'; import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store'; import { getKbResource } from './get_kb_resource'; +import { isV2KnowledgeBaseEnabled } 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 @@ -57,17 +58,21 @@ export const postKnowledgeBaseRoute = ( response ): Promise> => { const resp = buildResponse(response); - const assistantContext = await context.elasticAssistant; - const logger = assistantContext.logger; + const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); + const assistantContext = ctx.elasticAssistant; + const logger = ctx.elasticAssistant.logger; const telemetry = assistantContext.telemetry; const elserId = await getElser(); - const core = await context.core; + const core = ctx.core; const esClient = core.elasticsearch.client.asInternalUser; const soClient = core.savedObjects.getClient(); + // FF Check for V2 KB + const v2KnowledgeBaseEnabled = isV2KnowledgeBaseEnabled({ context: ctx, request }); + try { const knowledgeBaseDataClient = - await assistantContext.getAIAssistantKnowledgeBaseDataClient(); + await assistantContext.getAIAssistantKnowledgeBaseDataClient(v2KnowledgeBaseEnabled); if (!knowledgeBaseDataClient) { return response.custom({ body: { success: false }, statusCode: 500 }); } 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 b98ddfc6ac8f97..3d004994b32364 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,12 +81,15 @@ export class RequestContextFactory implements IRequestContextFactory { telemetry: core.analytics, - getAIAssistantKnowledgeBaseDataClient: memoize(() => { + // 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((v2KnowledgeBaseEnabled = false) => { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantKnowledgeBaseDataClient({ spaceId: getSpaceId(), logger: this.logger, currentUser, + v2KnowledgeBaseEnabled, }); }), diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 6581a2aa4c0f5e..6885b07a42c30f 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -119,7 +119,9 @@ export interface ElasticAssistantApiRequestHandlerContext { getSpaceId: () => string; getCurrentUser: () => AuthenticatedUser | null; getAIAssistantConversationsDataClient: () => Promise; - getAIAssistantKnowledgeBaseDataClient: () => Promise; + getAIAssistantKnowledgeBaseDataClient: ( + v2KnowledgeBaseEnabled?: boolean + ) => Promise; getAttackDiscoveryDataClient: () => Promise; getAIAssistantPromptsDataClient: () => Promise; getAIAssistantAnonymizationFieldsDataClient: () => Promise; diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index da2f19445b2992..c210253af04a45 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -47,6 +47,7 @@ "@kbn/security-plugin", "@kbn/apm-utils", "@kbn/std", + "@kbn/zod", ], "exclude": [ "target/**/*", diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts index 3e69305a9af641..7739de18857aa7 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts @@ -45,7 +45,7 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = { () => `KnowledgeBaseRetrievalToolParams:input\n ${JSON.stringify(input, null, 2)}` ); - const docs = await kbDataClient.getKnowledgeBaseDocuments({ + const docs = await kbDataClient.getKnowledgeBaseDocumentEntries({ query: input.query, kbResource: 'user', required: false, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts index 3dc1c19ce79f25..90f0ad0b96e6ee 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -9,6 +9,7 @@ import { DynamicStructuredTool } from '@langchain/core/tools'; import { z } from '@kbn/zod'; import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base'; +import { DocumentEntryType } from '@kbn/elastic-assistant-common'; import type { KnowledgeBaseEntryCreateProps } from '@kbn/elastic-assistant-common'; import { APP_UI_ID } from '../../../../common'; @@ -18,7 +19,7 @@ export interface KnowledgeBaseWriteToolParams extends AssistantToolParams { const toolDetails = { description: - "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, with no other text, and whether or not the entry is required.", + "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, a short UI friendly name for the entry, and whether or not the entry is required.", id: 'knowledge-base-write-tool', name: 'KnowledgeBaseWriteTool', }; @@ -39,25 +40,39 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { name: toolDetails.name, description: toolDetails.description, schema: z.object({ + name: z + .string() + .describe(`This is what the user will use to refer to the entry in the future.`), query: z.string().describe(`Summary of items/things to save in the knowledge base`), required: z .boolean() .describe( `Whether or not the entry is required to always be included in conversations. Is only true if the user explicitly asks for it to be required or always included in conversations, otherwise this is always false.` - ), + ) + .default(false), }), func: async (input, _, cbManager) => { logger.debug( () => `KnowledgeBaseWriteToolParams:input\n ${JSON.stringify(input, null, 2)}` ); - const knowledgeBaseEntry: KnowledgeBaseEntryCreateProps = { - metadata: { kbResource: 'user', source: 'conversation', required: input.required }, - text: input.query, - }; + // Backwards compatibility with v1 schema -- createKnowledgeBaseEntry() technically supports both for now + const knowledgeBaseEntry: KnowledgeBaseEntryCreateProps = + kbDataClient.isV2KnowledgeBaseEnabled + ? { + name: input.name, + kbResource: 'user', + source: 'conversation', + required: input.required, + text: input.query, + type: DocumentEntryType.value, + } + : ({ + metadata: { kbResource: 'user', source: 'conversation', required: input.required }, + text: input.query, + } as unknown as KnowledgeBaseEntryCreateProps); logger.debug(() => `knowledgeBaseEntry\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}`); - const resp = await kbDataClient.createKnowledgeBaseEntry({ knowledgeBaseEntry }); if (resp == null) { diff --git a/x-pack/test/security_solution_api_integration/package.json b/x-pack/test/security_solution_api_integration/package.json index f750078df62cea..726d6b61feed53 100644 --- a/x-pack/test/security_solution_api_integration/package.json +++ b/x-pack/test/security_solution_api_integration/package.json @@ -45,6 +45,13 @@ "intialize-server:explore": "node scripts/index.js server explore trial_license_complete_tier", "run-tests:explore": "node scripts/index.js runner explore trial_license_complete_tier", + "genai_kb_entries:server:serverless": "npm run initialize-server:genai:trial_complete knowledge_base/entries serverless", + "genai_kb_entries:runner:serverless": "npm run run-tests:genai:trial_complete knowledge_base/entries serverless serverlessEnv", + "genai_kb_entries:qa:serverless": "npm run run-tests:genai:trial_complete knowledge_base/entries serverless qaPeriodicEnv", + "genai_kb_entries:qa:serverless:release": "npm run run-tests:genai:trial_complete knowledge_base/entries serverless qaEnv", + "genai_kb_entries:server:ess": "npm run initialize-server:genai:trial_complete knowledge_base/entries ess", + "genai_kb_entries:runner:ess": "npm run run-tests:genai:trial_complete knowledge_base/entries ess essEnv", + "nlp_cleanup_task:complete:server:serverless": "npm run initialize-server:genai:trial_complete nlp_cleanup_task serverless", "nlp_cleanup_task:complete:runner:serverless": "npm run run-tests:genai:trial_complete nlp_cleanup_task serverless serverlessEnv", "nlp_cleanup_task:complete:qa:serverless": "npm run run-tests:genai:trial_complete nlp_cleanup_task serverless qaPeriodicEnv", diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/basic.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/basic.ts new file mode 100644 index 00000000000000..e31e36cf5f4bf3 --- /dev/null +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/basic.ts @@ -0,0 +1,60 @@ +/* + * 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 expect from 'expect'; +import { + DocumentEntryCreateFields, + DocumentEntryType, + IndexEntryCreateFields, + IndexEntryType, +} from '@kbn/elastic-assistant-common'; +import { FtrProviderContext } from '../../../../../ftr_provider_context'; +import { createEntry } from '../utils/create_entry'; + +const documentEntry: DocumentEntryCreateFields = { + name: 'Sample Document Entry', + type: DocumentEntryType.value, + required: false, + source: 'api', + kbResource: 'user', + namespace: 'default', + text: 'This is a sample document entry', + users: [], +}; + +const indexEntry: IndexEntryCreateFields = { + name: 'Sample Index Entry', + type: IndexEntryType.value, + namespace: 'default', + index: 'sample-index', + field: 'sample-field', + description: 'This is a sample index entry', + users: [], + queryDescription: 'Use sample-field to search in sample-index', +}; + +export default ({ getService }: FtrProviderContext) => { + const supertest = getService('supertest'); + const log = getService('log'); + + // TODO: Fill out tests + describe.skip('@ess @serverless Basic Security AI Assistant Knowledge Base Entries', () => { + describe('Create Entries', () => { + it('should create a new document entry', async () => { + const entry = await createEntry(supertest, log, documentEntry); + + expect(entry).toEqual(documentEntry); + }); + + it('should create a new index entry', async () => { + const entry = await createEntry(supertest, log, indexEntry); + + expect(entry).toEqual(indexEntry); + }); + }); + }); +}; diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts new file mode 100644 index 00000000000000..8b26b7b4652489 --- /dev/null +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/ess.config.ts @@ -0,0 +1,26 @@ +/* + * 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 { FtrConfigProviderContext } from '@kbn/test'; + +export default async function ({ readConfigFile }: FtrConfigProviderContext) { + const functionalConfig = await readConfigFile( + require.resolve('../../../../../../config/ess/config.base.trial') + ); + + return { + ...functionalConfig.getAll(), + kbnTestServer: { + ...functionalConfig.get('kbnTestServer'), + serverArgs: [...functionalConfig.get('kbnTestServer.serverArgs')], + }, + testFiles: [require.resolve('..')], + junit: { + reportName: 'GenAI - Knowledge Base Entries Tests - ESS Env - Trial License', + }, + }; +} diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/serverless.config.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/serverless.config.ts new file mode 100644 index 00000000000000..129f7243059cab --- /dev/null +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/configs/serverless.config.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 { createTestConfig } from '../../../../../../config/serverless/config.base'; + +export default createTestConfig({ + kbnTestServerArgs: [], + testFiles: [require.resolve('..')], + junit: { + reportName: 'GenAI - Knowledge Base Entries Tests - Serverless Env - Complete Tier', + }, +}); diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/index.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/index.ts new file mode 100644 index 00000000000000..a8f259ee4c8e28 --- /dev/null +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/index.ts @@ -0,0 +1,14 @@ +/* + * 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 { FtrProviderContext } from '../../../../../ftr_provider_context'; + +export default function ({ loadTestFile }: FtrProviderContext) { + describe('GenAI - Knowledge Base Entries APIs', function () { + loadTestFile(require.resolve('./basic')); + }); +} diff --git a/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/create_entry.ts b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/create_entry.ts new file mode 100644 index 00000000000000..86210d53c45e22 --- /dev/null +++ b/x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/create_entry.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 { ELASTIC_HTTP_VERSION_HEADER } from '@kbn/core-http-common'; +import type { ToolingLog } from '@kbn/tooling-log'; +import type SuperTest from 'supertest'; + +import { + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL, + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryResponse, +} from '@kbn/elastic-assistant-common'; +import { routeWithNamespace } from '../../../../../../common/utils/security_solution'; + +/** + * Creates a Knowledge Base Entry + * @param supertest The supertest deps + * @param log The tooling logger + * @param entry The entry to create + * @param namespace The Kibana Space to create the entry in (optional) + */ +export const createEntry = async ( + supertest: SuperTest.Agent, + log: ToolingLog, + entry: KnowledgeBaseEntryCreateProps, + namespace?: string +): Promise => { + const route = routeWithNamespace(ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL, namespace); + const response = await supertest + .post(route) + .set('kbn-xsrf', 'true') + .set(ELASTIC_HTTP_VERSION_HEADER, '1') + .send(entry); + if (response.status !== 200) { + throw new Error( + `Unexpected non 200 ok when attempting to create entry: ${JSON.stringify( + response.status + )},${JSON.stringify(response, null, 4)}` + ); + } else { + return response.body; + } +}; diff --git a/x-pack/test/security_solution_api_integration/tsconfig.json b/x-pack/test/security_solution_api_integration/tsconfig.json index 1f558e3c3f0516..82decfa5a6db34 100644 --- a/x-pack/test/security_solution_api_integration/tsconfig.json +++ b/x-pack/test/security_solution_api_integration/tsconfig.json @@ -47,6 +47,7 @@ "@kbn/utility-types", "@kbn/timelines-plugin", "@kbn/dev-cli-runner", + "@kbn/elastic-assistant-common", "@kbn/search-types", "@kbn/security-plugin", "@kbn/ftr-common-functional-ui-services",