[agent builder] ES|QL tools: interpolate query

Author: pgayvallet

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql/index.ts

@@ -7,3 +7,4 @@
 
 export { executeEsql, type EsqlResponse } from './execute_esql';
 export { extractEsqlQueries, esqlResponseToJson } from './misc';
+export { interpolateEsqlQuery } from './interpolate_query';

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql/interpolate_query.test.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 { interpolateEsqlQuery } from './interpolate_query';
+
+describe('interpolateEsqlQuery', () => {
+  it('should correctly interpolate string and number parameters', () => {
+    const template = 'FROM support_ticket | WHERE priority == ?priority AND assigne_id == ?user_id';
+    const params = { priority: 'high', user_id: 45 };
+    const expected = 'FROM support_ticket | WHERE priority == "high" AND assigne_id == 45';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should correctly interpolate boolean parameters', () => {
+    const template = 'FROM nodes | WHERE is_active == ?active AND is_primary == ?primary';
+    const params = { active: true, primary: false };
+    const expected = 'FROM nodes | WHERE is_active == true AND is_primary == false';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should replace all occurrences of a placeholder', () => {
+    const template = 'FROM logs | WHERE user.id == ?user_id OR target.id == ?user_id';
+    const params = { user_id: 101 };
+    const expected = 'FROM logs | WHERE user.id == 101 OR target.id == 101';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should handle an empty parameters object without changing the template', () => {
+    const template = 'FROM events | WHERE event.code == ?code';
+    const params = {};
+    expect(interpolateEsqlQuery(template, params)).toBe(template);
+  });
+
+  it('should ignore extra parameters that are not in the template', () => {
+    const template = 'FROM metrics | WHERE host == ?host';
+    const params = { host: 'server-alpha', unused_param: 'ignore' };
+    const expected = 'FROM metrics | WHERE host == "server-alpha"';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should not partially replace longer placeholders (word boundary test)', () => {
+    const template = 'FROM users | WHERE user_name == ?user_name AND user_id == ?user_id';
+    const params = { user: 'test', user_id: 99 };
+    const expected = 'FROM users | WHERE user_name == ?user_name AND user_id == 99';
+    // Only ?user_id should be replaced, ?user should not affect ?user_name
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should correctly handle an empty string parameter', () => {
+    const template = 'FROM products | WHERE category == ?category';
+    const params = { category: '' };
+    const expected = 'FROM products | WHERE category == ""';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should correctly handle the number 0 as a parameter', () => {
+    const template = 'FROM tasks | WHERE status_code == ?status';
+    const params = { status: 0 };
+    const expected = 'FROM tasks | WHERE status_code == 0';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+
+  it('should handle a mix of all data types in one query', () => {
+    const template =
+      'FROM telemetry | WHERE device_id == ?device AND temp > ?temp AND enabled == ?enabled';
+    const params = { device: 'sensor-3-14', temp: 25.5, enabled: true };
+    const expected =
+      'FROM telemetry | WHERE device_id == "sensor-3-14" AND temp > 25.5 AND enabled == true';
+    expect(interpolateEsqlQuery(template, params)).toBe(expected);
+  });
+});

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql/interpolate_query.ts

@@ -0,0 +1,35 @@
+/*
+ * 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.
+ */
+
+/**
+ * Interpolates parameters into a templated ESQL query string.
+ *
+ * @param template The ESQL query template with '?' placeholders (e.g., "?user_id").
+ * @param params An object where keys match the placeholder names and values are the data to insert.
+ * @returns The interpolated ESQL query string.
+ *
+ * **Important** This is meant as a workaround until a proper util gets exposed from `@kbn/esql-ast`,
+ *               and likely doesn't cover all edge cases.
+ */
+export const interpolateEsqlQuery = (template: string, params: Record<string, unknown>): string => {
+  let interpolatedQuery = template;
+
+  for (const key in params) {
+    if (Object.prototype.hasOwnProperty.call(params, key)) {
+      const value = params[key];
+      const placeholder = new RegExp(`\\?${key}\\b`, 'g');
+
+      // Format the value based on its type
+      const formattedValue = typeof value === 'string' ? `"${value}"` : String(value);
+
+      // Replace all occurrences of the placeholder
+      interpolatedQuery = interpolatedQuery.replace(placeholder, formattedValue);
+    }
+  }
+
+  return interpolatedQuery;
+};

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/index.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-export { esqlResponseToJson, extractEsqlQueries, executeEsql } from './esql';
+export { esqlResponseToJson, extractEsqlQueries, executeEsql, interpolateEsqlQuery } from './esql';
 export {
   flattenMapping,
   cleanupMapping,

x-pack/platform/plugins/shared/onechat/server/services/tools/persisted/tool_types/esql/to_tool_definition.ts

@@ -11,6 +11,7 @@ import { z } from '@kbn/zod';
 import { ToolResultType } from '@kbn/onechat-common/tools/tool_result';
 import type { FieldValue } from '@elastic/elasticsearch/lib/api/types';
 import { getToolResultId } from '@kbn/onechat-server/src/tools';
+import { interpolateEsqlQuery } from '@kbn/onechat-genai-utils/tools/utils';
 import type { ToolPersistedDefinition } from '../../client';
 import type { InternalToolDefinition } from '../../../tool_provider';
 
@@ -36,14 +37,23 @@ export function toToolDefinition<TSchema extends z.ZodObject<any> = z.ZodObject<
         params: paramArray as unknown as FieldValue[],
       });
 
+      // need the interpolated query to return in the results / to display in the UI
+      const interpolatedQuery = interpolateEsqlQuery(configuration.query, params);
+
       return {
         results: [
+          {
+            type: ToolResultType.query,
+            data: {
+              esql: interpolatedQuery,
+            },
+          },
           {
             tool_result_id: getToolResultId(),
             type: ToolResultType.tabularData,
             data: {
               source: 'esql',
-              query: configuration.query,
+              query: interpolatedQuery,
               columns: result.columns,
               values: result.values,
             },