Migrate mergeSampleDocumentsWithFieldCaps to mergeSampleDocumentsWithSchema, replacing field caps logic with ES|QL schema retrieval, updating related helpers, tests, and integrations.

Author: achyutjhunjhunwala

x-pack/platform/packages/shared/kbn-ai-tools/index.ts

@@ -9,7 +9,7 @@ export { describeDataset } from './src/tools/describe_dataset';
 
 export { formatDocumentAnalysis } from './src/tools/describe_dataset/format_document_analysis';
 
-export { mergeSampleDocumentsWithFieldCaps } from './src/tools/describe_dataset/merge_sample_documents_with_field_caps';
+export { mergeSampleDocumentsWithSchema } from './src/tools/describe_dataset/merge_sample_documents_with_schema';
 export {
   getSampleDocuments,
   getSampleDocumentsEsql,
@@ -30,5 +30,10 @@ export {
   P_VALUE_SIGNIFICANCE_HIGH,
   P_VALUE_SIGNIFICANCE_MEDIUM,
 } from './src/utils/p_value_to_label';
+export {
+  getEsqlColumnSchema,
+  type EsqlColumnSchema,
+  type GetEsqlColumnSchemaParams,
+} from './src/utils/get_esql_column_schema';
 
 export { executeAsEsqlAgent } from './src/tools/esql';

x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/get_diverse_sample_documents.test.ts

@@ -17,16 +17,10 @@ jest.mock('./get_sample_documents', () => ({
 const getSampleDocumentsEsqlMock = jest.mocked(getSampleDocumentsEsql);
 
 const createEsClient = () => {
-  const fieldCaps = jest.fn().mockResolvedValue({
-    fields: {
-      message: {},
-    },
-  });
   const query = jest.fn();
 
   return {
-    esClient: { fieldCaps, esql: { query } } as unknown as ElasticsearchClient,
-    fieldCaps,
+    esClient: { esql: { query } } as unknown as ElasticsearchClient,
     query,
   };
 };
@@ -40,6 +34,13 @@ const countResponse = (total: number) => ({
   values: [[total]],
 });
 
+const schemaResponse = (
+  columns: Array<{ name: string; type: string }> = [{ name: 'message', type: 'text' }]
+) => ({
+  columns,
+  values: [],
+});
+
 const pass1Response = (
   values: unknown[][] = [
     ['logs-a:doc-1', 10, 'error'],
@@ -76,6 +77,7 @@ describe('getDiverseSampleDocuments', () => {
   it('uses ES|QL count, categorize pass, and composite-key source fetch', async () => {
     const { esClient, query } = createEsClient();
     query
+      .mockResolvedValueOnce(schemaResponse())
       .mockResolvedValueOnce(countResponse(10))
       .mockResolvedValueOnce(pass1Response())
       .mockResolvedValueOnce(pass2Response());
@@ -90,11 +92,12 @@ describe('getDiverseSampleDocuments', () => {
       logger,
     });
 
-    expect(query.mock.calls[0][0].query).toBe('FROM logs-a, logs-b | STATS total = COUNT(*)');
-    expect(query.mock.calls[1][0].query).toBe(
+    expect(query.mock.calls[0][0].query).toBe('FROM logs-a, logs-b | LIMIT 0');
+    expect(query.mock.calls[1][0].query).toBe('FROM logs-a, logs-b | STATS total = COUNT(*)');
+    expect(query.mock.calls[2][0].query).toBe(
       'FROM logs-a, logs-b METADATA _index, _id | EVAL doc_key = CONCAT(_index, ":", _id) | STATS representative_key = TOP(doc_key, 1, "desc"), count = COUNT(*) BY pattern = CATEGORIZE(message) | SORT count DESC | LIMIT 3'
     );
-    expect(query.mock.calls[2][0].query).toBe(
+    expect(query.mock.calls[3][0].query).toBe(
       'FROM logs-a, logs-b METADATA _index, _id, _source | EVAL doc_key = CONCAT(_index, ":", _id) | WHERE doc_key IN ("logs-b:doc-2") | KEEP _index, _id, _source | LIMIT 1'
     );
     expect(result.hits).toEqual([
@@ -105,6 +108,7 @@ describe('getDiverseSampleDocuments', () => {
   it('adds SAMPLE when the population is large', async () => {
     const { esClient, query } = createEsClient();
     query
+      .mockResolvedValueOnce(schemaResponse())
       .mockResolvedValueOnce(countResponse(10_000_000))
       .mockResolvedValueOnce(pass1Response([['logs-a:doc-1', 10, 'error']]))
       .mockResolvedValueOnce(pass2Response([['logs-a', 'doc-1', { message: 'error one' }]]));
@@ -119,12 +123,12 @@ describe('getDiverseSampleDocuments', () => {
       logger,
     });
 
-    expect(query.mock.calls[1][0].query).toContain('| SAMPLE 0.01 |');
+    expect(query.mock.calls[2][0].query).toContain('| SAMPLE 0.01 |');
   });
 
   it('short-circuits when the count query returns zero', async () => {
     const { esClient, query } = createEsClient();
-    query.mockResolvedValueOnce(countResponse(0));
+    query.mockResolvedValueOnce(schemaResponse()).mockResolvedValueOnce(countResponse(0));
 
     const result = await getDiverseSampleDocuments({
       esClient,
@@ -136,14 +140,15 @@ describe('getDiverseSampleDocuments', () => {
       logger,
     });
 
-    expect(query).toHaveBeenCalledTimes(1);
+    expect(query).toHaveBeenCalledTimes(2);
     expect(result).toEqual({ hits: [] });
   });
 
   it('falls back to random ES|QL sampling when no message field exists', async () => {
-    const { esClient, fieldCaps, query } = createEsClient();
-    fieldCaps.mockResolvedValueOnce({ fields: {} });
-    query.mockResolvedValueOnce(countResponse(10));
+    const { esClient, query } = createEsClient();
+    query
+      .mockResolvedValueOnce(schemaResponse([{ name: 'host.name', type: 'keyword' }]))
+      .mockResolvedValueOnce(countResponse(10));
     getSampleDocumentsEsqlMock.mockResolvedValueOnce({
       hits: [{ _index: 'logs-a', _id: 'doc-1', _source: { event: 'one' } }],
       total: 1,
@@ -166,13 +171,42 @@ describe('getDiverseSampleDocuments', () => {
       end: 200,
       sampleSize: 1,
     });
-    expect(query).toHaveBeenCalledTimes(1);
+    expect(query).toHaveBeenCalledTimes(2);
     expect(result.hits).toEqual([{ _index: 'logs-a', _id: 'doc-1', _source: { event: 'one' } }]);
   });
 
+  it('uses body.text when it is the first available text field candidate', async () => {
+    const { esClient, query } = createEsClient();
+    query
+      .mockResolvedValueOnce(
+        schemaResponse([
+          { name: 'message', type: 'keyword' },
+          { name: 'body.text', type: 'text' },
+        ])
+      )
+      .mockResolvedValueOnce(countResponse(10))
+      .mockResolvedValueOnce(pass1Response([['logs-a:doc-1', 10, 'body pattern']]))
+      .mockResolvedValueOnce(
+        pass2Response([['logs-a', 'doc-1', { body: { text: 'body value' } }]])
+      );
+
+    await getDiverseSampleDocuments({
+      esClient,
+      index: 'logs-*',
+      start: 100,
+      end: 200,
+      size: 1,
+      offset: 0,
+      logger,
+    });
+
+    expect(query.mock.calls[2][0].query).toContain('CATEGORIZE(body.text)');
+  });
+
   it('drops patterns whose composite key is missing from pass 2', async () => {
     const { esClient, query } = createEsClient();
     query
+      .mockResolvedValueOnce(schemaResponse())
       .mockResolvedValueOnce(countResponse(10))
       .mockResolvedValueOnce(pass1Response())
       .mockResolvedValueOnce(pass2Response([['logs-a', 'doc-1', { message: 'error one' }]]));

x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/get_diverse_sample_documents.ts

@@ -11,6 +11,7 @@ import type { ElasticsearchClient } from '@kbn/core/server';
 import type { ESQLSearchResponse } from '@kbn/es-types';
 import { dateRangeQuery } from '@kbn/es-query';
 import type { Logger } from '@kbn/logging';
+import { getEsqlColumnSchema } from '../../utils/get_esql_column_schema';
 import { getSampleDocumentsEsql } from './get_sample_documents';
 
 const MESSAGE_FIELD_CANDIDATES = ['message', 'body.text'];
@@ -39,9 +40,8 @@ export async function getDiverseSampleDocuments({
   const filter = { bool: { filter: timeRangeFilter } };
   const indices = Array.isArray(index) ? index : [index];
 
-  // TODO: migrate this fieldCaps probe to ES|QL in https://github.com/elastic/streams-program/issues/1220.
   const [messageField, totalDocs] = await Promise.all([
-    detectMessageField({ esClient, index, timeRangeFilter }),
+    detectMessageField({ esClient, index, start, end }),
     runEsqlCount({ esClient, indices, filter }),
   ]);
 
@@ -119,27 +119,21 @@ export async function getDiverseSampleDocuments({
 async function detectMessageField({
   esClient,
   index,
-  timeRangeFilter,
+  start,
+  end,
 }: {
   esClient: ElasticsearchClient;
   index: string | string[];
-  timeRangeFilter: ReturnType<typeof dateRangeQuery>;
+  start: number;
+  end: number;
 }): Promise<string | undefined> {
-  const fieldCapsResponse = await esClient.fieldCaps({
-    index,
-    fields: MESSAGE_FIELD_CANDIDATES,
-    index_filter: {
-      bool: {
-        filter: timeRangeFilter,
-      },
-    },
-    types: ['text', 'match_only_text'],
-  });
-
-  const fieldsFound = Object.keys(fieldCapsResponse.fields);
+  const columns = await getEsqlColumnSchema({ esClient, index, start, end });
+  const textColumnNames = new Set(
+    columns.filter((column) => column.type === 'text').map((column) => column.name)
+  );
 
   for (const candidate of MESSAGE_FIELD_CANDIDATES) {
-    if (fieldsFound.includes(candidate)) {
+    if (textColumnNames.has(candidate)) {
       return candidate;
     }
   }

x-pack/platform/packages/shared/kbn-ai-tools/src/tools/describe_dataset/index.test.ts

@@ -0,0 +1,140 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ElasticsearchClient } from '@kbn/core/server';
+import { describeDataset } from '.';
+
+const createEsClient = () => {
+  const query = jest.fn();
+  return {
+    esClient: { esql: { query } } as unknown as ElasticsearchClient,
+    query,
+  };
+};
+
+const schemaResponse = ({
+  columns = [
+    { name: '@timestamp', type: 'date' },
+    { name: 'message', type: 'text' },
+    { name: 'host.name', type: 'keyword' },
+  ],
+}: {
+  columns?: Array<{ name: string; type: string; original_types?: string[] }>;
+} = {}) => ({
+  columns,
+  values: [],
+});
+
+const hitsResponse = ({
+  values = [
+    ['doc-1', { message: 'hello', host: { name: 'host-a' } }],
+    ['doc-2', { message: 'hello', host: { name: 'host-b' } }],
+  ],
+}: {
+  values?: unknown[][];
+} = {}) => ({
+  columns: [
+    { name: '_id', type: 'keyword' },
+    { name: '_source', type: 'object' },
+  ],
+  values,
+});
+
+const countResponse = (total: number) => ({
+  columns: [{ name: 'total', type: 'long' }],
+  values: [[total]],
+});
+
+describe('describeDataset', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('merges ES|QL schema columns with sampled documents', async () => {
+    const { esClient, query } = createEsClient();
+    query
+      .mockResolvedValueOnce(schemaResponse())
+      .mockResolvedValueOnce(hitsResponse())
+      .mockResolvedValueOnce(countResponse(2));
+
+    const analysis = await describeDataset({
+      esClient,
+      index: 'logs-*',
+      start: 100,
+      end: 200,
+    });
+
+    expect(query.mock.calls[0][0].query).toBe('FROM logs-* | LIMIT 0');
+    expect(query.mock.calls[1][0].query).toBe('FROM logs-* METADATA _id, _source | LIMIT 1000');
+    expect(query.mock.calls[2][0].query).toBe('FROM logs-* | STATS total = COUNT(*)');
+    expect(analysis.total).toBe(2);
+    expect(analysis.sampled).toBe(2);
+    expect(analysis.fields).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          name: 'message',
+          types: ['text'],
+          cardinality: 1,
+          documentsWithValue: 2,
+          values: [{ value: 'hello', count: 2 }],
+        }),
+        expect.objectContaining({
+          name: 'host.name',
+          types: ['keyword'],
+          cardinality: 2,
+          documentsWithValue: 2,
+        }),
+      ])
+    );
+  });
+
+  it('uses the population count instead of sampled hit count for total', async () => {
+    const { esClient, query } = createEsClient();
+    query
+      .mockResolvedValueOnce(schemaResponse())
+      .mockResolvedValueOnce(hitsResponse({ values: [['doc-1', { message: 'hello' }]] }))
+      .mockResolvedValueOnce(countResponse(15_000));
+
+    const analysis = await describeDataset({
+      esClient,
+      index: 'logs-*',
+      start: 100,
+      end: 200,
+    });
+
+    expect(analysis.sampled).toBe(1);
+    expect(analysis.total).toBe(15_000);
+  });
+
+  it('uses original types for unsupported cross-index columns', async () => {
+    const { esClient, query } = createEsClient();
+    query
+      .mockResolvedValueOnce(
+        schemaResponse({
+          columns: [{ name: 'host.name', type: 'unsupported', original_types: ['ip', 'keyword'] }],
+        })
+      )
+      .mockResolvedValueOnce(hitsResponse({ values: [['doc-1', { host: { name: '127.0.0.1' } }]] }))
+      .mockResolvedValueOnce(countResponse(1));
+
+    const analysis = await describeDataset({
+      esClient,
+      index: ['logs-a', 'logs-b'],
+      start: 100,
+      end: 200,
+    });
+
+    expect(analysis.fields).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          name: 'host.name',
+          types: ['ip', 'keyword'],
+        }),
+      ])
+    );
+  });
+});