github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-24 01:38:56 -04:00
Code Issues Projects Releases Packages Wiki Activity

[8.12] [Obs AI Assistant] Perform functions etc on the server (#172590) (#172824)

Browse source 
# Backport

This will backport the following commits from `main` to `8.12`:
- [[Obs AI Assistant] Perform functions etc on the server
(#172590)](https://github.com/elastic/kibana/pull/172590)

<!--- Backport version: 8.9.7 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Dario
Gieselaar","email":"dario.gieselaar@elastic.co"},"sourceCommit":{"committedDate":"2023-12-07T14:18:41Z","message":"[Obs
AI Assistant] Perform functions etc on the server
(#172590)\n\nCo-authored-by: kibanamachine
<42973632+kibanamachine@users.noreply.github.com>","sha":"56b36b9042b42c702a57568f1612af3d433d2df1","branchLabelMapping":{"^v8.13.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:enhancement","Team:APM","v8.12.0","apm:review","v8.13.0"],"number":172590,"url":"https://github.com/elastic/kibana/pull/172590","mergeCommit":{"message":"[Obs
AI Assistant] Perform functions etc on the server
(#172590)\n\nCo-authored-by: kibanamachine
<42973632+kibanamachine@users.noreply.github.com>","sha":"56b36b9042b42c702a57568f1612af3d433d2df1"}},"sourceBranch":"main","suggestedTargetBranches":["8.12"],"targetPullRequestStates":[{"branch":"8.12","label":"v8.12.0","labelRegex":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v8.13.0","labelRegex":"^v8.13.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/172590","number":172590,"mergeCommit":{"message":"[Obs
AI Assistant] Perform functions etc on the server
(#172590)\n\nCo-authored-by: kibanamachine
<42973632+kibanamachine@users.noreply.github.com>","sha":"56b36b9042b42c702a57568f1612af3d433d2df1"}}]}]
BACKPORT-->

Co-authored-by: Dario Gieselaar <dario.gieselaar@elastic.co>
This commit is contained in:
Kibana Machine 2023-12-07 10:41:47 -05:00 • committed by GitHub
parent 1806937483
commit 0db574a1c9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
76 changed files with 4590 additions and 2658 deletions
Download patch file Download diff file Expand all files Collapse all files

77
x-pack/plugins/apm/public/assistant_functions/get_apm_services_list.ts
View file

@ -1,77 +0,0 @@
/*
* 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 { i18n } from '@kbn/i18n';
import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
import { ServiceHealthStatus } from '../../common/service_health_status';
import { callApmApi } from '../services/rest/create_call_apm_api';
import { NON_EMPTY_STRING } from '../utils/non_empty_string_ref';
export function registerGetApmServicesListFunction({
registerFunction,
}: {
registerFunction: RegisterFunctionDefinition;
}) {
registerFunction(
{
name: 'get_apm_services_list',
contexts: ['apm'],
description: `Gets a list of services`,
descriptionForUser: i18n.translate(
'xpack.apm.observabilityAiAssistant.functions.registerGetApmServicesList.descriptionForUser',
{
defaultMessage: `Gets the list of monitored services, their health status, and alerts.`,
}
),
parameters: {
type: 'object',
additionalProperties: false,
properties: {
'service.environment': {
...NON_EMPTY_STRING,
description:
'Optionally filter the services by the environments that they are running in',
},
start: {
...NON_EMPTY_STRING,
description:
'The start of the time range, in Elasticsearch date math, like `now`.',
},
end: {
...NON_EMPTY_STRING,
description:
'The end of the time range, in Elasticsearch date math, like `now-24h`.',
},
healthStatus: {
type: 'array',
description: 'Filter service list by health status',
additionalProperties: false,
additionalItems: false,
items: {
type: 'string',
enum: [
ServiceHealthStatus.unknown,
ServiceHealthStatus.healthy,
ServiceHealthStatus.warning,
ServiceHealthStatus.critical,
],
},
},
},
required: ['start', 'end'],
} as const,
},
async ({ arguments: args }, signal) => {
return callApmApi('POST /internal/apm/assistant/get_services_list', {
signal,
params: {
body: args,
},
});
}
);
}

380
x-pack/plugins/apm/public/assistant_functions/get_apm_timeseries.tsx
View file

@ -5,294 +5,164 @@
* 2.0.
*/
import { EuiFlexGroup, EuiFlexItem, EuiText } from '@elastic/eui';
import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
import { useKibana } from '@kbn/kibana-react-plugin/public';
import type {
RegisterRenderFunctionDefinition,
RenderFunction,
} from '@kbn/observability-ai-assistant-plugin/public/types';
import { groupBy } from 'lodash';
import React from 'react';
import { i18n } from '@kbn/i18n';
import { useKibana } from '@kbn/kibana-react-plugin/public';
import { FETCH_STATUS } from '../hooks/use_fetcher';
import { callApmApi } from '../services/rest/create_call_apm_api';
import { getTimeZone } from '../components/shared/charts/helper/timezone';
import { TimeseriesChart } from '../components/shared/charts/timeseries_chart';
import { ChartPointerEventContextProvider } from '../context/chart_pointer_event/chart_pointer_event_context';
import { ApmThemeProvider } from '../components/routing/app_root';
import { Coordinate, TimeSeries } from '../../typings/timeseries';
import {
ChartType,
getTimeSeriesColor,
} from '../components/shared/charts/helper/get_timeseries_color';
import { LatencyAggregationType } from '../../common/latency_aggregation_types';
import {
asPercent,
asTransactionRate,
getDurationFormatter,
} from '../../common/utils/formatters';
import type {
GetApmTimeseriesFunctionArguments,
GetApmTimeseriesFunctionResponse,
} from '../../server/assistant_functions/get_apm_timeseries';
import { Coordinate, TimeSeries } from '../../typings/timeseries';
import { ApmThemeProvider } from '../components/routing/app_root';
import {
ChartType,
getTimeSeriesColor,
} from '../components/shared/charts/helper/get_timeseries_color';
import { getTimeZone } from '../components/shared/charts/helper/timezone';
import { TimeseriesChart } from '../components/shared/charts/timeseries_chart';
import {
getMaxY,
getResponseTimeTickFormatter,
} from '../components/shared/charts/transaction_charts/helper';
import { NON_EMPTY_STRING } from '../utils/non_empty_string_ref';
import { ChartPointerEventContextProvider } from '../context/chart_pointer_event/chart_pointer_event_context';
import { FETCH_STATUS } from '../hooks/use_fetcher';
export function registerGetApmTimeseriesFunction({
registerFunction,
registerRenderFunction,
}: {
registerFunction: RegisterFunctionDefinition;
registerRenderFunction: RegisterRenderFunctionDefinition;
}) {
registerFunction(
{
contexts: ['apm'],
name: 'get_apm_timeseries',
descriptionForUser: i18n.translate(
'xpack.apm.observabilityAiAssistant.functions.registerGetApmTimeseries.descriptionForUser',
{
defaultMessage: `Display different APM metrics, like throughput, failure rate, or latency, for any service or all services, or any or all of its dependencies, both as a timeseries and as a single statistic. Additionally, the function will return any changes, such as spikes, step and trend changes, or dips. You can also use it to compare data by requesting two different time ranges, or for instance two different service versions`,
}
),
description: `Visualise and analyse different APM metrics, like throughput, failure rate, or latency, for any service or all services, or any or all of its dependencies, both as a timeseries and as a single statistic. A visualisation will be displayed above your reply - DO NOT attempt to display or generate an image yourself, or any other placeholder. Additionally, the function will return any changes, such as spikes, step and trend changes, or dips. You can also use it to compare data by requesting two different time ranges, or for instance two different service versions.`,
parameters: {
type: 'object',
properties: {
start: {
type: 'string',
description:
'The start of the time range, in Elasticsearch date math, like `now`.',
},
end: {
type: 'string',
description:
'The end of the time range, in Elasticsearch date math, like `now-24h`.',
},
stats: {
type: 'array',
items: {
type: 'object',
properties: {
timeseries: {
description: 'The metric to be displayed',
oneOf: [
{
type: 'object',
properties: {
name: {
type: 'string',
enum: [
'transaction_throughput',
'transaction_failure_rate',
],
},
'transaction.type': {
type: 'string',
description: 'The transaction type',
},
},
required: ['name'],
},
{
type: 'object',
properties: {
name: {
type: 'string',
enum: [
'exit_span_throughput',
'exit_span_failure_rate',
'exit_span_latency',
],
},
'span.destination.service.resource': {
type: 'string',
description:
'The name of the downstream dependency for the service',
},
},
required: ['name'],
},
{
type: 'object',
properties: {
name: {
type: 'string',
const: 'error_event_rate',
},
},
required: ['name'],
},
{
type: 'object',
properties: {
name: {
type: 'string',
const: 'transaction_latency',
},
'transaction.type': {
type: 'string',
},
function: {
type: 'string',
enum: ['avg', 'p95', 'p99'],
},
},
required: ['name', 'function'],
},
],
},
'service.name': {
...NON_EMPTY_STRING,
description: 'The name of the service',
},
'service.environment': {
description:
'The environment that the service is running in. If undefined, all environments will be included. Only use this if you have confirmed the environment that the service is running in.',
},
filter: {
type: 'string',
description:
'a KQL query to filter the data by. If no filter should be applied, leave it empty.',
},
title: {
type: 'string',
description:
'A unique, human readable, concise title for this specific group series.',
},
offset: {
type: 'string',
description:
'The offset. Right: 15m. 8h. 1d. Wrong: -15m. -8h. -1d.',
},
},
required: ['service.name', 'timeseries', 'title'],
},
},
},
required: ['stats', 'start', 'end'],
} as const,
},
async ({ arguments: { stats, start, end } }, signal) => {
const response = await callApmApi(
'POST /internal/apm/assistant/get_apm_timeseries',
{
signal,
params: {
body: { stats: stats as any, start, end },
},
}
);
registerRenderFunction('get_apm_timeseries', (parameters) => {
const { response } = parameters as Parameters<
RenderFunction<
GetApmTimeseriesFunctionArguments,
GetApmTimeseriesFunctionResponse
>
>[0];
return response;
},
({ arguments: args, response }) => {
const groupedSeries = groupBy(response.data, (series) => series.group);
const groupedSeries = groupBy(response.data, (series) => series.group);
const {
services: { uiSettings },
} = useKibana();
const {
services: { uiSettings },
} = useKibana();
const timeZone = getTimeZone(uiSettings);
const timeZone = getTimeZone(uiSettings);
return (
<ChartPointerEventContextProvider>
<ApmThemeProvider>
<EuiFlexGroup direction="column">
{Object.values(groupedSeries).map((groupSeries) => {
const groupId = groupSeries[0].group;
return (
<ChartPointerEventContextProvider>
<ApmThemeProvider>
<EuiFlexGroup direction="column">
{Object.values(groupedSeries).map((groupSeries) => {
const groupId = groupSeries[0].group;
const maxY = getMaxY(groupSeries);
const latencyFormatter = getDurationFormatter(maxY, 10, 1000);
const maxY = getMaxY(groupSeries);
const latencyFormatter = getDurationFormatter(maxY, 10, 1000);
let yLabelFormat: (value: number) => string;
let yLabelFormat: (value: number) => string;
const firstStat = groupSeries[0].stat;
const firstStat = groupSeries[0].stat;
switch (firstStat.timeseries.name) {
case 'transaction_throughput':
case 'exit_span_throughput':
case 'error_event_rate':
yLabelFormat = asTransactionRate;
break;
switch (firstStat.timeseries.name) {
case 'transaction_throughput':
case 'exit_span_throughput':
case 'error_event_rate':
yLabelFormat = asTransactionRate;
break;
case 'transaction_latency':
case 'exit_span_latency':
yLabelFormat =
getResponseTimeTickFormatter(latencyFormatter);
break;
case 'transaction_latency':
case 'exit_span_latency':
yLabelFormat = getResponseTimeTickFormatter(latencyFormatter);
break;
case 'transaction_failure_rate':
case 'exit_span_failure_rate':
yLabelFormat = (y) => asPercent(y || 0, 100);
break;
}
case 'transaction_failure_rate':
case 'exit_span_failure_rate':
yLabelFormat = (y) => asPercent(y || 0, 100);
break;
}
const timeseries: Array<TimeSeries<Coordinate>> =
groupSeries.map((series): TimeSeries<Coordinate> => {
let chartType: ChartType;
const timeseries: Array<TimeSeries<Coordinate>> = groupSeries.map(
(series): TimeSeries<Coordinate> => {
let chartType: ChartType;
const data = series.data;
const data = series.data;
switch (series.stat.timeseries.name) {
case 'transaction_throughput':
case 'exit_span_throughput':
chartType = ChartType.THROUGHPUT;
break;
switch (series.stat.timeseries.name) {
case 'transaction_throughput':
case 'exit_span_throughput':
chartType = ChartType.THROUGHPUT;
break;
case 'transaction_failure_rate':
case 'exit_span_failure_rate':
chartType = ChartType.FAILED_TRANSACTION_RATE;
break;
case 'transaction_failure_rate':
case 'exit_span_failure_rate':
chartType = ChartType.FAILED_TRANSACTION_RATE;
break;
case 'transaction_latency':
if (
series.stat.timeseries.function ===
LatencyAggregationType.p99
) {
chartType = ChartType.LATENCY_P99;
} else if (
series.stat.timeseries.function ===
LatencyAggregationType.p95
) {
chartType = ChartType.LATENCY_P95;
} else {
chartType = ChartType.LATENCY_AVG;
}
break;
case 'exit_span_latency':
case 'transaction_latency':
if (
series.stat.timeseries.function ===
LatencyAggregationType.p99
) {
chartType = ChartType.LATENCY_P99;
} else if (
series.stat.timeseries.function ===
LatencyAggregationType.p95
) {
chartType = ChartType.LATENCY_P95;
} else {
chartType = ChartType.LATENCY_AVG;
break;
}
break;
case 'error_event_rate':
chartType = ChartType.ERROR_OCCURRENCES;
break;
}
case 'exit_span_latency':
chartType = ChartType.LATENCY_AVG;
break;
return {
title: series.id,
type: 'line',
color: getTimeSeriesColor(chartType!).currentPeriodColor,
data,
};
});
case 'error_event_rate':
chartType = ChartType.ERROR_OCCURRENCES;
break;
}
return (
<EuiFlexItem grow={false} key={groupId}>
<EuiFlexGroup direction="column" gutterSize="s">
<EuiFlexItem>
<EuiText size="m">{groupId}</EuiText>
<TimeseriesChart
comparisonEnabled={false}
fetchStatus={FETCH_STATUS.SUCCESS}
id={groupId}
timeZone={timeZone}
timeseries={timeseries}
yLabelFormat={yLabelFormat!}
/>
</EuiFlexItem>
</EuiFlexGroup>
</EuiFlexItem>
);
})}
</EuiFlexGroup>
</ApmThemeProvider>
</ChartPointerEventContextProvider>
);
}
);
return {
title: series.id,
type: 'line',
color: getTimeSeriesColor(chartType!).currentPeriodColor,
data,
};
}
);
return (
<EuiFlexItem grow={false} key={groupId}>
<EuiFlexGroup direction="column" gutterSize="s">
<EuiFlexItem>
<EuiText size="m">{groupId}</EuiText>
<TimeseriesChart
comparisonEnabled={false}
fetchStatus={FETCH_STATUS.SUCCESS}
id={groupId}
timeZone={timeZone}
timeseries={timeseries}
yLabelFormat={yLabelFormat!}
/>
</EuiFlexItem>
</EuiFlexGroup>
</EuiFlexItem>
);
})}
</EuiFlexGroup>
</ApmThemeProvider>
</ChartPointerEventContextProvider>
);
});
}

145
x-pack/plugins/apm/public/assistant_functions/index.ts
View file

@ -5,152 +5,15 @@
* 2.0.
*/
import type { CoreStart } from '@kbn/core/public';
import type {
RegisterContextDefinition,
RegisterFunctionDefinition,
} from '@kbn/observability-ai-assistant-plugin/common/types';
import type { ApmPluginStartDeps } from '../plugin';
import {
createCallApmApi,
callApmApi,
} from '../services/rest/create_call_apm_api';
import { registerGetApmCorrelationsFunction } from './get_apm_correlations';
import { registerGetApmDownstreamDependenciesFunction } from './get_apm_downstream_dependencies';
import { registerGetApmErrorDocumentFunction } from './get_apm_error_document';
import { registerGetApmServicesListFunction } from './get_apm_services_list';
import { registerGetApmServiceSummaryFunction } from './get_apm_service_summary';
import { RegisterRenderFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/public/types';
import { registerGetApmTimeseriesFunction } from './get_apm_timeseries';
export async function registerAssistantFunctions({
pluginsStart,
coreStart,
registerContext,
registerFunction,
signal,
registerRenderFunction,
}: {
pluginsStart: ApmPluginStartDeps;
coreStart: CoreStart;
registerContext: RegisterContextDefinition;
registerFunction: RegisterFunctionDefinition;
signal: AbortSignal;
registerRenderFunction: RegisterRenderFunctionDefinition;
}) {
createCallApmApi(coreStart);
const response = await callApmApi('GET /internal/apm/has_data', {
signal,
});
if (!response.hasData) {
return;
}
registerGetApmTimeseriesFunction({
registerFunction,
});
registerGetApmErrorDocumentFunction({
registerFunction,
});
registerGetApmCorrelationsFunction({
registerFunction,
});
registerGetApmDownstreamDependenciesFunction({
registerFunction,
});
registerGetApmServiceSummaryFunction({
registerFunction,
});
registerGetApmServicesListFunction({
registerFunction,
});
registerContext({
name: 'apm',
description: `
When analyzing APM data, prefer the APM specific functions over the generic Lens,
Elasticsearch or Kibana ones, unless those are explicitly requested by the user.
When requesting metrics for a service, make sure you also know what environment
it is running in. Metrics aggregated over multiple environments are useless.
There are four important data types in Elastic APM. Each of them have the
following fields:
- service.name: the name of the service
- service.node.name: the id of the service instance (often the hostname)
- service.environment: the environment (often production, development)
- agent.name: the name of the agent (go, java, etc)
The four data types are transactions, exit spans, error events, and application
metrics.
Transactions have three metrics: throughput, failure rate, and latency. The
fields are:
- transaction.type: often request or page-load (the main transaction types),
but can also be worker, or route-change.
- transaction.name: The name of the transaction group, often something like
'GET /api/product/:productId'
- transaction.result: The result. Used to capture HTTP response codes
(2xx,3xx,4xx,5xx) for request transactions.
- event.outcome: whether the transaction was succesful or not. success,
failure, or unknown.
Exit spans have three metrics: throughput, failure rate and latency. The fields
are:
- span.type: db, external
- span.subtype: the type of database (redis, postgres) or protocol (http, grpc)
- span.destination.service.resource: the address of the destination of the call
- event.outcome: whether the transaction was succesful or not. success,
failure, or unknown.
Error events have one metric, error event rate. The fields are:
- error.grouping_name: a human readable keyword that identifies the error group
For transaction metrics we also collect anomalies. These are scored 0 (low) to
100 (critical).
For root cause analysis, locate a change point in the relevant metrics for a
service or downstream dependency. You can locate a change point by using a
sliding window, e.g. start with a small time range, like 30m, and make it
bigger until you identify a change point. It's very important to identify a
change point. If you don't have a change point, ask the user for next steps.
You can also use an anomaly or a deployment as a change point. Then, compare
data before the change with data after the change. You can either use the
groupBy parameter in get_apm_chart to get the most occuring values in a certain
data set, or you can use correlations to see for which field and value the
frequency has changed when comparing the foreground set to the background set.
This is useful when comparing data from before the change point with after the
change point. For instance, you might see a specific error pop up more often
after the change point.
When comparing anomalies and changes in timeseries, first, zoom in to a smaller
time window, at least 30 minutes before and 30 minutes after the change
occured. E.g., if the anomaly occured at 2023-07-05T08:15:00.000Z, request a
time window that starts at 2023-07-05T07:45:00.000Z and ends at
2023-07-05T08:45:00.000Z. When comparing changes in different timeseries and
anomalies to determine a correlation, make sure to compare the timestamps. If
in doubt, rate the likelihood of them being related, given the time difference,
between 1 and 10. If below 5, assume it's not related. Mention this likelihood
(and the time difference) to the user.
Your goal is to help the user determine the root cause of an issue quickly and
transparently. If you see a change or
anomaly in a metric for a service, try to find similar changes in the metrics
for the traffic to its downstream dependencies, by comparing transaction
metrics to span metrics. To inspect the traffic from one service to a
downstream dependency, first get the downstream dependencies for a service,
then get the span metrics from that service (\`service.name\`) to its
downstream dependency (\`span.destination.service.resource\`). For instance,
for an anomaly in throughput, first inspect \`transaction_throughput\` for
\`service.name\`. Then, inspect \`exit_span_throughput\` for its downstream
dependencies, by grouping by \`span.destination.service.resource\`. Repeat this
process over the next service its downstream dependencies until you identify a
root cause. If you can not find any similar changes, use correlations or
grouping to find attributes that could be causes for the change.`,
registerRenderFunction,
});
}

8
x-pack/plugins/apm/public/plugin.ts
View file

@ -428,15 +428,11 @@ export class ApmPlugin implements Plugin<ApmPluginSetup, ApmPluginStart> {
const { fleet } = plugins;
plugins.observabilityAIAssistant.service.register(
async ({ signal, registerContext, registerFunction }) => {
async ({ registerRenderFunction }) => {
const mod = await import('./assistant_functions');
mod.registerAssistantFunctions({
coreStart: core,
pluginsStart: plugins,
registerContext,
registerFunction,
signal,
registerRenderFunction,
});
}
);

21
x-pack/plugins/apm/public/assistant_functions/get_apm_correlations.ts → x-pack/plugins/apm/server/assistant_functions/get_apm_correlations.ts
View file

@ -6,15 +6,14 @@
*/
import { i18n } from '@kbn/i18n';
import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
import { FunctionRegistrationParameters } from '.';
import { CorrelationsEventType } from '../../common/assistant/constants';
import { callApmApi } from '../services/rest/create_call_apm_api';
import { getApmCorrelationValues } from '../routes/assistant_functions/get_apm_correlation_values';
export function registerGetApmCorrelationsFunction({
apmEventClient,
registerFunction,
}: {
registerFunction: RegisterFunctionDefinition;
}) {
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'get_apm_correlations',
@ -113,12 +112,12 @@ export function registerGetApmCorrelationsFunction({
} as const,
},
async ({ arguments: args }, signal) => {
return callApmApi('POST /internal/apm/assistant/get_correlation_values', {
signal,
params: {
body: args,
},
});
return {
content: await getApmCorrelationValues({
arguments: args as any,
apmEventClient,
}),
};
}
);
}

24
x-pack/plugins/apm/public/assistant_functions/get_apm_downstream_dependencies.ts → x-pack/plugins/apm/server/assistant_functions/get_apm_downstream_dependencies.ts
View file

@ -6,14 +6,13 @@
*/
import { i18n } from '@kbn/i18n';
import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
import { callApmApi } from '../services/rest/create_call_apm_api';
import type { FunctionRegistrationParameters } from '.';
import { getAssistantDownstreamDependencies } from '../routes/assistant_functions/get_apm_downstream_dependencies';
export function registerGetApmDownstreamDependenciesFunction({
apmEventClient,
registerFunction,
}: {
registerFunction: RegisterFunctionDefinition;
}) {
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'get_apm_downstream_dependencies',
@ -57,15 +56,12 @@ export function registerGetApmDownstreamDependenciesFunction({
} as const,
},
async ({ arguments: args }, signal) => {
return callApmApi(
'GET /internal/apm/assistant/get_downstream_dependencies',
{
signal,
params: {
query: args,
},
}
);
return {
content: await getAssistantDownstreamDependencies({
arguments: args,
apmEventClient,
}),
};
}
);
}

21
x-pack/plugins/apm/public/assistant_functions/get_apm_error_document.ts → x-pack/plugins/apm/server/assistant_functions/get_apm_error_document.ts
View file

@ -6,14 +6,13 @@
*/
import { i18n } from '@kbn/i18n';
import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
import { callApmApi } from '../services/rest/create_call_apm_api';
import type { FunctionRegistrationParameters } from '.';
import { getApmErrorDocument } from '../routes/assistant_functions/get_apm_error_document';
export function registerGetApmErrorDocumentFunction({
apmEventClient,
registerFunction,
}: {
registerFunction: RegisterFunctionDefinition;
}) {
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'get_apm_error_document',
@ -55,12 +54,12 @@ export function registerGetApmErrorDocumentFunction({
} as const,
},
async ({ arguments: args }, signal) => {
return callApmApi('GET /internal/apm/assistant/get_error_document', {
signal,
params: {
query: args,
},
});
return {
content: await getApmErrorDocument({
apmEventClient,
arguments: args,
}),
};
}
);
}

46
x-pack/plugins/apm/public/assistant_functions/get_apm_service_summary.ts → x-pack/plugins/apm/server/assistant_functions/get_apm_service_summary.ts
View file

@ -5,16 +5,19 @@
* 2.0.
*/
import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
import { i18n } from '@kbn/i18n';
import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
import { callApmApi } from '../services/rest/create_call_apm_api';
import type { FunctionRegistrationParameters } from '.';
import { getApmAlertsClient } from '../lib/helpers/get_apm_alerts_client';
import { getMlClient } from '../lib/helpers/get_ml_client';
import { getApmServiceSummary } from '../routes/assistant_functions/get_apm_service_summary';
import { NON_EMPTY_STRING } from '../utils/non_empty_string_ref';
export function registerGetApmServiceSummaryFunction({
resources,
apmEventClient,
registerFunction,
}: {
registerFunction: RegisterFunctionDefinition;
}) {
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'get_apm_service_summary',
@ -58,12 +61,33 @@ alerts and anomalies.`,
} as const,
},
async ({ arguments: args }, signal) => {
return callApmApi('GET /internal/apm/assistant/get_service_summary', {
signal,
params: {
query: args,
},
});
const { context, request, plugins, logger } = resources;
const [annotationsClient, esClient, apmAlertsClient, mlClient] =
await Promise.all([
plugins.observability.setup.getScopedAnnotationsClient(
context,
request
),
context.core.then(
(coreContext): ElasticsearchClient =>
coreContext.elasticsearch.client.asCurrentUser
),
getApmAlertsClient(resources),
getMlClient(resources),
]);
return {
content: await getApmServiceSummary({
apmEventClient,
annotationsClient,
esClient,
apmAlertsClient,
mlClient,
logger,
arguments: args,
}),
};
}
);
}

137
x-pack/plugins/apm/server/assistant_functions/get_apm_services_list.ts Normal file
View file

@ -0,0 +1,137 @@
/*
* 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 datemath from '@elastic/datemath';
import { i18n } from '@kbn/i18n';
import { FunctionRegistrationParameters } from '.';
import { ApmDocumentType } from '../../common/document_type';
import { ENVIRONMENT_ALL } from '../../common/environment_filter_values';
import { RollupInterval } from '../../common/rollup';
import { ServiceHealthStatus } from '../../common/service_health_status';
import { getApmAlertsClient } from '../lib/helpers/get_apm_alerts_client';
import { getMlClient } from '../lib/helpers/get_ml_client';
import { getRandomSampler } from '../lib/helpers/get_random_sampler';
import { getServicesItems } from '../routes/services/get_services/get_services_items';
import { NON_EMPTY_STRING } from '../utils/non_empty_string_ref';
export interface ApmServicesListItem {
'service.name': string;
'agent.name'?: string;
'transaction.type'?: string;
alertsCount: number;
healthStatus: ServiceHealthStatus;
'service.environment'?: string[];
}
export function registerGetApmServicesListFunction({
apmEventClient,
resources,
registerFunction,
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'get_apm_services_list',
contexts: ['apm'],
description: `Gets a list of services`,
descriptionForUser: i18n.translate(
'xpack.apm.observabilityAiAssistant.functions.registerGetApmServicesList.descriptionForUser',
{
defaultMessage: `Gets the list of monitored services, their health status, and alerts.`,
}
),
parameters: {
type: 'object',
additionalProperties: false,
properties: {
'service.environment': {
...NON_EMPTY_STRING,
description:
'Optionally filter the services by the environments that they are running in',
},
start: {
...NON_EMPTY_STRING,
description:
'The start of the time range, in Elasticsearch date math, like `now`.',
},
end: {
...NON_EMPTY_STRING,
description:
'The end of the time range, in Elasticsearch date math, like `now-24h`.',
},
healthStatus: {
type: 'array',
description: 'Filter service list by health status',
additionalProperties: false,
additionalItems: false,
items: {
type: 'string',
enum: [
ServiceHealthStatus.unknown,
ServiceHealthStatus.healthy,
ServiceHealthStatus.warning,
ServiceHealthStatus.critical,
],
},
},
},
required: ['start', 'end'],
} as const,
},
async ({ arguments: args }, signal) => {
const { healthStatus } = args;
const [apmAlertsClient, mlClient, randomSampler] = await Promise.all([
getApmAlertsClient(resources),
getMlClient(resources),
getRandomSampler({
security: resources.plugins.security,
probability: 1,
request: resources.request,
}),
]);
const start = datemath.parse(args.start)?.valueOf()!;
const end = datemath.parse(args.end)?.valueOf()!;
const serviceItems = await getServicesItems({
apmAlertsClient,
apmEventClient,
documentType: ApmDocumentType.TransactionMetric,
start,
end,
environment: args['service.environment'] || ENVIRONMENT_ALL.value,
kuery: '',
logger: resources.logger,
randomSampler,
rollupInterval: RollupInterval.OneMinute,
serviceGroup: null,
mlClient,
useDurationSummary: false,
});
let mappedItems = serviceItems.items.map((item): ApmServicesListItem => {
return {
'service.name': item.serviceName,
'agent.name': item.agentName,
alertsCount: item.alertsCount ?? 0,
healthStatus: item.healthStatus ?? ServiceHealthStatus.unknown,
'service.environment': item.environments,
'transaction.type': item.transactionType,
};
});
if (healthStatus && healthStatus.length) {
mappedItems = mappedItems.filter((item): boolean =>
healthStatus.includes(item.healthStatus)
);
}
return {
content: mappedItems,
};
}
);
}

174
x-pack/plugins/apm/server/assistant_functions/get_apm_timeseries.ts Normal file
View file

@ -0,0 +1,174 @@
/*
* 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 { i18n } from '@kbn/i18n';
import { FromSchema } from 'json-schema-to-ts';
import { omit } from 'lodash';
import { FunctionRegistrationParameters } from '.';
import {
ApmTimeseries,
getApmTimeseries,
} from '../routes/assistant_functions/get_apm_timeseries';
import { NON_EMPTY_STRING } from '../utils/non_empty_string_ref';
const parameters = {
type: 'object',
properties: {
start: {
type: 'string',
description:
'The start of the time range, in Elasticsearch date math, like `now`.',
},
end: {
type: 'string',
description:
'The end of the time range, in Elasticsearch date math, like `now-24h`.',
},
stats: {
type: 'array',
items: {
type: 'object',
properties: {
timeseries: {
description: 'The metric to be displayed',
oneOf: [
{
type: 'object',
properties: {
name: {
type: 'string',
enum: [
'transaction_throughput',
'transaction_failure_rate',
],
},
'transaction.type': {
type: 'string',
description: 'The transaction type',
},
},
required: ['name'],
},
{
type: 'object',
properties: {
name: {
type: 'string',
enum: [
'exit_span_throughput',
'exit_span_failure_rate',
'exit_span_latency',
],
},
'span.destination.service.resource': {
type: 'string',
description:
'The name of the downstream dependency for the service',
},
},
required: ['name'],
},
{
type: 'object',
properties: {
name: {
type: 'string',
const: 'error_event_rate',
},
},
required: ['name'],
},
{
type: 'object',
properties: {
name: {
type: 'string',
const: 'transaction_latency',
},
'transaction.type': {
type: 'string',
},
function: {
type: 'string',
enum: ['avg', 'p95', 'p99'],
},
},
required: ['name', 'function'],
},
],
},
'service.name': {
...NON_EMPTY_STRING,
description: 'The name of the service',
},
'service.environment': {
description:
'The environment that the service is running in. If undefined, all environments will be included. Only use this if you have confirmed the environment that the service is running in.',
},
filter: {
type: 'string',
description:
'a KQL query to filter the data by. If no filter should be applied, leave it empty.',
},
title: {
type: 'string',
description:
'A unique, human readable, concise title for this specific group series.',
},
offset: {
type: 'string',
description:
'The offset. Right: 15m. 8h. 1d. Wrong: -15m. -8h. -1d.',
},
},
required: ['service.name', 'timeseries', 'title'],
},
},
},
required: ['stats', 'start', 'end'],
} as const;
export function registerGetApmTimeseriesFunction({
apmEventClient,
registerFunction,
}: FunctionRegistrationParameters) {
registerFunction(
{
contexts: ['apm'],
name: 'get_apm_timeseries',
descriptionForUser: i18n.translate(
'xpack.apm.observabilityAiAssistant.functions.registerGetApmTimeseries.descriptionForUser',
{
defaultMessage: `Display different APM metrics, like throughput, failure rate, or latency, for any service or all services, or any or all of its dependencies, both as a timeseries and as a single statistic. Additionally, the function will return any changes, such as spikes, step and trend changes, or dips. You can also use it to compare data by requesting two different time ranges, or for instance two different service versions`,
}
),
description: `Visualise and analyse different APM metrics, like throughput, failure rate, or latency, for any service or all services, or any or all of its dependencies, both as a timeseries and as a single statistic. A visualisation will be displayed above your reply - DO NOT attempt to display or generate an image yourself, or any other placeholder. Additionally, the function will return any changes, such as spikes, step and trend changes, or dips. You can also use it to compare data by requesting two different time ranges, or for instance two different service versions.`,
parameters,
},
async (
{ arguments: args },
signal
): Promise<GetApmTimeseriesFunctionResponse> => {
const timeseries = await getApmTimeseries({
apmEventClient,
arguments: args as any,
});
return {
content: timeseries.map(
(series): Omit<ApmTimeseries, 'data'> => omit(series, 'data')
),
data: timeseries,
};
}
);
}
export type GetApmTimeseriesFunctionArguments = FromSchema<typeof parameters>;
export interface GetApmTimeseriesFunctionResponse {
content: Array<Omit<ApmTimeseries, 'data'>>;
data: ApmTimeseries[];
}

186
x-pack/plugins/apm/server/assistant_functions/index.ts Normal file
View file

@ -0,0 +1,186 @@
/*
* 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 { CoreSetup } from '@kbn/core-lifecycle-server';
import type { Logger } from '@kbn/logging';
import type {
ChatRegistrationFunction,
RegisterFunction,
} from '@kbn/observability-ai-assistant-plugin/server/service/types';
import type { IRuleDataClient } from '@kbn/rule-registry-plugin/server';
import type { APMConfig } from '..';
import type { ApmFeatureFlags } from '../../common/apm_feature_flags';
import { APMEventClient } from '../lib/helpers/create_es_client/create_apm_event_client';
import { getApmEventClient } from '../lib/helpers/get_apm_event_client';
import type { APMRouteHandlerResources } from '../routes/apm_routes/register_apm_server_routes';
import { hasHistoricalAgentData } from '../routes/historical_data/has_historical_agent_data';
import { registerGetApmCorrelationsFunction } from './get_apm_correlations';
import { registerGetApmDownstreamDependenciesFunction } from './get_apm_downstream_dependencies';
import { registerGetApmErrorDocumentFunction } from './get_apm_error_document';
import { registerGetApmServicesListFunction } from './get_apm_services_list';
import { registerGetApmServiceSummaryFunction } from './get_apm_service_summary';
import { registerGetApmTimeseriesFunction } from './get_apm_timeseries';
export interface FunctionRegistrationParameters {
apmEventClient: APMEventClient;
registerFunction: RegisterFunction;
resources: APMRouteHandlerResources;
}
export function registerAssistantFunctions({
coreSetup,
config,
featureFlags,
logger,
kibanaVersion,
ruleDataClient,
plugins,
}: {
coreSetup: CoreSetup;
config: APMConfig;
featureFlags: ApmFeatureFlags;
logger: Logger;
kibanaVersion: string;
ruleDataClient: IRuleDataClient;
plugins: APMRouteHandlerResources['plugins'];
}): ChatRegistrationFunction {
return async ({ resources, registerContext, registerFunction }) => {
const apmRouteHandlerResources: APMRouteHandlerResources = {
context: resources.context,
request: resources.request,
core: {
setup: coreSetup,
start: () =>
coreSetup.getStartServices().then(([coreStart]) => coreStart),
},
params: {
query: {
_inspect: false,
},
},
config,
featureFlags,
logger,
kibanaVersion,
ruleDataClient,
plugins,
getApmIndices: async () => {
const coreContext = await resources.context.core;
const apmIndices = await plugins.apmDataAccess.setup.getApmIndices(
coreContext.savedObjects.client
);
return apmIndices;
},
};
const apmEventClient = await getApmEventClient(apmRouteHandlerResources);
const hasData = await hasHistoricalAgentData(apmEventClient);
if (!hasData) {
return;
}
const parameters: FunctionRegistrationParameters = {
resources: apmRouteHandlerResources,
apmEventClient,
registerFunction,
};
registerGetApmServicesListFunction(parameters);
registerGetApmServiceSummaryFunction(parameters);
registerGetApmErrorDocumentFunction(parameters);
registerGetApmDownstreamDependenciesFunction(parameters);
registerGetApmCorrelationsFunction(parameters);
registerGetApmTimeseriesFunction(parameters);
registerContext({
name: 'apm',
description: `
When analyzing APM data, prefer the APM specific functions over the generic Lens,
Elasticsearch or Kibana ones, unless those are explicitly requested by the user.
When requesting metrics for a service, make sure you also know what environment
it is running in. Metrics aggregated over multiple environments are useless.
There are four important data types in Elastic APM. Each of them have the
following fields:
- service.name: the name of the service
- service.node.name: the id of the service instance (often the hostname)
- service.environment: the environment (often production, development)
- agent.name: the name of the agent (go, java, etc)
The four data types are transactions, exit spans, error events, and application
metrics.
Transactions have three metrics: throughput, failure rate, and latency. The
fields are:
- transaction.type: often request or page-load (the main transaction types),
but can also be worker, or route-change.
- transaction.name: The name of the transaction group, often something like
'GET /api/product/:productId'
- transaction.result: The result. Used to capture HTTP response codes
(2xx,3xx,4xx,5xx) for request transactions.
- event.outcome: whether the transaction was succesful or not. success,
failure, or unknown.
Exit spans have three metrics: throughput, failure rate and latency. The fields
are:
- span.type: db, external
- span.subtype: the type of database (redis, postgres) or protocol (http, grpc)
- span.destination.service.resource: the address of the destination of the call
- event.outcome: whether the transaction was succesful or not. success,
failure, or unknown.
Error events have one metric, error event rate. The fields are:
- error.grouping_name: a human readable keyword that identifies the error group
For transaction metrics we also collect anomalies. These are scored 0 (low) to
100 (critical).
For root cause analysis, locate a change point in the relevant metrics for a
service or downstream dependency. You can locate a change point by using a
sliding window, e.g. start with a small time range, like 30m, and make it
bigger until you identify a change point. It's very important to identify a
change point. If you don't have a change point, ask the user for next steps.
You can also use an anomaly or a deployment as a change point. Then, compare
data before the change with data after the change. You can either use the
groupBy parameter in get_apm_chart to get the most occuring values in a certain
data set, or you can use correlations to see for which field and value the
frequency has changed when comparing the foreground set to the background set.
This is useful when comparing data from before the change point with after the
change point. For instance, you might see a specific error pop up more often
after the change point.
When comparing anomalies and changes in timeseries, first, zoom in to a smaller
time window, at least 30 minutes before and 30 minutes after the change
occured. E.g., if the anomaly occured at 2023-07-05T08:15:00.000Z, request a
time window that starts at 2023-07-05T07:45:00.000Z and ends at
2023-07-05T08:45:00.000Z. When comparing changes in different timeseries and
anomalies to determine a correlation, make sure to compare the timestamps. If
in doubt, rate the likelihood of them being related, given the time difference,
between 1 and 10. If below 5, assume it's not related. Mention this likelihood
(and the time difference) to the user.
Your goal is to help the user determine the root cause of an issue quickly and
transparently. If you see a change or
anomaly in a metric for a service, try to find similar changes in the metrics
for the traffic to its downstream dependencies, by comparing transaction
metrics to span metrics. To inspect the traffic from one service to a
downstream dependency, first get the downstream dependencies for a service,
then get the span metrics from that service (\`service.name\`) to its
downstream dependency (\`span.destination.service.resource\`). For instance,
for an anomaly in throughput, first inspect \`transaction_throughput\` for
\`service.name\`. Then, inspect \`exit_span_throughput\` for its downstream
dependencies, by grouping by \`span.destination.service.resource\`. Repeat this
process over the next service its downstream dependencies until you identify a
root cause. If you can not find any similar changes, use correlations or
grouping to find attributes that could be causes for the change.`,
});
};
}

17
x-pack/plugins/apm/server/plugin.ts
View file

@ -49,6 +49,7 @@ import { scheduleSourceMapMigration } from './routes/source_maps/schedule_source
import { createApmSourceMapIndexTemplate } from './routes/source_maps/create_apm_source_map_index_template';
import { addApiKeysToEveryPackagePolicyIfMissing } from './routes/fleet/api_keys/add_api_keys_to_policies_if_missing';
import { apmTutorialCustomIntegration } from '../common/tutorial/tutorials';
import { registerAssistantFunctions } from './assistant_functions';
export class APMPlugin
implements
@ -167,6 +168,8 @@ export class APMPlugin
APM_SERVER_FEATURE_ID
);
const kibanaVersion = this.initContext.env.packageInfo.version;
registerRoutes({
core: {
setup: core,
@ -179,7 +182,7 @@ export class APMPlugin
ruleDataClient,
plugins: resourcePlugins,
telemetryUsageCounter,
kibanaVersion: this.initContext.env.packageInfo.version,
kibanaVersion,
});
const { getApmIndices } = plugins.apmDataAccess;
@ -230,6 +233,18 @@ export class APMPlugin
this.logger?.error(e);
});
plugins.observabilityAIAssistant.service.register(
registerAssistantFunctions({
config: this.currentConfig!,
coreSetup: core,
featureFlags: this.currentConfig!.featureFlags,
kibanaVersion,
logger: this.logger.get('assistant'),
plugins: resourcePlugins,
ruleDataClient,
})
);
return { config$ };
}

130
x-pack/plugins/apm/server/routes/assistant_functions/route.ts
View file

@ -4,21 +4,13 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import datemath from '@elastic/datemath';
import { ElasticsearchClient } from '@kbn/core/server';
import * as t from 'io-ts';
import { omit } from 'lodash';
import { ApmDocumentType } from '../../../common/document_type';
import { ENVIRONMENT_ALL } from '../../../common/environment_filter_values';
import { RollupInterval } from '../../../common/rollup';
import { ServiceHealthStatus } from '../../../common/service_health_status';
import type { APMError } from '../../../typings/es_schemas/ui/apm_error';
import { getApmAlertsClient } from '../../lib/helpers/get_apm_alerts_client';
import { getApmEventClient } from '../../lib/helpers/get_apm_event_client';
import { getMlClient } from '../../lib/helpers/get_ml_client';
import { getRandomSampler } from '../../lib/helpers/get_random_sampler';
import { createApmServerRoute } from '../apm_routes/create_apm_server_route';
import { getServicesItems } from '../services/get_services/get_services_items';
import {
CorrelationValue,
correlationValuesRouteRt,
@ -29,7 +21,6 @@ import {
getAssistantDownstreamDependencies,
type APMDownstreamDependency,
} from './get_apm_downstream_dependencies';
import { errorRouteRt, getApmErrorDocument } from './get_apm_error_document';
import {
getApmServiceSummary,
serviceSummaryRouteRt,
@ -167,130 +158,9 @@ const getApmCorrelationValuesRoute = createApmServerRoute({
},
});
const getApmErrorDocRoute = createApmServerRoute({
endpoint: 'GET /internal/apm/assistant/get_error_document',
params: t.type({
query: errorRouteRt,
}),
options: {
tags: ['access:apm'],
},
handler: async (
resources
): Promise<{ content: Array<Partial<APMError>> }> => {
const { params } = resources;
const apmEventClient = await getApmEventClient(resources);
const { query } = params;
return {
content: await getApmErrorDocument({
apmEventClient,
arguments: query,
}),
};
},
});
export interface ApmServicesListItem {
'service.name': string;
'agent.name'?: string;
'transaction.type'?: string;
alertsCount: number;
healthStatus: ServiceHealthStatus;
'service.environment'?: string[];
}
type ApmServicesListContent = ApmServicesListItem[];
const getApmServicesListRoute = createApmServerRoute({
endpoint: 'POST /internal/apm/assistant/get_services_list',
params: t.type({
body: t.intersection([
t.type({
start: t.string,
end: t.string,
}),
t.partial({
'service.environment': t.string,
healthStatus: t.array(
t.union([
t.literal(ServiceHealthStatus.unknown),
t.literal(ServiceHealthStatus.healthy),
t.literal(ServiceHealthStatus.warning),
t.literal(ServiceHealthStatus.critical),
])
),
}),
]),
}),
options: {
tags: ['access:apm'],
},
handler: async (resources): Promise<{ content: ApmServicesListContent }> => {
const { params } = resources;
const { body } = params;
const { healthStatus } = body;
const [apmEventClient, apmAlertsClient, mlClient, randomSampler] =
await Promise.all([
getApmEventClient(resources),
getApmAlertsClient(resources),
getMlClient(resources),
getRandomSampler({
security: resources.plugins.security,
probability: 1,
request: resources.request,
}),
]);
const start = datemath.parse(body.start)?.valueOf()!;
const end = datemath.parse(body.end)?.valueOf()!;
const serviceItems = await getServicesItems({
apmAlertsClient,
apmEventClient,
documentType: ApmDocumentType.TransactionMetric,
start,
end,
environment: body['service.environment'] || ENVIRONMENT_ALL.value,
kuery: '',
logger: resources.logger,
randomSampler,
rollupInterval: RollupInterval.OneMinute,
serviceGroup: null,
mlClient,
useDurationSummary: false,
});
let mappedItems = serviceItems.items.map((item): ApmServicesListItem => {
return {
'service.name': item.serviceName,
'agent.name': item.agentName,
alertsCount: item.alertsCount ?? 0,
healthStatus: item.healthStatus ?? ServiceHealthStatus.unknown,
'service.environment': item.environments,
'transaction.type': item.transactionType,
};
});
if (healthStatus && healthStatus.length) {
mappedItems = mappedItems.filter((item): boolean =>
healthStatus.includes(item.healthStatus)
);
}
return {
content: mappedItems,
};
},
});
export const assistantRouteRepository = {
...getApmTimeSeriesRoute,
...getApmServiceSummaryRoute,
...getApmErrorDocRoute,
...getApmCorrelationValuesRoute,
...getDownstreamDependenciesRoute,
...getApmServicesListRoute,
};

8
x-pack/plugins/apm/server/types.ts
View file

@ -65,6 +65,10 @@ import {
ProfilingDataAccessPluginSetup,
ProfilingDataAccessPluginStart,
} from '@kbn/profiling-data-access-plugin/server';
import type {
ObservabilityAIAssistantPluginSetup,
ObservabilityAIAssistantPluginStart,
} from '@kbn/observability-ai-assistant-plugin/server';
import { APMConfig } from '.';
export interface APMPluginSetup {
@ -82,7 +86,7 @@ export interface APMPluginSetupDependencies {
metricsDataAccess: MetricsDataPluginSetup;
dataViews: {};
share: SharePluginSetup;
observabilityAIAssistant: ObservabilityAIAssistantPluginSetup;
// optional dependencies
actions?: ActionsPlugin['setup'];
alerting?: AlertingPlugin['setup'];
@ -108,7 +112,7 @@ export interface APMPluginStartDependencies {
metricsDataAccess: MetricsDataPluginSetup;
dataViews: DataViewsServerPluginStart;
share: undefined;
observabilityAIAssistant: ObservabilityAIAssistantPluginStart;
// optional dependencies
actions?: ActionsPlugin['start'];
alerting?: AlertingPlugin['start'];

0
x-pack/plugins/apm/public/utils/non_empty_string_ref.ts → x-pack/plugins/apm/server/utils/non_empty_string_ref.ts
View file

3
x-pack/plugins/apm/tsconfig.json
View file

@ -106,7 +106,8 @@
"@kbn/custom-icons",
"@kbn/elastic-agent-utils",
"@kbn/shared-ux-link-redirect-app",
"@kbn/observability-get-padded-alert-time-range-util"
"@kbn/observability-get-padded-alert-time-range-util",
"@kbn/core-lifecycle-server"
],
"exclude": ["target/**/*"]
}

109
x-pack/plugins/observability_ai_assistant/common/conversation_complete.ts Normal file
View file

@ -0,0 +1,109 @@
/*
* 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.
*/
/* eslint-disable max-classes-per-file*/
import { i18n } from '@kbn/i18n';
import { Message } from './types';
export enum StreamingChatResponseEventType {
ChatCompletionChunk = 'chatCompletionChunk',
ConversationCreate = 'conversationCreate',
ConversationUpdate = 'conversationUpdate',
MessageAdd = 'messageAdd',
ConversationCompletionError = 'conversationCompletionError',
}
type StreamingChatResponseEventBase<
TEventType extends StreamingChatResponseEventType,
TData extends {}
> = {
type: TEventType;
} & TData;
type ChatCompletionChunkEvent = StreamingChatResponseEventBase<
StreamingChatResponseEventType.ChatCompletionChunk,
{
id: string;
message: {
content?: string;
function_call?: {
name?: string;
arguments?: string;
};
};
}
>;
export type ConversationCreateEvent = StreamingChatResponseEventBase<
StreamingChatResponseEventType.ConversationCreate,
{
conversation: {
id: string;
title: string;
last_updated: string;
};
}
>;
export type ConversationUpdateEvent = StreamingChatResponseEventBase<
StreamingChatResponseEventType.ConversationUpdate,
{
conversation: {
id: string;
title: string;
last_updated: string;
};
}
>;
export type MessageAddEvent = StreamingChatResponseEventBase<
StreamingChatResponseEventType.MessageAdd,
{ message: Message; id: string }
>;
export type ConversationCompletionErrorEvent = StreamingChatResponseEventBase<
StreamingChatResponseEventType.ConversationCompletionError,
{ error: { message: string; stack?: string; code?: ChatCompletionErrorCode } }
>;
export type StreamingChatResponseEvent =
| ChatCompletionChunkEvent
| ConversationCreateEvent
| ConversationUpdateEvent
| MessageAddEvent
| ConversationCompletionErrorEvent;
export enum ChatCompletionErrorCode {
InternalError = 'internalError',
NotFound = 'notFound',
}
export class ConversationCompletionError extends Error {
code: ChatCompletionErrorCode;
constructor(code: ChatCompletionErrorCode, message: string) {
super(message);
this.code = code;
}
}
export class ConversationNotFoundError extends ConversationCompletionError {
constructor() {
super(
ChatCompletionErrorCode.NotFound,
i18n.translate(
'xpack.observabilityAiAssistant.conversationCompletionError.conversationNotFound',
{
defaultMessage: 'Conversation not found',
}
)
);
}
}
export function isChatCompletionError(error: Error): error is ConversationCompletionError {
return error instanceof ConversationCompletionError;
}

122
x-pack/plugins/observability_ai_assistant/common/functions/lens.tsx Normal file
View file

@ -0,0 +1,122 @@
/*
* 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 { FromSchema } from 'json-schema-to-ts';
import { FIELD_FORMAT_IDS } from '@kbn/field-formats-plugin/common';
export enum SeriesType {
Bar = 'bar',
Line = 'line',
Area = 'area',
BarStacked = 'bar_stacked',
AreaStacked = 'area_stacked',
BarHorizontal = 'bar_horizontal',
BarPercentageStacked = 'bar_percentage_stacked',
AreaPercentageStacked = 'area_percentage_stacked',
BarHorizontalPercentageStacked = 'bar_horizontal_percentage_stacked',
}
export const lensFunctionDefinition = {
name: 'lens',
contexts: ['core'],
description:
"Use this function to create custom visualizations, using Lens, that can be saved to dashboards. This function does not return data to the assistant, it only shows it to the user. When using this function, make sure to use the recall function to get more information about how to use it, with how you want to use it. Make sure the query also contains information about the user's request. The visualisation is displayed to the user above your reply, DO NOT try to generate or display an image yourself.",
descriptionForUser:
'Use this function to create custom visualizations, using Lens, that can be saved to dashboards.',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
layers: {
type: 'array',
items: {
type: 'object',
additionalProperties: false,
properties: {
label: {
type: 'string',
},
formula: {
type: 'string',
description:
'The formula for calculating the value, e.g. sum(my_field_name). Query the knowledge base to get more information about the syntax and available formulas.',
},
filter: {
type: 'string',
description: 'A KQL query that will be used as a filter for the series',
},
format: {
type: 'object',
additionalProperties: false,
properties: {
id: {
type: 'string',
description:
'How to format the value. When using duration, make sure the value is seconds OR is converted to seconds using math functions. Ask the user for clarification in which unit the value is stored, or derive it from the field name.',
enum: [
FIELD_FORMAT_IDS.BYTES,
FIELD_FORMAT_IDS.CURRENCY,
FIELD_FORMAT_IDS.DURATION,
FIELD_FORMAT_IDS.NUMBER,
FIELD_FORMAT_IDS.PERCENT,
FIELD_FORMAT_IDS.STRING,
],
},
},
required: ['id'],
},
},
required: ['label', 'formula', 'format'],
},
},
timeField: {
type: 'string',
default: '@timefield',
description:
'time field to use for XY chart. Use @timefield if its available on the index.',
},
breakdown: {
type: 'object',
additionalProperties: false,
properties: {
field: {
type: 'string',
},
},
required: ['field'],
},
indexPattern: {
type: 'string',
},
seriesType: {
type: 'string',
enum: [
SeriesType.Area,
SeriesType.AreaPercentageStacked,
SeriesType.AreaStacked,
SeriesType.Bar,
SeriesType.BarHorizontal,
SeriesType.BarHorizontalPercentageStacked,
SeriesType.BarPercentageStacked,
SeriesType.BarStacked,
SeriesType.Line,
],
},
start: {
type: 'string',
description: 'The start of the time range, in Elasticsearch datemath',
},
end: {
type: 'string',
description: 'The end of the time range, in Elasticsearch datemath',
},
},
required: ['layers', 'indexPattern', 'start', 'end', 'timeField'],
} as const,
};
export type LensFunctionArguments = FromSchema<typeof lensFunctionDefinition['parameters']>;

53
x-pack/plugins/observability_ai_assistant/common/types.ts
View file

@ -5,10 +5,20 @@
* 2.0.
*/
import type { FromSchema } from 'json-schema-to-ts';
import type { JSONSchema } from 'json-schema-to-ts';
import React from 'react';
import { Observable } from 'rxjs';
import type {
CreateChatCompletionResponse,
CreateChatCompletionResponseChoicesInner,
} from 'openai';
import type { Observable } from 'rxjs';
export type CreateChatCompletionResponseChunk = Omit<CreateChatCompletionResponse, 'choices'> & {
choices: Array<
Omit<CreateChatCompletionResponseChoicesInner, 'message'> & {
delta: { content?: string; function_call?: { name?: string; arguments?: string } };
}
>;
};
export enum MessageRole {
System = 'system',
@ -89,12 +99,12 @@ export interface ContextDefinition {
description: string;
}
type FunctionResponse =
export type FunctionResponse =
| {
content?: any;
data?: any;
}
| Observable<PendingMessage>;
| Observable<CreateChatCompletionResponseChunk>;
export enum FunctionVisibility {
System = 'system',
@ -102,7 +112,9 @@ export enum FunctionVisibility {
All = 'all',
}
interface FunctionOptions<TParameters extends CompatibleJSONSchema = CompatibleJSONSchema> {
export interface FunctionDefinition<
TParameters extends CompatibleJSONSchema = CompatibleJSONSchema
> {
name: string;
description: string;
visibility?: FunctionVisibility;
@ -111,36 +123,7 @@ interface FunctionOptions<TParameters extends CompatibleJSONSchema = CompatibleJ
contexts: string[];
}
type RespondFunction<TArguments, TResponse extends FunctionResponse> = (
options: { arguments: TArguments; messages: Message[]; connectorId: string },
signal: AbortSignal
) => Promise<TResponse>;
type RenderFunction<TArguments, TResponse extends FunctionResponse> = (options: {
arguments: TArguments;
response: TResponse;
}) => React.ReactNode;
export interface FunctionDefinition {
options: FunctionOptions;
respond: (
options: { arguments: any; messages: Message[]; connectorId: string },
signal: AbortSignal
) => Promise<FunctionResponse>;
render?: RenderFunction<any, any>;
}
export type RegisterContextDefinition = (options: ContextDefinition) => void;
export type RegisterFunctionDefinition = <
TParameters extends CompatibleJSONSchema,
TResponse extends FunctionResponse,
TArguments = FromSchema<TParameters>
>(
options: FunctionOptions<TParameters>,
respond: RespondFunction<TArguments, TResponse>,
render?: RenderFunction<TArguments, TResponse>
) => void;
export type ContextRegistry = Map<string, ContextDefinition>;
export type FunctionRegistry = Map<string, FunctionDefinition>;

34
x-pack/plugins/observability_ai_assistant/common/utils/concatenate_openai_chunks.ts Normal file
View file

@ -0,0 +1,34 @@
/*
* 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 { cloneDeep } from 'lodash';
import { type Observable, scan } from 'rxjs';
import { CreateChatCompletionResponseChunk, MessageRole } from '../types';
export const concatenateOpenAiChunks =
() => (source: Observable<CreateChatCompletionResponseChunk>) =>
source.pipe(
scan(
(acc, { choices }) => {
acc.message.content += choices[0].delta.content ?? '';
acc.message.function_call.name += choices[0].delta.function_call?.name ?? '';
acc.message.function_call.arguments += choices[0].delta.function_call?.arguments ?? '';
return cloneDeep(acc);
},
{
message: {
content: '',
function_call: {
name: '',
arguments: '',
trigger: MessageRole.Assistant as const,
},
role: MessageRole.Assistant,
},
}
)
);

29
x-pack/plugins/observability_ai_assistant/common/utils/filter_function_definitions.ts Normal file
View file

@ -0,0 +1,29 @@
/*
* 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 { FunctionDefinition } from '../types';
export function filterFunctionDefinitions({
contexts,
filter,
definitions,
}: {
contexts?: string[];
filter?: string;
definitions: FunctionDefinition[];
}) {
return contexts || filter
? definitions.filter((fn) => {
const matchesContext =
!contexts || fn.contexts.some((context) => contexts.includes(context));
const matchesFilter =
!filter || fn.name.includes(filter) || fn.description.includes(filter);
return matchesContext && matchesFilter;
})
: definitions;
}

45
x-pack/plugins/observability_ai_assistant/common/utils/process_openai_stream.ts Normal file
View file

@ -0,0 +1,45 @@
/*
* 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.
*/
/* eslint-disable max-classes-per-file*/
import { filter, map, Observable, tap } from 'rxjs';
import type { CreateChatCompletionResponseChunk } from '../types';
class TokenLimitReachedError extends Error {
constructor() {
super(`Token limit reached`);
}
}
class ServerError extends Error {}
export function processOpenAiStream() {
return (source: Observable<string>): Observable<CreateChatCompletionResponseChunk> =>
source.pipe(
map((line) => line.substring(6)),
filter((line) => !!line && line !== '[DONE]'),
map(
(line) =>
JSON.parse(line) as CreateChatCompletionResponseChunk | { error: { message: string } }
),
tap((line) => {
if ('error' in line) {
throw new ServerError(line.error.message);
}
if (
'choices' in line &&
line.choices.length &&
line.choices[0].finish_reason === 'length'
) {
throw new TokenLimitReachedError();
}
}),
filter(
(line): line is CreateChatCompletionResponseChunk =>
'object' in line && line.object === 'chat.completion.chunk'
)
);
}

3
x-pack/plugins/observability_ai_assistant/jest.config.js
View file

@ -11,5 +11,8 @@ module.exports = {
roots: ['<rootDir>/x-pack/plugins/observability_ai_assistant'],
setupFiles: ['<rootDir>/x-pack/plugins/observability_ai_assistant/.storybook/jest_setup.js'],
collectCoverage: true,
collectCoverageFrom: [
'<rootDir>/x-pack/plugins/observability_ai_assistant/{common,public,server}/**/*.{js,ts,tsx}',
],
coverageReporters: ['html'],
};

3
x-pack/plugins/observability_ai_assistant/kibana.jsonc
View file

@ -8,6 +8,7 @@
"browser": true,
"configPath": ["xpack", "observabilityAIAssistant"],
"requiredPlugins": [
"alerting",
"actions",
"dataViews",
"features",
@ -21,7 +22,7 @@
"triggersActionsUi",
"dataViews"
],
"requiredBundles": ["fieldFormats", "kibanaReact", "kibanaUtils"],
"requiredBundles": [ "kibanaReact", "kibanaUtils"],
"optionalPlugins": [],
"extraPublicDirs": []
}

2
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx
View file

@ -73,7 +73,7 @@ export function ChatBody({
connectorsManagementHref: string;
currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
startedFrom?: StartedFrom;
onConversationUpdate: (conversation: Conversation) => void;
onConversationUpdate: (conversation: { conversation: Conversation['conversation'] }) => void;
}) {
const license = useLicense();
const hasCorrectLicense = license?.hasAtLeast('enterprise');

8
x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
View file

@ -186,12 +186,12 @@ function mapFunctions({
selectedFunctionName: string | undefined;
}) {
return functions
.filter((func) => func.options.visibility !== FunctionVisibility.System)
.filter((func) => func.visibility !== FunctionVisibility.System)
.map((func) => ({
label: func.options.name,
searchableLabel: func.options.descriptionForUser || func.options.description,
label: func.name,
searchableLabel: func.descriptionForUser || func.description,
checked:
func.options.name === selectedFunctionName
func.name === selectedFunctionName
? ('on' as EuiSelectableOptionCheckedType)
: ('off' as EuiSelectableOptionCheckedType),
}));

1
x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx
View file

@ -45,6 +45,7 @@ function ChatContent({
chatService,
connectorId,
initialMessages,
persist: false,
});
const lastAssistantResponse = last(

86
x-pack/plugins/observability_ai_assistant/public/functions/alerts.ts
View file

@ -1,86 +0,0 @@
/*
* 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 { RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
const DEFAULT_FEATURE_IDS = [
'apm',
'infrastructure',
'logs',
'uptime',
'slo',
'observability',
] as const;
export function registerAlertsFunction({
service,
registerFunction,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
}) {
registerFunction(
{
name: 'alerts',
contexts: ['core'],
description:
'Get alerts for Observability. Display the response in tabular format if appropriate.',
descriptionForUser: 'Get alerts for Observability',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
featureIds: {
type: 'array',
additionalItems: false,
items: {
type: 'string',
enum: DEFAULT_FEATURE_IDS,
},
description:
'The Observability apps for which to retrieve alerts. By default it will return alerts for all apps.',
},
start: {
type: 'string',
description: 'The start of the time range, in Elasticsearch date math, like `now`.',
},
end: {
type: 'string',
description: 'The end of the time range, in Elasticsearch date math, like `now-24h`.',
},
filter: {
type: 'string',
description:
'a KQL query to filter the data by. If no filter should be applied, leave it empty.',
},
includeRecovered: {
type: 'boolean',
description:
'Whether to include recovered/closed alerts. Defaults to false, which means only active alerts will be returned',
},
},
required: ['start', 'end'],
} as const,
},
({ arguments: { start, end, featureIds, filter, includeRecovered } }, signal) => {
return service.callApi('POST /internal/observability_ai_assistant/functions/alerts', {
params: {
body: {
start,
end,
featureIds:
featureIds && featureIds.length > 0 ? featureIds : DEFAULT_FEATURE_IDS.concat(),
filter,
includeRecovered,
},
},
signal,
});
}
);
}

153
x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
View file

@ -1,153 +0,0 @@
/*
* 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 { chunk, groupBy, uniq } from 'lodash';
import { CreateChatCompletionResponse } from 'openai';
import { FunctionVisibility, MessageRole, RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
export function registerGetDatasetInfoFunction({
service,
registerFunction,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
}) {
registerFunction(
{
name: 'get_dataset_info',
contexts: ['core'],
visibility: FunctionVisibility.System,
description: `Use this function to get information about indices/datasets available and the fields available on them.
providing empty string as index name will retrieve all indices
else list of all fields for the given index will be given. if no fields are returned this means no indices were matched by provided index pattern.
wildcards can be part of index name.`,
descriptionForUser:
'This function allows the assistant to get information about available indices and their fields.',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
index: {
type: 'string',
description:
'index pattern the user is interested in or empty string to get information about all available indices',
},
},
required: ['index'],
} as const,
},
async ({ arguments: { index }, messages, connectorId }, signal) => {
const response = await service.callApi(
'POST /internal/observability_ai_assistant/functions/get_dataset_info',
{
params: {
body: {
index,
},
},
signal,
}
);
const allFields = response.fields;
const fieldNames = uniq(allFields.map((field) => field.name));
const groupedFields = groupBy(allFields, (field) => field.name);
const relevantFields = await Promise.all(
chunk(fieldNames, 500).map(async (fieldsInChunk) => {
const chunkResponse = (await service.callApi(
'POST /internal/observability_ai_assistant/chat',
{
signal,
params: {
query: {
stream: false,
},
body: {
connectorId,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: `You are a helpful assistant for Elastic Observability.
Your task is to create a list of field names that are relevant
to the conversation, using ONLY the list of fields and
types provided in the last user message. DO NOT UNDER ANY
CIRCUMSTANCES include fields not mentioned in this list.`,
},
},
...messages.slice(1),
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: `This is the list:
${fieldsInChunk.join('\n')}`,
},
},
],
functions: [
{
name: 'fields',
description: 'The fields you consider relevant to the conversation',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
fields: {
type: 'array',
additionalProperties: false,
addditionalItems: false,
items: {
type: 'string',
additionalProperties: false,
addditionalItems: false,
},
},
},
required: ['fields'],
},
},
],
functionCall: 'fields',
},
},
}
)) as CreateChatCompletionResponse;
return chunkResponse.choices[0].message?.function_call?.arguments
? (
JSON.parse(chunkResponse.choices[0].message?.function_call?.arguments) as {
fields: string[];
}
).fields
.filter((field) => fieldNames.includes(field))
.map((field) => {
const fieldDescriptors = groupedFields[field];
return `${field}:${fieldDescriptors
.map((descriptor) => descriptor.type)
.join(',')}`;
})
: [chunkResponse.choices[0].message?.content ?? ''];
})
);
return {
content: {
indices: response.indices,
fields: relevantFields.flat(),
},
};
}
);
}

85
x-pack/plugins/observability_ai_assistant/public/functions/index.ts
View file

@ -5,88 +5,21 @@
* 2.0.
*/
import dedent from 'dedent';
import type { CoreStart } from '@kbn/core/public';
import type { RegisterContextDefinition, RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantPluginStartDependencies } from '../types';
import type { ObservabilityAIAssistantService } from '../types';
import { registerElasticsearchFunction } from './elasticsearch';
import { registerKibanaFunction } from './kibana';
import { registerLensFunction } from './lens';
import { registerRecallFunction } from './recall';
import { registerGetDatasetInfoFunction } from './get_dataset_info';
import { registerSummarizationFunction } from './summarize';
import { registerAlertsFunction } from './alerts';
import { registerEsqlFunction } from './esql';
import type {
ObservabilityAIAssistantPluginStartDependencies,
ObservabilityAIAssistantService,
RegisterRenderFunctionDefinition,
} from '../types';
import { registerLensRenderFunction } from './lens';
export async function registerFunctions({
registerFunction,
registerContext,
registerRenderFunction,
service,
pluginsStart,
coreStart,
signal,
}: {
registerFunction: RegisterFunctionDefinition;
registerContext: RegisterContextDefinition;
registerRenderFunction: RegisterRenderFunctionDefinition;
service: ObservabilityAIAssistantService;
pluginsStart: ObservabilityAIAssistantPluginStartDependencies;
coreStart: CoreStart;
signal: AbortSignal;
}) {
return service
.callApi('GET /internal/observability_ai_assistant/kb/status', {
signal,
})
.then((response) => {
const isReady = response.ready;
let description = dedent(
`You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities.
It's very important to not assume what the user is meaning. Ask them for clarification if needed.
If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
In KQL, escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
/\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
If multiple functions are suitable, use the most specific and easy one. E.g., when the user asks to visualise APM data, use the APM functions (if available) rather than Lens.
If a function call fails, DO NOT UNDER ANY CIRCUMSTANCES execute it again. Ask the user for guidance and offer them options.
Note that ES|QL (the Elasticsearch query language, which is NOT Elasticsearch SQL, but a new piped language) is the preferred query language.
If the user asks about a query, or ES|QL, always call the "esql" function. DO NOT UNDER ANY CIRCUMSTANCES generate ES|QL queries yourself. Even if the "recall" function was used before that, follow it up with the "esql" function.`
);
if (isReady) {
description += `You can use the "summarize" functions to store new information you have learned in a knowledge database. Once you have established that you did not know the answer to a question, and the user gave you this information, it's important that you create a summarisation of what you have learned and store it in the knowledge database. Don't create a new summarization if you see a similar summarization in the conversation, instead, update the existing one by re-using its ID.
Additionally, you can use the "recall" function to retrieve relevant information from the knowledge database.
`;
description += `Here are principles you MUST adhere to, in order:
- DO NOT make any assumptions about where and how users have stored their data. ALWAYS first call get_dataset_info function with empty string to get information about available indices. Once you know about available indices you MUST use this function again to get a list of available fields for specific index. If user provides an index name make sure its a valid index first before using it to retrieve the field list by calling this function with an empty string!
`;
registerSummarizationFunction({ service, registerFunction });
registerRecallFunction({ service, registerFunction });
registerLensFunction({ service, pluginsStart, registerFunction });
} else {
description += `You do not have a working memory. Don't try to recall information via the "recall" function. If the user expects you to remember the previous conversations, tell them they can set up the knowledge base. A banner is available at the top of the conversation to set this up.`;
}
registerElasticsearchFunction({ service, registerFunction });
registerEsqlFunction({ service, registerFunction });
registerKibanaFunction({ service, registerFunction, coreStart });
registerAlertsFunction({ service, registerFunction });
registerGetDatasetInfoFunction({ service, registerFunction });
registerContext({
name: 'core',
description: dedent(description),
});
});
registerLensRenderFunction({ service, pluginsStart, registerRenderFunction });
}

123
x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
View file

@ -6,17 +6,18 @@
*/
import { EuiButton, EuiFlexGroup, EuiFlexItem, EuiLoadingSpinner } from '@elastic/eui';
import type { DataViewsServicePublic } from '@kbn/data-views-plugin/public/types';
import { FIELD_FORMAT_IDS } from '@kbn/field-formats-plugin/common';
import { i18n } from '@kbn/i18n';
import { LensAttributesBuilder, XYChart, XYDataLayer } from '@kbn/lens-embeddable-utils';
import type { LensEmbeddableInput, LensPublicStart } from '@kbn/lens-plugin/public';
import React, { useState } from 'react';
import useAsync from 'react-use/lib/useAsync';
import { i18n } from '@kbn/i18n';
import { Assign } from 'utility-types';
import type { RegisterFunctionDefinition } from '../../common/types';
import type { LensFunctionArguments } from '../../common/functions/lens';
import type {
ObservabilityAIAssistantPluginStartDependencies,
ObservabilityAIAssistantService,
RegisterRenderFunctionDefinition,
RenderFunction,
} from '../types';
export enum SeriesType {
@ -137,120 +138,20 @@ function Lens({
);
}
export function registerLensFunction({
export function registerLensRenderFunction({
service,
registerFunction,
registerRenderFunction,
pluginsStart,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
registerRenderFunction: RegisterRenderFunctionDefinition;
pluginsStart: ObservabilityAIAssistantPluginStartDependencies;
}) {
registerFunction(
{
name: 'lens',
contexts: ['core'],
description:
"Use this function to create custom visualizations, using Lens, that can be saved to dashboards. This function does not return data to the assistant, it only shows it to the user. When using this function, make sure to use the recall function to get more information about how to use it, with how you want to use it. Make sure the query also contains information about the user's request. The visualisation is displayed to the user above your reply, DO NOT try to generate or display an image yourself.",
descriptionForUser:
'Use this function to create custom visualizations, using Lens, that can be saved to dashboards.',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
layers: {
type: 'array',
items: {
type: 'object',
additionalProperties: false,
properties: {
label: {
type: 'string',
},
formula: {
type: 'string',
description:
'The formula for calculating the value, e.g. sum(my_field_name). Query the knowledge base to get more information about the syntax and available formulas.',
},
filter: {
type: 'string',
description: 'A KQL query that will be used as a filter for the series',
},
format: {
type: 'object',
additionalProperties: false,
properties: {
id: {
type: 'string',
description:
'How to format the value. When using duration, make sure the value is seconds OR is converted to seconds using math functions. Ask the user for clarification in which unit the value is stored, or derive it from the field name.',
enum: [
FIELD_FORMAT_IDS.BYTES,
FIELD_FORMAT_IDS.CURRENCY,
FIELD_FORMAT_IDS.DURATION,
FIELD_FORMAT_IDS.NUMBER,
FIELD_FORMAT_IDS.PERCENT,
FIELD_FORMAT_IDS.STRING,
],
},
},
required: ['id'],
},
},
required: ['label', 'formula', 'format'],
},
},
timeField: {
type: 'string',
default: '@timefield',
description:
'time field to use for XY chart. Use @timefield if its available on the index.',
},
breakdown: {
type: 'object',
additionalProperties: false,
properties: {
field: {
type: 'string',
},
},
required: ['field'],
},
indexPattern: {
type: 'string',
},
seriesType: {
type: 'string',
enum: [
SeriesType.Area,
SeriesType.AreaPercentageStacked,
SeriesType.AreaStacked,
SeriesType.Bar,
SeriesType.BarHorizontal,
SeriesType.BarHorizontalPercentageStacked,
SeriesType.BarPercentageStacked,
SeriesType.BarStacked,
SeriesType.Line,
],
},
start: {
type: 'string',
description: 'The start of the time range, in Elasticsearch datemath',
},
end: {
type: 'string',
description: 'The end of the time range, in Elasticsearch datemath',
},
},
required: ['layers', 'indexPattern', 'start', 'end', 'timeField'],
} as const,
},
async () => {
return {
content: {},
};
},
({ arguments: { layers, indexPattern, breakdown, seriesType, start, end, timeField } }) => {
registerRenderFunction(
'lens',
({
arguments: { layers, indexPattern, breakdown, seriesType, start, end, timeField },
}: Parameters<RenderFunction<LensFunctionArguments, {}>>[0]) => {
const xyDataLayer = new XYDataLayer({
data: layers.map((layer) => ({
type: 'formula',

311
x-pack/plugins/observability_ai_assistant/public/hooks/use_chat.test.ts
View file

@ -8,20 +8,26 @@ import type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import { type RenderHookResult, renderHook, act } from '@testing-library/react-hooks';
import { Observable, Subject } from 'rxjs';
import { MessageRole } from '../../common';
import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
import {
ChatCompletionErrorCode,
ConversationCompletionError,
StreamingChatResponseEvent,
StreamingChatResponseEventType,
} from '../../common/conversation_complete';
import type { ObservabilityAIAssistantChatService } from '../types';
import { type UseChatResult, useChat, type UseChatProps, ChatState } from './use_chat';
import * as useKibanaModule from './use_kibana';
type MockedChatService = DeeplyMockedKeys<ObservabilityAIAssistantChatService>;
const mockChatService: MockedChatService = {
chat: jest.fn(),
complete: jest.fn(),
analytics: {
optIn: jest.fn(),
reportEvent: jest.fn(),
telemetryCounter$: new Observable() as any,
},
chat: jest.fn(),
executeFunction: jest.fn(),
getContexts: jest.fn().mockReturnValue([{ name: 'core', description: '' }]),
getFunctions: jest.fn().mockReturnValue([]),
hasFunction: jest.fn().mockReturnValue(false),
@ -63,6 +69,7 @@ describe('useChat', () => {
},
},
],
persist: false,
} as UseChatProps,
});
});
@ -80,7 +87,7 @@ describe('useChat', () => {
});
describe('when calling next()', () => {
let subject: Subject<PendingMessage>;
let subject: Subject<StreamingChatResponseEvent>;
beforeEach(() => {
hookResult = renderHook(useChat, {
@ -88,12 +95,13 @@ describe('useChat', () => {
connectorId: 'my-connector',
chatService: mockChatService,
initialMessages: [],
persist: false,
} as UseChatProps,
});
subject = new Subject();
mockChatService.chat.mockReturnValueOnce(subject);
mockChatService.complete.mockReturnValueOnce(subject);
act(() => {
hookResult.result.current.next([
@ -118,11 +126,23 @@ describe('useChat', () => {
act(() => {
hookResult.result.current.next([]);
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message-id',
message: {
role: MessageRole.User,
content: 'goodbye',
},
});
subject.next({
type: StreamingChatResponseEventType.MessageAdd,
id: 'my-message-id',
message: {
'@timestamp': new Date().toISOString(),
message: {
content: 'goodbye',
role: MessageRole.Assistant,
},
},
});
subject.complete();
});
});
@ -141,9 +161,10 @@ describe('useChat', () => {
it('updates the returned messages', () => {
act(() => {
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message-id',
message: {
content: 'good',
role: MessageRole.Assistant,
},
});
});
@ -156,15 +177,28 @@ describe('useChat', () => {
it('updates the returned messages and the loading state', () => {
act(() => {
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message-id',
message: {
content: 'good',
role: MessageRole.Assistant,
},
});
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message-id',
message: {
content: 'goodbye',
role: MessageRole.Assistant,
content: 'bye',
},
});
subject.next({
type: StreamingChatResponseEventType.MessageAdd,
id: 'my-message-id',
message: {
'@timestamp': new Date().toISOString(),
message: {
content: 'goodbye',
role: MessageRole.Assistant,
},
},
});
subject.complete();
@ -179,13 +213,13 @@ describe('useChat', () => {
beforeEach(() => {
act(() => {
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message-id',
message: {
content: 'good',
role: MessageRole.Assistant,
},
aborted: true,
});
subject.complete();
hookResult.result.current.stop();
});
});
@ -203,13 +237,15 @@ describe('useChat', () => {
beforeEach(() => {
act(() => {
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message-id',
message: {
content: 'good',
role: MessageRole.Assistant,
},
error: new Error('foo'),
});
subject.complete();
subject.error(
new ConversationCompletionError(ChatCompletionErrorCode.InternalError, 'foo')
);
});
});
@ -222,248 +258,5 @@ describe('useChat', () => {
expect(addErrorMock).toHaveBeenCalled();
});
});
describe('after the LLM responds with a function call', () => {
let resolve: (data: any) => void;
let reject: (error: Error) => void;
beforeEach(() => {
mockChatService.executeFunction.mockResolvedValueOnce(
new Promise((...args) => {
resolve = args[0];
reject = args[1];
})
);
act(() => {
subject.next({
message: {
content: '',
role: MessageRole.Assistant,
function_call: {
name: 'my_function',
arguments: JSON.stringify({ foo: 'bar' }),
trigger: MessageRole.Assistant,
},
},
});
subject.complete();
});
});
it('the chat state stays loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
it('adds a message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(3);
expect(messages[2]).toEqual({
'@timestamp': expect.any(String),
message: {
content: '',
function_call: {
arguments: JSON.stringify({ foo: 'bar' }),
name: 'my_function',
trigger: MessageRole.Assistant,
},
role: MessageRole.Assistant,
},
});
});
describe('the function call succeeds', () => {
beforeEach(async () => {
subject = new Subject();
mockChatService.chat.mockReturnValueOnce(subject);
await act(async () => {
resolve({ content: { foo: 'bar' }, data: { bar: 'foo' } });
});
});
it('adds a message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(4);
expect(messages[3]).toEqual({
'@timestamp': expect.any(String),
message: {
content: JSON.stringify({ foo: 'bar' }),
data: JSON.stringify({ bar: 'foo' }),
name: 'my_function',
role: MessageRole.User,
},
});
});
it('keeps the chat state in loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
it('sends the function call back to the LLM for a response', () => {
expect(mockChatService.chat).toHaveBeenCalledTimes(2);
expect(mockChatService.chat).toHaveBeenLastCalledWith({
connectorId: 'my-connector',
messages: hookResult.result.current.messages,
});
});
});
describe('the function call fails', () => {
beforeEach(async () => {
subject = new Subject();
mockChatService.chat.mockReturnValue(subject);
await act(async () => {
reject(new Error('connection error'));
});
});
it('keeps the chat state in loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
it('adds a message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(4);
expect(messages[3]).toEqual({
'@timestamp': expect.any(String),
message: {
content: JSON.stringify({
message: 'Error: connection error',
error: {},
}),
name: 'my_function',
role: MessageRole.User,
},
});
});
it('does not show an error toast', () => {
expect(addErrorMock).not.toHaveBeenCalled();
});
it('sends the function call back to the LLM for a response', () => {
expect(mockChatService.chat).toHaveBeenCalledTimes(2);
expect(mockChatService.chat).toHaveBeenLastCalledWith({
connectorId: 'my-connector',
messages: hookResult.result.current.messages,
});
});
});
describe('stop() is called', () => {
beforeEach(() => {
act(() => {
hookResult.result.current.stop();
});
});
it('sets the chatState to aborted', () => {
expect(hookResult.result.current.state).toBe(ChatState.Aborted);
});
it('has called the abort controller', () => {
const signal = mockChatService.executeFunction.mock.calls[0][0].signal;
expect(signal.aborted).toBe(true);
});
it('is not updated after the promise is rejected', () => {
const numRenders = hookResult.result.all.length;
act(() => {
reject(new Error('Request aborted'));
});
expect(numRenders).toBe(hookResult.result.all.length);
});
it('removes all subscribers', () => {
expect(subject.observed).toBe(false);
});
});
describe('setMessages() is called', () => {});
});
});
describe('when calling next() with the recall function available', () => {
let subject: Subject<PendingMessage>;
beforeEach(async () => {
hookResult = renderHook(useChat, {
initialProps: {
connectorId: 'my-connector',
chatService: mockChatService,
initialMessages: [],
} as UseChatProps,
});
subject = new Subject();
mockChatService.hasFunction.mockReturnValue(true);
mockChatService.executeFunction.mockResolvedValueOnce({
content: [
{
id: 'my_document',
text: 'My text',
},
],
});
mockChatService.chat.mockReturnValueOnce(subject);
await act(async () => {
hookResult.result.current.next([
...hookResult.result.current.messages,
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'hello',
},
},
]);
});
});
it('adds a user message and a recall function request', () => {
expect(hookResult.result.current.messages[1].message.content).toBe('hello');
expect(hookResult.result.current.messages[2].message.function_call?.name).toBe('recall');
expect(hookResult.result.current.messages[2].message.content).toBe('');
expect(hookResult.result.current.messages[2].message.function_call?.arguments).toBe(
JSON.stringify({ queries: [], contexts: [] })
);
expect(hookResult.result.current.messages[3].message.name).toBe('recall');
expect(hookResult.result.current.messages[3].message.content).toBe(
JSON.stringify([
{
id: 'my_document',
text: 'My text',
},
])
);
});
it('executes the recall function', () => {
expect(mockChatService.executeFunction).toHaveBeenCalled();
expect(mockChatService.executeFunction).toHaveBeenCalledWith({
signal: expect.any(AbortSignal),
connectorId: 'my-connector',
args: JSON.stringify({ queries: [], contexts: [] }),
name: 'recall',
messages: [...hookResult.result.current.messages.slice(0, -1)],
});
});
it('sends the user message, function request and recall response to the LLM', () => {
expect(mockChatService.chat).toHaveBeenCalledWith({
connectorId: 'my-connector',
messages: [...hookResult.result.current.messages],
});
});
});
});

286
x-pack/plugins/observability_ai_assistant/public/hooks/use_chat.ts
View file

@ -6,12 +6,16 @@
*/
import { i18n } from '@kbn/i18n';
import { last } from 'lodash';
import { merge } from 'lodash';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { isObservable } from 'rxjs';
import { type Message, MessageRole } from '../../common';
import { MessageRole, type Message } from '../../common';
import {
ConversationCreateEvent,
ConversationUpdateEvent,
StreamingChatResponseEventType,
} from '../../common/conversation_complete';
import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
import type { ObservabilityAIAssistantChatService } from '../types';
import { useKibana } from './use_kibana';
import { useOnce } from './use_once';
@ -22,6 +26,13 @@ export enum ChatState {
Aborted = 'aborted',
}
function getWithSystemMessage(messages: Message[], systemMessage: Message) {
return [
systemMessage,
...messages.filter((message) => message.message.role !== MessageRole.System),
];
}
export interface UseChatResult {
messages: Message[];
setMessages: (messages: Message[]) => void;
@ -32,16 +43,22 @@ export interface UseChatResult {
export interface UseChatProps {
initialMessages: Message[];
initialConversationId?: string;
chatService: ObservabilityAIAssistantChatService;
connectorId?: string;
persist: boolean;
onConversationUpdate?: (event: ConversationCreateEvent | ConversationUpdateEvent) => void;
onChatComplete?: (messages: Message[]) => void;
}
export function useChat({
initialMessages,
initialConversationId: initialConversationIdFromProps,
chatService,
connectorId,
onConversationUpdate,
onChatComplete,
persist,
}: UseChatProps): UseChatResult {
const [chatState, setChatState] = useState(ChatState.Ready);
@ -51,9 +68,11 @@ export function useChat({
useOnce(initialMessages);
const initialConversationId = useOnce(initialConversationIdFromProps);
const [messages, setMessages] = useState<Message[]>(initialMessages);
const [pendingMessage, setPendingMessage] = useState<PendingMessage>();
const [pendingMessages, setPendingMessages] = useState<Message[]>();
const abortControllerRef = useRef(new AbortController());
@ -62,13 +81,27 @@ export function useChat({
} = useKibana();
const onChatCompleteRef = useRef(onChatComplete);
onChatCompleteRef.current = onChatComplete;
const onConversationUpdateRef = useRef(onConversationUpdate);
onConversationUpdateRef.current = onConversationUpdate;
const handleSignalAbort = useCallback(() => {
setChatState(ChatState.Aborted);
}, []);
const handleError = useCallback(
(error: Error) => {
notifications.toasts.addError(error, {
title: i18n.translate('xpack.observabilityAiAssistant.failedToLoadResponse', {
defaultMessage: 'Failed to load response from the AI Assistant',
}),
});
setChatState(ChatState.Error);
},
[notifications.toasts]
);
const next = useCallback(
async (nextMessages: Message[]) => {
// make sure we ignore any aborts for the previous signal
@ -77,173 +110,134 @@ export function useChat({
// cancel running requests
abortControllerRef.current.abort();
const lastMessage = last(nextMessages);
abortControllerRef.current = new AbortController();
const allMessages = [
systemMessage,
...nextMessages.filter((message) => message.message.role !== MessageRole.System),
];
setPendingMessages([]);
setMessages(nextMessages);
setMessages(allMessages);
if (!lastMessage || !connectorId) {
if (!connectorId || !nextMessages.length) {
setChatState(ChatState.Ready);
onChatCompleteRef.current?.(nextMessages);
return;
}
const isUserMessage = lastMessage.message.role === MessageRole.User;
const functionCall = lastMessage.message.function_call;
const isAssistantMessageWithFunctionRequest =
lastMessage.message.role === MessageRole.Assistant && functionCall && !!functionCall.name;
const isFunctionResult = isUserMessage && !!lastMessage.message.name;
const isRecallFunctionAvailable = chatService.hasFunction('recall');
if (!isUserMessage && !isAssistantMessageWithFunctionRequest) {
setChatState(ChatState.Ready);
onChatCompleteRef.current?.(nextMessages);
return;
}
const abortController = (abortControllerRef.current = new AbortController());
abortController.signal.addEventListener('abort', handleSignalAbort);
setChatState(ChatState.Loading);
if (isUserMessage && !isFunctionResult && isRecallFunctionAvailable) {
const allMessagesWithRecall = allMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'recall',
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.Assistant,
},
},
});
next(allMessagesWithRecall);
return;
const next$ = chatService.complete({
connectorId,
messages: getWithSystemMessage(nextMessages, systemMessage),
persist,
signal: abortControllerRef.current.signal,
conversationId: initialConversationId,
});
function getPendingMessages() {
return [
...completedMessages,
...(pendingMessage
? [
merge(
{
message: {
role: MessageRole.Assistant,
function_call: { trigger: MessageRole.Assistant as const },
},
},
pendingMessage
),
]
: []),
];
}
function handleError(error: Error) {
setChatState(ChatState.Error);
notifications.toasts.addError(error, {
title: i18n.translate('xpack.observabilityAiAssistant.failedToLoadResponse', {
defaultMessage: 'Failed to load response from the AI Assistant',
}),
});
}
const completedMessages: Message[] = [];
const response = isAssistantMessageWithFunctionRequest
? await chatService
.executeFunction({
name: functionCall.name,
signal: abortController.signal,
args: functionCall.arguments,
connectorId,
messages: allMessages,
})
.catch((error) => {
return {
content: {
message: error.toString(),
error,
},
data: undefined,
};
})
: chatService.chat({
messages: allMessages,
connectorId,
});
let pendingMessage:
| {
'@timestamp': string;
message: { content: string; function_call: { name: string; arguments: string } };
}
| undefined;
if (abortController.signal.aborted) {
return;
}
const subscription = next$.subscribe({
next: (event) => {
switch (event.type) {
case StreamingChatResponseEventType.ChatCompletionChunk:
if (!pendingMessage) {
pendingMessage = {
'@timestamp': new Date().toISOString(),
message: {
content: event.message.content || '',
function_call: {
name: event.message.function_call?.name || '',
arguments: event.message.function_call?.arguments || '',
},
},
};
} else {
pendingMessage.message.content += event.message.content || '';
pendingMessage.message.function_call.name +=
event.message.function_call?.name || '';
pendingMessage.message.function_call.arguments +=
event.message.function_call?.arguments || '';
}
break;
if (isObservable(response)) {
let localPendingMessage: PendingMessage = {
message: {
content: '',
role: MessageRole.User,
},
};
case StreamingChatResponseEventType.MessageAdd:
pendingMessage = undefined;
completedMessages.push(event.message);
break;
const subscription = response.subscribe({
next: (nextPendingMessage) => {
localPendingMessage = nextPendingMessage;
setPendingMessage(nextPendingMessage);
},
complete: () => {
setPendingMessage(undefined);
const allMessagesWithResolved = allMessages.concat({
message: {
...localPendingMessage.message,
},
'@timestamp': new Date().toISOString(),
});
if (localPendingMessage.aborted) {
setChatState(ChatState.Aborted);
setMessages(allMessagesWithResolved);
} else if (localPendingMessage.error) {
handleError(localPendingMessage.error);
setMessages(allMessagesWithResolved);
} else {
next(allMessagesWithResolved);
}
},
error: (error) => {
handleError(error);
},
});
case StreamingChatResponseEventType.ConversationCreate:
case StreamingChatResponseEventType.ConversationUpdate:
onConversationUpdateRef.current?.(event);
break;
}
setPendingMessages(getPendingMessages());
},
complete: () => {
setChatState(ChatState.Ready);
const completed = nextMessages.concat(completedMessages);
setMessages(completed);
setPendingMessages([]);
onChatCompleteRef.current?.(completed);
},
error: (error) => {
setPendingMessages([]);
setMessages(nextMessages.concat(getPendingMessages()));
handleError(error);
},
});
abortController.signal.addEventListener('abort', () => {
subscription.unsubscribe();
});
} else {
const allMessagesWithFunctionReply = allMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
name: functionCall!.name,
role: MessageRole.User,
content: JSON.stringify(response.content),
data: JSON.stringify(response.data),
},
});
next(allMessagesWithFunctionReply);
}
abortControllerRef.current.signal.addEventListener('abort', () => {
handleSignalAbort();
subscription.unsubscribe();
});
},
[connectorId, chatService, handleSignalAbort, notifications.toasts, systemMessage]
[
connectorId,
chatService,
handleSignalAbort,
systemMessage,
handleError,
persist,
initialConversationId,
]
);
useEffect(() => {
const controller = abortControllerRef.current;
return () => {
abortControllerRef.current.abort();
controller.abort();
};
}, []);
const memoizedMessages = useMemo(() => {
const includingSystemMessage = [
systemMessage,
...messages.filter((message) => message.message.role !== MessageRole.System),
];
return pendingMessage
? includingSystemMessage.concat({
...pendingMessage,
'@timestamp': new Date().toISOString(),
})
: includingSystemMessage;
}, [systemMessage, messages, pendingMessage]);
return getWithSystemMessage(messages.concat(pendingMessages ?? []), systemMessage);
}, [systemMessage, messages, pendingMessages]);
const setMessagesWithAbort = useCallback((nextMessages: Message[]) => {
abortControllerRef.current.abort();
setPendingMessage(undefined);
setPendingMessages([]);
setChatState(ChatState.Ready);
setMessages(nextMessages);
}, []);

272
x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.test.tsx
View file

@ -4,28 +4,32 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import React from 'react';
import {
useConversation,
type UseConversationProps,
type UseConversationResult,
} from './use_conversation';
import type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import {
act,
renderHook,
type RenderHookResult,
type WrapperComponent,
} from '@testing-library/react-hooks';
import type { ObservabilityAIAssistantService, PendingMessage } from '../types';
import type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
import * as useKibanaModule from './use_kibana';
import { Message, MessageRole } from '../../common';
import { ChatState } from './use_chat';
import { createMockChatService } from '../service/create_mock_chat_service';
import { merge } from 'lodash';
import React from 'react';
import { Subject } from 'rxjs';
import { MessageRole } from '../../common';
import {
StreamingChatResponseEvent,
StreamingChatResponseEventType,
} from '../../common/conversation_complete';
import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
import { EMPTY_CONVERSATION_TITLE } from '../i18n';
import { merge, omit } from 'lodash';
import { createMockChatService } from '../service/create_mock_chat_service';
import type { ObservabilityAIAssistantService } from '../types';
import { ChatState } from './use_chat';
import {
useConversation,
type UseConversationProps,
type UseConversationResult,
} from './use_conversation';
import * as useKibanaModule from './use_kibana';
let hookResult: RenderHookResult<UseConversationProps, UseConversationResult>;
@ -270,8 +274,9 @@ describe('useConversation', () => {
});
});
describe('when chat completes without an initial conversation id', () => {
let subject: Subject<PendingMessage>;
describe('when chat completes', () => {
const subject: Subject<StreamingChatResponseEvent> = new Subject();
let onConversationUpdate: jest.Mock;
const expectedMessages = [
{
'@timestamp': expect.any(String),
@ -310,7 +315,6 @@ describe('useConversation', () => {
},
];
beforeEach(() => {
subject = new Subject();
mockService.callApi.mockImplementation(async (endpoint, request) =>
merge(
{
@ -323,6 +327,8 @@ describe('useConversation', () => {
)
);
onConversationUpdate = jest.fn();
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
@ -343,224 +349,66 @@ describe('useConversation', () => {
},
},
],
onConversationUpdate,
},
wrapper,
});
mockChatService.chat.mockImplementationOnce(() => {
mockChatService.complete.mockImplementationOnce(() => {
return subject;
});
act(() => {
hookResult.result.current.next(
hookResult.result.current.messages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello again',
},
})
);
});
});
describe('when chat completes with an error', () => {
describe('and the conversation is created or updated', () => {
beforeEach(async () => {
mockService.callApi.mockClear();
act(() => {
await act(async () => {
hookResult.result.current.next(
hookResult.result.current.messages.concat({
'@timestamp': new Date().toISOString(),
message: {
content: 'Hello again',
role: MessageRole.User,
},
})
);
subject.next({
type: StreamingChatResponseEventType.ChatCompletionChunk,
id: 'my-message',
message: {
role: MessageRole.Assistant,
content: 'Goodbye',
},
error: new Error(),
});
subject.complete();
});
await act(async () => {});
});
it('does not store the conversation', () => {
expect(mockService.callApi).not.toHaveBeenCalled();
});
});
describe('when chat completes without an error', () => {
beforeEach(async () => {
act(() => {
subject.next({
type: StreamingChatResponseEventType.MessageAdd,
id: 'my-message',
message: {
role: MessageRole.Assistant,
content: 'Goodbye again',
'@timestamp': new Date().toISOString(),
message: {
content: 'Goodbye',
role: MessageRole.Assistant,
},
},
});
subject.next({
type: StreamingChatResponseEventType.ConversationUpdate,
conversation: {
id: 'my-conversation-id',
title: 'My title',
last_updated: new Date().toISOString(),
},
});
subject.complete();
});
await act(async () => {});
});
it('the conversation is created including the initial messages', async () => {
expect(mockService.callApi.mock.calls[0]).toEqual([
'POST /internal/observability_ai_assistant/conversation',
{
params: {
body: {
conversation: {
'@timestamp': expect.any(String),
conversation: {
title: EMPTY_CONVERSATION_TITLE,
},
messages: expectedMessages,
labels: {},
numeric_labels: {},
public: false,
},
},
},
signal: null,
it('calls the onConversationUpdate hook', () => {
expect(onConversationUpdate).toHaveBeenCalledWith({
conversation: {
id: 'my-conversation-id',
last_updated: expect.any(String),
title: 'My title',
},
]);
expect(hookResult.result.current.conversation.error).toBeUndefined();
expect(hookResult.result.current.messages).toEqual(expectedMessages);
});
});
});
describe('when chat completes with an initial conversation id', () => {
let subject: Subject<PendingMessage>;
const initialMessages: Message[] = [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: '',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'user',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: 'assistant',
},
},
];
beforeEach(async () => {
mockService.callApi.mockImplementation(async (endpoint, request) => ({
'@timestamp': new Date().toISOString(),
conversation: {
id: 'my-conversation-id',
title: EMPTY_CONVERSATION_TITLE,
},
labels: {},
numeric_labels: {},
public: false,
messages: initialMessages,
}));
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialConversationId: 'my-conversation-id',
},
wrapper,
});
await act(async () => {});
});
it('the conversation is loadeded', async () => {
expect(mockService.callApi.mock.calls[0]).toEqual([
'GET /internal/observability_ai_assistant/conversation/{conversationId}',
{
signal: expect.anything(),
params: {
path: {
conversationId: 'my-conversation-id',
},
},
},
]);
expect(hookResult.result.current.messages).toEqual(
initialMessages.map((msg) => ({ ...msg, '@timestamp': expect.any(String) }))
);
});
describe('after chat completes', () => {
const nextUserMessage: Message = {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello again',
},
};
const nextAssistantMessage: Message = {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: 'Goodbye again',
},
};
beforeEach(async () => {
mockService.callApi.mockClear();
subject = new Subject();
mockChatService.chat.mockImplementationOnce(() => {
return subject;
});
act(() => {
hookResult.result.current.next(
hookResult.result.current.messages.concat(nextUserMessage)
);
subject.next(omit(nextAssistantMessage, '@timestamp'));
subject.complete();
});
await act(async () => {});
});
it('saves the updated message', () => {
expect(mockService.callApi.mock.calls[0]).toEqual([
'PUT /internal/observability_ai_assistant/conversation/{conversationId}',
{
params: {
path: {
conversationId: 'my-conversation-id',
},
body: {
conversation: {
'@timestamp': expect.any(String),
conversation: {
title: EMPTY_CONVERSATION_TITLE,
id: 'my-conversation-id',
},
messages: initialMessages
.concat([nextUserMessage, nextAssistantMessage])
.map((msg) => ({ ...msg, '@timestamp': expect.any(String) })),
labels: {},
numeric_labels: {},
public: false,
},
},
},
signal: null,
},
]);
});
});
});

72
x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts
View file

@ -38,7 +38,7 @@ export interface UseConversationProps {
initialTitle?: string;
chatService: ObservabilityAIAssistantChatService;
connectorId: string | undefined;
onConversationUpdate?: (conversation: Conversation) => void;
onConversationUpdate?: (conversation: { conversation: Conversation['conversation'] }) => void;
}
export type UseConversationResult = {
@ -101,76 +101,16 @@ export function useConversation({
});
};
const save = (nextMessages: Message[]) => {
const conversationObject = conversation.value!;
const nextConversationObject = merge({}, omit(conversationObject, 'messages'), {
messages: nextMessages,
});
return (
displayedConversationId
? update(
merge(
{ conversation: { id: displayedConversationId } },
nextConversationObject
) as Conversation
)
: service
.callApi(`POST /internal/observability_ai_assistant/conversation`, {
signal: null,
params: {
body: {
conversation: nextConversationObject,
},
},
})
.then((nextConversation) => {
setDisplayedConversationId(nextConversation.conversation.id);
if (connectorId) {
service
.callApi(
`PUT /internal/observability_ai_assistant/conversation/{conversationId}/auto_title`,
{
signal: null,
params: {
path: {
conversationId: nextConversation.conversation.id,
},
body: {
connectorId,
},
},
}
)
.then(() => {
onConversationUpdate?.(nextConversation);
return conversation.refresh();
});
}
return nextConversation;
})
.catch((err) => {
notifications.toasts.addError(err, {
title: i18n.translate('xpack.observabilityAiAssistant.errorCreatingConversation', {
defaultMessage: 'Could not create conversation',
}),
});
throw err;
})
).then((nextConversation) => {
onConversationUpdate?.(nextConversation);
return nextConversation;
});
};
const { next, messages, setMessages, state, stop } = useChat({
initialMessages,
initialConversationId,
chatService,
connectorId,
onChatComplete: (nextMessages) => {
save(nextMessages);
onConversationUpdate: (event) => {
setDisplayedConversationId(event.conversation.id);
onConversationUpdate?.({ conversation: event.conversation });
},
persist: true,
});
const [displayedConversationId, setDisplayedConversationId] = useState(initialConversationId);

10
x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts
View file

@ -23,21 +23,19 @@ export const useJsonEditorModel = ({
}) => {
const chatService = useObservabilityAIAssistantChatService();
const functionDefinition = chatService
.getFunctions()
.find((func) => func.options.name === functionName);
const functionDefinition = chatService.getFunctions().find((func) => func.name === functionName);
return useMemo(() => {
if (!functionDefinition) {
return {};
}
const schema = { ...functionDefinition.options.parameters };
const schema = { ...functionDefinition.parameters };
const initialJsonString = initialJson
? initialJson
: functionDefinition.options.parameters.properties
? JSON.stringify(createInitializedObject(functionDefinition.options.parameters), null, 4)
: functionDefinition.parameters.properties
? JSON.stringify(createInitializedObject(functionDefinition.parameters), null, 4)
: '';
languages.json.jsonDefaults.setDiagnosticsOptions({

12
x-pack/plugins/observability_ai_assistant/public/plugin.tsx
View file

@ -107,19 +107,19 @@ export class ObservabilityAIAssistantPlugin
shareStart: pluginsStart.share,
}));
service.register(async ({ signal, registerContext, registerFunction }) => {
service.register(async ({ registerRenderFunction }) => {
const mod = await import('./functions');
return mod.registerFunctions({
service,
signal,
pluginsStart,
coreStart,
registerContext,
registerFunction,
registerRenderFunction,
});
});
return { service, useGenAIConnectors: () => useGenAIConnectorsWithoutContext(service) };
return {
service,
useGenAIConnectors: () => useGenAIConnectorsWithoutContext(service),
};
}
}

6
x-pack/plugins/observability_ai_assistant/public/service/create_chat_service.test.ts
View file

@ -39,6 +39,12 @@ describe('createChatService', () => {
}
beforeEach(async () => {
clientSpy.mockImplementationOnce(async () => {
return {
functionDefinitions: [],
contextDefinitions: [],
};
});
service = await createChatService({
analytics: {
optIn: () => {},

289
x-pack/plugins/observability_ai_assistant/public/service/create_chat_service.ts
View file

@ -4,8 +4,6 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
/* eslint-disable max-classes-per-file*/
import { Validator, type Schema, type OutputUnit } from '@cfworker/json-schema';
import { AnalyticsServiceStart, HttpResponse } from '@kbn/core/public';
import { AbortError } from '@kbn/kibana-utils-plugin/common';
@ -16,45 +14,82 @@ import {
catchError,
concatMap,
delay,
filter as rxJsFilter,
finalize,
map,
of,
scan,
shareReplay,
tap,
Subject,
timestamp,
map,
tap,
} from 'rxjs';
import {
ContextRegistry,
FunctionRegistry,
ChatCompletionErrorCode,
ConversationCompletionError,
type StreamingChatResponseEvent,
StreamingChatResponseEventType,
} from '../../common/conversation_complete';
import {
FunctionVisibility,
Message,
MessageRole,
type RegisterContextDefinition,
type RegisterFunctionDefinition,
type FunctionRegistry,
type FunctionResponse,
type Message,
} from '../../common/types';
import { ObservabilityAIAssistantAPIClient } from '../api';
import { filterFunctionDefinitions } from '../../common/utils/filter_function_definitions';
import { processOpenAiStream } from '../../common/utils/process_openai_stream';
import type { ObservabilityAIAssistantAPIClient } from '../api';
import type {
AssistantRegistrationFunction,
CreateChatCompletionResponseChunk,
ChatRegistrationRenderFunction,
ObservabilityAIAssistantChatService,
PendingMessage,
RenderFunction,
} from '../types';
import { readableStreamReaderIntoObservable } from '../utils/readable_stream_reader_into_observable';
class TokenLimitReachedError extends Error {
constructor() {
super(`Token limit reached`);
}
}
const MIN_DELAY = 35;
class ServerError extends Error {}
function toObservable(response: HttpResponse<IncomingMessage>) {
const status = response.response?.status;
export class FunctionArgsValidationError extends Error {
constructor(public readonly errors: OutputUnit[]) {
super('Function arguments are invalid');
if (!status || status >= 400) {
throw new Error(response.response?.statusText || 'Unexpected error');
}
const reader = response.response.body?.getReader();
if (!reader) {
throw new Error('Could not get reader from response');
}
return readableStreamReaderIntoObservable(reader).pipe(
// append a timestamp of when each value was emitted
timestamp(),
// use the previous timestamp to calculate a target
// timestamp for emitting the next value
scan((acc, value) => {
const lastTimestamp = acc.timestamp || 0;
const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp);
return {
timestamp: emitAt,
value: value.value,
};
}),
// add the delay based on the elapsed time
// using concatMap(of(value).pipe(delay(50))
// leads to browser issues because timers
// are throttled when the tab is not active
concatMap((value) => {
const now = Date.now();
const delayFor = value.timestamp - now;
if (delayFor <= 0) {
return of(value.value);
}
return of(value.value).pipe(delay(delayFor));
})
);
}
export async function createChatService({
@ -65,73 +100,41 @@ export async function createChatService({
}: {
analytics: AnalyticsServiceStart;
signal: AbortSignal;
registrations: AssistantRegistrationFunction[];
registrations: ChatRegistrationRenderFunction[];
client: ObservabilityAIAssistantAPIClient;
}): Promise<ObservabilityAIAssistantChatService> {
const contextRegistry: ContextRegistry = new Map();
const functionRegistry: FunctionRegistry = new Map();
const validators = new Map<string, Validator>();
const renderFunctionRegistry: Map<string, RenderFunction<unknown, FunctionResponse>> = new Map();
const registerContext: RegisterContextDefinition = (context) => {
contextRegistry.set(context.name, context);
const [{ functionDefinitions, contextDefinitions }] = await Promise.all([
client('GET /internal/observability_ai_assistant/functions', {
signal: setupAbortSignal,
}),
...registrations.map((registration) => {
return registration({
registerRenderFunction: (name, renderFn) => {
renderFunctionRegistry.set(name, renderFn);
},
});
}),
]);
functionDefinitions.forEach((fn) => {
functionRegistry.set(fn.name, fn);
});
const getFunctions = (options?: { contexts?: string[]; filter?: string }) => {
return filterFunctionDefinitions({
...options,
definitions: functionDefinitions,
});
};
const registerFunction: RegisterFunctionDefinition = (def, respond, render) => {
validators.set(def.name, new Validator(def.parameters as Schema, '2020-12', true));
functionRegistry.set(def.name, { options: def, respond, render });
};
const getContexts: ObservabilityAIAssistantChatService['getContexts'] = () => {
return Array.from(contextRegistry.values());
};
const getFunctions: ObservabilityAIAssistantChatService['getFunctions'] = ({
contexts,
filter,
} = {}) => {
const allFunctions = Array.from(functionRegistry.values());
return contexts || filter
? allFunctions.filter((fn) => {
const matchesContext =
!contexts || fn.options.contexts.some((context) => contexts.includes(context));
const matchesFilter =
!filter || fn.options.name.includes(filter) || fn.options.description.includes(filter);
return matchesContext && matchesFilter;
})
: allFunctions;
};
await Promise.all(
registrations.map((fn) => fn({ signal: setupAbortSignal, registerContext, registerFunction }))
);
function validate(name: string, parameters: unknown) {
const validator = validators.get(name)!;
const result = validator.validate(parameters);
if (!result.valid) {
throw new FunctionArgsValidationError(result.errors);
}
}
return {
analytics,
executeFunction: async ({ name, args, signal, messages, connectorId }) => {
const fn = functionRegistry.get(name);
if (!fn) {
throw new Error(`Function ${name} not found`);
}
const parsedArguments = args ? JSON.parse(args) : {};
validate(name, parsedArguments);
return await fn.respond({ arguments: parsedArguments, messages, connectorId }, signal);
},
renderFunction: (name, args, response) => {
const fn = functionRegistry.get(name);
const fn = renderFunctionRegistry.get(name);
if (!fn) {
throw new Error(`Function ${name} not found`);
@ -144,15 +147,57 @@ export async function createChatService({
data: JSON.parse(response.data ?? '{}'),
};
return fn.render?.({ response: parsedResponse, arguments: parsedArguments });
return fn?.({ response: parsedResponse, arguments: parsedArguments });
},
getContexts,
getContexts: () => contextDefinitions,
getFunctions,
hasFunction: (name: string) => {
return !!getFunctions().find((fn) => fn.options.name === name);
return functionRegistry.has(name);
},
hasRenderFunction: (name: string) => {
return !!getFunctions().find((fn) => fn.options.name === name)?.render;
return renderFunctionRegistry.has(name);
},
complete({ connectorId, messages, conversationId, persist, signal }) {
const subject = new Subject<StreamingChatResponseEvent>();
client('POST /internal/observability_ai_assistant/chat/complete', {
params: {
body: {
messages,
connectorId,
conversationId,
persist,
},
},
signal,
asResponse: true,
rawResponse: true,
})
.then((_response) => {
const response = _response as unknown as HttpResponse<IncomingMessage>;
const response$ = toObservable(response)
.pipe(
map((line) => JSON.parse(line) as StreamingChatResponseEvent),
tap((event) => {
if (event.type === StreamingChatResponseEventType.ConversationCompletionError) {
const code = event.error.code ?? ChatCompletionErrorCode.InternalError;
const message = event.error.message;
throw new ConversationCompletionError(code, message);
}
})
)
.subscribe(subject);
signal.addEventListener('abort', () => {
response$.unsubscribe();
});
})
.catch((err) => {
subject.error(err);
subject.complete();
});
return subject;
},
chat({
connectorId,
@ -184,8 +229,8 @@ export async function createChatService({
callFunctions === 'none'
? []
: functions
.filter((fn) => fn.options.visibility !== FunctionVisibility.User)
.map((fn) => pick(fn.options, 'name', 'description', 'parameters')),
.filter((fn) => fn.visibility !== FunctionVisibility.User)
.map((fn) => pick(fn, 'name', 'description', 'parameters')),
},
},
signal: controller.signal,
@ -195,51 +240,9 @@ export async function createChatService({
.then((_response) => {
const response = _response as unknown as HttpResponse<IncomingMessage>;
const status = response.response?.status;
if (!status || status >= 400) {
throw new Error(response.response?.statusText || 'Unexpected error');
}
const reader = response.response.body?.getReader();
if (!reader) {
throw new Error('Could not get reader from response');
}
const subscription = readableStreamReaderIntoObservable(reader)
const subscription = toObservable(response)
.pipe(
// lines start with 'data: '
map((line) => line.substring(6)),
// a message completes with the line '[DONE]'
rxJsFilter((line) => !!line && line !== '[DONE]'),
// parse the JSON, add the type
map(
(line) =>
JSON.parse(line) as
| CreateChatCompletionResponseChunk
| { error: { message: string } }
),
// validate the message. in some cases OpenAI
// will throw halfway through the message
tap((line) => {
if ('error' in line) {
throw new ServerError(line.error.message);
}
}),
// there also might be some metadata that we need
// to exclude
rxJsFilter(
(line): line is CreateChatCompletionResponseChunk =>
'object' in line && line.object === 'chat.completion.chunk'
),
// this is how OpenAI signals that the context window
// limit has been exceeded
tap((line) => {
if (line.choices[0].finish_reason === 'length') {
throw new TokenLimitReachedError();
}
}),
processOpenAiStream(),
// merge the messages
scan(
(acc, { choices }) => {
@ -301,8 +304,6 @@ export async function createChatService({
subject.complete();
});
const MIN_DELAY = 35;
const pendingMessages$ = subject.pipe(
// make sure the request is only triggered once,
// even with multiple subscribers
@ -311,32 +312,6 @@ export async function createChatService({
// abort the running request
finalize(() => {
controller.abort();
}),
// append a timestamp of when each value was emitted
timestamp(),
// use the previous timestamp to calculate a target
// timestamp for emitting the next value
scan((acc, value) => {
const lastTimestamp = acc.timestamp || 0;
const emitAt = Math.max(lastTimestamp + MIN_DELAY, value.timestamp);
return {
timestamp: emitAt,
value: value.value,
};
}),
// add the delay based on the elapsed time
// using concatMap(of(value).pipe(delay(50))
// leads to browser issues because timers
// are throttled when the tab is not active
concatMap((value) => {
const now = Date.now();
const delayFor = value.timestamp - now;
if (delayFor <= 0) {
return of(value.value);
}
return of(value.value).pipe(delay(delayFor));
})
);

6
x-pack/plugins/observability_ai_assistant/public/service/create_mock_chat_service.ts
View file

@ -5,7 +5,7 @@
* 2.0.
*/
import { TelemetryCounter } from '@kbn/analytics-client';
import type { TelemetryCounter } from '@kbn/analytics-client';
import type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import { Observable } from 'rxjs';
import type { ObservabilityAIAssistantChatService } from '../types';
@ -14,13 +14,13 @@ type MockedChatService = DeeplyMockedKeys<ObservabilityAIAssistantChatService>;
export const createMockChatService = (): MockedChatService => {
const mockChatService: MockedChatService = {
chat: jest.fn(),
complete: jest.fn(),
analytics: {
optIn: jest.fn(),
reportEvent: jest.fn(),
telemetryCounter$: new Observable<TelemetryCounter>() as any,
},
chat: jest.fn(),
executeFunction: jest.fn(),
getContexts: jest.fn().mockReturnValue([{ name: 'core', description: '' }]),
getFunctions: jest.fn().mockReturnValue([]),
hasFunction: jest.fn().mockReturnValue(false),

5
x-pack/plugins/observability_ai_assistant/public/service/create_service.ts
View file

@ -10,7 +10,7 @@ import type { LicensingPluginStart } from '@kbn/licensing-plugin/public';
import type { SecurityPluginStart } from '@kbn/security-plugin/public';
import type { SharePluginStart } from '@kbn/share-plugin/public';
import { createCallObservabilityAIAssistantAPI } from '../api';
import type { AssistantRegistrationFunction, ObservabilityAIAssistantService } from '../types';
import type { ChatRegistrationRenderFunction, ObservabilityAIAssistantService } from '../types';
export function createService({
analytics,
@ -29,7 +29,7 @@ export function createService({
}): ObservabilityAIAssistantService {
const client = createCallObservabilityAIAssistantAPI(coreStart);
const registrations: AssistantRegistrationFunction[] = [];
const registrations: ChatRegistrationRenderFunction[] = [];
return {
isEnabled: () => {
@ -42,7 +42,6 @@ export function createService({
const mod = await import('./create_chat_service');
return await mod.createChatService({ analytics, client, signal, registrations });
},
callApi: client,
getCurrentUser: () => securityStart.authc.getCurrentUser(),
getLicense: () => licenseStart.license$,

55
x-pack/plugins/observability_ai_assistant/public/types.ts
View file

@ -20,11 +20,6 @@ import type {
TriggersAndActionsUIPublicPluginSetup,
TriggersAndActionsUIPublicPluginStart,
} from '@kbn/triggers-actions-ui-plugin/public';
import type { Serializable } from '@kbn/utility-types';
import type {
CreateChatCompletionResponse,
CreateChatCompletionResponseChoicesInner,
} from 'openai';
import type { Observable } from 'rxjs';
import type { LensPublicSetup, LensPublicStart } from '@kbn/lens-plugin/public';
import type {
@ -36,23 +31,17 @@ import type { SharePluginStart } from '@kbn/share-plugin/public';
import type {
ContextDefinition,
FunctionDefinition,
FunctionResponse,
Message,
RegisterContextDefinition,
RegisterFunctionDefinition,
} from '../common/types';
import type { ObservabilityAIAssistantAPIClient } from './api';
import type { PendingMessage } from '../common/types';
import { UseGenAIConnectorsResult } from './hooks/use_genai_connectors';
import type { StreamingChatResponseEvent } from '../common/conversation_complete';
import type { UseGenAIConnectorsResult } from './hooks/use_genai_connectors';
/* eslint-disable @typescript-eslint/no-empty-interface*/
export type CreateChatCompletionResponseChunk = Omit<CreateChatCompletionResponse, 'choices'> & {
choices: Array<
Omit<CreateChatCompletionResponseChoicesInner, 'message'> & {
delta: { content?: string; function_call?: { name?: string; arguments?: string } };
}
>;
};
export type { CreateChatCompletionResponseChunk } from '../common/types';
export interface ObservabilityAIAssistantChatService {
analytics: AnalyticsServiceStart;
@ -61,17 +50,17 @@ export interface ObservabilityAIAssistantChatService {
connectorId: string;
function?: 'none' | 'auto';
}) => Observable<PendingMessage>;
complete: (options: {
messages: Message[];
connectorId: string;
persist: boolean;
conversationId?: string;
signal: AbortSignal;
}) => Observable<StreamingChatResponseEvent>;
getContexts: () => ContextDefinition[];
getFunctions: (options?: { contexts?: string[]; filter?: string }) => FunctionDefinition[];
hasFunction: (name: string) => boolean;
hasRenderFunction: (name: string) => boolean;
executeFunction: ({}: {
name: string;
args: string | undefined;
messages: Message[];
signal: AbortSignal;
connectorId: string;
}) => Promise<{ content?: Serializable; data?: Serializable } | Observable<PendingMessage>>;
renderFunction: (
name: string,
args: string | undefined,
@ -79,12 +68,6 @@ export interface ObservabilityAIAssistantChatService {
) => React.ReactNode;
}
export type AssistantRegistrationFunction = ({}: {
signal: AbortSignal;
registerFunction: RegisterFunctionDefinition;
registerContext: RegisterContextDefinition;
}) => Promise<void>;
export interface ObservabilityAIAssistantService {
isEnabled: () => boolean;
callApi: ObservabilityAIAssistantAPIClient;
@ -92,9 +75,23 @@ export interface ObservabilityAIAssistantService {
getLicense: () => Observable<ILicense>;
getLicenseManagementLocator: () => SharePluginStart;
start: ({}: { signal: AbortSignal }) => Promise<ObservabilityAIAssistantChatService>;
register: (fn: AssistantRegistrationFunction) => void;
register: (fn: ChatRegistrationRenderFunction) => void;
}
export type RenderFunction<TArguments, TResponse extends FunctionResponse> = (options: {
arguments: TArguments;
response: TResponse;
}) => React.ReactNode;
export type RegisterRenderFunctionDefinition<
TFunctionArguments = any,
TFunctionResponse extends FunctionResponse = FunctionResponse
> = (name: string, render: RenderFunction<TFunctionArguments, TFunctionResponse>) => void;
export type ChatRegistrationRenderFunction = ({}: {
registerRenderFunction: RegisterRenderFunctionDefinition;
}) => Promise<void>;
export interface ObservabilityAIAssistantPluginStart {
service: ObservabilityAIAssistantService;
useGenAIConnectors: () => UseGenAIConnectorsResult;

59
x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
View file

@ -7,7 +7,12 @@
import { merge, uniqueId } from 'lodash';
import { DeepPartial } from 'utility-types';
import { MessageRole, Conversation, FunctionDefinition, Message } from '../../common/types';
import {
type Conversation,
type FunctionDefinition,
type Message,
MessageRole,
} from '../../common/types';
import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
type BuildMessageProps = DeepPartial<Message> & {
@ -121,28 +126,25 @@ export function buildConversation(params?: Partial<Conversation>) {
export function buildFunction(): FunctionDefinition {
return {
options: {
name: 'elasticsearch',
contexts: ['core'],
description: 'Call Elasticsearch APIs on behalf of the user',
descriptionForUser: 'Call Elasticsearch APIs on behalf of the user',
parameters: {
type: 'object',
properties: {
method: {
type: 'string',
description: 'The HTTP method of the Elasticsearch endpoint',
enum: ['GET', 'PUT', 'POST', 'DELETE', 'PATCH'] as const,
},
path: {
type: 'string',
description: 'The path of the Elasticsearch endpoint, including query parameters',
},
name: 'elasticsearch',
contexts: ['core'],
description: 'Call Elasticsearch APIs on behalf of the user',
descriptionForUser: 'Call Elasticsearch APIs on behalf of the user',
parameters: {
type: 'object',
properties: {
method: {
type: 'string',
description: 'The HTTP method of the Elasticsearch endpoint',
enum: ['GET', 'PUT', 'POST', 'DELETE', 'PATCH'] as const,
},
path: {
type: 'string',
description: 'The path of the Elasticsearch endpoint, including query parameters',
},
required: ['method' as const, 'path' as const],
},
required: ['method' as const, 'path' as const],
},
respond: async (options: { arguments: any }, signal: AbortSignal) => ({}),
};
}
@ -150,16 +152,13 @@ export const buildFunctionElasticsearch = buildFunction;
export function buildFunctionServiceSummary(): FunctionDefinition {
return {
options: {
name: 'get_service_summary',
contexts: ['core'],
description:
'Gets a summary of a single service, including: the language, service version, deployments, infrastructure, alerting, etc. ',
descriptionForUser: 'Get a summary for a single service.',
parameters: {
type: 'object',
},
name: 'get_service_summary',
contexts: ['core'],
description:
'Gets a summary of a single service, including: the language, service version, deployments, infrastructure, alerting, etc. ',
descriptionForUser: 'Get a summary for a single service.',
parameters: {
type: 'object',
},
respond: async (options: { arguments: any }, signal: AbortSignal) => ({}),
};
}

2
x-pack/plugins/observability_ai_assistant/public/utils/create_initialized_object.ts
View file

@ -7,7 +7,7 @@
import { FunctionDefinition } from '../../common/types';
type Params = FunctionDefinition['options']['parameters'];
type Params = FunctionDefinition['parameters'];
export function createInitializedObject(parameters: Params) {
const emptyObject: Record<string, string | any> = {};

31
x-pack/plugins/observability_ai_assistant/public/utils/storybook_decorator.tsx
View file

@ -4,22 +4,22 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import React, { ComponentType } from 'react';
import { Observable } from 'rxjs';
import { i18n } from '@kbn/i18n';
import { KibanaContextProvider } from '@kbn/kibana-react-plugin/public';
import type { Serializable } from '@kbn/utility-types';
import type { AuthenticatedUser } from '@kbn/security-plugin/common';
import type { SharePluginStart } from '@kbn/share-plugin/public';
import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
import React, { ComponentType } from 'react';
import { Observable } from 'rxjs';
import type { StreamingChatResponseEvent } from '../../common/conversation_complete';
import { ObservabilityAIAssistantAPIClient } from '../api';
import type { Message } from '../../common';
import { ObservabilityAIAssistantChatServiceProvider } from '../context/observability_ai_assistant_chat_service_provider';
import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
import type {
ObservabilityAIAssistantChatService,
ObservabilityAIAssistantService,
PendingMessage,
} from '../types';
import { buildFunctionElasticsearch, buildFunctionServiceSummary } from './builders';
import { ObservabilityAIAssistantChatServiceProvider } from '../context/observability_ai_assistant_chat_service_provider';
const chatService: ObservabilityAIAssistantChatService = {
analytics: {
@ -27,18 +27,17 @@ const chatService: ObservabilityAIAssistantChatService = {
reportEvent: () => {},
telemetryCounter$: new Observable(),
},
chat: (options: { messages: Message[]; connectorId: string }) => new Observable<PendingMessage>(),
chat: (options) => new Observable<PendingMessage>(),
complete: (options) => new Observable<StreamingChatResponseEvent>(),
getContexts: () => [],
getFunctions: () => [buildFunctionElasticsearch(), buildFunctionServiceSummary()],
executeFunction: async ({}: {
name: string;
args: string | undefined;
messages: Message[];
signal: AbortSignal;
}): Promise<{ content?: Serializable; data?: Serializable }> => ({}),
renderFunction: (name: string, args: string | undefined, response: {}) => (
// eslint-disable-next-line @kbn/i18n/strings_should_be_translated_with_i18n
<div>Hello! {name}</div>
renderFunction: (name) => (
<div>
{i18n.translate('xpack.observabilityAiAssistant.chatService.div.helloLabel', {
defaultMessage: 'Hello',
})}
{name}
</div>
),
hasFunction: () => true,
hasRenderFunction: () => true,

153
x-pack/plugins/observability_ai_assistant/server/functions/alerts.ts Normal file
View file

@ -0,0 +1,153 @@
/*
* 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 datemath from '@elastic/datemath';
import { fromKueryExpression, toElasticsearchQuery } from '@kbn/es-query';
import { ParsedTechnicalFields } from '@kbn/rule-registry-plugin/common';
import {
ALERT_STATUS,
ALERT_STATUS_ACTIVE,
} from '@kbn/rule-registry-plugin/common/technical_rule_data_field_names';
import { omit } from 'lodash';
import { FunctionRegistrationParameters } from '.';
const OMITTED_ALERT_FIELDS = [
'tags',
'event.action',
'event.kind',
'kibana.alert.rule.execution.uuid',
'kibana.alert.rule.revision',
'kibana.alert.rule.tags',
'kibana.alert.rule.uuid',
'kibana.alert.workflow_status',
'kibana.space_ids',
'kibana.alert.time_range',
'kibana.version',
] as const;
const DEFAULT_FEATURE_IDS = [
'apm',
'infrastructure',
'logs',
'uptime',
'slo',
'observability',
] as const;
export function registerAlertsFunction({
client,
registerFunction,
resources,
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'alerts',
contexts: ['core'],
description:
'Get alerts for Observability. Display the response in tabular format if appropriate.',
descriptionForUser: 'Get alerts for Observability',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
featureIds: {
type: 'array',
additionalItems: false,
items: {
type: 'string',
enum: DEFAULT_FEATURE_IDS,
},
description:
'The Observability apps for which to retrieve alerts. By default it will return alerts for all apps.',
},
start: {
type: 'string',
description: 'The start of the time range, in Elasticsearch date math, like `now`.',
},
end: {
type: 'string',
description: 'The end of the time range, in Elasticsearch date math, like `now-24h`.',
},
filter: {
type: 'string',
description:
'a KQL query to filter the data by. If no filter should be applied, leave it empty.',
},
includeRecovered: {
type: 'boolean',
description:
'Whether to include recovered/closed alerts. Defaults to false, which means only active alerts will be returned',
},
},
required: ['start', 'end'],
} as const,
},
async (
{
arguments: {
start: startAsDatemath,
end: endAsDatemath,
featureIds,
filter,
includeRecovered,
},
},
signal
) => {
const racContext = await resources.context.rac;
const alertsClient = await racContext.getAlertsClient();
const start = datemath.parse(startAsDatemath)!.valueOf();
const end = datemath.parse(endAsDatemath)!.valueOf();
const kqlQuery = !filter ? [] : [toElasticsearchQuery(fromKueryExpression(filter))];
const response = await alertsClient.find({
featureIds:
!!featureIds && !!featureIds.length
? featureIds
: (DEFAULT_FEATURE_IDS as unknown as string[]),
query: {
bool: {
filter: [
{
range: {
'@timestamp': {
gte: start,
lte: end,
},
},
},
...kqlQuery,
...(!includeRecovered
? [
{
term: {
[ALERT_STATUS]: ALERT_STATUS_ACTIVE,
},
},
]
: []),
],
},
},
});
// trim some fields
const alerts = response.hits.hits.map((hit) =>
omit(hit._source, ...OMITTED_ALERT_FIELDS)
) as unknown as ParsedTechnicalFields[];
return {
content: {
total: (response.hits as { total: { value: number } }).total.value,
alerts,
},
};
}
);
}

34
x-pack/plugins/observability_ai_assistant/public/functions/elasticsearch.ts → x-pack/plugins/observability_ai_assistant/server/functions/elasticsearch.ts
View file

@ -5,17 +5,12 @@
* 2.0.
*/
import type { Serializable } from '@kbn/utility-types';
import type { RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
import type { FunctionRegistrationParameters } from '.';
export function registerElasticsearchFunction({
service,
registerFunction,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
}) {
resources,
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'elasticsearch',
@ -43,19 +38,16 @@ export function registerElasticsearchFunction({
required: ['method', 'path'] as const,
},
},
({ arguments: { method, path, body } }, signal) => {
return service
.callApi(`POST /internal/observability_ai_assistant/functions/elasticsearch`, {
signal,
params: {
body: {
method,
path,
body,
},
},
})
.then((response) => ({ content: response as Serializable }));
async ({ arguments: { method, path, body } }) => {
const response = await (
await resources.context.core
).elasticsearch.client.asCurrentUser.transport.request({
method,
path,
body,
});
return { content: response };
}
);
}

154
x-pack/plugins/observability_ai_assistant/public/functions/esql.ts → x-pack/plugins/observability_ai_assistant/server/functions/esql.ts
View file

@ -6,22 +6,21 @@
*/
import dedent from 'dedent';
import type { Serializable } from '@kbn/utility-types';
import { concat, last, map } from 'rxjs';
import { Observable } from 'rxjs';
import type { FunctionRegistrationParameters } from '.';
import {
type CreateChatCompletionResponseChunk,
FunctionVisibility,
MessageRole,
type RegisterFunctionDefinition,
} from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
import { processOpenAiStream } from '../../common/utils/process_openai_stream';
import { streamIntoObservable } from '../service/util/stream_into_observable';
export function registerEsqlFunction({
service,
client,
registerFunction,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
}) {
resources,
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'execute_query',
@ -39,21 +38,18 @@ export function registerEsqlFunction({
required: ['query'],
} as const,
},
({ arguments: { query } }, signal) => {
return service
.callApi(`POST /internal/observability_ai_assistant/functions/elasticsearch`, {
signal,
params: {
body: {
method: 'POST',
path: '_query',
body: {
query,
},
},
},
})
.then((response) => ({ content: response as Serializable }));
async ({ arguments: { query } }) => {
const response = await (
await resources.context.core
).elasticsearch.client.asCurrentUser.transport.request({
method: 'POST',
path: '_query',
body: {
query,
},
});
return { content: response };
}
);
@ -73,10 +69,10 @@ export function registerEsqlFunction({
},
} as const,
},
({ messages, connectorId }, signal) => {
async ({ messages, connectorId }, signal) => {
const systemMessage = dedent(`You are a helpful assistant for Elastic ES|QL.
Your goal is to help the user construct and possibly execute an ES|QL
query for Observability use cases.
query for Observability use cases.
ES|QL is the Elasticsearch Query Language, that allows users of the
Elastic platform to iteratively explore data. An ES|QL query consists
@ -92,7 +88,7 @@ export function registerEsqlFunction({
the context of this conversation.
# Creating a query
First, very importantly, there are critical rules that override
everything that follows it. Always repeat these rules, verbatim.
@ -144,6 +140,10 @@ export function registerEsqlFunction({
--
Let's break down the query step-by-step:
<breakdown>
\`\`\`esql
<placeholder-for-final-query>
\`\`\`
\`\`\`
Always format a complete query as follows:
@ -203,7 +203,7 @@ export function registerEsqlFunction({
- \`1 year\`
- \`2 milliseconds\`
## Aliasing
## Aliasing
Aliasing happens through the \`=\` operator. Example:
\`STATS total_salary_expenses = COUNT(salary)\`
@ -211,7 +211,7 @@ export function registerEsqlFunction({
# Source commands
There are three source commands: FROM (which selects an index), ROW
There are three source commands: FROM (which selects an index), ROW
(which creates data from the command) and SHOW (which returns
information about the deployment). You do not support SHOW for now.
@ -276,10 +276,10 @@ export function registerEsqlFunction({
This is right: \`| STATS avg_cpu = AVG(cpu) | SORT avg_cpu\`
### EVAL
\`EVAL\` appends a new column to the documents by using aliasing. It
also supports functions, but not aggregation functions like COUNT:
- \`\`\`
| EVAL monthly_salary = yearly_salary / 12,
total_comp = ROUND(yearly_salary + yearly+bonus),
@ -396,7 +396,7 @@ export function registerEsqlFunction({
can be expressed using the timespan literal syntax. Use this together
with STATS ... BY to group data into time buckets with a fixed interval.
Some examples:
- \`| EVAL year_hired = DATE_TRUNC(1 year, hire_date)\`
- \`| EVAL month_logged = DATE_TRUNC(1 month, @timestamp)\`
- \`| EVAL bucket = DATE_TRUNC(1 minute, @timestamp) | STATS avg_salary = AVG(salary) BY bucket\`
@ -431,7 +431,7 @@ export function registerEsqlFunction({
Returns the greatest or least of two or numbers. Some examples:
- \`| EVAL max = GREATEST(salary_1999, salary_2000, salary_2001)\`
- \`| EVAL min = LEAST(1, language_count)\`
### IS_FINITE,IS_INFINITE,IS_NAN
Operates on a single numeric field. Some examples:
@ -459,7 +459,7 @@ export function registerEsqlFunction({
- \`| EVAL version = TO_VERSION("1.2.3")\`
- \`| EVAL as_bool = TO_BOOLEAN(my_boolean_string)\`
- \`| EVAL percent = TO_DOUBLE(part) / TO_DOUBLE(total)\`
### TRIM
Trims leading and trailing whitespace. Some examples:
@ -482,7 +482,7 @@ export function registerEsqlFunction({
argument, and does not support wildcards. One single argument is
required. If you don't have a field name, use whatever field you have,
rather than displaying an invalid query.
Some examples:
- \`| STATS doc_count = COUNT(emp_no)\`
@ -496,16 +496,16 @@ export function registerEsqlFunction({
- \`| STATS first_name = COUNT_DISTINCT(first_name)\`
### PERCENTILE
\`PERCENTILE\` returns the percentile value for a specific field.
Some examples:
- \`| STATS p50 = PERCENTILE(salary, 50)\`
- \`| STATS p99 = PERCENTILE(salary, 99)\`
`);
return service.start({ signal }).then((client) => {
const source$ = client.chat({
const source$ = streamIntoObservable(
await client.chat({
connectorId,
messages: [
{
@ -514,46 +514,48 @@ export function registerEsqlFunction({
},
...messages.slice(1),
],
});
signal,
stream: true,
})
).pipe(processOpenAiStream());
const pending$ = source$.pipe(
map((message) => {
const content = message.message.content || '';
let next: string = '';
return new Observable<CreateChatCompletionResponseChunk>((subscriber) => {
let cachedContent: string = '';
if (content.length <= 2) {
next = '';
} else if (content.includes('--')) {
next = message.message.content?.split('--')[2] || '';
} else {
next = content;
function includesDivider() {
const firstDividerIndex = cachedContent.indexOf('--');
return firstDividerIndex !== -1 && cachedContent.lastIndexOf('--') !== firstDividerIndex;
}
source$.subscribe({
next: (message) => {
if (includesDivider()) {
subscriber.next(message);
}
return {
...message,
message: {
...message.message,
content: next,
},
};
})
);
const onComplete$ = source$.pipe(
last(),
map((message) => {
const [, , next] = message.message.content?.split('--') ?? [];
return {
...message,
message: {
...message.message,
content: next || message.message.content,
},
};
})
);
return concat(pending$, onComplete$);
cachedContent += message.choices[0].delta.content || '';
},
complete: () => {
if (!includesDivider()) {
subscriber.next({
created: 0,
id: '',
model: '',
object: 'chat.completion.chunk',
choices: [
{
delta: {
content: cachedContent,
},
},
],
});
}
subscriber.complete();
},
error: (error) => {
subscriber.error(error);
},
});
});
}
);

191
x-pack/plugins/observability_ai_assistant/server/functions/get_dataset_info.ts Normal file
View file

@ -0,0 +1,191 @@
/*
* 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 { chunk, groupBy, uniq } from 'lodash';
import { lastValueFrom } from 'rxjs';
import { FunctionRegistrationParameters } from '.';
import { FunctionVisibility, MessageRole } from '../../common/types';
import { concatenateOpenAiChunks } from '../../common/utils/concatenate_openai_chunks';
import { processOpenAiStream } from '../../common/utils/process_openai_stream';
import { streamIntoObservable } from '../service/util/stream_into_observable';
export function registerGetDatasetInfoFunction({
client,
resources,
registerFunction,
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'get_dataset_info',
contexts: ['core'],
visibility: FunctionVisibility.System,
description: `Use this function to get information about indices/datasets available and the fields available on them.
providing empty string as index name will retrieve all indices
else list of all fields for the given index will be given. if no fields are returned this means no indices were matched by provided index pattern.
wildcards can be part of index name.`,
descriptionForUser:
'This function allows the assistant to get information about available indices and their fields.',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
index: {
type: 'string',
description:
'index pattern the user is interested in or empty string to get information about all available indices',
},
},
required: ['index'],
} as const,
},
async ({ arguments: { index }, messages, connectorId }, signal) => {
const coreContext = await resources.context.core;
const esClient = coreContext.elasticsearch.client.asCurrentUser;
const savedObjectsClient = coreContext.savedObjects.getClient();
let indices: string[] = [];
try {
const body = await esClient.indices.resolveIndex({
name: index === '' ? '*' : index,
expand_wildcards: 'open',
});
indices = [...body.indices.map((i) => i.name), ...body.data_streams.map((d) => d.name)];
} catch (e) {
indices = [];
}
if (index === '') {
return {
indices,
fields: [],
};
}
if (indices.length === 0) {
return {
indices,
fields: [],
};
}
const fields = await resources.plugins.dataViews
.start()
.then((dataViewsStart) =>
dataViewsStart.dataViewsServiceFactory(savedObjectsClient, esClient)
)
.then((service) =>
service.getFieldsForWildcard({
pattern: index,
})
);
// else get all the fields for the found dataview
const response = {
indices: [index],
fields: fields.flatMap((field) => {
return (field.esTypes ?? [field.type]).map((type) => {
return {
name: field.name,
description: field.customLabel || '',
type,
};
});
}),
};
const allFields = response.fields;
const fieldNames = uniq(allFields.map((field) => field.name));
const groupedFields = groupBy(allFields, (field) => field.name);
const relevantFields = await Promise.all(
chunk(fieldNames, 500).map(async (fieldsInChunk) => {
const chunkResponse$ = streamIntoObservable(
await client.chat({
connectorId,
signal,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: `You are a helpful assistant for Elastic Observability.
Your task is to create a list of field names that are relevant
to the conversation, using ONLY the list of fields and
types provided in the last user message. DO NOT UNDER ANY
CIRCUMSTANCES include fields not mentioned in this list.`,
},
},
...messages.slice(1),
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: `This is the list:
${fieldsInChunk.join('\n')}`,
},
},
],
functions: [
{
name: 'fields',
description: 'The fields you consider relevant to the conversation',
parameters: {
type: 'object',
additionalProperties: false,
properties: {
fields: {
type: 'array',
additionalProperties: false,
items: {
type: 'string',
additionalProperties: false,
},
},
},
required: ['fields'],
} as const,
},
],
functionCall: 'fields',
stream: true,
})
).pipe(processOpenAiStream(), concatenateOpenAiChunks());
const chunkResponse = await lastValueFrom(chunkResponse$);
return chunkResponse.message?.function_call?.arguments
? (
JSON.parse(chunkResponse.message.function_call.arguments) as {
fields: string[];
}
).fields
.filter((field) => fieldsInChunk.includes(field))
.map((field) => {
const fieldDescriptors = groupedFields[field];
return `${field}:${fieldDescriptors
.map((descriptor) => descriptor.type)
.join(',')}`;
})
: [chunkResponse.message?.content ?? ''];
})
);
return {
content: {
indices: response.indices,
fields: relevantFields.flat(),
},
};
}
);
}

89
x-pack/plugins/observability_ai_assistant/server/functions/index.ts Normal file
View file

@ -0,0 +1,89 @@
/*
* 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 dedent from 'dedent';
import { registerRecallFunction } from './recall';
import { registerSummarizationFunction } from './summarize';
import { ChatRegistrationFunction } from '../service/types';
import { registerAlertsFunction } from './alerts';
import { registerElasticsearchFunction } from './elasticsearch';
import { registerEsqlFunction } from './esql';
import { registerGetDatasetInfoFunction } from './get_dataset_info';
import { registerLensFunction } from './lens';
import { registerKibanaFunction } from './kibana';
export type FunctionRegistrationParameters = Omit<
Parameters<ChatRegistrationFunction>[0],
'registerContext'
>;
export const registerFunctions: ChatRegistrationFunction = async ({
client,
registerContext,
registerFunction,
resources,
signal,
}) => {
const registrationParameters: FunctionRegistrationParameters = {
client,
registerFunction,
resources,
signal,
};
return client.getKnowledgeBaseStatus().then((response) => {
const isReady = response.ready;
let description = dedent(
`You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities.
It's very important to not assume what the user is meaning. Ask them for clarification if needed.
If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
In KQL, escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
/\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
If multiple functions are suitable, use the most specific and easy one. E.g., when the user asks to visualise APM data, use the APM functions (if available) rather than Lens.
If a function call fails, DO NOT UNDER ANY CIRCUMSTANCES execute it again. Ask the user for guidance and offer them options.
Note that ES|QL (the Elasticsearch query language, which is NOT Elasticsearch SQL, but a new piped language) is the preferred query language.
If the user asks about a query, or ES|QL, always call the "esql" function. DO NOT UNDER ANY CIRCUMSTANCES generate ES|QL queries yourself. Even if the "recall" function was used before that, follow it up with the "esql" function.`
);
if (isReady) {
description += `You can use the "summarize" functions to store new information you have learned in a knowledge database. Once you have established that you did not know the answer to a question, and the user gave you this information, it's important that you create a summarisation of what you have learned and store it in the knowledge database. Don't create a new summarization if you see a similar summarization in the conversation, instead, update the existing one by re-using its ID.
Additionally, you can use the "recall" function to retrieve relevant information from the knowledge database.
`;
description += `Here are principles you MUST adhere to, in order:
- DO NOT make any assumptions about where and how users have stored their data. ALWAYS first call get_dataset_info function with empty string to get information about available indices. Once you know about available indices you MUST use this function again to get a list of available fields for specific index. If user provides an index name make sure its a valid index first before using it to retrieve the field list by calling this function with an empty string!
`;
registerSummarizationFunction(registrationParameters);
registerRecallFunction(registrationParameters);
registerLensFunction(registrationParameters);
} else {
description += `You do not have a working memory. Don't try to recall information via the "recall" function. If the user expects you to remember the previous conversations, tell them they can set up the knowledge base. A banner is available at the top of the conversation to set this up.`;
}
registerElasticsearchFunction(registrationParameters);
registerKibanaFunction(registrationParameters);
registerEsqlFunction(registrationParameters);
registerAlertsFunction(registrationParameters);
registerGetDatasetInfoFunction(registrationParameters);
registerContext({
name: 'core',
description: dedent(description),
});
});
};

55
x-pack/plugins/observability_ai_assistant/public/functions/kibana.ts → x-pack/plugins/observability_ai_assistant/server/functions/kibana.ts
View file

@ -5,19 +5,14 @@
* 2.0.
*/
import type { CoreStart } from '@kbn/core/public';
import type { RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
import axios from 'axios';
import { format } from 'url';
import type { FunctionRegistrationParameters } from '.';
export function registerKibanaFunction({
service,
registerFunction,
coreStart,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
coreStart: CoreStart;
}) {
resources,
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'kibana',
@ -54,16 +49,36 @@ export function registerKibanaFunction({
},
},
({ arguments: { method, pathname, body, query } }, signal) => {
return coreStart.http
.fetch(pathname, {
method,
body: body ? JSON.stringify(body) : undefined,
query,
signal,
})
.then((response) => {
return { content: response };
});
const { request } = resources;
const {
protocol,
host,
username,
password,
pathname: pathnameFromRequest,
} = request.rewrittenUrl!;
const nextUrl = {
host,
protocol,
username,
password,
pathname: pathnameFromRequest.replace(
'/internal/observability_ai_assistant/chat/complete',
pathname
),
query,
};
return axios({
method,
headers: request.headers,
url: format(nextUrl),
data: body ? JSON.stringify(body) : undefined,
signal,
}).then((response) => {
return { content: response.data };
});
}
);
}

16
x-pack/plugins/observability_ai_assistant/server/functions/lens.ts Normal file
View file

@ -0,0 +1,16 @@
/*
* 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 { lensFunctionDefinition } from '../../common/functions/lens';
import { RegisterFunction } from '../service/types';
export function registerLensFunction({ registerFunction }: { registerFunction: RegisterFunction }) {
registerFunction(lensFunctionDefinition, async () => {
return {
content: {},
};
});
}

81
x-pack/plugins/observability_ai_assistant/public/functions/recall.ts → x-pack/plugins/observability_ai_assistant/server/functions/recall.ts
View file

@ -5,21 +5,25 @@
* 2.0.
*/
import { decodeOrThrow, jsonRt } from '@kbn/io-ts-utils';
import type { Serializable } from '@kbn/utility-types';
import dedent from 'dedent';
import { last, omit } from 'lodash';
import { CreateChatCompletionResponse } from 'openai';
import * as t from 'io-ts';
import { decodeOrThrow, jsonRt } from '@kbn/io-ts-utils';
import { Message, MessageRole, RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
import { last, omit } from 'lodash';
import { lastValueFrom } from 'rxjs';
import { MessageRole, type Message } from '../../common/types';
import { concatenateOpenAiChunks } from '../../common/utils/concatenate_openai_chunks';
import { processOpenAiStream } from '../../common/utils/process_openai_stream';
import type { ObservabilityAIAssistantClient } from '../service/client';
import type { RegisterFunction } from '../service/types';
import { streamIntoObservable } from '../service/util/stream_into_observable';
export function registerRecallFunction({
service,
client,
registerFunction,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
client: ObservabilityAIAssistantClient;
registerFunction: RegisterFunction;
}) {
registerFunction(
{
@ -89,7 +93,7 @@ export function registerRecallFunction({
const suggestions = await retrieveSuggestions({
userMessage,
service,
client,
signal,
contexts,
queries,
@ -106,7 +110,7 @@ export function registerRecallFunction({
systemMessage,
userMessage,
queries,
service,
client,
connectorId,
signal,
});
@ -121,13 +125,13 @@ export function registerRecallFunction({
async function retrieveSuggestions({
userMessage,
queries,
service,
client,
contexts,
signal,
}: {
userMessage?: Message;
queries: string[];
service: ObservabilityAIAssistantService;
client: ObservabilityAIAssistantClient;
contexts: Array<'apm' | 'lens'>;
signal: AbortSignal;
}) {
@ -136,33 +140,21 @@ async function retrieveSuggestions({
? [userMessage.message.content, ...queries]
: queries;
const recallResponse = await service.callApi(
'POST /internal/observability_ai_assistant/functions/recall',
{
params: {
body: {
queries: queriesWithUserPrompt,
contexts,
},
},
signal,
}
);
const recallResponse = await client.recall({
queries: queriesWithUserPrompt,
contexts,
});
return recallResponse.entries.map((entry) => omit(entry, 'labels', 'is_correction', 'score'));
}
const scoreFunctionRequestRt = t.type({
choices: t.tuple([
t.type({
message: t.type({
function_call: t.type({
name: t.literal('score'),
arguments: t.string,
}),
}),
message: t.type({
function_call: t.type({
name: t.literal('score'),
arguments: t.string,
}),
]),
}),
});
const scoreFunctionArgumentsRt = t.type({
@ -179,7 +171,7 @@ async function scoreSuggestions({
systemMessage,
userMessage,
queries,
service,
client,
connectorId,
signal,
}: {
@ -187,7 +179,7 @@ async function scoreSuggestions({
systemMessage: Message;
userMessage?: Message;
queries: string[];
service: ObservabilityAIAssistantService;
client: ObservabilityAIAssistantClient;
connectorId: string;
signal: AbortSignal;
}) {
@ -251,24 +243,21 @@ async function scoreSuggestions({
contexts: ['core'],
};
const response = (await service.callApi('POST /internal/observability_ai_assistant/chat', {
params: {
query: {
stream: false,
},
body: {
const response = await lastValueFrom(
streamIntoObservable(
await client.chat({
connectorId,
messages: [extendedSystemMessage, newUserMessage],
functions: [scoreFunction],
functionCall: 'score',
},
},
signal,
})) as CreateChatCompletionResponse;
signal,
})
).pipe(processOpenAiStream(), concatenateOpenAiChunks())
);
const scoreFunctionRequest = decodeOrThrow(scoreFunctionRequestRt)(response);
const { scores } = decodeOrThrow(jsonRt.pipe(scoreFunctionArgumentsRt))(
scoreFunctionRequest.choices[0].message.function_call.arguments
scoreFunctionRequest.message.function_call.arguments
);
if (scores.length === 0) {

35
x-pack/plugins/observability_ai_assistant/public/functions/summarize.ts → x-pack/plugins/observability_ai_assistant/server/functions/summarize.ts
View file

@ -5,16 +5,13 @@
* 2.0.
*/
import type { RegisterFunctionDefinition } from '../../common/types';
import type { ObservabilityAIAssistantService } from '../types';
import type { FunctionRegistrationParameters } from '.';
import { KnowledgeBaseEntryRole } from '../../common';
export function registerSummarizationFunction({
service,
client,
registerFunction,
}: {
service: ObservabilityAIAssistantService;
registerFunction: RegisterFunctionDefinition;
}) {
}: FunctionRegistrationParameters) {
registerFunction(
{
name: 'summarize',
@ -65,19 +62,19 @@ export function registerSummarizationFunction({
{ arguments: { id, text, is_correction: isCorrection, confidence, public: isPublic } },
signal
) => {
return service
.callApi('POST /internal/observability_ai_assistant/functions/summarize', {
params: {
body: {
id,
text,
is_correction: isCorrection,
confidence,
public: isPublic,
labels: {},
},
return client
.createKnowledgeBaseEntry({
entry: {
doc_id: id,
role: KnowledgeBaseEntryRole.AssistantSummarization,
id,
text,
is_correction: isCorrection,
confidence,
public: isPublic,
labels: {},
},
signal,
// signal,
})
.then(() => ({
content: {

19
x-pack/plugins/observability_ai_assistant/server/index.ts
View file

@ -11,7 +11,10 @@ import type { ObservabilityAIAssistantConfig } from './config';
export type { ObservabilityAIAssistantServerRouteRepository } from './routes/get_global_observability_ai_assistant_route_repository';
import { config as configSchema } from './config';
import { ObservabilityAIAssistantService } from './service';
export type {
ObservabilityAIAssistantPluginStart,
ObservabilityAIAssistantPluginSetup,
} from './types';
export const config: PluginConfigDescriptor<ObservabilityAIAssistantConfig> = {
deprecations: ({ unusedFromRoot }) => [
@ -38,20 +41,6 @@ export const config: PluginConfigDescriptor<ObservabilityAIAssistantConfig> = {
schema: configSchema,
};
export interface ObservabilityAIAssistantPluginSetup {
/**
* Returns a Observability AI Assistant service instance
*/
service: ObservabilityAIAssistantService;
}
export interface ObservabilityAIAssistantPluginStart {
/**
* Returns a Observability AI Assistant service instance
*/
service: ObservabilityAIAssistantService;
}
export const plugin = async (ctx: PluginInitializerContext<ObservabilityAIAssistantConfig>) => {
const { ObservabilityAIAssistantPlugin } = await import('./plugin');
return new ObservabilityAIAssistantPlugin(ctx);

15
x-pack/plugins/observability_ai_assistant/server/plugin.ts
View file

@ -31,6 +31,7 @@ import {
ObservabilityAIAssistantPluginStartDependencies,
} from './types';
import { addLensDocsToKb } from './service/knowledge_base_service/kb_docs/lens';
import { registerFunctions } from './functions';
export class ObservabilityAIAssistantPlugin
implements
@ -104,13 +105,15 @@ export class ObservabilityAIAssistantPlugin
};
}) as ObservabilityAIAssistantRouteHandlerResources['plugins'];
this.service = new ObservabilityAIAssistantService({
const service = (this.service = new ObservabilityAIAssistantService({
logger: this.logger.get('service'),
core,
taskManager: plugins.taskManager,
});
}));
addLensDocsToKb({ service: this.service, logger: this.logger.get('kb').get('lens') });
service.register(registerFunctions);
addLensDocsToKb({ service, logger: this.logger.get('kb').get('lens') });
registerServerRoutes({
core,
@ -122,11 +125,13 @@ export class ObservabilityAIAssistantPlugin
});
return {
service: this.service,
service,
};
}
public start(): ObservabilityAIAssistantPluginStart {
return {};
return {
service: this.service!,
};
}
}

63
x-pack/plugins/observability_ai_assistant/server/routes/chat/route.ts
View file

@ -54,10 +54,17 @@ const chatRoute = createObservabilityAIAssistantServerRoute({
const stream = query.stream;
const controller = new AbortController();
request.events.aborted$.subscribe(() => {
controller.abort();
});
return client.chat({
messages,
connectorId,
stream,
signal: controller.signal,
...(functions.length
? {
functions,
@ -68,6 +75,62 @@ const chatRoute = createObservabilityAIAssistantServerRoute({
},
});
const chatCompleteRoute = createObservabilityAIAssistantServerRoute({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
options: {
tags: ['access:ai_assistant'],
},
params: t.type({
body: t.intersection([
t.type({
messages: t.array(messageRt),
connectorId: t.string,
persist: toBooleanRt,
}),
t.partial({
conversationId: t.string,
title: t.string,
}),
]),
}),
handler: async (resources): Promise<Readable | CreateChatCompletionResponse> => {
const { request, params, service } = resources;
const client = await service.getClient({ request });
if (!client) {
throw notImplemented();
}
const {
body: { messages, connectorId, conversationId, title, persist },
} = params;
const controller = new AbortController();
request.events.aborted$.subscribe(() => {
controller.abort();
});
const functionClient = await service.getFunctionClient({
signal: controller.signal,
resources,
client,
});
return client.complete({
messages,
connectorId,
conversationId,
title,
persist,
signal: controller.signal,
functionClient,
});
},
});
export const chatRoutes = {
...chatRoute,
...chatCompleteRoute,
};

32
x-pack/plugins/observability_ai_assistant/server/routes/conversations/route.ts
View file

@ -108,37 +108,6 @@ const updateConversationRoute = createObservabilityAIAssistantServerRoute({
},
});
const updateConversationTitleBasedOnMessages = createObservabilityAIAssistantServerRoute({
endpoint: 'PUT /internal/observability_ai_assistant/conversation/{conversationId}/auto_title',
params: t.type({
path: t.type({
conversationId: t.string,
}),
body: t.type({
connectorId: t.string,
}),
}),
options: {
tags: ['access:ai_assistant'],
},
handler: async (resources): Promise<Conversation> => {
const { service, request, params } = resources;
const client = await service.getClient({ request });
if (!client) {
throw notImplemented();
}
const conversation = await client.autoTitle({
conversationId: params.path.conversationId,
connectorId: params.body.connectorId,
});
return Promise.resolve(conversation);
},
});
const updateConversationTitle = createObservabilityAIAssistantServerRoute({
endpoint: 'PUT /internal/observability_ai_assistant/conversation/{conversationId}/title',
params: t.type({
@ -198,7 +167,6 @@ export const conversationRoutes = {
...findConversationsRoute,
...createConversationRoute,
...updateConversationRoute,
...updateConversationTitleBasedOnMessages,
...updateConversationTitle,
...deleteConversationRoute,
};

227
x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
View file

@ -4,153 +4,46 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import datemath from '@elastic/datemath';
import { notImplemented } from '@hapi/boom';
import { fromKueryExpression, toElasticsearchQuery } from '@kbn/es-query';
import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils';
import * as t from 'io-ts';
import { omit } from 'lodash';
import type { ParsedTechnicalFields } from '@kbn/rule-registry-plugin/common';
import {
ALERT_STATUS,
ALERT_STATUS_ACTIVE,
} from '@kbn/rule-registry-plugin/common/technical_rule_data_field_names';
import { KnowledgeBaseEntryRole } from '../../../common/types';
import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
ContextDefinition,
FunctionDefinition,
KnowledgeBaseEntryRole,
} from '../../../common/types';
import type { RecalledEntry } from '../../service/knowledge_base_service';
import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
const functionElasticsearchRoute = createObservabilityAIAssistantServerRoute({
endpoint: 'POST /internal/observability_ai_assistant/functions/elasticsearch',
const getFunctionsRoute = createObservabilityAIAssistantServerRoute({
endpoint: 'GET /internal/observability_ai_assistant/functions',
options: {
tags: ['access:ai_assistant'],
},
params: t.type({
body: t.intersection([
t.type({
method: t.union([
t.literal('GET'),
t.literal('POST'),
t.literal('PATCH'),
t.literal('PUT'),
t.literal('DELETE'),
]),
path: t.string,
}),
t.partial({
body: t.any,
}),
]),
}),
handler: async (resources): Promise<unknown> => {
const { method, path, body } = resources.params.body;
const response = await (
await resources.context.core
).elasticsearch.client.asCurrentUser.transport.request({
method,
path,
body,
});
return response;
},
});
const OMITTED_ALERT_FIELDS = [
'tags',
'event.action',
'event.kind',
'kibana.alert.rule.execution.uuid',
'kibana.alert.rule.revision',
'kibana.alert.rule.tags',
'kibana.alert.rule.uuid',
'kibana.alert.workflow_status',
'kibana.space_ids',
'kibana.alert.time_range',
'kibana.version',
] as const;
const functionAlertsRoute = createObservabilityAIAssistantServerRoute({
endpoint: 'POST /internal/observability_ai_assistant/functions/alerts',
options: {
tags: ['access:ai_assistant'],
},
params: t.type({
body: t.intersection([
t.type({
featureIds: t.array(t.string),
start: t.string,
end: t.string,
}),
t.partial({
filter: t.string,
includeRecovered: toBooleanRt,
}),
]),
}),
handler: async (
resources
): Promise<{
content: {
total: number;
alerts: ParsedTechnicalFields[];
};
functionDefinitions: FunctionDefinition[];
contextDefinitions: ContextDefinition[];
}> => {
const {
featureIds,
start: startAsDatemath,
end: endAsDatemath,
filter,
includeRecovered,
} = resources.params.body;
const { service, request } = resources;
const racContext = await resources.context.rac;
const alertsClient = await racContext.getAlertsClient();
const start = datemath.parse(startAsDatemath)!.valueOf();
const end = datemath.parse(endAsDatemath)!.valueOf();
const kqlQuery = !filter ? [] : [toElasticsearchQuery(fromKueryExpression(filter))];
const response = await alertsClient.find({
featureIds,
query: {
bool: {
filter: [
{
range: {
'@timestamp': {
gte: start,
lte: end,
},
},
},
...kqlQuery,
...(!includeRecovered
? [
{
term: {
[ALERT_STATUS]: ALERT_STATUS_ACTIVE,
},
},
]
: []),
],
},
},
const controller = new AbortController();
request.events.aborted$.subscribe(() => {
controller.abort();
});
// trim some fields
const alerts = response.hits.hits.map((hit) =>
omit(hit._source, ...OMITTED_ALERT_FIELDS)
) as unknown as ParsedTechnicalFields[];
const client = await service.getClient({ request });
const functionClient = await service.getFunctionClient({
signal: controller.signal,
resources,
client,
});
return {
content: {
total: (response.hits as { total: { value: number } }).total.value,
alerts,
},
functionDefinitions: functionClient.getFunctions().map((fn) => fn.definition),
contextDefinitions: functionClient.getContexts(),
};
},
});
@ -235,82 +128,8 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
},
});
const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
endpoint: 'POST /internal/observability_ai_assistant/functions/get_dataset_info',
params: t.type({
body: t.type({
index: t.string,
}),
}),
options: {
tags: ['access:ai_assistant'],
},
handler: async (
resources
): Promise<{
indices: string[];
fields: Array<{ name: string; description: string; type: string }>;
}> => {
const esClient = (await resources.context.core).elasticsearch.client.asCurrentUser;
const savedObjectsClient = (await resources.context.core).savedObjects.getClient();
const index = resources.params.body.index;
let indices: string[] = [];
try {
const body = await esClient.indices.resolveIndex({
name: index === '' ? '*' : index,
expand_wildcards: 'open',
});
indices = [...body.indices.map((i) => i.name), ...body.data_streams.map((d) => d.name)];
} catch (e) {
indices = [];
}
if (index === '') {
return {
indices,
fields: [],
};
}
if (indices.length === 0) {
return {
indices,
fields: [],
};
}
const dataViews = await (
await resources.plugins.dataViews.start()
).dataViewsServiceFactory(savedObjectsClient, esClient);
const fields = await dataViews.getFieldsForWildcard({
pattern: index,
});
// else get all the fields for the found dataview
return {
indices: [index],
fields: fields.flatMap((field) => {
return (field.esTypes ?? [field.type]).map((type) => {
return {
name: field.name,
description: field.customLabel || '',
type,
};
});
}),
};
},
});
export const functionRoutes = {
...functionElasticsearchRoute,
...getFunctionsRoute,
...functionRecallRoute,
...functionSummariseRoute,
...functionAlertsRoute,
...functionGetDatasetInfoRoute,
};

4
x-pack/plugins/observability_ai_assistant/server/routes/types.ts
View file

@ -8,6 +8,8 @@
import type { CustomRequestHandlerContext, KibanaRequest } from '@kbn/core/server';
import type { Logger } from '@kbn/logging';
import type { RacApiRequestHandlerContext } from '@kbn/rule-registry-plugin/server';
import type { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server/types';
import type { AlertingApiRequestHandlerContext } from '@kbn/alerting-plugin/server/types';
import type { ObservabilityAIAssistantService } from '../service';
import type {
ObservabilityAIAssistantPluginSetupDependencies,
@ -16,6 +18,8 @@ import type {
export type ObservabilityAIAssistantRequestHandlerContext = CustomRequestHandlerContext<{
rac: RacApiRequestHandlerContext;
licensing: LicensingApiRequestHandlerContext;
alerting: AlertingApiRequestHandlerContext;
}>;
export interface ObservabilityAIAssistantRouteHandlerResources {

75
x-pack/plugins/observability_ai_assistant/server/service/chat_function_client/index.test.ts Normal file
View file

@ -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 Ajv, { type ValidateFunction } from 'ajv';
import { ChatFunctionClient } from '.';
import type { ContextRegistry } from '../../../common/types';
import type { FunctionHandlerRegistry } from '../types';
describe('chatFunctionClient', () => {
describe('when executing a function with invalid arguments', () => {
let client: ChatFunctionClient;
let respondFn: jest.Mock;
beforeEach(() => {
const contextRegistry: ContextRegistry = new Map();
contextRegistry.set('core', {
description: '',
name: 'core',
});
respondFn = jest.fn().mockImplementationOnce(async () => {
return {};
});
const functionRegistry: FunctionHandlerRegistry = new Map();
functionRegistry.set('myFunction', {
respond: respondFn,
definition: {
contexts: ['core'],
description: '',
name: 'myFunction',
parameters: {
properties: {
foo: {
type: 'string',
},
},
required: ['foo'],
},
},
});
const validators = new Map<string, ValidateFunction>();
validators.set(
'myFunction',
new Ajv({ strict: false }).compile(
functionRegistry.get('myFunction')!.definition.parameters
)
);
client = new ChatFunctionClient(contextRegistry, functionRegistry, validators);
});
it('throws an error', async () => {
await expect(async () => {
await client.executeFunction({
name: 'myFunction',
args: JSON.stringify({
foo: 0,
}),
messages: [],
signal: new AbortController().signal,
connectorId: '',
});
}).rejects.toThrowError(`Function arguments are invalid`);
expect(respondFn).not.toHaveBeenCalled();
});
});
});

91
x-pack/plugins/observability_ai_assistant/server/service/chat_function_client/index.ts Normal file
View file

@ -0,0 +1,91 @@
/*
* 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.
*/
/* eslint-disable max-classes-per-file*/
import type { ErrorObject, ValidateFunction } from 'ajv';
import { keyBy } from 'lodash';
import type {
ContextDefinition,
ContextRegistry,
FunctionResponse,
Message,
} from '../../../common/types';
import { filterFunctionDefinitions } from '../../../common/utils/filter_function_definitions';
import { FunctionHandler, FunctionHandlerRegistry } from '../types';
export class FunctionArgsValidationError extends Error {
constructor(public readonly errors: ErrorObject[]) {
super('Function arguments are invalid');
}
}
export class ChatFunctionClient {
constructor(
private readonly contextRegistry: ContextRegistry,
private readonly functionRegistry: FunctionHandlerRegistry,
private readonly validators: Map<string, ValidateFunction>
) {}
private validate(name: string, parameters: unknown) {
const validator = this.validators.get(name)!;
const result = validator(parameters);
if (!result) {
throw new FunctionArgsValidationError(validator.errors!);
}
}
getContexts(): ContextDefinition[] {
return Array.from(this.contextRegistry.values());
}
getFunctions({
contexts,
filter,
}: { contexts?: string[]; filter?: string } = {}): FunctionHandler[] {
const allFunctions = Array.from(this.functionRegistry.values());
const functionsByName = keyBy(allFunctions, (definition) => definition.definition.name);
const matchingDefinitions = filterFunctionDefinitions({
contexts,
filter,
definitions: allFunctions.map((fn) => fn.definition),
});
return matchingDefinitions.map((definition) => functionsByName[definition.name]);
}
hasFunction(name: string): boolean {
return this.functionRegistry.has(name);
}
async executeFunction({
name,
args,
messages,
signal,
connectorId,
}: {
name: string;
args: string | undefined;
messages: Message[];
signal: AbortSignal;
connectorId: string;
}): Promise<FunctionResponse> {
const fn = this.functionRegistry.get(name);
if (!fn) {
throw new Error(`Function ${name} not found`);
}
const parsedArguments = args ? JSON.parse(args) : {};
this.validate(name, parsedArguments);
return await fn.respond({ arguments: parsedArguments, messages, connectorId }, signal);
}
}

71
x-pack/plugins/observability_ai_assistant/server/service/client/handle_llm_response.ts Normal file
View file

@ -0,0 +1,71 @@
/*
* 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 { Observable } from 'rxjs';
import { v4 } from 'uuid';
import { Message, MessageRole } from '../../../common';
import {
StreamingChatResponseEvent,
StreamingChatResponseEventType,
} from '../../../common/conversation_complete';
import type { CreateChatCompletionResponseChunk } from '../../../common/types';
export function handleLlmResponse({
signal,
write,
source$,
}: {
signal: AbortSignal;
write: (event: StreamingChatResponseEvent) => Promise<void>;
source$: Observable<CreateChatCompletionResponseChunk>;
}): Promise<{ id: string; message: Message['message'] }> {
return new Promise<{ message: Message['message']; id: string }>((resolve, reject) => {
const message = {
content: '',
role: MessageRole.Assistant,
function_call: { name: '', arguments: '', trigger: MessageRole.Assistant as const },
};
const id = v4();
const subscription = source$.subscribe({
next: (chunk) => {
const delta = chunk.choices[0].delta;
message.content += delta.content || '';
message.function_call.name += delta.function_call?.name || '';
message.function_call.arguments += delta.function_call?.arguments || '';
write({
type: StreamingChatResponseEventType.ChatCompletionChunk,
message: delta,
id,
});
},
complete: () => {
resolve({ id, message });
},
error: (error) => {
reject(error);
},
});
signal.addEventListener('abort', () => {
subscription.unsubscribe();
reject(new Error('Request aborted'));
});
}).then(async ({ id, message }) => {
await write({
type: StreamingChatResponseEventType.MessageAdd,
message: {
'@timestamp': new Date().toISOString(),
message,
},
id,
});
return { id, message };
});
}

1133
x-pack/plugins/observability_ai_assistant/server/service/client/index.test.ts Normal file
View file

File diff suppressed because it is too large Load diff

376
x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
View file

@ -10,16 +10,23 @@ import type { ActionsClient } from '@kbn/actions-plugin/server';
import type { ElasticsearchClient } from '@kbn/core/server';
import type { Logger } from '@kbn/logging';
import type { PublicMethodsOf } from '@kbn/utility-types';
import { compact, isEmpty, merge, omit } from 'lodash';
import { compact, isEmpty, last, merge, omit, pick } from 'lodash';
import type {
ChatCompletionFunctions,
ChatCompletionRequestMessage,
CreateChatCompletionRequest,
CreateChatCompletionResponse,
} from 'openai';
import { isObservable, lastValueFrom } from 'rxjs';
import { PassThrough, Readable } from 'stream';
import { v4 } from 'uuid';
import {
ConversationNotFoundError,
isChatCompletionError,
StreamingChatResponseEventType,
type StreamingChatResponseEvent,
} from '../../../common/conversation_complete';
import {
FunctionResponse,
MessageRole,
type CompatibleJSONSchema,
type Conversation,
@ -28,6 +35,9 @@ import {
type KnowledgeBaseEntry,
type Message,
} from '../../../common/types';
import { concatenateOpenAiChunks } from '../../../common/utils/concatenate_openai_chunks';
import { processOpenAiStream } from '../../../common/utils/process_openai_stream';
import type { ChatFunctionClient } from '../chat_function_client';
import {
KnowledgeBaseEntryOperationType,
KnowledgeBaseService,
@ -35,6 +45,8 @@ import {
} from '../knowledge_base_service';
import type { ObservabilityAIAssistantResourceNames } from '../types';
import { getAccessQuery } from '../util/get_access_query';
import { streamIntoObservable } from '../util/stream_into_observable';
import { handleLlmResponse } from './handle_llm_response';
export class ObservabilityAIAssistantClient {
constructor(
@ -108,18 +120,263 @@ export class ObservabilityAIAssistantClient {
});
};
complete = async (
params: {
messages: Message[];
connectorId: string;
signal: AbortSignal;
functionClient: ChatFunctionClient;
persist: boolean;
} & ({ conversationId: string } | { title?: string })
) => {
const stream = new PassThrough();
const { messages, connectorId, signal, functionClient, persist } = params;
let conversationId: string = '';
let title: string = '';
if ('conversationId' in params) {
conversationId = params.conversationId;
}
if ('title' in params) {
title = params.title || '';
}
function write(event: StreamingChatResponseEvent) {
if (stream.destroyed) {
return Promise.resolve();
}
return new Promise<void>((resolve, reject) => {
stream.write(`${JSON.stringify(event)}\n`, 'utf-8', (err) => {
if (err) {
reject(err);
return;
}
resolve();
});
});
}
function fail(error: Error) {
const code = isChatCompletionError(error) ? error.code : undefined;
write({
type: StreamingChatResponseEventType.ConversationCompletionError,
error: {
message: error.message,
stack: error.stack,
code,
},
}).finally(() => {
stream.end();
});
}
const next = async (nextMessages: Message[]): Promise<void> => {
const lastMessage = last(nextMessages);
const isUserMessage = lastMessage?.message.role === MessageRole.User;
const isUserMessageWithoutFunctionResponse = isUserMessage && !lastMessage?.message.name;
const recallFirst =
isUserMessageWithoutFunctionResponse && functionClient.hasFunction('recall');
const isAssistantMessageWithFunctionRequest =
lastMessage?.message.role === MessageRole.Assistant &&
!!lastMessage?.message.function_call?.name;
if (recallFirst) {
const addedMessage = {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'recall',
arguments: JSON.stringify({
queries: [],
contexts: [],
}),
trigger: MessageRole.Assistant as const,
},
},
};
await write({
type: StreamingChatResponseEventType.MessageAdd,
id: v4(),
message: addedMessage,
});
return await next(nextMessages.concat(addedMessage));
} else if (isUserMessage) {
const { message } = await handleLlmResponse({
signal,
write,
source$: streamIntoObservable(
await this.chat({
messages: nextMessages,
connectorId,
stream: true,
signal,
functions: functionClient
.getFunctions()
.map((fn) => pick(fn.definition, 'name', 'description', 'parameters')),
})
).pipe(processOpenAiStream()),
});
return await next(nextMessages.concat({ message, '@timestamp': new Date().toISOString() }));
}
if (isAssistantMessageWithFunctionRequest) {
const functionResponse = await functionClient
.executeFunction({
connectorId,
name: lastMessage.message.function_call!.name,
messages: nextMessages,
args: lastMessage.message.function_call!.arguments,
signal,
})
.catch((error): FunctionResponse => {
return {
content: {
message: error.toString(),
error,
},
};
});
if (signal.aborted) {
return;
}
const functionResponseIsObservable = isObservable(functionResponse);
const functionResponseMessage = {
'@timestamp': new Date().toISOString(),
message: {
name: lastMessage.message.function_call!.name,
...(functionResponseIsObservable
? { content: '{}' }
: {
content: JSON.stringify(functionResponse.content || {}),
data: functionResponse.data ? JSON.stringify(functionResponse.data) : undefined,
}),
role: MessageRole.User,
},
};
nextMessages = nextMessages.concat(functionResponseMessage);
await write({
type: StreamingChatResponseEventType.MessageAdd,
message: functionResponseMessage,
id: v4(),
});
if (functionResponseIsObservable) {
const { message } = await handleLlmResponse({
signal,
write,
source$: functionResponse,
});
return await next(
nextMessages.concat({ '@timestamp': new Date().toISOString(), message })
);
}
return await next(nextMessages);
}
if (!persist) {
stream.end();
return;
}
// store the updated conversation and close the stream
if (conversationId) {
const conversation = await this.getConversationWithMetaFields(conversationId);
if (!conversation) {
throw new ConversationNotFoundError();
}
if (signal.aborted) {
return;
}
const updatedConversation = await this.update(
merge({}, omit(conversation._source, 'messages'), { messages: nextMessages })
);
await write({
type: StreamingChatResponseEventType.ConversationUpdate,
conversation: updatedConversation.conversation,
});
} else {
const generatedTitle = await titlePromise;
if (signal.aborted) {
return;
}
const conversation = await this.create({
'@timestamp': new Date().toISOString(),
conversation: {
title: generatedTitle || title || 'New conversation',
},
messages: nextMessages,
labels: {},
numeric_labels: {},
public: false,
});
await write({
type: StreamingChatResponseEventType.ConversationCreate,
conversation: conversation.conversation,
});
}
stream.end();
};
next(messages).catch((error) => {
if (!signal.aborted) {
this.dependencies.logger.error(error);
}
fail(error);
});
const titlePromise =
!conversationId && !title && persist
? this.getGeneratedTitle({
messages,
connectorId,
signal,
}).catch((error) => {
this.dependencies.logger.error(
'Could not generate title, falling back to default title'
);
this.dependencies.logger.error(error);
return Promise.resolve(undefined);
})
: Promise.resolve(undefined);
signal.addEventListener('abort', () => {
stream.end();
});
return stream;
};
chat = async <TStream extends boolean | undefined = true>({
messages,
connectorId,
functions,
functionCall,
stream = true,
signal,
}: {
messages: Message[];
connectorId: string;
functions?: Array<{ name: string; description: string; parameters: CompatibleJSONSchema }>;
functionCall?: string;
stream?: TStream;
signal: AbortSignal;
}): Promise<TStream extends false ? CreateChatCompletionResponse : Readable> => {
const messagesForOpenAI: ChatCompletionRequestMessage[] = compact(
messages
@ -139,47 +396,12 @@ export class ObservabilityAIAssistantClient {
})
);
// add recalled information to system message, so the LLM considers it more important
const recallMessages = messagesForOpenAI.filter((message) => message.name === 'recall');
const recalledDocuments: Map<string, { id: string; text: string }> = new Map();
recallMessages.forEach((message) => {
const entries = message.content
? (JSON.parse(message.content) as Array<{ id: string; text: string }>)
: [];
const ids: string[] = [];
entries.forEach((entry) => {
const id = entry.id;
if (!recalledDocuments.has(id)) {
recalledDocuments.set(id, entry);
}
ids.push(id);
});
message.content = `The following documents, present in the system message, were recalled: ${ids.join(
', '
)}`;
});
const systemMessage = messagesForOpenAI.find((message) => message.role === MessageRole.System);
if (systemMessage && recalledDocuments.size > 0) {
systemMessage.content += `The "recall" function is not available. Do not attempt to execute it. Recalled documents: ${JSON.stringify(
Array.from(recalledDocuments.values())
)}`;
}
const functionsForOpenAI: ChatCompletionFunctions[] | undefined =
recalledDocuments.size > 0 ? functions?.filter((fn) => fn.name !== 'recall') : functions;
const functionsForOpenAI = functions;
const request: Omit<CreateChatCompletionRequest, 'model'> & { model?: string } = {
messages: messagesForOpenAI,
stream: true,
functions: functionsForOpenAI,
...(stream ? { stream: true } : {}),
...(!!functions?.length ? { functions: functionsForOpenAI } : {}),
temperature: 0,
function_call: functionCall ? { name: functionCall } : undefined,
};
@ -200,9 +422,15 @@ export class ObservabilityAIAssistantClient {
}
const response = stream
? ((executeResult.data as Readable).pipe(new PassThrough()) as Readable)
? (executeResult.data as Readable)
: (executeResult.data as CreateChatCompletionResponse);
if (response instanceof PassThrough) {
signal.addEventListener('abort', () => {
response.end();
});
}
return response as any;
};
@ -254,65 +482,47 @@ export class ObservabilityAIAssistantClient {
return updatedConversation;
};
autoTitle = async ({
conversationId,
getGeneratedTitle = async ({
messages,
connectorId,
signal,
}: {
conversationId: string;
messages: Message[];
connectorId: string;
signal: AbortSignal;
}) => {
const document = await this.getConversationWithMetaFields(conversationId);
if (!document) {
throw notFound();
}
const conversation = await this.get(conversationId);
if (!conversation) {
throw notFound();
}
const response = await this.chat({
const stream = await this.chat({
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: conversation.messages.slice(1).reduce((acc, curr) => {
role: MessageRole.User,
content: messages.slice(1).reduce((acc, curr) => {
return `${acc} ${curr.message.role}: ${curr.message.content}`;
}, 'You are a helpful assistant for Elastic Observability. Assume the following message is the start of a conversation between you and a user; give this conversation a title based on this content: '),
}, 'You are a helpful assistant for Elastic Observability. Assume the following message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Here is the content:'),
},
},
],
connectorId,
stream: false,
stream: true,
signal,
});
if ('object' in response && response.object === 'chat.completion') {
const input =
response.choices[0].message?.content || `Conversation on ${conversation['@timestamp']}`;
const response = await lastValueFrom(
streamIntoObservable(stream).pipe(processOpenAiStream(), concatenateOpenAiChunks())
);
// This regular expression captures a string enclosed in single or double quotes.
// It extracts the string content without the quotes.
// Example matches:
// - "Hello, World!" => Captures: Hello, World!
// - 'Another Example' => Captures: Another Example
// - JustTextWithoutQuotes => Captures: JustTextWithoutQuotes
const match = input.match(/^["']?([^"']+)["']?$/);
const title = match ? match[1] : input;
const input = response.message?.content || '';
const updatedConversation: Conversation = merge(
{},
conversation,
{ conversation: { title } },
this.getConversationUpdateValues(new Date().toISOString())
);
await this.setTitle({ conversationId, title });
return updatedConversation;
}
return conversation;
// This regular expression captures a string enclosed in single or double quotes.
// It extracts the string content without the quotes.
// Example matches:
// - "Hello, World!" => Captures: Hello, World!
// - 'Another Example' => Captures: Another Example
// - JustTextWithoutQuotes => Captures: JustTextWithoutQuotes
const match = input.match(/^["']?([^"']+)["']?$/);
const title = match ? match[1] : input;
return title;
};
setTitle = async ({ conversationId, title }: { conversationId: string; title: string }) => {

116
x-pack/plugins/observability_ai_assistant/server/service/index.ts
View file

@ -12,20 +12,60 @@ import type { CoreSetup, CoreStart, KibanaRequest, Logger } from '@kbn/core/serv
import type { SecurityPluginStart } from '@kbn/security-plugin/server';
import { getSpaceIdFromPath } from '@kbn/spaces-plugin/common';
import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
import Ajv, { type ValidateFunction } from 'ajv';
import { once } from 'lodash';
import { KnowledgeBaseEntryRole } from '../../common/types';
import {
ContextRegistry,
KnowledgeBaseEntryRole,
RegisterContextDefinition,
} from '../../common/types';
import type { ObservabilityAIAssistantPluginStartDependencies } from '../types';
import { ChatFunctionClient } from './chat_function_client';
import { ObservabilityAIAssistantClient } from './client';
import { conversationComponentTemplate } from './conversation_component_template';
import { kbComponentTemplate } from './kb_component_template';
import { KnowledgeBaseEntryOperationType, KnowledgeBaseService } from './knowledge_base_service';
import type { ObservabilityAIAssistantResourceNames } from './types';
import type {
ChatRegistrationFunction,
FunctionHandlerRegistry,
ObservabilityAIAssistantResourceNames,
RegisterFunction,
RespondFunctionResources,
} from './types';
import { splitKbText } from './util/split_kb_text';
const ajv = new Ajv({
strict: false,
});
function getResourceName(resource: string) {
return `.kibana-observability-ai-assistant-${resource}`;
}
export function createResourceNamesMap() {
return {
componentTemplate: {
conversations: getResourceName('component-template-conversations'),
kb: getResourceName('component-template-kb'),
},
aliases: {
conversations: getResourceName('conversations'),
kb: getResourceName('kb'),
},
indexPatterns: {
conversations: getResourceName('conversations*'),
kb: getResourceName('kb*'),
},
indexTemplate: {
conversations: getResourceName('index-template-conversations'),
kb: getResourceName('index-template-kb'),
},
pipelines: {
kb: getResourceName('kb-ingest-pipeline'),
},
};
}
export const ELSER_MODEL_ID = '.elser_model_2';
export const INDEX_QUEUED_DOCUMENTS_TASK_ID = 'observabilityAIAssistant:indexQueuedDocumentsTask';
@ -46,27 +86,9 @@ export class ObservabilityAIAssistantService {
private readonly logger: Logger;
private kbService?: KnowledgeBaseService;
private readonly resourceNames: ObservabilityAIAssistantResourceNames = {
componentTemplate: {
conversations: getResourceName('component-template-conversations'),
kb: getResourceName('component-template-kb'),
},
aliases: {
conversations: getResourceName('conversations'),
kb: getResourceName('kb'),
},
indexPatterns: {
conversations: getResourceName('conversations*'),
kb: getResourceName('kb*'),
},
indexTemplate: {
conversations: getResourceName('index-template-conversations'),
kb: getResourceName('index-template-kb'),
},
pipelines: {
kb: getResourceName('kb-ingest-pipeline'),
},
};
private readonly resourceNames: ObservabilityAIAssistantResourceNames = createResourceNamesMap();
private readonly registrations: ChatRegistrationFunction[] = [];
constructor({
logger,
@ -100,6 +122,12 @@ export class ObservabilityAIAssistantService {
});
}
getKnowledgeBaseStatus() {
return this.init().then(() => {
return this.kbService!.status();
});
}
init = once(async () => {
try {
const [coreStart, pluginsStart] = await this.core.getStartServices();
@ -224,13 +252,18 @@ export class ObservabilityAIAssistantService {
}: {
request: KibanaRequest;
}): Promise<ObservabilityAIAssistantClient> {
const controller = new AbortController();
request.events.aborted$.subscribe(() => {
controller.abort();
});
const [_, [coreStart, plugins]] = await Promise.all([
this.init(),
this.core.getStartServices() as Promise<
[CoreStart, { security: SecurityPluginStart; actions: ActionsPluginStart }, unknown]
>,
]);
const user = plugins.security.authc.getCurrentUser(request);
if (!user) {
@ -255,6 +288,37 @@ export class ObservabilityAIAssistantService {
});
}
async getFunctionClient({
signal,
resources,
client,
}: {
signal: AbortSignal;
resources: RespondFunctionResources;
client: ObservabilityAIAssistantClient;
}): Promise<ChatFunctionClient> {
const contextRegistry: ContextRegistry = new Map();
const functionHandlerRegistry: FunctionHandlerRegistry = new Map();
const validators = new Map<string, ValidateFunction>();
const registerContext: RegisterContextDefinition = (context) => {
contextRegistry.set(context.name, context);
};
const registerFunction: RegisterFunction = (definition, respond) => {
validators.set(definition.name, ajv.compile(definition.parameters));
functionHandlerRegistry.set(definition.name, { definition, respond });
};
await Promise.all(
this.registrations.map((fn) =>
fn({ signal, registerContext, registerFunction, resources, client })
)
);
return new ChatFunctionClient(contextRegistry, functionHandlerRegistry, validators);
}
addToKnowledgeBase(entries: KnowledgeBaseEntryRequest[]): void {
this.init()
.then(() => {
@ -308,4 +372,8 @@ export class ObservabilityAIAssistantService {
})
);
}
register(fn: ChatRegistrationFunction) {
this.registrations.push(fn);
}
}

48
x-pack/plugins/observability_ai_assistant/server/service/types.ts
View file

@ -5,6 +5,54 @@
* 2.0.
*/
import type { FromSchema } from 'json-schema-to-ts';
import type {
CompatibleJSONSchema,
FunctionDefinition,
FunctionResponse,
Message,
RegisterContextDefinition,
} from '../../common/types';
import type { ObservabilityAIAssistantRouteHandlerResources } from '../routes/types';
import type { ObservabilityAIAssistantClient } from './client';
export type RespondFunctionResources = Pick<
ObservabilityAIAssistantRouteHandlerResources,
'context' | 'logger' | 'plugins' | 'request'
>;
type RespondFunction<TArguments, TResponse extends FunctionResponse> = (
options: {
arguments: TArguments;
messages: Message[];
connectorId: string;
},
signal: AbortSignal
) => Promise<TResponse>;
export interface FunctionHandler {
definition: FunctionDefinition;
respond: RespondFunction<any, FunctionResponse>;
}
export type RegisterFunction = <
TParameters extends CompatibleJSONSchema = any,
TResponse extends FunctionResponse = any,
TArguments = FromSchema<TParameters>
>(
definition: FunctionDefinition<TParameters>,
respond: RespondFunction<TArguments, TResponse>
) => void;
export type FunctionHandlerRegistry = Map<string, FunctionHandler>;
export type ChatRegistrationFunction = ({}: {
signal: AbortSignal;
resources: RespondFunctionResources;
client: ObservabilityAIAssistantClient;
registerFunction: RegisterFunction;
registerContext: RegisterContextDefinition;
}) => Promise<void>;
export interface ObservabilityAIAssistantResourceNames {
componentTemplate: {
conversations: string;

24
x-pack/plugins/observability_ai_assistant/server/service/util/stream_into_observable.ts Normal file
View file

@ -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.
*/
import { concatMap, filter, from, map, Observable } from 'rxjs';
import type { Readable } from 'stream';
export function streamIntoObservable(readable: Readable): Observable<string> {
let lineBuffer = '';
return from(readable).pipe(
map((chunk: Buffer) => chunk.toString('utf-8')),
map((part) => {
const lines = (lineBuffer + part).split('\n');
lineBuffer = lines.pop() || ''; // Keep the last incomplete line for the next chunk
return lines;
}),
concatMap((lines) => lines),
filter((line) => line.trim() !== '')
);
}

18
x-pack/plugins/observability_ai_assistant/server/types.ts
View file

@ -18,10 +18,22 @@ import type {
TaskManagerStartContract,
} from '@kbn/task-manager-plugin/server';
import { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
import { ObservabilityAIAssistantService } from './service';
export interface ObservabilityAIAssistantPluginSetup {
/**
* Returns a Observability AI Assistant service instance
*/
service: ObservabilityAIAssistantService;
}
export interface ObservabilityAIAssistantPluginStart {
/**
* Returns a Observability AI Assistant service instance
*/
service: ObservabilityAIAssistantService;
}
/* eslint-disable @typescript-eslint/no-empty-interface*/
export interface ObservabilityAIAssistantPluginStart {}
export interface ObservabilityAIAssistantPluginSetup {}
export interface ObservabilityAIAssistantPluginSetupDependencies {
actions: ActionsPluginSetup;
security: SecurityPluginSetup;

1
x-pack/plugins/translations/translations/fr-FR.json
View file

@ -29672,7 +29672,6 @@
"xpack.observabilityAiAssistant.conversationStartTitle": "a démarré une conversation",
"xpack.observabilityAiAssistant.couldNotFindConversationTitle": "Conversation introuvable",
"xpack.observabilityAiAssistant.emptyConversationTitle": "Nouvelle conversation",
"xpack.observabilityAiAssistant.errorCreatingConversation": "Impossible de créer une conversation",
"xpack.observabilityAiAssistant.errorSettingUpKnowledgeBase": "Impossible de configurer la base de connaissances",
"xpack.observabilityAiAssistant.errorUpdatingConversation": "Impossible de mettre à jour la conversation",
"xpack.observabilityAiAssistant.experimentalTitle": "Version d'évaluation technique",

1
x-pack/plugins/translations/translations/ja-JP.json
View file

@ -29672,7 +29672,6 @@
"xpack.observabilityAiAssistant.conversationStartTitle": "会話を開始しました",
"xpack.observabilityAiAssistant.couldNotFindConversationTitle": "会話が見つかりません",
"xpack.observabilityAiAssistant.emptyConversationTitle": "新しい会話",
"xpack.observabilityAiAssistant.errorCreatingConversation": "会話を作成できませんでした",
"xpack.observabilityAiAssistant.errorSettingUpKnowledgeBase": "ナレッジベースをセットアップできませんでした",
"xpack.observabilityAiAssistant.errorUpdatingConversation": "会話を更新できませんでした",
"xpack.observabilityAiAssistant.experimentalTitle": "テクニカルプレビュー",

1
x-pack/plugins/translations/translations/zh-CN.json
View file

@ -29669,7 +29669,6 @@
"xpack.observabilityAiAssistant.conversationStartTitle": "已开始对话",
"xpack.observabilityAiAssistant.couldNotFindConversationTitle": "未找到对话",
"xpack.observabilityAiAssistant.emptyConversationTitle": "新对话",
"xpack.observabilityAiAssistant.errorCreatingConversation": "无法创建对话",
"xpack.observabilityAiAssistant.errorSettingUpKnowledgeBase": "无法设置知识库",
"xpack.observabilityAiAssistant.errorUpdatingConversation": "无法更新对话",
"xpack.observabilityAiAssistant.experimentalTitle": "技术预览",

147
x-pack/test/observability_ai_assistant_api_integration/common/create_llm_proxy.ts Normal file
View file

@ -0,0 +1,147 @@
/*
* 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 getPort from 'get-port';
import http, { type Server } from 'http';
import { once, pull } from 'lodash';
import { createOpenAiChunk } from './create_openai_chunk';
type Request = http.IncomingMessage;
type Response = http.ServerResponse<http.IncomingMessage> & { req: http.IncomingMessage };
type RequestHandler = (request: Request, response: Response) => void;
interface RequestInterceptor {
name: string;
when: (body: string) => boolean;
}
export interface LlmResponseSimulator {
status: (code: number) => Promise<void>;
next: (
msg:
| string
| {
content?: string;
function_call?: { name: string; arguments: string };
}
) => Promise<void>;
error: (error: Error) => Promise<void>;
complete: () => Promise<void>;
write: (chunk: string) => Promise<void>;
}
export class LlmProxy {
server: Server;
interceptors: Array<RequestInterceptor & { handle: RequestHandler }> = [];
constructor(private readonly port: number) {
this.server = http
.createServer(async (request, response) => {
const interceptors = this.interceptors.concat();
const body = await new Promise<string>((resolve, reject) => {
let concatenated = '';
request.on('data', (chunk) => {
concatenated += chunk.toString();
});
request.on('close', () => {
resolve(concatenated);
});
});
while (interceptors.length) {
const interceptor = interceptors.shift()!;
if (interceptor.when(body)) {
pull(this.interceptors, interceptor);
interceptor.handle(request, response);
return;
}
}
throw new Error('No interceptors found to handle request');
})
.listen(port);
}
getPort() {
return this.port;
}
clear() {
this.interceptors.length = 0;
}
close() {
this.server.close();
}
intercept(
name: string,
when: RequestInterceptor['when']
): {
waitForIntercept: () => Promise<LlmResponseSimulator>;
} {
const waitForInterceptPromise = Promise.race([
new Promise<LlmResponseSimulator>((outerResolve, outerReject) => {
this.interceptors.push({
name,
when,
handle: (request, response) => {
function write(chunk: string) {
return new Promise<void>((resolve) => response.write(chunk, () => resolve()));
}
function end() {
return new Promise<void>((resolve) => response.end(resolve));
}
const simulator: LlmResponseSimulator = {
status: once(async (status: number) => {
response.writeHead(status, {
'Content-Type': 'text/event-stream',
'Cache-Control': 'no-cache',
Connection: 'keep-alive',
});
}),
next: (msg) => {
const chunk = createOpenAiChunk(msg);
return write(`data: ${JSON.stringify(chunk)}\n`);
},
write: (chunk: string) => {
return write(chunk);
},
complete: async () => {
await write('data: [DONE]');
await end();
},
error: async (error) => {
await write(`data: ${JSON.stringify({ error })}`);
await end();
},
};
outerResolve(simulator);
},
});
}),
new Promise<LlmResponseSimulator>((_, reject) => {
setTimeout(() => reject(new Error('Operation timed out')), 5000);
}),
]);
return {
waitForIntercept: () => waitForInterceptPromise,
};
}
}
export async function createLlmProxy() {
const port = await getPort({ port: getPort.makeRange(9000, 9100) });
return new LlmProxy(port);
}

27
x-pack/test/observability_ai_assistant_api_integration/common/create_openai_chunk.ts Normal file
View file

@ -0,0 +1,27 @@
/*
* 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 { CreateChatCompletionResponseChunk } from '@kbn/observability-ai-assistant-plugin/common/types';
import { v4 } from 'uuid';
export function createOpenAiChunk(
msg: string | { content?: string; function_call?: { name: string; arguments?: string } }
): CreateChatCompletionResponseChunk {
msg = typeof msg === 'string' ? { content: msg } : msg;
return {
id: v4(),
object: 'chat.completion.chunk',
created: 0,
model: 'gpt-4',
choices: [
{
delta: msg,
},
],
};
}

256
x-pack/test/observability_ai_assistant_api_integration/tests/complete/complete.spec.ts Normal file
View file

@ -0,0 +1,256 @@
/*
* 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 { Response } from 'supertest';
import { MessageRole, type Message } from '@kbn/observability-ai-assistant-plugin/common';
import { omit } from 'lodash';
import { PassThrough } from 'stream';
import expect from '@kbn/expect';
import {
ConversationCreateEvent,
StreamingChatResponseEvent,
StreamingChatResponseEventType,
} from '@kbn/observability-ai-assistant-plugin/common/conversation_complete';
import { CreateChatCompletionRequest } from 'openai';
import { createLlmProxy, LlmProxy } from '../../common/create_llm_proxy';
import { createOpenAiChunk } from '../../common/create_openai_chunk';
import { FtrProviderContext } from '../../common/ftr_provider_context';
export default function ApiTest({ getService }: FtrProviderContext) {
const supertest = getService('supertest');
const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantAPIClient');
const COMPLETE_API_URL = `/internal/observability_ai_assistant/chat/complete`;
const messages: Message[] = [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'You are a helpful assistant',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Good morning!',
},
},
];
describe('Complete', () => {
let proxy: LlmProxy;
let connectorId: string;
before(async () => {
proxy = await createLlmProxy();
const response = await supertest
.post('/api/actions/connector')
.set('kbn-xsrf', 'foo')
.send({
name: 'OpenAI',
connector_type_id: '.gen-ai',
config: {
apiProvider: 'OpenAI',
apiUrl: `http://localhost:${proxy.getPort()}`,
},
secrets: {
apiKey: 'my-api-key',
},
})
.expect(200);
connectorId = response.body.id;
});
after(async () => {
await supertest
.delete(`/api/actions/connector/${connectorId}`)
.set('kbn-xsrf', 'foo')
.expect(204);
proxy.close();
});
it('returns a streaming response from the server', async () => {
const interceptor = proxy.intercept('conversation', () => true);
const receivedChunks: any[] = [];
const passThrough = new PassThrough();
supertest
.post(COMPLETE_API_URL)
.set('kbn-xsrf', 'foo')
.send({
messages,
connectorId,
persist: false,
})
.pipe(passThrough);
passThrough.on('data', (chunk) => {
receivedChunks.push(chunk.toString());
});
const simulator = await interceptor.waitForIntercept();
await simulator.status(200);
const chunk = JSON.stringify(createOpenAiChunk('Hello'));
await simulator.write(`data: ${chunk.substring(0, 10)}`);
await simulator.write(`${chunk.substring(10)}\n`);
await simulator.complete();
await new Promise<void>((resolve) => passThrough.on('end', () => resolve()));
const parsedChunks = receivedChunks
.join('')
.split('\n')
.map((line) => line.trim())
.filter(Boolean)
.map((line) => JSON.parse(line) as StreamingChatResponseEvent);
expect(parsedChunks.length).to.be(2);
expect(omit(parsedChunks[0], 'id')).to.eql({
type: StreamingChatResponseEventType.ChatCompletionChunk,
message: {
content: 'Hello',
},
});
expect(omit(parsedChunks[1], 'id', 'message.@timestamp')).to.eql({
type: StreamingChatResponseEventType.MessageAdd,
message: {
message: {
content: 'Hello',
role: MessageRole.Assistant,
function_call: {
name: '',
arguments: '',
trigger: MessageRole.Assistant,
},
},
},
});
});
describe('when creating a new conversation', async () => {
let lines: StreamingChatResponseEvent[];
before(async () => {
const titleInterceptor = proxy.intercept(
'title',
(body) => (JSON.parse(body) as CreateChatCompletionRequest).messages.length === 1
);
const conversationInterceptor = proxy.intercept(
'conversation',
(body) => (JSON.parse(body) as CreateChatCompletionRequest).messages.length !== 1
);
const responsePromise = new Promise<Response>((resolve, reject) => {
supertest
.post(COMPLETE_API_URL)
.set('kbn-xsrf', 'foo')
.send({
messages,
connectorId,
persist: true,
})
.end((err, response) => {
if (err) {
return reject(err);
}
return resolve(response);
});
});
const [conversationSimulator, titleSimulator] = await Promise.all([
conversationInterceptor.waitForIntercept(),
titleInterceptor.waitForIntercept(),
]);
await titleSimulator.status(200);
await titleSimulator.next('My generated title');
await titleSimulator.complete();
await conversationSimulator.status(200);
await conversationSimulator.next('Hello');
await conversationSimulator.next(' again');
await conversationSimulator.complete();
const response = await responsePromise;
lines = String(response.body)
.split('\n')
.map((line) => line.trim())
.filter(Boolean)
.map((line) => JSON.parse(line) as StreamingChatResponseEvent);
});
it('creates a new conversation', async () => {
expect(omit(lines[0], 'id')).to.eql({
type: StreamingChatResponseEventType.ChatCompletionChunk,
message: {
content: 'Hello',
},
});
expect(omit(lines[1], 'id')).to.eql({
type: StreamingChatResponseEventType.ChatCompletionChunk,
message: {
content: ' again',
},
});
expect(omit(lines[2], 'id', 'message.@timestamp')).to.eql({
type: StreamingChatResponseEventType.MessageAdd,
message: {
message: {
content: 'Hello again',
function_call: {
arguments: '',
name: '',
trigger: MessageRole.Assistant,
},
role: MessageRole.Assistant,
},
},
});
expect(omit(lines[3], 'conversation.id', 'conversation.last_updated')).to.eql({
type: StreamingChatResponseEventType.ConversationCreate,
conversation: {
title: 'My generated title',
},
});
});
after(async () => {
const createdConversationId = lines.filter(
(line): line is ConversationCreateEvent =>
line.type === StreamingChatResponseEventType.ConversationCreate
)[0]?.conversation.id;
await observabilityAIAssistantAPIClient
.writeUser({
endpoint: 'DELETE /internal/observability_ai_assistant/conversation/{conversationId}',
params: {
path: {
conversationId: createdConversationId,
},
},
})
.expect(200);
});
});
// todo
it.skip('updates an existing conversation', async () => {});
// todo
it.skip('executes a function', async () => {});
});
}

46
x-pack/test/observability_ai_assistant_api_integration/tests/functions/elasticsearch.spec.ts
View file

@ -1,46 +0,0 @@
/*
* 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 expect from '@kbn/expect';
import type { FtrProviderContext } from '../../common/ftr_provider_context';
export default function ApiTest({ getService }: FtrProviderContext) {
const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantAPIClient');
describe('Functions: elasticsearch', () => {
it('executes a search request', async () => {
const response = await observabilityAIAssistantAPIClient
.readUser({
endpoint: 'POST /internal/observability_ai_assistant/functions/elasticsearch',
params: {
body: {
method: 'GET',
path: '_all/_search',
body: {
query: {
bool: {
filter: [
{
term: {
matches_no_docs: 'true',
},
},
],
},
},
track_total_hits: false,
},
},
},
})
.expect(200);
expect((response.body as any).hits.hits).to.eql([]);
expect((response.body as any).hits.total).to.eql(undefined);
});
});
}