diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1cd9dd3dc14d7..bb8b6216a55b5 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -862,6 +862,7 @@ x-pack/platform/packages/shared/index-lifecycle-management/index_lifecycle_manag x-pack/platform/packages/shared/index-management/index_management_shared_types @elastic/kibana-management x-pack/platform/packages/shared/ingest-pipelines @elastic/kibana-management x-pack/platform/packages/shared/kbn-ai-assistant @elastic/search-kibana @elastic/obs-ai-assistant +x-pack/platform/packages/shared/kbn-ai-tools @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-ai-tools-cli @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-alerting-comparators @elastic/response-ops x-pack/platform/packages/shared/kbn-apm-types @elastic/obs-ux-infra_services-team @@ -875,6 +876,7 @@ x-pack/platform/packages/shared/kbn-evals @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-event-stacktrace @elastic/obs-ux-infra_services-team @elastic/obs-ux-logs-team x-pack/platform/packages/shared/kbn-inference-cli @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-inference-endpoint-ui-common @elastic/appex-ai-infra +x-pack/platform/packages/shared/kbn-inference-prompt-utils @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-inference-tracing @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-inference-tracing-config @elastic/appex-ai-infra x-pack/platform/packages/shared/kbn-key-value-metadata-table @elastic/obs-ux-infra_services-team @elastic/obs-ux-logs-team diff --git a/package.json b/package.json index 54824585588e1..a0708e83abfa4 100644 --- a/package.json +++ b/package.json @@ -177,6 +177,7 @@ "@kbn/ai-assistant-icon": "link:x-pack/platform/packages/shared/ai-assistant/icon", "@kbn/ai-assistant-management-plugin": "link:src/platform/plugins/shared/ai_assistant_management/selection", "@kbn/ai-security-labs-content": "link:x-pack/solutions/security/packages/ai-security-labs-content", + "@kbn/ai-tools": "link:x-pack/platform/packages/shared/kbn-ai-tools", "@kbn/aiops-change-point-detection": "link:x-pack/platform/packages/private/ml/aiops_change_point_detection", "@kbn/aiops-common": "link:x-pack/platform/packages/shared/ml/aiops_common", "@kbn/aiops-components": "link:x-pack/platform/packages/private/ml/aiops_components", @@ -617,6 +618,7 @@ "@kbn/inference-endpoint-ui-common": "link:x-pack/platform/packages/shared/kbn-inference-endpoint-ui-common", "@kbn/inference-langchain": "link:x-pack/platform/packages/shared/ai-infra/inference-langchain", "@kbn/inference-plugin": "link:x-pack/platform/plugins/shared/inference", + "@kbn/inference-prompt-utils": "link:x-pack/platform/packages/shared/kbn-inference-prompt-utils", "@kbn/inference-tracing": "link:x-pack/platform/packages/shared/kbn-inference-tracing", "@kbn/inference-tracing-config": "link:x-pack/platform/packages/shared/kbn-inference-tracing-config", "@kbn/infra-forge": "link:x-pack/platform/packages/private/kbn-infra-forge", diff --git a/tsconfig.base.json b/tsconfig.base.json index f7af272ef9fc4..eb83a7707e55f 100644 --- a/tsconfig.base.json +++ b/tsconfig.base.json @@ -26,6 +26,8 @@ "@kbn/ai-assistant-management-plugin/*": ["src/platform/plugins/shared/ai_assistant_management/selection/*"], "@kbn/ai-security-labs-content": ["x-pack/solutions/security/packages/ai-security-labs-content"], "@kbn/ai-security-labs-content/*": ["x-pack/solutions/security/packages/ai-security-labs-content/*"], + "@kbn/ai-tools": ["x-pack/platform/packages/shared/kbn-ai-tools"], + "@kbn/ai-tools/*": ["x-pack/platform/packages/shared/kbn-ai-tools/*"], "@kbn/ai-tools-cli": ["x-pack/platform/packages/shared/kbn-ai-tools-cli"], "@kbn/ai-tools-cli/*": ["x-pack/platform/packages/shared/kbn-ai-tools-cli/*"], "@kbn/aiops-change-point-detection": ["x-pack/platform/packages/private/ml/aiops_change_point_detection"], @@ -1146,6 +1148,8 @@ "@kbn/inference-langchain/*": ["x-pack/platform/packages/shared/ai-infra/inference-langchain/*"], "@kbn/inference-plugin": ["x-pack/platform/plugins/shared/inference"], "@kbn/inference-plugin/*": ["x-pack/platform/plugins/shared/inference/*"], + "@kbn/inference-prompt-utils": ["x-pack/platform/packages/shared/kbn-inference-prompt-utils"], + "@kbn/inference-prompt-utils/*": ["x-pack/platform/packages/shared/kbn-inference-prompt-utils/*"], "@kbn/inference-tracing": ["x-pack/platform/packages/shared/kbn-inference-tracing"], "@kbn/inference-tracing/*": ["x-pack/platform/packages/shared/kbn-inference-tracing/*"], "@kbn/inference-tracing-config": ["x-pack/platform/packages/shared/kbn-inference-tracing-config"], diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts index ad1b66e367e4e..d7cee77d72a26 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts @@ -26,6 +26,7 @@ export { type ToolCallsOf, type ToolCallbacksOf, type ToolCall, + type ToolCallback, type ToolDefinition, type ToolOptions, type FunctionCallingMode, diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts index 674019e0c0dce..42529e9918e1c 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts @@ -42,6 +42,7 @@ export { type ToolSchema, type ToolSchemaType, type FromToolSchema } from './too export { ToolChoiceType, type ToolCallbacksOf, + type ToolCallback, type ToolOptions, type ToolDefinition, type ToolCall, diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/tools.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/tools.ts index 8f485df6fcc5e..93890bf03117c 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/tools.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/tools.ts @@ -28,6 +28,8 @@ type ToolCallbacksOfTools | undefi } : never; +export type ToolCallback = (toolCall: ToolCall) => Promise; + export type ToolCallbacksOf = TToolOptions extends { tools?: Record; } diff --git a/x-pack/platform/packages/shared/kbn-ai-tools-cli/recipes/esql.ts b/x-pack/platform/packages/shared/kbn-ai-tools-cli/recipes/esql.ts new file mode 100644 index 0000000000000..fd5b6a0654e73 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools-cli/recipes/esql.ts @@ -0,0 +1,41 @@ +/* + * 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 { runRecipe } from '@kbn/inference-cli'; +import { executeAsEsqlAgent } from '@kbn/ai-tools'; +import moment from 'moment'; +import { inspect } from 'util'; + +runRecipe( + { + name: 'answer_esql', + flags: { + string: ['prompt'], + help: ` + --prompt The user prompt for generating ES|QL + `, + }, + }, + async ({ inferenceClient, kibanaClient, flags, esClient, logger, log, signal }) => { + const now = moment(); + + const end = now.valueOf(); + + const start = now.clone().subtract(1, 'days').valueOf(); + + const response = await executeAsEsqlAgent({ + start, + end, + esClient, + inferenceClient, + logger, + prompt: String(flags.prompt), + signal, + }); + + log.info(inspect(response, { depth: null })); + } +); diff --git a/x-pack/platform/packages/shared/kbn-ai-tools-cli/tsconfig.json b/x-pack/platform/packages/shared/kbn-ai-tools-cli/tsconfig.json index 800fac2a27df3..6f7ce09b08fd4 100644 --- a/x-pack/platform/packages/shared/kbn-ai-tools-cli/tsconfig.json +++ b/x-pack/platform/packages/shared/kbn-ai-tools-cli/tsconfig.json @@ -17,6 +17,8 @@ "@kbn/core", "@kbn/cache-cli", "@kbn/dev-cli-runner", - "@kbn/kibana-api-cli" + "@kbn/kibana-api-cli", + "@kbn/inference-cli", + "@kbn/ai-tools" ] } diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/README.md b/x-pack/platform/packages/shared/kbn-ai-tools/README.md new file mode 100644 index 0000000000000..520e78123f617 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/README.md @@ -0,0 +1,3 @@ +# @kbn/ai-tools + +A collection of GenAI & ML related (broadly defined) tools, such `describeDataset` which summarizes a data stream or index by sampling data and pre-processing it. diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/index.ts b/x-pack/platform/packages/shared/kbn-ai-tools/index.ts new file mode 100644 index 0000000000000..1962cd74e4115 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/index.ts @@ -0,0 +1,10 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export { describeDataset } from './src/tools/describe_dataset'; +export { sortAndTruncateAnalyzedFields } from './src/tools/describe_dataset/sort_and_truncate_analyzed_fields'; +export { executeAsEsqlAgent } from './src/tools/esql'; diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/jest.config.js b/x-pack/platform/packages/shared/kbn-ai-tools/jest.config.js new file mode 100644 index 0000000000000..447b93496bd89 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/jest.config.js @@ -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. + */ + +module.exports = { + preset: '@kbn/test/jest_node', + rootDir: '../../../../..', + roots: ['/x-pack/platform/packages/shared/kbn-ai-tools'], +}; diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/kibana.jsonc b/x-pack/platform/packages/shared/kbn-ai-tools/kibana.jsonc new file mode 100644 index 0000000000000..9827674a40244 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/kibana.jsonc @@ -0,0 +1,7 @@ +{ + "type": "shared-server", + "id": "@kbn/ai-tools", + "owner": "@elastic/appex-ai-infra", + "group": "platform", + "visibility": "shared" +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/package.json b/x-pack/platform/packages/shared/kbn-ai-tools/package.json new file mode 100644 index 0000000000000..c1ede973f9179 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/package.json @@ -0,0 +1,6 @@ +{ + "name": "@kbn/ai-tools", + "private": true, + "version": "1.0.0", + "license": "Elastic License 2.0" +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/document_analysis.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/document_analysis.ts new file mode 100644 index 0000000000000..be896571ca217 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/document_analysis.ts @@ -0,0 +1,24 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export interface DocumentAnalysis { + total: number; + sampled: number; + fields: Array<{ + name: string; + types: string[]; + cardinality: number | null; + values: Array; + empty: boolean; + }>; +} + +export interface TruncatedDocumentAnalysis { + fields: string[]; + total: number; + sampled: number; +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/get_sample_documents.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/get_sample_documents.ts new file mode 100644 index 0000000000000..cd992261468c3 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/get_sample_documents.ts @@ -0,0 +1,75 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { ElasticsearchClient } from '@kbn/core/server'; +import { QueryDslFieldAndFormat, SearchHit } from '@elastic/elasticsearch/lib/api/types'; +import { kqlQuery, rangeQuery } from './queries'; + +export function getSampleDocuments({ + esClient, + index, + start, + end, + kql, + size = 1000, + fields = [ + { + field: '*', + include_unmapped: true, + }, + ], + _source = false, + timeout = '5s', +}: { + esClient: ElasticsearchClient; + index: string; + start: number; + end: number; + kql?: string; + size?: number; + fields?: Array; + _source?: boolean; + timeout?: string; +}) { + return esClient + .search>({ + index, + size, + track_total_hits: true, + timeout, + query: { + bool: { + must: [...kqlQuery(kql), ...rangeQuery(start, end)], + should: [ + { + function_score: { + functions: [ + { + random_score: {}, + }, + ], + }, + }, + ], + }, + }, + sort: { + _score: { + order: 'desc', + }, + }, + _source, + fields, + }) + .then((response) => ({ + hits: response.hits.hits as Array>>, + total: + typeof response.hits.total === 'number' + ? response.hits.total! + : response.hits.total?.value!, + })); +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/index.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/index.ts new file mode 100644 index 0000000000000..68bc7be32c140 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/index.ts @@ -0,0 +1,53 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { getSampleDocuments } from './get_sample_documents'; +import { mergeSampleDocumentsWithFieldCaps } from './merge_sample_documents_with_field_caps'; +import { rangeQuery } from './queries'; + +export async function describeDataset({ + esClient, + start, + end, + index, + kql, +}: { + esClient: ElasticsearchClient; + start: number; + end: number; + index: string; + kql?: string; +}) { + const [fieldCaps, hits] = await Promise.all([ + esClient.fieldCaps({ + index, + fields: '*', + index_filter: { + bool: { + filter: rangeQuery(start, end), + }, + }, + }), + getSampleDocuments({ + esClient, + index, + start, + end, + kql, + }), + ]); + + const total = hits.total; + + const analysis = mergeSampleDocumentsWithFieldCaps({ + hits: hits.hits, + total: total ?? 0, + fieldCaps, + }); + + return analysis; +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/merge_sample_documents_with_field_caps.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/merge_sample_documents_with_field_caps.ts new file mode 100644 index 0000000000000..976e89fe0a5cc --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/merge_sample_documents_with_field_caps.ts @@ -0,0 +1,99 @@ +/* + * 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 { castArray, sortBy, uniq } from 'lodash'; +import { FieldCapsResponse, SearchHit } from '@elastic/elasticsearch/lib/api/types'; +import { getFlattenedObject } from '@kbn/std'; +import type { DocumentAnalysis } from './document_analysis'; + +export function mergeSampleDocumentsWithFieldCaps({ + total, + hits, + fieldCaps, +}: { + total: number; + hits: SearchHit[]; + fieldCaps: FieldCapsResponse; +}): DocumentAnalysis { + const nonEmptyFields = new Set(); + const fieldValues = new Map>(); + + const samples = hits.map((hit) => ({ + ...hit.fields, + ...getFlattenedObject(hit._source ?? {}), + })); + + const specs = Object.entries(fieldCaps.fields).map(([name, capabilities]) => { + return { name, esTypes: Object.keys(capabilities) }; + }); + + const typesByFields = new Map( + specs.map(({ name, esTypes }) => { + return [name, esTypes ?? []]; + }) + ); + + for (const document of samples) { + Object.keys(document).forEach((field) => { + if (!nonEmptyFields.has(field)) { + nonEmptyFields.add(field); + } + + if (!typesByFields.has(field)) { + typesByFields.set(field, []); + } + + const values = castArray(document[field]); + + const currentFieldValues = fieldValues.get(field) ?? []; + + values.forEach((value) => { + if (typeof value === 'string' || typeof value === 'number' || typeof value === 'boolean') { + currentFieldValues.push(value); + } + }); + + fieldValues.set(field, currentFieldValues); + }); + } + + const fields = Array.from(typesByFields.entries()).flatMap(([name, types]) => { + const values = fieldValues.get(name); + + const countByValues = new Map(); + + values?.forEach((value) => { + const currentCount = countByValues.get(value) ?? 0; + countByValues.set(value, currentCount + 1); + }); + + const sortedValues = sortBy( + Array.from(countByValues.entries()).map(([value, count]) => { + return { + value, + count, + }; + }), + 'count', + 'desc' + ); + + return { + name, + types, + empty: !nonEmptyFields.has(name), + cardinality: countByValues.size || null, + values: uniq(sortedValues.flatMap(({ value }) => value)), + }; + }); + + return { + total, + sampled: samples.length, + fields, + }; +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/queries.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/queries.ts new file mode 100644 index 0000000000000..92ad67a2639aa --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/queries.ts @@ -0,0 +1,40 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { estypes } from '@elastic/elasticsearch'; + +export function rangeQuery( + start?: number, + end?: number, + field = '@timestamp' +): estypes.QueryDslQueryContainer[] { + return [ + { + range: { + [field]: { + gte: start, + lte: end, + format: 'epoch_millis', + }, + }, + }, + ]; +} + +export function kqlQuery(kuery?: string): estypes.QueryDslQueryContainer[] { + if (!kuery) { + return []; + } + + return [ + { + kql: { + query: kuery, + }, + } as estypes.QueryDslQueryContainer, + ]; +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/sort_and_truncate_analyzed_fields.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/sort_and_truncate_analyzed_fields.ts new file mode 100644 index 0000000000000..0a6bdfd4ce627 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/sort_and_truncate_analyzed_fields.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 { partition, shuffle } from 'lodash'; +import { truncateList } from '@kbn/inference-common'; +import type { DocumentAnalysis, TruncatedDocumentAnalysis } from './document_analysis'; + +export function sortAndTruncateAnalyzedFields( + analysis: DocumentAnalysis, + options: { dropEmpty?: boolean } = {} +): TruncatedDocumentAnalysis { + const { dropEmpty = false } = options; + const { fields, ...meta } = analysis; + const [nonEmptyFields, emptyFields] = partition(analysis.fields, (field) => !field.empty); + + // randomize field selection to get a somewhat more illustrative set of fields when + // the # of fields exceeds the threshold, instead of alphabetically sorted + // additionally, prefer non-empty fields over empty fields + const sortedFields = [...shuffle(nonEmptyFields), ...shuffle(emptyFields)]; + + const filteredFields = dropEmpty ? sortedFields.filter((field) => !field.empty) : fields; + + return { + ...meta, + fields: truncateList( + filteredFields.map((field) => { + const types = field.types.join(',') || '(unnmapped)'; + let label = `${field.name}:${types}`; + + if (field.empty) { + return `${field.name} (empty)`; + } + + label += ` - ${field.cardinality} distinct values`; + + if (field.name === '@timestamp' || field.name === 'event.ingested') { + return `${label}`; + } + + const shortValues = field.values.filter((value) => { + return String(value).length <= 1024; + }); + + if (shortValues.length) { + return `${label} (${truncateList( + shortValues.map((value) => '`' + value + '`'), + field.types.includes('text') || field.types.includes('match_only_text') ? 2 : 10 + ).join(', ')})`; + } + + return label; + }), + 500 + ).sort(), + }; +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/content_prompt_template.text b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/content_prompt_template.text new file mode 100644 index 0000000000000..fb181b7858720 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/content_prompt_template.text @@ -0,0 +1,7 @@ +**ES|QL instructions:** + +`{{{esql_system_prompt}}}` + +**User prompt:** + +{{{prompt}}} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/esql_task_description.text b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/esql_task_description.text new file mode 100644 index 0000000000000..3fa755750149c --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/esql_task_description.text @@ -0,0 +1,450 @@ +/** +* FOR THE PROMPT-GENERATION LLM +* +* Make sure the reasoning includes examples of ES|QL queries, errors and how to correct them. +* Also have some examples of hypothetical tools being available, like `visualize_esql`. +* Additionally, the LLM MUST understand that any answer BEFORE the `complete` call is not shown +* to the user. Therefore, if a reasoning steps include the answer, the LLM MUST call complete and +* repeat the answer in a user-facing manner. +*/ + +Your task is to answer any ES|QL related questions, including generating and validating ES|QL queries, based on a user prompt. There are at least tools available (use JSON when writing examples) you can use to retrieve information and validate your answer: + +- `list_datasets(name=string[])`: this returns a list of indices, data streams and aliases, similar to _resolve/_index. + +- `describe_dataset(index=string, kql=?string)`: this returns an aggregated analysis of a data set, based on field-value pairs that are available in n sampled documents. Here's an example of what is returned: + +```json +{ + "total": 653236, + "sampled": 1000, + "fields": [ + "@timestamp:date - 996 distinct values", + "filepath:(unmapped): - 1 distinct values (`my-service.log`)", + "host.name:keyword - 3 distinct values (`host3`, `host2`, `host1`)", + "..." + ] +} +``` + +- `get_documentation(commands=string[], functions=string[])`: returns documentation for one or more ES|QL commands or functions + +- `validate_queries(queries=string[])`: validate one or more ES|QL queries. This will also return validation errors for both syntax and planning, e.g.: + +- `run_queries(queries=string[])`: run one or more validated ES|QL queries and retrieve the results. + +``` +{ + "error": { + "root_cause": [ + { + "type": "verification_exception", + "reason": "Found 1 problem\nline 2:5: Unknown index" + } + ] + } +} +``` + +or + +``` +{ + "error": { + "root_cause": [ + { + "type": "parsing_exception", + "reason": "line 3:9: mismatched input 'LIIT' expecting {'dissect', 'drop', 'enrich', 'eval', 'grok', 'keep', 'limit', 'mv_expand', 'rename', 'sort', 'stats', 'where', 'lookup'}" + } + ], + "type": "parsing_exception", + "reason": "line 3:9: mismatched input 'LIIT' expecting {'dissect', 'drop', 'enrich', 'eval', 'grok', 'keep', 'limit', 'mv_expand', 'rename', 'sort', 'stats', 'where', 'lookup'}", + "caused_by": { + "type": "input_mismatch_exception", + "reason": null + } + }, + "status": 400 +} +``` + +You must first gather enough information using these tools. After you have completed your research you must call either `complete`. Your next reply is the response to the user's prompt. + +There also might be additional tools available, but you can only call them after calling `complete` as a last step. However, you must take these tools into account when planning, validating and reasoning. If there are no other tool calls, use the last message to respond to the user's question. + +/** +* FOR THE PROMPT-GENERATION LLM +* +* What follows below is the ES|QL system prompt. It should be included in +* the rendered system prompt as {{{esql_system_prompt}}}. Use it to rewrite +* the task description, but don't cite it verbatim, as it will be included +* anyway. +*/ + +==== START OF ES|QL SYSTEM PROMPT ==== + +You are a helpful assistant for generating and executing ES|QL queries. +Your goal is to help the user construct an ES|QL query for their data. + +VERY IMPORTANT: When writing ES|QL queries, make sure to ONLY use commands, functions +and operators listed in the current documentation. + +# Limitations + +- ES|QL currently does not support pagination. +- A query will never return more than 10000 rows. + +# Syntax + +An ES|QL query is composed of a source command followed by a series +of processing commands, separated by a pipe character: |. + +For example: + + | + | + +## Source commands + +Source commands select a data source. + +There are three source commands: +- FROM: Selects one or multiple indices, data streams or aliases to use as source. +- ROW: Produces a row with one or more columns with values that you specify. +- SHOW: returns information about the deployment. + +## Processing commands + +ES|QL processing commands change an input table by adding, removing, or +changing rows and columns. + +The following processing commands are available: + +- DISSECT: extracts structured data out of a string, using a dissect pattern +- DROP: drops one or more columns +- ENRICH: adds data from existing indices as new columns +- EVAL: adds a new column with calculated values, using various type of functions +- GROK: extracts structured data out of a string, using a grok pattern +- KEEP: keeps one or more columns, drop the ones that are not kept +- LIMIT: returns the first n number of rows. The maximum value for this is 10000 +- MV_EXPAND: expands multi-value columns into a single row per value +- RENAME: renames a column +- STATS ... BY: groups rows according to a common value and calculates + one or more aggregated values over the grouped rows. STATS supports aggregation + function and can group using grouping functions. +- SORT: sorts the row in a table by a column. Expressions are not supported. +- WHERE: Filters rows based on a boolean condition. WHERE supports the same functions as EVAL. + +## Functions and operators + +### Grouping functions + +BUCKET: Creates groups of values out of a datetime or numeric input +CATEGORIZE: Organize textual data into groups of similar format + +### Aggregation functions + +AVG: calculates the average of a numeric field +COUNT: returns the total number of input values +COUNT_DISTINCT: return the number of distinct values in a field +MAX: calculates the maximum value of a field +MEDIAN: calculates the median value of a numeric field +MEDIAN_ABSOLUTE_DEVIATION: calculates the median absolute deviation of a numeric field +MIN: calculates the minimum value of a field +PERCENTILE: calculates a specified percentile of a numeric field +STD_DEV: calculates the standard deviation of a numeric field +SUM: calculates the total sum of a numeric expression +TOP: collects the top values for a specified field +VALUES: returns all values in a group as a multivalued field +WEIGHTED_AVG: calculates the weighted average of a numeric expression + +### Conditional functions and expressions + +Conditional functions return one of their arguments by evaluating in an if-else manner + +CASE: accepts pairs of conditions and values and returns the value that belongs to the first condition that evaluates to true +COALESCE: returns the first non-null argument from the list of provided arguments +GREATEST: returns the maximum value from multiple columns +LEAST: returns the smallest value from multiple columns + +### Search functions + +Search functions perform full-text search against the data + +MATCH: execute a match query on a specified field (tech preview) +QSTR: performs a Lucene query string query (tech preview) + +### Date-time functions + +DATE_DIFF: calculates the difference between two timestamps in a given unit +DATE_EXTRACT: extract a specific part of a date +DATE_FORMAT: returns a string representation of a date using the provided format +DATE_PARSE: convert a date string into a date +DATE_TRUNC: rounds down a date to the nearest specified interval +NOW: returns the current date and time + +### Mathematical functions + +ABS: returns the absolute value of a number +ACOS: returns the arccosine of a number +ASIN: returns the arcsine of a number +ATAN: returns the arctangent of a number +ATAN2: returns the angle from the positive x-axis to a point (x, y) +CBRT: calculates the cube root of a given number +CEIL: rounds a number up to the nearest integer +COS: returns the cosine of a given angle +COSH: returns the hyperbolic cosine of a given angle +E: returns Euler's number +EXP: returns the value of Euler's number raised to the power of a given number +FLOOR: rounds a number down to the nearest integer +HYPOT: calculate the hypotenuse of two numbers +LOG: calculates the logarithm of a given value to a specified base +LOG10: calculates the logarithm of a value to base 10 +PI: returns the mathematical constant Pi +POW: calculates the value of a base raised to the power of an exponent +ROUND: rounds a numeric value to a specified number of decimal +SIGNUM: returns the sign of a given number +SIN: calculates the sine of a given angle +SINH: calculates the hyperbolic sine of a given angle +SQRT: calculates the square root of a given number +TAN: calculates the tangent of a given angle +TANH: calculates the hyperbolic tangent of a given angle +TAU: returns the mathematical constant τ (tau) + +### String functions + +BIT_LENGTH: calculates the bit length of a string +BYTE_LENGTH: calculates the byte length of a string +CONCAT: combines two or more strings into one +ENDS_WITH: checks if a given string ends with a specified suffix +FROM_BASE64: decodes a base64 string +HASH: computes the hash of a given input using a specified algorithm +LEFT: extracts a specified number of characters from the start of a string +LENGTH: calculates the character length of a given string +LOCATE: returns the position of a specified substring within a string +LTRIM: remove leading whitespaces from a string +REPEAT: generates a string by repeating a specified string a certain number of times +REPLACE: substitutes any match of a regular expression within a string with a replacement string +REVERSE: reverses a string +RIGHT: extracts a specified number of characters from the end of a string +RTRIM: remove trailing whitespaces from a string +SPACE: creates a string composed of a specific number of spaces +SPLIT: split a single valued string into multiple strings based on a delimiter +STARTS_WITH: checks if a given string begins with another specified string +SUBSTRING: extracts a portion of a string +TO_BASE64: encodes a string to a base64 +TO_LOWER: converts a string to lowercase +TO_UPPER: converts a string to uppercase +TRIM: removes leading and trailing whitespaces from a string + +### IP Functions + +CIDR_MATCH: checks if an IP address falls within specified network blocks +IP_PREFIX: truncates an IP address to a specified prefix length + +### Type conversion functions + +TO_BOOLEAN +TO_CARTESIANPOINT +TO_CARTESIANSHAPE +TO_DATETIME (prefer DATE_PARSE to convert strings to datetime) +TO_DATEPERIOD +TO_DEGREES +TO_DOUBLE +TO_GEOPOINT +TO_GEOSHAPE +TO_INTEGER +TO_IP +TO_LONG +TO_RADIANS +TO_STRING +TO_TIMEDURATION +TO_UNSIGNED_LONG +TO_VERSION + +### Multivalue functions + +Multivalue function are used to manipulate and transform multi-value fields. + +MV_APPEND: concatenates the values of two multi-value fields +MV_AVG: returns the average of all values in a multivalued field +MV_CONCAT: transforms a multivalued string expression into a single valued string +MV_COUNT: counts the total number of values in a multivalued expression +MV_DEDUPE: eliminates duplicate values from a multivalued field +MV_FIRST: returns the first value of a multivalued field +MV_LAST: returns the last value of a multivalued field +MV_MAX: returns the max value of a multivalued field +MV_MEDIAN: returns the median value of a multivalued field +MV_MEDIAN_ABSOLUTE_DEVIATION: returns the median absolute deviation of a multivalued field +MV_MIN: returns the min value of a multivalued field +MV_PERCENTILE: returns the specified percentile of a multivalued field +MV_SLIDE: extract a subset of a multivalued field using specified start and end index values +MV_SORT: sorts a multivalued field in lexicographical order. +MV_SUM: returns the sum of all values of a multivalued field +MV_ZIP: combines the values from two multivalued fields with a specified delimiter + +### Spacial functions + +ST_CONTAINS: checks if the first specified geometry encompasses the second one +ST_DISJOINT: checks if two geometries or geometry columns are disjoint +ST_DISTANCE: calculates the distance between two points +ST_ENVELOPE: calculates the minimum bounding box for the provided geometry +ST_INTERSECTS: checks if two geometries intersect +ST_WITHIN: checks if the first geometry is located within the second geometry +ST_X: extracts the x coordinate from a given point +ST_XMAX: extracts the maximum value of the x coordinates from a geometry +ST_XMIN: extracts the minimum value of the x coordinates from a geometry +ST_Y: extracts the y coordinate from a given point +ST_YMAX: extracts the maximum value of the y coordinates from a geometry +ST_YMIN: extracts the minimum value of the y coordinates from a geometry + +### Spacial aggregations functions + +ST_EXTENT_AGG: calculates the spatial extent over a field that has a geometry type +ST_CENTROID_AGG: calculates the spatial centroid over a spatial point geometry field + +### Operators + +Binary operators: ==, !=, <, <=, >, >=, +, -, *, /, % +Logical operators: AND, OR, NOT +Predicates: IS NULL, IS NOT NULL +Unary operators: - +IN: test if a field or expression is in a list of literals +LIKE: filter data based on string patterns using wildcards +RLIKE: filter data based on string patterns using regular expressions +Cast (`::`): provides a convenient alternative syntax to the `TO_` conversion functions + +# Usage examples + +Here are some examples of ES|QL queries: + +**Returns the 10 latest errors from the logs** +```esql +FROM logs +| WHERE level == "ERROR" +| SORT @timestamp DESC +| LIMIT 10 +``` + +**Returns the title and description of last month's blog articles** +```esql +FROM blogposts +| WHERE published > NOW() - 1 month +| KEEP title, description +| SORT title +``` + +**Returns the number of employees from the "NL" country using STATS** +```esql +FROM employees +| WHERE country == "NL" +| STATS COUNT(*) +``` + +**Returns the number of order for each month over last year** +```esql +FROM orders +| WHERE order_date > NOW() - 1 year +| STATS count = COUNT(*) BY date_bucket = BUCKET(order_date, 1 month) +``` + +**Extracting structured data from logs using DISSECT** +```esql +FROM postgres-logs* +// messages are similar to "2023-01-23T12:15:00.000Z - some text - 127.0.0.1" +| DISSECT message "%{date} - %{msg} - %{ip}" +// keep columns created by the dissect command +| KEEP date, msg, ip +// evaluate date from string representation +| EVAL date = DATE_PARSE("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'", date) +``` + +**Find contributors which first name starts with "b", sort them by number of commits and +then returns their first and last names for the top 5** +```esql +FROM commits +| WHERE TO_LOWER(first_name) LIKE "b*" +| STATS doc_count = COUNT(*) by first_name, last_name +| SORT doc_count DESC +| KEEP first_name, last_name +| LIMIT 5 +``` + +**Returning average salary per hire date split in 20 buckets using BUCKET** +```esql +FROM employees +| WHERE hire_date >= "1985-01-01T00:00:00Z" AND hire_date < "1986-01-01T00:00:00Z" +| STATS avg_salary = AVG(salary) BY date_bucket = BUCKET(hire_date, 20, "1985-01-01T00:00:00Z", "1986-01-01T00:00:00Z") +| SORT bucket +``` + +**Returning number of employees grouped by buckets of salary** +```esql +FROM employees +| WHERE hire_date >= "1985-01-01T00:00:00Z" AND hire_date < "1986-01-01T00:00:00Z" +| STATS c = COUNT(1) BY b = BUCKET(salary, 5000.) +| SORT b +``` + +**returns total and recent hire counts plus ratio break down by country** +```esql +FROM employees +// insert a boolean column using case for conditional evaluation +| EVAL is_recent_hire = CASE(hire_date <= "2023-01-01T00:00:00Z", 1, 0) +// using stats with multiple grouping expressions +| STATS total_recent_hires = SUM(is_recent_hire), total_hires = COUNT(*) BY country +// evaluate the recent hiring rate by country based on the previous grouping expressions +| EVAL recent_hiring_rate = total_recent_hires / total_hires +``` + +**computes failure ratios from logs** +```esql +FROM logs-* +| WHERE @timestamp <= NOW() - 24 hours +// convert a keyword field into a numeric field to aggregate over it +| EVAL is_5xx = CASE(http.response.status_code >= 500, 1, 0) +// count total events and failed events to calculate a rate +| STATS total_events = COUNT(*), total_failures = SUM(is_5xx) BY host.hostname, bucket = BUCKET(@timestamp, 1 hour) +// evaluate the failure ratio +| EVAL failure_rate_per_host = total_failures / total_events +// drops the temporary columns +| DROP total_events, total_failures +``` + +**Returning the number of logs grouped by level over the past 24h** +```esql +FROM logs-* +| WHERE @timestamp <= NOW() - 24 hours +| STATS count = COUNT(*) BY log.level +| SORT count DESC +``` + +**Returning all first names for each first letter** +```esql +FROM employees +// evaluate first letter +| EVAL first_letter = SUBSTRING(first_name, 0, 1) +// group all first_name into a multivalued field, break down by first_letter +| STATS first_name = MV_SORT(VALUES(first_name)) BY first_letter +| SORT first_letter +``` + +**Retrieving the min, max and average value from a multivalued field** +```esql +FROM bag_of_numbers +| EVAL min = MV_MIN(numbers), max = MV_MAX(numbers), avg = MV_AVG(numbers) +| KEEP bad_id, min, max, avg +``` + +**Converts a date string into datetime using DATE_PARSE** +```esql +FROM personal_info +// birth_date is a text field storing date with the "yyyy-MM-dd" format +| EVAL birth=DATE_PARSE("yyyy-MM-dd", birth_date) +| KEEP user_name, birth +| SORT birth +``` + +==== END OF ES|QL SYSTEM PROMPT ==== + +=== END OF TASK DESCRIPTION === diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/index.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/index.ts new file mode 100644 index 0000000000000..d8a00c3137f31 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/index.ts @@ -0,0 +1,195 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { errors } from '@elastic/elasticsearch'; +import { ElasticsearchClient, Logger } from '@kbn/core/server'; +import { + BoundInferenceClient, + PromptResponse, + ToolCallbacksOf, + ToolDefinition, + ToolOptions, + truncateList, +} from '@kbn/inference-common'; +import { PromptCompositeResponse, PromptOptions } from '@kbn/inference-common/src/prompt/api'; +import { EsqlDocumentBase, runAndValidateEsqlQuery } from '@kbn/inference-plugin/server'; +import { executeAsReasoningAgent } from '@kbn/inference-prompt-utils'; +import { omit, once } from 'lodash'; +import moment from 'moment'; +import { describeDataset, sortAndTruncateAnalyzedFields } from '../../..'; +import { EsqlPrompt } from './prompt'; + +const loadEsqlDocBase = once(() => EsqlDocumentBase.load()); + +export async function executeAsEsqlAgent | undefined>( + options: { + inferenceClient: BoundInferenceClient; + esClient: ElasticsearchClient; + logger: Logger; + start?: number; + end?: number; + signal: AbortSignal; + prompt: string; + tools?: TTools; + } & (TTools extends Record + ? { toolCallbacks: ToolCallbacksOf<{ tools: TTools }> } + : {}) +): PromptCompositeResponse & { tools: TTools; stream: false }>; + +export async function executeAsEsqlAgent({ + inferenceClient, + esClient, + start, + end, + signal, + prompt, + tools, + toolCallbacks, +}: { + inferenceClient: BoundInferenceClient; + esClient: ElasticsearchClient; + start?: number; + end?: number; + signal: AbortSignal; + prompt: string; + tools?: Record; + toolCallbacks?: ToolCallbacksOf; +}): Promise { + const docBase = await loadEsqlDocBase(); + + async function runEsqlQuery(query: string) { + return await runAndValidateEsqlQuery({ + query, + client: esClient, + }).then((response) => { + if (response.error || response.errorMessages?.length) { + return { + error: + response.error && response.error instanceof errors.ResponseError + ? omit(response.error, 'meta') + : response.error, + errorMessages: response.errorMessages, + }; + } + + return { + columns: response.columns, + rows: response.rows, + }; + }); + } + + const assistantReply = await executeAsReasoningAgent({ + inferenceClient, + prompt: EsqlPrompt, + abortSignal: signal, + tools, + toolCallbacks: { + ...toolCallbacks, + list_datasets: async (toolCall) => { + return esClient.indices + .resolveIndex({ + name: toolCall.function.arguments.name.flatMap((index) => index.split(',')), + allow_no_indices: true, + }) + .then((response) => { + return { + ...response, + data_streams: response.data_streams.map((dataStream) => { + return { + name: dataStream.name, + timestamp_field: dataStream.timestamp_field, + }; + }), + }; + }); + }, + describe_dataset: async (toolCall) => { + const analysis = await describeDataset({ + esClient, + index: toolCall.function.arguments.index, + kql: toolCall.function.arguments.kql, + start: start ?? moment().subtract(24, 'hours').valueOf(), + end: Date.now(), + }); + + return { + analysis: sortAndTruncateAnalyzedFields(analysis), + }; + }, + get_documentation: async (toolCall) => { + return docBase.getDocumentation( + toolCall.function.arguments.commands.concat(toolCall.function.arguments.functions), + { generateMissingKeywordDoc: true } + ); + }, + run_queries: async (toolCall) => { + const results = await Promise.all( + toolCall.function.arguments.queries.map(async (query) => { + const response = await runEsqlQuery(query); + + const cols = response.columns ?? []; + const docs = + response.rows?.map((row) => { + const doc: Record = {}; + row.forEach((value, idx) => { + const col = cols[idx]; + if (value !== null) { + doc[col.name] = value; + } + }); + }) ?? []; + + return { + query, + response: { + docs: truncateList(docs, 50), + }, + }; + }) + ); + + return { + queries: results, + }; + }, + validate_queries: async (toolCall) => { + const results = await Promise.all( + toolCall.function.arguments.queries.map(async (query) => { + return { + query, + validation: await runEsqlQuery(query + ' | LIMIT 0').then((response) => { + if ('error' in response) { + return { + valid: false, + ...response, + }; + } + + const cols = truncateList(response.columns?.map((col) => col.name) ?? [], 10); + return { + valid: true, + ...(cols.length ? { columns: cols } : {}), + }; + }), + }; + }) + ); + + return { + results, + }; + }, + }, + input: { + prompt, + esql_system_prompt: docBase.getSystemMessage(), + }, + }); + + return assistantReply; +} diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/prompt.ts b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/prompt.ts new file mode 100644 index 0000000000000..755a680597503 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/prompt.ts @@ -0,0 +1,120 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { z } from '@kbn/zod'; +import { createPrompt } from '@kbn/inference-common'; +import systemPromptTemplate from './system_prompt_template.text'; +import contentPromptTemplate from './content_prompt_template.text'; + +export const EsqlPrompt = createPrompt({ + name: 'esql_prompt', + description: 'Answer ES|QL related questions', + input: z.object({ + prompt: z.string(), + esql_system_prompt: z.string(), + }), +}) + .version({ + system: { + mustache: { + template: systemPromptTemplate, + }, + }, + template: { + mustache: { + template: contentPromptTemplate, + }, + }, + temperature: 0.25, + tools: { + get_documentation: { + description: 'Get documentation about specific ES|QL commands or functions', + schema: { + type: 'object', + properties: { + commands: { + type: 'array', + items: { + type: 'string', + }, + }, + functions: { + type: 'array', + items: { + type: 'string', + }, + }, + }, + required: ['commands', 'functions'], + }, + }, + validate_queries: { + description: 'Validate one or more ES|QL queries for syntax errors and/or mapping issues', + schema: { + type: 'object', + properties: { + queries: { + type: 'array', + items: { + type: 'string', + }, + }, + }, + required: ['queries'], + }, + }, + run_queries: { + description: 'Run one or more validated ES|QL queries and retrieve the results', + schema: { + type: 'object', + properties: { + queries: { + type: 'array', + items: { + type: 'string', + }, + }, + }, + required: ['queries'], + }, + }, + list_datasets: { + description: + 'List datasets (index, data stream, aliases) based on a name or pattern, similar to _resolve/_index', + schema: { + type: 'object', + properties: { + name: { + type: 'array', + items: { + type: 'string', + }, + }, + }, + required: ['name'], + }, + }, + describe_dataset: { + description: `Get dataset description via sampling of documents`, + schema: { + type: 'object', + properties: { + index: { + type: 'string', + description: 'Index, data stream or index pattern you want to analyze', + }, + kql: { + type: 'string', + description: 'KQL for filtering the data', + }, + }, + required: ['index'], + }, + }, + } as const, + }) + .get(); diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/system_prompt_template.text b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/system_prompt_template.text new file mode 100644 index 0000000000000..b0ffc5c7b6737 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/system_prompt_template.text @@ -0,0 +1,251 @@ +## 1. Purpose + +You are an **expert ES|QL Query Assistant**. Your purpose is to help users by generating, validating, and explaining ES|QL (Elasticsearch Query Language) queries. You will achieve this by methodically using a specialized set of tools to explore the available data, construct a valid query, and then present a final, accurate answer to the user. + +Your workflow is a strict loop: +1. **Gather context** with a specific ES|QL tool. +2. **Think in the clear** via a structured **Reasoning Monologue** after *every* tool response. +3. Repeat Steps 1-2 until you have enough information, then produce one final answer. + +--- + +## 2. Goal & Success Criteria + +### Goal +Your primary goal is to accurately answer the user's question by providing a correct and well-formed ES|QL query, along with any necessary explanations or results. + +### Success criteria +* **Accuracy:** The final ES|QL query must be syntactically correct and semantically valid for the user's data schema. +* **Relevance:** The query must directly address the user's request. +* **Efficiency:** You should aim to determine the correct query with a logical and minimal number of tool calls. +* **Clarity:** The final answer, provided after calling `complete()`, should be clear, user-friendly, and explain the query if necessary. + +--- + +## 3. Available Tools + +| Tool | Function | Notes | +| :--- | :--- | :--- | +| `list_datasets()` | Returns a list of available indices, data streams, and aliases. | Call this first to see what data is available. | +| `describe_dataset(index)` | Analyzes a dataset's schema and fields from a sample of documents. | Essential for discovering field names and types. | +| `get_documentation(commands=[], functions=[])` | Retrieves official documentation for ES|QL commands and functions. | Use this to verify syntax or understand functionality. | +| `validate_queries(queries=[])` | Validates the syntax and semantics of one or more ES|QL queries without running them. | Always use this before `run_queries`. | +| `run_queries(queries=[])` | Executes one or more validated ES|QL queries and returns the results. | Use this to get the final data for the user. | +| `reason()` | **Begin a Reasoning Monologue.** | Outputs your private thoughts. Must use sentinel tags (see §4). | +| `complete()` | Declare readiness to answer. | Ends the tool loop and triggers the **Definitive Output**. | + +*Note: Additional tools, such as `visualize_esql`, might be available. You can mention these as possibilities in your final answer, but you can only call them after `complete()` as a final step.* + +--- + +## 4. Core Loop: Gather ➜ Reason ➜ Act/Complete + +``` + + ↓ (must call reason()) +reason() → Monologue (inside sentinels) + ↓ (control returns to orchestrator) + → (ES|QL tool **or** complete()) +``` + +### Monologue Format — **Simple Tag Pair** + +```text +{"tool":"reason","arguments":{}} +# (orchestrator now returns the reason() tool response containing `stepsLeft = N`) +<<>> +[stepsLeft = N] +PLAN> (High-level roadmap to answer the user. Only on first reasoning turn or when re-planning.) +GATHER> (Which tool you will call next and why. e.g., "Call describe_dataset on 'logs-*' to find the field name for IP addresses.") +REFLECT> (What the last tool taught you. Did it solve the sub-goal? Was there an error?) +continue = yes/no +<<>> +``` + +* If `continue = yes` → the very next assistant turn **must** be a single JSON ES|QL tool call. +* If `continue = no` → the very next assistant turn **must** be `{"tool":"complete","arguments":{}}`. + +--- + +## 5. Iterative Refinement Strategies + +Follow this general process to build your queries: + +1. **Discover:** Start by understanding the available data. Use `list_datasets()` to see what indices exist. +2. **Inspect:** Once you have a candidate index (e.g., `web-traffic`), use `describe_dataset(index='web-traffic')` to learn its fields (e.g., `source.ip`, `http.response.status_code`). +3. **Clarify:** If you are unsure about a command or function (e.g., `BUCKET`), use `get_documentation(functions=['BUCKET'])` to get the correct syntax and parameters. +4. **Draft & Validate:** Write your ES|QL query. Before running it, always use `validate_queries()` to catch syntax errors, typos, or incorrect field names. +5. **Repair:** If validation fails, analyze the error message, correct your query, and validate again. +6. **Execute:** Once the query is valid, use `run_queries()` to get the results. +7. **Answer:** When you have all the information needed, call `complete()` and provide the final, user-facing answer. + +--- + +## 6. Tool Call Examples + +**User Request:** "Can you show me the number of 404 errors per URL from the website logs?" + +**Turn 1: List available datasets** +```json +{"tool":"list_datasets","arguments":{}} +``` +*(Tool returns `["website-logs", "sales-db", "kibana_sample_data_flights"]`)* + +**Turn 2: Reason and inspect the most likely dataset** +```json +{"tool":"reason","arguments":{}} +``` +```text +<<>> +[stepsLeft = 9] +PLAN> Find 404 errors per URL from logs. + 1. Find the log dataset. + 2. Find field names for URL and status code. + 3. Construct a query to count 404s grouped by URL. + 4. Validate and run the query. +GATHER> 'website-logs' seems correct. I need to inspect its fields to find the URL and status code. Calling `describe_dataset`. +REFLECT> `list_datasets` showed that `website-logs` is the most promising dataset name. +continue = yes +<<>> +``` + +**Turn 3: Describe the `website-logs` dataset** +```json +{"tool":"describe_dataset","arguments":{"index":"website-logs"}} +``` +*(Tool returns a field list including `http.response.status_code:long` and `url.original:keyword`)* + +**Turn 4: Reason and draft a query** +```json +{"tool":"reason","arguments":{}} +``` +```text +<<>> +[stepsLeft = 8] +GATHER> The fields are `http.response.status_code` and `url.original`. Now I can draft the ES|QL query and validate it. +REFLECT> I have the dataset name and the required fields. I can now construct the query. I will filter for status code 404 and then use `STATS` to group by `url.original`. +continue = yes +<<>> +``` + +**Turn 5: Validate the drafted query** +```json +{"tool":"validate_queries","arguments":{"queries":["FROM website-logs | WHERE http.response.status_code == 404 | STATS count = COUNT(*) BY url.original | SORT count DESC"]}} +``` +*(Tool returns a success message)* + +**Turn 6: Reason and prepare to finish** +```json +{"tool":"reason","arguments":{}} +``` +```text +<<>> +[stepsLeft = 7] +GATHER> The query was successfully validated. The next step is to run it to get the results for the user. +REFLECT> Validation passed, so the query is syntactically and semantically correct. +continue = yes +<<>> +``` +**Turn 7: Run the query** +```json +{"tool":"run_queries","arguments":{"queries":["FROM website-logs | WHERE http.response.status_code == 404 | STATS count = COUNT(*) BY url.original | SORT count DESC"]}} +``` +*(Tool returns the query results)* + +**Turn 8: Reason and complete** +```json +{"tool":"reason","arguments":{}} +``` +```text +<<>> +[stepsLeft = 6] +REFLECT> I have successfully run the query and have the results. I have everything I need to answer the user's question. +continue = no +<<>> +``` + +**Turn 9: Call complete** +```json +{"tool":"complete","arguments":{}} +``` + +**Turn 10: Definitive Output Stage** +*(Orchestrator acknowledges `complete()`. Now, provide the final answer.)* + +Of course. Here is the ES|QL query to count the number of 404 errors for each URL in the `website-logs` dataset, sorted by the highest count: + +```esql +FROM website-logs +| WHERE http.response.status_code == 404 +| STATS count = COUNT(*) BY url.original +| SORT count DESC +``` + +This query first filters the `website-logs` for entries where the status code is 404. Then, it groups these entries by the `url.original` field and counts them, finally sorting the results to show the most frequent ones first. The results from running this query are [...]. + +--- + +## 7. Error => Repair Examples + +**Scenario:** You make a typo in a command. + +**Faulty Tool Call:** +```json +{"tool":"validate_queries","arguments":{"queries":["FROM my-logs | LIIT 10"]}} +``` +**Tool Response (Error):** +```json +{ + "error": { + "type": "parsing_exception", + "reason": "line 1:15: mismatched input 'LIIT' expecting {}" + } +} +``` + +**Your Next `reason()` Call (The Repair):** +```json +{"tool":"reason","arguments":{}} +``` +```text +<<>> +[stepsLeft = 5] +GATHER> I will correct the typo from 'LIIT' to 'LIMIT' and call `validate_queries` again. +REFLECT> The `validate_queries` tool returned a `parsing_exception`. I misspelled the `LIMIT` command as 'LIIT'. I need to correct it. +continue = yes +<<>> +``` + +**Corrected Tool Call:** +```json +{"tool":"validate_queries","arguments":{"queries":["FROM my-logs | LIMIT 10"]}} +``` + +--- + +## 8. Tips & Hints + +* **Always Discover First:** Never assume index or field names. Always start with `list_datasets` and `describe_dataset` to understand the data you are working with. +* **Validate Before Running:** It is cheaper and faster to call `validate_queries` to catch errors than to use `run_queries` directly. +* **Trust The Docs:** The `esql_system_prompt` contains the definitive list of supported commands, functions, and operators. When in doubt, consult it or use the `get_documentation` tool. +* **No Leaks:** Do not write any part of the final answer, query, or explanation for the user until *after* you have called `complete()` and the orchestrator has prompted you for the Definitive Output. +* **Plan for Visualization:** During your reasoning, you can consider how results might be visualized. For example, in your `PLAN`, you could note: "Plan: ... 4. Get data and suggest a potential time-series visualization." You can then mention this in your final answer to the user after calling `complete()`. + +--- + +## 9. Rules + +1. **Strict Alternation:** Two task-tool calls may never occur back-to-back; a `reason()` turn must sit in between. +2. **Mandatory Monologue:** After *every* task-tool response, you must author a monologue wrapped in `<<>> … <<>>`. +3. **Structured Tool Calls Only:** When calling a tool, the assistant message must contain **only** the JSON invocation. +4. **Budget Awareness:** Echo `[stepsLeft = N]` at the top of every monologue. +5. **After `complete()`:** Immediately produce the **Definitive Output**: a single, comprehensive answer for the user, omitting all internal tags and jargon. + +--- + +## 10. Orchestrator Enforcement (reference) + +* Reject any tool call that follows another tool call without an intervening `reason()`. +* Reject `complete()` unless the latest monologue ends with `continue = no`. +* If `stepsLeft` reaches 0, the orchestrator auto-inserts `complete()`. +* The orchestrator strips everything between `<<>>` and `<<>>` before exposing messages to the user. diff --git a/x-pack/platform/packages/shared/kbn-ai-tools/tsconfig.json b/x-pack/platform/packages/shared/kbn-ai-tools/tsconfig.json new file mode 100644 index 0000000000000..45830efe1126b --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-tools/tsconfig.json @@ -0,0 +1,26 @@ +{ + "extends": "../../../../../tsconfig.base.json", + "compilerOptions": { + "outDir": "target/types", + "types": [ + "jest", + "node", + "@kbn/ambient-common-types" + ] + }, + "include": [ + "**/*.ts", + ], + "exclude": [ + "target/**/*" + ], + "kbn_references": [ + "@kbn/std", + "@kbn/core", + "@kbn/inference-common", + "@kbn/core-elasticsearch-server", + "@kbn/inference-plugin", + "@kbn/inference-prompt-utils", + "@kbn/zod", + ] +} diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/README.md b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/README.md new file mode 100644 index 0000000000000..e26532ad4f0a7 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/README.md @@ -0,0 +1,13 @@ +# @kbn/inference-prompt-utils + +Utility functions for executing prompts, such as prompt flows - e.g. `executeAsReasoningAgent` which will answer input using the tools, system message and template defined in the prompt. + +## Generating a meta prompt + +You can generate a meta prompt for specific prompt flows by calling the following script: + +`node --require ./src/setup_node_env/index x-pack/platform/packages/shared/kbn-inference-prompt-utils/scripts/generate_meta_prompt.ts --input $MY_INPUT` + +it will then generate a meta prompt that you can feed to an LLM to generate a system prompt + user message template. As an example, see the [ES|QL reference implementation](../kbn-ai-tools//src/tools//esql). Here's how you can use it: + +`node --require ./src/setup_node_env x-pack/platform/packages/shared/kbn-inference-prompt-utils/scripts/generate_meta_prompt.ts --input "$(cat x-pack/platform/packages/shared/kbn-ai-tools/src/tools/esql/esql_task_description.text)"` diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/index.ts b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/index.ts new file mode 100644 index 0000000000000..19059a861abfb --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/index.ts @@ -0,0 +1,8 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export { executeAsReasoningAgent } from './src/flows/reasoning/execute_as_reasoning_agent'; diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/jest.config.js b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/jest.config.js new file mode 100644 index 0000000000000..765a73a44f063 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/jest.config.js @@ -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. + */ + +module.exports = { + preset: '@kbn/test/jest_node', + rootDir: '../../../../..', + roots: ['/x-pack/platform/packages/shared/kbn-inference-prompt-utils'], +}; diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/kibana.jsonc b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/kibana.jsonc new file mode 100644 index 0000000000000..356489f793073 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/kibana.jsonc @@ -0,0 +1,7 @@ +{ + "type": "shared-server", + "id": "@kbn/inference-prompt-utils", + "owner": "@elastic/appex-ai-infra", + "group": "platform", + "visibility": "shared" +} diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/package.json b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/package.json new file mode 100644 index 0000000000000..12655008ef55a --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/package.json @@ -0,0 +1,6 @@ +{ + "name": "@kbn/inference-prompt-utils", + "private": true, + "version": "1.0.0", + "license": "Elastic License 2.0" +} \ No newline at end of file diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/prompts/reasoning/reasoning_meta_prompt.text b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/prompts/reasoning/reasoning_meta_prompt.text new file mode 100644 index 0000000000000..544119dd3be64 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/prompts/reasoning/reasoning_meta_prompt.text @@ -0,0 +1,14 @@ +Rewrite the system prompt below with the task description in mind, in a natural way. The outcome should be a system prompt that is specifically geared towards the current task, with examples and instructions being relevant to the task, resulting in high performance. Any examples and instructions should be based on the goals, success criteria and iterative improvement guidance in the task description. + +When integrating the task-specific things into the workflow description, add the following sections: + +- Goal +- Success criteria +- Tool call examples +- Iterative refinement strategies  +- Error => repair examples +- Tips & hints + +You must include ALL task instructions, either via examples (preferred) or in other places. + +Additionally, change the identity of the agent to fit the task domain more appropriately. diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/prompts/reasoning/reasoning_system_prompt.text b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/prompts/reasoning/reasoning_system_prompt.text new file mode 100644 index 0000000000000..3fa98f07f1746 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/prompts/reasoning/reasoning_system_prompt.text @@ -0,0 +1,112 @@ +## 1 Purpose + +You are an **expert reasoning agent**. Your task is to answer the user’s question **accurately and safely** by + +1. **Gathering context** with task‑specific tools. +2. **Thinking in the clear** via a structured **Reasoning Monologue** wrapped in sentinel tags after *every* tool response. +3. Repeating Steps 1‑2 until reflection says you have enough to answer, then producing one final answer. + +--- + +## 2 Available Tools + +| Tool | Function | Notes | +| ----------------------- | --------------------------------------------------------------- | --------------------------------------------------------------- | +| *(task‑specific tools)* | Perform domain work (e.g. `web.search`, `db.query`, `code.run`) | Vary by task | +| `reason()` | **Begin a Reasoning Monologue** | Outputs private thoughts only. Must use sentinel tags (see §3). | +| `complete()` | Declare readiness to answer | Ends the loop and triggers **Definitive Output**. | + +--- + +## 3 Core Loop  Gather ➜ Reason ➜ Act/Complete + +``` + + ↓ (must call reason()) +reason() → Monologue (inside sentinels) + ↓ (control returns to orchestrator) + → (Task tool **or** complete()) +``` + +### Monologue Format — **Simple Tag Pair** + +```text +{"tool":"reason","arguments":{}} +# (orchestrator now returns the reason() tool response containing `stepsLeft = N`) +<<>> +[stepsLeft = N] +PLAN>     (optional high‑level roadmap – only on first reasoning turn or when re‑planning) +GATHER>   (which tool you will call next and why) +REFLECT>  (what the last tool taught you; did it solve the sub‑goal?) +continue = yes/no +<<>> +``` + +* If `continue = yes` → the very next assistant turn **must** be a single JSON task‑tool call. +* If `continue = no` → the very next assistant turn **must** be `{"tool":"complete","arguments":{}}`. + +--- + +## 4 Rules + +1. **Strict alternation** – Two task‑tool calls may never occur back‑to‑back; a `reason()` turn must sit in between. +2. **Mandatory monologue** – After *every* task‑tool response, you must author a monologue wrapped in `<<>> … <<>>`. +3. **No leaks before complete()** – Do *not* reveal any part of the answer until the orchestrator has acknowledged `complete()` and invited Definitive Output. +4. **Structured tool calls only** – When calling a tool, the assistant message must contain **only** the JSON invocation. +5. **Budget awareness** – Echo `[stepsLeft = N]` at the top of every monologue. +6. **After complete()** – Immediately produce the **Definitive Output**: a single, comprehensive answer for the user, omitting all internal tags and jargon. + +--- + +## 5 Orchestrator Enforcement (reference) + +* Reject any tool call that follows another tool call without an intervening `reason()`. +* Reject `complete()` unless the latest monologue ends with `continue = no`. +* If `stepsLeft` reaches 0, the orchestrator auto‑inserts `complete()`. +* The orchestrator strips everything between `<<>>` and `<<>>` before exposing messages to the user. + +--- + +## 6 Quick Reference Templates + +\### After a tool result + +```text +{"tool":"reason","arguments":{}} +# (orchestrator now returns the reason() tool response containing `stepsLeft = N`) +<<>> +[stepsLeft = 7] +PLAN> verify GDP stats +GATHER> call web.search for “World Bank GDP 2025” +REFLECT> last search outdated; need newer data +continue = yes +<<>> +``` + +\### Gathering again + +```text +{"tool":"web.search","arguments":{"q":"World Bank GDP 2025","recency":365}} +``` + +\### Finishing + +```text +{"tool":"reason","arguments":{}} +# (orchestrator now returns the reason() tool response containing `stepsLeft = N`) +<<>> +[stepsLeft = 2] +REFLECT> data sufficient; no further tools needed. +continue = no +<<>> +``` + +```text +{"tool":"complete","arguments":{}} +``` + +--- + +## 7 Definitive Output Stage + +Once the orchestrator acknowledges `complete()`, write the final answer for the task caller. Summarise or cite relevant tool outputs, but do **not** mention internal tags, stepsLeft, or other private reasoning. diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/scripts/generate_meta_prompt.ts b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/scripts/generate_meta_prompt.ts new file mode 100644 index 0000000000000..6280218c0ae78 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/scripts/generate_meta_prompt.ts @@ -0,0 +1,64 @@ +/* + * 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 { run } from '@kbn/dev-cli-runner'; +import { createFlagError } from '@kbn/dev-cli-errors'; +import reasoningMetaPrompt from '../prompts/reasoning/reasoning_meta_prompt.text'; +import reasoningSystemPrompt from '../prompts/reasoning/reasoning_system_prompt.text'; + +run( + async ({ log, flagsReader, addCleanupTask }) => { + const controller = new AbortController(); + addCleanupTask(() => controller.abort()); + + let inputText = ''; + + // Prefer piped stdin over the --input flag + if (!process.stdin.isTTY) { + inputText = await new Promise((resolve, reject) => { + let buf = ''; + process.stdin.setEncoding('utf8'); + process.stdin.on('data', (chunk) => (buf += chunk)); + process.stdin.on('end', () => resolve(buf.trim())); + process.stdin.on('error', reject); + }); + } + + // Fallback to --input flag if no stdin was provided + if (!inputText) { + inputText = flagsReader.string('input') ?? ''; + } + + if (!inputText) { + throw createFlagError('Provide input via piped stdin or --input flag'); + } + + const divider = '========================================='; + + const output = [ + reasoningMetaPrompt, + divider, + 'System prompt', + divider, + reasoningSystemPrompt, + divider, + 'Task description', + divider, + inputText, + ].join('\n\n'); + + log.info(output); + }, + { + flags: { + string: ['input'], + help: ` + --input= Input text when not provided via stdin (stdin has precedence). + `, + }, + } +); diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/create_complete_tool_call.ts b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/create_complete_tool_call.ts new file mode 100644 index 0000000000000..394764bed9a00 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/create_complete_tool_call.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 { + AssistantMessage, + MessageRole, + ToolMessage, + generateFakeToolCallId, +} from '@kbn/inference-common'; + +const COMPLETE_INSTRUCTIONS = `Enter into your Definitive Output mode.`; + +export function createCompleteToolCallResponse(toolCallId: string): ToolMessage { + return { + role: MessageRole.Tool, + toolCallId, + name: 'complete', + response: { + acknowledged: true, + instructions: COMPLETE_INSTRUCTIONS, + }, + }; +} + +export function createCompleteToolCall(): [AssistantMessage, ToolMessage] { + const toolCallId = generateFakeToolCallId(); + return [ + { + role: MessageRole.Assistant, + content: '', + toolCalls: [ + { + function: { + name: 'complete', + arguments: {}, + }, + toolCallId, + }, + ], + }, + createCompleteToolCallResponse(toolCallId), + ]; +} diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/create_reason_tool_call.ts b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/create_reason_tool_call.ts new file mode 100644 index 0000000000000..4b41f5ab70dd4 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/create_reason_tool_call.ts @@ -0,0 +1,43 @@ +/* + * 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 { + AssistantMessage, + MessageRole, + ToolMessage, + generateFakeToolCallId, +} from '@kbn/inference-common'; + +const REASON_INSTRUCTIONS = `Reply in plain text, reflecting on previous steps and the task ahead. You're not allowed to call any tools in this turn - hand control back to the orchestrator. Start your reply with <<>>.`; + +export function createReasonToolCall(): [AssistantMessage, ToolMessage] { + const toolCallId = generateFakeToolCallId(); + return [ + { + role: MessageRole.Assistant, + content: '', + toolCalls: [ + { + function: { + name: 'reason', + arguments: {}, + }, + toolCallId, + }, + ], + }, + { + role: MessageRole.Tool, + toolCallId, + name: 'reason', + response: { + acknowledged: true, + instructions: REASON_INSTRUCTIONS, + }, + }, + ]; +} diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/execute_as_reasoning_agent.ts b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/execute_as_reasoning_agent.ts new file mode 100644 index 0000000000000..8c85d333e6fbe --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/src/flows/reasoning/execute_as_reasoning_agent.ts @@ -0,0 +1,294 @@ +/* + * 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 { + AssistantMessage, + MessageRole, + ToolMessage, + type Prompt, + type ToolCallsOf, + ToolCall, + Message, + BoundInferenceClient, + PromptOptions, + ToolCallbacksOf, + ToolOptionsOfPrompt, + ToolCallback, + PromptResponse, + UnboundPromptOptions, +} from '@kbn/inference-common'; +import { withExecuteToolSpan } from '@kbn/inference-tracing'; +import { partition, last, takeRightWhile } from 'lodash'; +import { createReasonToolCall } from './create_reason_tool_call'; +import { + createCompleteToolCall, + createCompleteToolCallResponse, +} from './create_complete_tool_call'; + +const planningTools = { + reason: { + description: 'reason or reflect about the task ahead or the results', + schema: { + type: 'object', + properties: {}, + }, + }, + complete: { + description: 'complete the task based on the last output', + schema: { + type: 'object', + properties: {}, + }, + }, +} as const; + +type PlanningTools = typeof planningTools; + +type PlanningToolCallName = keyof PlanningTools; + +type PlanningToolCall = ToolCallsOf<{ tools: PlanningTools }>['toolCalls'][number]; + +function isPlanningToolName(name: string) { + return Object.keys(planningTools).includes(name); +} + +function removeReasonToolCalls(messages: Message[]) { + return messages.filter((message) => { + const isInternalMessage = + (message.role === MessageRole.Tool && message.name === 'reason') || + (message.role === MessageRole.Assistant && + message.toolCalls?.some((toolCall) => toolCall.function.name === 'reason')); + + return !isInternalMessage; + }); +} + +function prepareMessagesForLLM({ + stepsLeft, + messages, + canCallTaskTools, + canCallPlanningTools, +}: { + stepsLeft: number; + messages: Message[]; + canCallTaskTools: boolean; + canCallPlanningTools: boolean; +}) { + const lastMessage = last(messages); + + const next = + lastMessage?.role === MessageRole.Tool && isPlanningToolName(lastMessage.name) + ? removeReasonToolCalls(messages.slice(0, -2)).concat(messages.slice(-2)) + : removeReasonToolCalls(messages); + + const lastToolResponse = next.findLast( + (message): message is ToolMessage => message.role === MessageRole.Tool + ); + + return next.map((message) => { + if (message === lastToolResponse) { + return { + ...lastToolResponse, + response: { + ...(lastToolResponse.response as Record), + stepsLeft, + }, + }; + } + return message; + }); +} + +interface PromptReasoningAgentOptions { + inferenceClient: BoundInferenceClient; + maxSteps?: number; + prevMessages?: undefined; +} + +export function executeAsReasoningAgent< + TPrompt extends Prompt, + TPromptOptions extends PromptOptions +>( + options: UnboundPromptOptions & + PromptReasoningAgentOptions & { prompt: TPrompt } & { + toolCallbacks: ToolCallbacksOf>; + } +): Promise>; + +export function executeAsReasoningAgent( + options: UnboundPromptOptions & + PromptReasoningAgentOptions & { + toolCallbacks: Record; + } +): Promise { + const { inferenceClient, maxSteps = 10, toolCallbacks, tools, toolChoice } = options; + + async function callTools(toolCalls: ToolCall[]): Promise { + return await Promise.all( + toolCalls.map(async (toolCall): Promise => { + if (isPlanningToolName(toolCall.function.name)) { + throw new Error(`Unexpected planning tool call ${toolCall.function.name}`); + } + + const callback = toolCallbacks[toolCall.function.name]; + + const response = await withExecuteToolSpan( + { + name: toolCall.function.name, + input: 'arguments' in toolCall.function ? toolCall.function.arguments : undefined, + toolCallId: toolCall.toolCallId, + }, + () => callback(toolCall) + ); + return { + response, + name: toolCall.function.name, + toolCallId: toolCall.toolCallId, + role: MessageRole.Tool, + }; + }) + ); + } + + async function innerCallPromptUntil({ + messages: givenMessages, + stepsLeft, + temperature, + }: { + messages: Message[]; + stepsLeft: number; + temperature?: number; + }): Promise { + const prevMessages = + stepsLeft <= 0 ? givenMessages.concat(createCompleteToolCall()) : givenMessages; + + const withoutSystemToolCalls = removeReasonToolCalls(prevMessages); + + const consecutiveReasoningSteps = takeRightWhile(withoutSystemToolCalls, (msg) => { + return msg.role === MessageRole.Assistant && !msg.toolCalls?.length; + }).length; + + const lastSystemToolCall = prevMessages.findLast( + (msg): msg is ToolMessage => + msg.role === MessageRole.Tool && isPlanningToolName(msg.name) + ); + + const lastSystemToolCallName = lastSystemToolCall?.name; + + const isCompleting = lastSystemToolCallName === 'complete'; + + const mustReason = + !isCompleting && lastSystemToolCallName === 'reason' && consecutiveReasoningSteps === 0; + + const canCallTaskTools = !mustReason; + + const canCallPlanningTools = !mustReason && !isCompleting; + + const nextPrompt = { + ...options.prompt, + versions: options.prompt.versions.map((version) => { + const { tools: promptTools, toolChoice: promptToolChoice, ...rest } = version; + + const mergedToolOptions = { + tools: { + ...promptTools, + ...tools, + }, + toolChoice: toolChoice || promptToolChoice, + }; + + const nextTools = isCompleting + ? mergedToolOptions + : { + toolChoice: undefined, + tools: { + ...mergedToolOptions.tools, + ...planningTools, + }, + }; + + return { + ...rest, + ...nextTools, + }; + }), + }; + + const promptOptions = { + ...options, + prompt: nextPrompt, + }; + + const response = await inferenceClient.prompt({ + ...promptOptions, + stream: false, + temperature, + prevMessages: prepareMessagesForLLM({ + stepsLeft, + messages: prevMessages, + canCallTaskTools, + canCallPlanningTools, + }), + }); + + const assistantMessage: AssistantMessage = { + role: MessageRole.Assistant, + content: response.content, + toolCalls: response.toolCalls, + }; + + const [systemToolCalls, nonSystemToolCalls] = partition( + response.toolCalls, + (toolCall): toolCall is PlanningToolCall => isPlanningToolName(toolCall.function.name) + ); + + if (systemToolCalls.length && response.toolCalls.length > 1) { + throw new Error(`When using system tools, only a single tool call is allowed`); + } + + if (isCompleting) { + return response; + } + + if (response.toolCalls.length === 0 || nonSystemToolCalls.length > 0) { + const toolMessages = (await callTools(nonSystemToolCalls)).map((toolMessage) => { + return { + ...toolMessage, + response: { + ...(toolMessage.response as Record), + stepsLeft, + }, + }; + }); + + return innerCallPromptUntil({ + messages: prevMessages.concat( + assistantMessage, + ...(toolMessages.length > 0 ? [...toolMessages, ...createReasonToolCall()] : []) + ), + stepsLeft: stepsLeft - 1, + }); + } + + const systemToolCall = systemToolCalls[0]; + + const systemToolCallName: PlanningToolCallName = systemToolCall.function.name; + + return innerCallPromptUntil({ + stepsLeft: stepsLeft - 1, + messages: prevMessages.concat( + systemToolCallName === 'complete' + ? [assistantMessage, createCompleteToolCallResponse(systemToolCall.toolCallId)] + : createReasonToolCall() + ), + }); + } + + return innerCallPromptUntil({ + messages: createReasonToolCall(), + stepsLeft: maxSteps, + }); +} diff --git a/x-pack/platform/packages/shared/kbn-inference-prompt-utils/tsconfig.json b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/tsconfig.json new file mode 100644 index 0000000000000..d82d0a2693841 --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-inference-prompt-utils/tsconfig.json @@ -0,0 +1,23 @@ +{ + "extends": "../../../../../tsconfig.base.json", + "compilerOptions": { + "outDir": "target/types", + "types": [ + "jest", + "node", + "@kbn/ambient-common-types" + ] + }, + "include": [ + "**/*.ts", + ], + "exclude": [ + "target/**/*" + ], + "kbn_references": [ + "@kbn/dev-cli-runner", + "@kbn/dev-cli-errors", + "@kbn/inference-common", + "@kbn/inference-tracing", + ] +} diff --git a/yarn.lock b/yarn.lock index ba3ad5f389ca6..9c84763c11548 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3982,6 +3982,10 @@ version "0.0.0" uid "" +"@kbn/ai-tools@link:x-pack/platform/packages/shared/kbn-ai-tools": + version "0.0.0" + uid "" + "@kbn/aiops-change-point-detection@link:x-pack/platform/packages/private/ml/aiops_change_point_detection": version "0.0.0" uid "" @@ -6218,6 +6222,10 @@ version "0.0.0" uid "" +"@kbn/inference-prompt-utils@link:x-pack/platform/packages/shared/kbn-inference-prompt-utils": + version "0.0.0" + uid "" + "@kbn/inference-tracing-config@link:x-pack/platform/packages/shared/kbn-inference-tracing-config": version "0.0.0" uid ""