[Obs AI Assistant] Write ES|QL docs with LLM (elastic#183173)

Uses the LLM to re-write docs from the `built-docs` repo, that have a lot of weird technical artifacts, to improve the output from the LLM. (cherry picked from commit 4fc13a4) # Conflicts: # x-pack/plugins/observability_solution/observability_ai_assistant_app/scripts/evaluation/kibana_client.ts # x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts
dgieselaar · Jun 10, 2024 · 0acb55e · 0acb55e
1 parent 7a22f1f
commit 0acb55e
Show file tree

Hide file tree

Showing 138 changed files with 3,545 additions and 2,897 deletions.
diff --git a/...k/plugins/observability_solution/observability_ai_assistant_app/scripts/evaluation/cli.ts b/...k/plugins/observability_solution/observability_ai_assistant_app/scripts/evaluation/cli.ts
@@ -13,9 +13,29 @@ export enum EvaluateWith {
   other = 'other',
 }
 
-export function options(y: Argv) {
-  const config = readKibanaConfig();
+const config = readKibanaConfig();
+
+export const kibanaOption = {
+  describe: 'Where Kibana is running',
+  string: true as const,
+  default: process.env.KIBANA_HOST || 'http://localhost:5601',
+};
+export const elasticsearchOption = {
+  alias: 'es',
+  describe: 'Where Elasticsearch is running',
+  string: true as const,
+  default: format({
+    ...parse(config['elasticsearch.hosts']),
+    auth: `${config['elasticsearch.username']}:${config['elasticsearch.password']}`,
+  }),
+};
+
+export const connectorIdOption = {
+  describe: 'The ID of the connector',
+  string: true as const,
+};
 
+export function options(y: Argv) {
   return y
     .option('files', {
       string: true as const,
@@ -27,30 +47,15 @@ export function options(y: Argv) {
       array: false,
       describe: 'A string or regex to filter scenarios by',
     })
-    .option('kibana', {
-      describe: 'Where Kibana is running',
-      string: true,
-      default: process.env.KIBANA_HOST || 'http://localhost:5601',
-    })
+    .option('kibana', kibanaOption)
     .option('spaceId', {
       describe:
         'The space to use. If space is set, conversations will only be cleared for that spaceId',
       string: true,
       array: false,
     })
-    .option('elasticsearch', {
-      alias: 'es',
-      describe: 'Where Elasticsearch is running',
-      string: true,
-      default: format({
-        ...parse(config['elasticsearch.hosts']),
-        auth: `${config['elasticsearch.username']}:${config['elasticsearch.password']}`,
-      }),
-    })
-    .option('connectorId', {
-      describe: 'The ID of the connector',
-      string: true,
-    })
+    .option('elasticsearch', elasticsearchOption)
+    .option('connectorId', connectorIdOption)
     .option('persist', {
       describe:
         'Whether the conversations should be stored. Adding this will generate a link at which the conversation can be opened.',

diff --git a/...ns/observability_solution/observability_ai_assistant_app/scripts/evaluation/evaluation.ts b/...ns/observability_solution/observability_ai_assistant_app/scripts/evaluation/evaluation.ts
@@ -8,7 +8,6 @@
 import { Client } from '@elastic/elasticsearch';
 import { run } from '@kbn/dev-cli-runner';
 import * as fastGlob from 'fast-glob';
-import inquirer from 'inquirer';
 import yargs from 'yargs';
 import chalk from 'chalk';
 import { castArray, omit } from 'lodash';
@@ -18,48 +17,14 @@ import Path from 'path';
 import * as table from 'table';
 import { TableUserConfig } from 'table';
 import { format, parse } from 'url';
-import { ToolingLog } from '@kbn/tooling-log';
 import { MessageRole } from '@kbn/observability-ai-assistant-plugin/common';
 import { EvaluateWith, options } from './cli';
 import { getServiceUrls } from './get_service_urls';
 import { KibanaClient } from './kibana_client';
 import { initServices } from './services';
 import { setupSynthtrace } from './setup_synthtrace';
 import { EvaluationResult } from './types';
-
-async function selectConnector({
-  connectors,
-  preferredId,
-  log,
-  message = 'Select a connector',
-}: {
-  connectors: Awaited<ReturnType<KibanaClient['getConnectors']>>;
-  preferredId?: string;
-  log: ToolingLog;
-  message?: string;
-}) {
-  let connector = connectors.find((item) => item.id === preferredId);
-
-  if (!connector && preferredId) {
-    log.warning(`Could not find connector ${preferredId}`);
-  }
-
-  if (!connector && connectors.length === 1) {
-    connector = connectors[0];
-    log.debug('Using the only connector found');
-  } else if (!connector) {
-    const connectorChoice = await inquirer.prompt({
-      type: 'list',
-      name: 'connector',
-      message,
-      choices: connectors.map((item) => ({ name: `${item.name} (${item.id})`, value: item.id })),
-    });
-
-    connector = connectors.find((item) => item.id === connectorChoice.connector)!;
-  }
-
-  return connector;
-}
+import { selectConnector } from './select_connector';
 
 function runEvaluations() {
   yargs(process.argv.slice(2))

diff --git a/...ervability_solution/observability_ai_assistant_app/scripts/evaluation/get_service_urls.ts b/...ervability_solution/observability_ai_assistant_app/scripts/evaluation/get_service_urls.ts
@@ -117,10 +117,6 @@ export async function getServiceUrls({
     elasticsearch = 'http://127.0.0.1:9200';
   }
 
-  if (!elasticsearch) {
-    throw new Error('Could not determine an Elasticsearch target');
-  }
-
   const parsedTarget = parse(elasticsearch);
 
   let auth = parsedTarget.auth;

diff --git a/...observability_solution/observability_ai_assistant_app/scripts/evaluation/kibana_client.ts b/...observability_solution/observability_ai_assistant_app/scripts/evaluation/kibana_client.ts
@@ -45,10 +45,7 @@ import {
 } from 'rxjs';
 import { format, parse, UrlObject } from 'url';
 import { inspect } from 'util';
-import type {
-  ObservabilityAIAssistantAPIClientRequestParamsOf,
-  APIReturnType,
-} from '@kbn/observability-ai-assistant-plugin/public';
+import type { ObservabilityAIAssistantAPIClientRequestParamsOf } from '@kbn/observability-ai-assistant-plugin/public';
 import { EvaluationResult } from './types';
 
 // eslint-disable-next-line spaced-comment
@@ -171,7 +168,7 @@ export class KibanaClient {
     connectorId: string;
     evaluationConnectorId: string;
     persist: boolean;
-    suite: Mocha.Suite;
+    suite?: Mocha.Suite;
   }): ChatClient {
     function getMessages(message: string | Array<Message['message']>): Array<Message['message']> {
       if (typeof message === 'string') {
@@ -187,34 +184,25 @@ export class KibanaClient {
 
     const that = this;
 
-    async function getFunctions() {
-      const {
-        data: { functionDefinitions },
-      }: AxiosResponse<APIReturnType<'GET /internal/observability_ai_assistant/functions'>> =
-        await that.axios.get(
-          that.getUrl({ pathname: '/internal/observability_ai_assistant/functions' })
-        );
-
-      return { functionDefinitions };
-    }
-
     let currentTitle: string = '';
 
-    suite.beforeEach(function () {
-      const currentTest: Mocha.Test = this.currentTest;
-      const titles: string[] = [];
-      titles.push(this.currentTest.title);
-      let parent = currentTest.parent;
-      while (parent) {
-        titles.push(parent.title);
-        parent = parent.parent;
-      }
-      currentTitle = titles.reverse().join(' ');
-    });
+    if (suite) {
+      suite.beforeEach(function () {
+        const currentTest: Mocha.Test = this.currentTest;
+        const titles: string[] = [];
+        titles.push(this.currentTest.title);
+        let parent = currentTest.parent;
+        while (parent) {
+          titles.push(parent.title);
+          parent = parent.parent;
+        }
+        currentTitle = titles.reverse().join(' ');
+      });
 
-    suite.afterEach(function () {
-      currentTitle = '';
-    });
+      suite.afterEach(function () {
+        currentTitle = '';
+      });
+    }
 
     const onResultCallbacks: Array<{
       callback: (result: EvaluationResult) => void;
@@ -246,13 +234,12 @@ export class KibanaClient {
                     {
                       message: error.message,
                       status: error.status,
-                      response: error.response?.data,
                     },
                     { depth: 10 }
                   )
                 );
               } else {
-                that.log.error(inspect(error, { depth: 10 }));
+                that.log.error(inspect(error, { depth: 5 }));
               }
 
               if (
@@ -329,14 +316,13 @@ export class KibanaClient {
 
     return {
       chat: async (message) => {
-        const { functionDefinitions } = await getFunctions();
         const messages = [
           ...getMessages(message).map((msg) => ({
             message: msg,
             '@timestamp': new Date().toISOString(),
           })),
         ];
-        return chat('chat', { messages, functions: functionDefinitions });
+        return chat('chat', { messages, functions: [] });
       },
       complete: async (...args) => {
         that.log.info(`Complete`);

diff --git a/...y_solution/observability_ai_assistant_app/scripts/evaluation/scenarios/esql/index.spec.ts b/...y_solution/observability_ai_assistant_app/scripts/evaluation/scenarios/esql/index.spec.ts
@@ -29,7 +29,7 @@ async function evaluateEsqlQuery({
     ...(expected
       ? [
           `Returns a ES|QL query that is functionally equivalent to:      
-      ${expected}`,
+      ${expected}. It's OK if column names are slightly different, as long as the expected end result is the same.`,
         ]
       : []),
     ...(execute
@@ -90,7 +90,7 @@ describe('ES|QL query generation', () => {
       it('top 10 unique domains', async () => {
         await evaluateEsqlQuery({
           question:
-            'For standard Elastic ECS compliant packetbeat data view, shows the top 10 unique destination.domain with more docs',
+            'For standard Elastic ECS compliant packetbeat data view, show me the top 10 unique destination.domain with the most docs',
           expected: `FROM packetbeat-*
           | STATS doc_count = COUNT(*) BY destination.domain
           | SORT doc_count DESC

diff --git a/...ervability_solution/observability_ai_assistant_app/scripts/evaluation/select_connector.ts b/...ervability_solution/observability_ai_assistant_app/scripts/evaluation/select_connector.ts
@@ -0,0 +1,44 @@
+/*
+ * 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 inquirer from 'inquirer';
+import { ToolingLog } from '@kbn/tooling-log';
+import { KibanaClient } from './kibana_client';
+
+export async function selectConnector({
+  connectors,
+  preferredId,
+  log,
+  message = 'Select a connector',
+}: {
+  connectors: Awaited<ReturnType<KibanaClient['getConnectors']>>;
+  preferredId?: string;
+  log: ToolingLog;
+  message?: string;
+}) {
+  let connector = connectors.find((item) => item.id === preferredId);
+
+  if (!connector && preferredId) {
+    log.warning(`Could not find connector ${preferredId}`);
+  }
+
+  if (!connector && connectors.length === 1) {
+    connector = connectors[0];
+    log.debug('Using the only connector found');
+  } else if (!connector) {
+    const connectorChoice = await inquirer.prompt({
+      type: 'list',
+      name: 'connector',
+      message,
+      choices: connectors.map((item) => ({ name: `${item.name} (${item.id})`, value: item.id })),
+    });
+
+    connector = connectors.find((item) => item.id === connectorChoice.connector)!;
+  }
+
+  return connector;
+}
diff --git a/...ty_solution/observability_ai_assistant_app/scripts/load_esql_docs/format_esql_examples.ts b/...ty_solution/observability_ai_assistant_app/scripts/load_esql_docs/format_esql_examples.ts