[Inference] Prompts: reasoning mode + ES|QL prompt

.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

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",

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"],

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,

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,

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/tools.ts

@@ -28,6 +28,8 @@ type ToolCallbacksOfTools<TTools extends Record<string, ToolDefinition> | undefi
       }
     : never;
 
+export type ToolCallback = (toolCall: ToolCall) => Promise<ToolMessage['response']>;
+
 export type ToolCallbacksOf<TToolOptions extends ToolOptions> = TToolOptions extends {
   tools?: Record<string, ToolDefinition>;
 }

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 }));
+  }
+);

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"
   ]
 }

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.

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';

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: ['<rootDir>/x-pack/platform/packages/shared/kbn-ai-tools'],
+};

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"
+}

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"
+}

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<string | number | boolean>;
+    empty: boolean;
+  }>;
+}
+
+export interface TruncatedDocumentAnalysis {
+  fields: string[];
+  total: number;
+  sampled: number;
+}

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<QueryDslFieldAndFormat | string>;
+  _source?: boolean;
+  timeout?: string;
+}) {
+  return esClient
+    .search<Record<string, any>>({
+      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<SearchHit<Record<string, any>>>,
+      total:
+        typeof response.hits.total === 'number'
+          ? response.hits.total!
+          : response.hits.total?.value!,
+    }));
+}

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;
+}