Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Jun 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,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
Expand Down Expand Up @@ -193,7 +190,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') {
Expand All @@ -209,36 +206,27 @@ 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 = '';
let firstSuiteName: 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(' ');
firstSuiteName = titles.filter((item) => item !== '')[0];
});
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(' ');
firstSuiteName = titles.filter((item) => item !== '')[0];
});

suite.afterEach(function () {
currentTitle = '';
});
suite.afterEach(function () {
currentTitle = '';
});
}

const onResultCallbacks: Array<{
callback: (result: EvaluationResult) => void;
Expand Down Expand Up @@ -278,13 +266,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 (error.message.includes('Status code: 429')) {
Expand Down Expand Up @@ -359,14 +346,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`);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
}

This file was deleted.

Loading
Loading